From d9f137f23b783af13da32968642bc72dc908784e Mon Sep 17 00:00:00 2001 From: Intron7 Date: Mon, 7 Sep 2026 18:47:29 +0200 Subject: [PATCH 1/4] add tile minor axis sweep kernel Signed-off-by: Intron7 --- CMakeLists.txt | 18 +- docs/release-notes/0.17.0.md | 1 + src/rapids_singlecell/_cuda/__init__.py | 16 +- src/rapids_singlecell/_cuda/aggr/aggr.cu | 44 +- .../_cuda/aggr/kernels_aggr.cuh | 74 ++ src/rapids_singlecell/_cuda/hvg/hvg.cu | 74 ++ .../_cuda/ligrec/kernels_ligrec.cuh | 107 +-- src/rapids_singlecell/_cuda/ligrec/ligrec.cu | 66 +- .../_cuda/mean_var/kernels_mv.cuh | 56 +- .../_cuda/mean_var/mean_var.cu | 69 +- src/rapids_singlecell/_cuda/minor_tiles.cuh | 644 ++++++++++++++++++ .../_cuda/nanmean/kernels_nanmean.cuh | 85 ++- .../_cuda/nanmean/nanmean.cu | 50 +- src/rapids_singlecell/_cuda/nb_types.h | 37 + .../_cuda/pr/kernels_pr_hvg.cuh | 23 - src/rapids_singlecell/_cuda/pr/pr.cu | 29 +- src/rapids_singlecell/_cuda/qc/kernels_qc.cuh | 133 ++-- src/rapids_singlecell/_cuda/qc/qc.cu | 133 ++-- .../_cuda/qc_dask/kernels_qcd.cuh | 37 +- .../_cuda/qc_dask/qc_kernels_dask.cu | 38 +- .../_cuda/spca/kernels_spca.cuh | 12 - src/rapids_singlecell/_cuda/spca/spca.cu | 38 +- src/rapids_singlecell/_utils/_sparse_rows.py | 25 + src/rapids_singlecell/get/_aggregated.py | 9 +- .../preprocessing/_hvg/_pearson_residuals.py | 5 +- .../preprocessing/_hvg/_seurat_v3.py | 30 +- src/rapids_singlecell/preprocessing/_qc.py | 17 +- .../preprocessing/_sparse_pca/_helper.py | 6 +- src/rapids_singlecell/preprocessing/_utils.py | 41 +- src/rapids_singlecell/squidpy_gpu/_ligrec.py | 10 +- src/rapids_singlecell/tools/_utils.py | 56 +- 31 files changed, 1469 insertions(+), 514 deletions(-) create mode 100644 src/rapids_singlecell/_cuda/minor_tiles.cuh create mode 100644 src/rapids_singlecell/_utils/_sparse_rows.py diff --git a/CMakeLists.txt b/CMakeLists.txt index a4855d792..c2f3e6ee1 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -80,33 +80,33 @@ endfunction() if (RSC_BUILD_EXTENSIONS) # CUDA modules - add_nb_cuda_module(_mean_var_cuda src/rapids_singlecell/_cuda/mean_var/mean_var.cu) + add_rmm_cuda_module(_mean_var_cuda src/rapids_singlecell/_cuda/mean_var/mean_var.cu) add_nb_cuda_module(_sparse2dense_cuda src/rapids_singlecell/_cuda/sparse2dense/sparse2dense.cu) add_nb_cuda_module(_jaccard_cuda src/rapids_singlecell/_cuda/jaccard/jaccard.cu) add_nb_cuda_module(_scale_cuda src/rapids_singlecell/_cuda/scale/scale.cu) - add_nb_cuda_module(_qc_cuda src/rapids_singlecell/_cuda/qc/qc.cu) - add_nb_cuda_module(_qc_dask_cuda src/rapids_singlecell/_cuda/qc_dask/qc_kernels_dask.cu) + add_rmm_cuda_module(_qc_cuda src/rapids_singlecell/_cuda/qc/qc.cu) + add_rmm_cuda_module(_qc_dask_cuda src/rapids_singlecell/_cuda/qc_dask/qc_kernels_dask.cu) add_nb_cuda_module(_bbknn_cuda src/rapids_singlecell/_cuda/bbknn/bbknn.cu) add_nb_cuda_module(_norm_cuda src/rapids_singlecell/_cuda/norm/norm.cu) add_nb_cuda_module(_gmm_cuda src/rapids_singlecell/_cuda/gmm/gmm.cu) target_link_libraries(_gmm_cuda PRIVATE CUDA::cublas) target_link_libraries(_gmm_cuda PRIVATE CUDA::cusolver) add_nb_cuda_module(_mixscale_cuda src/rapids_singlecell/_cuda/mixscale/mixscale.cu) - add_nb_cuda_module(_pr_cuda src/rapids_singlecell/_cuda/pr/pr.cu) + add_rmm_cuda_module(_pr_cuda src/rapids_singlecell/_cuda/pr/pr.cu) add_nb_cuda_module(_nn_descent_cuda src/rapids_singlecell/_cuda/nn_descent/nn_descent.cu) add_nb_cuda_module(_aucell_cuda src/rapids_singlecell/_cuda/aucell/aucell.cu) - add_nb_cuda_module(_nanmean_cuda src/rapids_singlecell/_cuda/nanmean/nanmean.cu) + add_rmm_cuda_module(_nanmean_cuda src/rapids_singlecell/_cuda/nanmean/nanmean.cu) add_nb_cuda_module(_autocorr_cuda src/rapids_singlecell/_cuda/autocorr/autocorr.cu) add_nb_cuda_module(_cooc_cuda src/rapids_singlecell/_cuda/cooc/cooc.cu) - add_nb_cuda_module(_aggr_cuda src/rapids_singlecell/_cuda/aggr/aggr.cu) - add_nb_cuda_module(_spca_cuda src/rapids_singlecell/_cuda/spca/spca.cu) - add_nb_cuda_module(_ligrec_cuda src/rapids_singlecell/_cuda/ligrec/ligrec.cu) + add_rmm_cuda_module(_aggr_cuda src/rapids_singlecell/_cuda/aggr/aggr.cu) + add_rmm_cuda_module(_spca_cuda src/rapids_singlecell/_cuda/spca/spca.cu) + add_rmm_cuda_module(_ligrec_cuda src/rapids_singlecell/_cuda/ligrec/ligrec.cu) add_nb_cuda_module(_pv_cuda src/rapids_singlecell/_cuda/pv/pv.cu) add_nb_cuda_module(_edistance_cuda src/rapids_singlecell/_cuda/edistance/edistance.cu) add_nb_cuda_module(_sinkhorn_cuda src/rapids_singlecell/_cuda/sinkhorn/sinkhorn.cu) add_nb_cuda_module(_guide_assignment_cuda src/rapids_singlecell/_cuda/guide_assignment/guide_assignment.cu) add_nb_cuda_module(_pseudobulk_cuda src/rapids_singlecell/_cuda/pseudobulk/pseudobulk.cu) - add_nb_cuda_module(_hvg_cuda src/rapids_singlecell/_cuda/hvg/hvg.cu) + add_rmm_cuda_module(_hvg_cuda src/rapids_singlecell/_cuda/hvg/hvg.cu) add_nb_cuda_module(_kde_cuda src/rapids_singlecell/_cuda/kde/kde.cu) add_rmm_cuda_module(_wilcoxon_cuda src/rapids_singlecell/_cuda/wilcoxon/wilcoxon.cu) add_rmm_cuda_module(_wilcoxon_sparse_cuda src/rapids_singlecell/_cuda/wilcoxon/wilcoxon_sparse.cu) diff --git a/docs/release-notes/0.17.0.md b/docs/release-notes/0.17.0.md index 6c673e71e..63eaafa69 100644 --- a/docs/release-notes/0.17.0.md +++ b/docs/release-notes/0.17.0.md @@ -37,6 +37,7 @@ ``` * Speed up {func}`~rapids_singlecell.tl.umap` and cut its peak memory by no longer moving the representation to the GPU. {pr}`768` {smaller}`S Dicks` * Keep the dataset on device for unbatched ``all_neighbors`` builds in {func}`~rapids_singlecell.pp.neighbors` {pr}`769` {smaller}`S Dicks` +* Speed up the sparse column reductions behind {func}`~rapids_singlecell.pp.calculate_qc_metrics`, {func}`~rapids_singlecell.pp.highly_variable_genes`, {func}`~rapids_singlecell.get.aggregate`, {func}`~rapids_singlecell.gr.ligrec` and the mean/variance helpers 2 to 4x with a shared-memory tile sweep {pr}`792` {smaller}`S Dicks` ```{rubric} Misc ``` diff --git a/src/rapids_singlecell/_cuda/__init__.py b/src/rapids_singlecell/_cuda/__init__.py index d109ecfbd..2f8855127 100644 --- a/src/rapids_singlecell/_cuda/__init__.py +++ b/src/rapids_singlecell/_cuda/__init__.py @@ -72,7 +72,21 @@ def _preload_rapids_runtime_libs() -> None: # CuPy-backed allocator injected here, so temporaries land on the caller's # current device resource (RMM pool / UVM aware) without linking librmm. _SCRATCH_MODULES = frozenset( - {"_wilcoxon_cuda", "_wilcoxon_sparse_cuda", "_rank_stream_cuda"} + { + "_wilcoxon_cuda", + "_wilcoxon_sparse_cuda", + "_rank_stream_cuda", + # minor-axis tile sweep (minor_tiles.cuh): flag, grouping, block ranges + "_mean_var_cuda", + "_nanmean_cuda", + "_qc_cuda", + "_qc_dask_cuda", + "_hvg_cuda", + "_pr_cuda", + "_spca_cuda", + "_aggr_cuda", + "_ligrec_cuda", + } ) _scratch_allocator = None diff --git a/src/rapids_singlecell/_cuda/aggr/aggr.cu b/src/rapids_singlecell/_cuda/aggr/aggr.cu index bc1b0a164..c62c30096 100644 --- a/src/rapids_singlecell/_cuda/aggr/aggr.cu +++ b/src/rapids_singlecell/_cuda/aggr/aggr.cu @@ -45,27 +45,29 @@ constexpr int BLOCK_SIZE_DENSE = 256; "provided"); \ } +// CSC keeps the per-gene block kernel (genes are the compressed axis). CSR +// runs the grouped tile sweep and returns whether unsorted rows were detected. template -static inline void launch_sparse_aggr(bool is_csc, const IdxT* indptr, - const IdxT* index, const T* data, - double* out_sum, double* out_count, - double* out_sqsum, const int* cats, - const bool* mask, size_t n_cells, - size_t n_genes, cudaStream_t stream) { - dim3 block(BLOCK_SIZE_SPARSE); +static inline bool launch_sparse_aggr( + bool is_csc, const IdxT* indptr, const IdxT* index, const T* data, + double* out_sum, double* out_count, double* out_sqsum, const int* cats, + const bool* mask, size_t n_cells, size_t n_genes, int n_groups, + long long nnz, bool assume_unsorted, cudaStream_t stream) { if (is_csc) { + dim3 block(BLOCK_SIZE_SPARSE); dim3 grid((unsigned)n_genes); csc_aggr_kernel<<>>( indptr, index, data, out_sum, out_count, out_sqsum, cats, mask, n_cells, n_genes); CUDA_CHECK_LAST_ERROR(csc_aggr_kernel); - } else { - dim3 grid((unsigned)n_cells); - csr_aggr_kernel<<>>( - indptr, index, data, out_sum, out_count, out_sqsum, cats, mask, - n_cells, n_genes); - CUDA_CHECK_LAST_ERROR(csr_aggr_kernel); + return false; } + RmmScratchPool pool; + const GroupedRows groups = + build_grouped_rows(pool, cats, mask, (int)n_cells, n_groups, stream); + AggrOp op{data, out_sum, out_count, out_sqsum, n_genes, 0}; + return minor_reduce(indptr, index, op, (int)n_cells, (int)n_genes, + nnz, assume_unsorted, stream, &groups); } template @@ -124,23 +126,30 @@ void def_sparse_aggr(nb::module_& m) { std::optional> out_sqsum, gpu_array_c cats, gpu_array_c mask, size_t n_cells, size_t n_genes, - bool is_csc, std::uintptr_t stream) { + bool is_csc, bool assume_unsorted, std::uintptr_t stream) { double* ps = out_sum ? out_sum->data() : nullptr; double* pc = out_count ? out_count->data() : nullptr; double* pq = out_sqsum ? out_sqsum->data() : nullptr; int active = (ps ? AGGR_SUM : 0) | (pc ? AGGR_COUNT : 0) | (pq ? AGGR_SQSUM : 0); + const int n_groups = out_sum ? (int)out_sum->shape(0) + : out_count ? (int)out_count->shape(0) + : out_sqsum ? (int)out_sqsum->shape(0) + : 0; + bool unsorted = false; #define LAUNCH(M) \ - launch_sparse_aggr( \ + unsorted = launch_sparse_aggr( \ is_csc, indptr.data(), index.data(), data.data(), ps, pc, pq, \ - cats.data(), mask.data(), n_cells, n_genes, (cudaStream_t)stream) + cats.data(), mask.data(), n_cells, n_genes, n_groups, \ + (long long)data.shape(0), assume_unsorted, (cudaStream_t)stream) AGGR_DISPATCH_MASK(active, LAUNCH); #undef LAUNCH + return unsorted; }, "indptr"_a, "index"_a, "data"_a, nb::kw_only(), "out_sum"_a = nb::none(), "out_count"_a = nb::none(), "out_sqsum"_a = nb::none(), "cats"_a, "mask"_a, "n_cells"_a, - "n_genes"_a, "is_csc"_a, "stream"_a = 0); + "n_genes"_a, "is_csc"_a, "assume_unsorted"_a = false, "stream"_a = 0); } template @@ -237,4 +246,5 @@ void register_bindings(nb::module_& m) { NB_MODULE(_aggr_cuda, m) { REGISTER_GPU_BINDINGS(register_bindings, m); + register_scratch_allocator(m); } diff --git a/src/rapids_singlecell/_cuda/aggr/kernels_aggr.cuh b/src/rapids_singlecell/_cuda/aggr/kernels_aggr.cuh index dcd187ba9..94458d921 100644 --- a/src/rapids_singlecell/_cuda/aggr/kernels_aggr.cuh +++ b/src/rapids_singlecell/_cuda/aggr/kernels_aggr.cuh @@ -1,6 +1,7 @@ #pragma once #include +#include "../minor_tiles.cuh" // Compile-time selector for which raw accumulators a kernel writes. Combined as // a bitmask so each kernel instantiation emits only the atomicAdds (and only @@ -11,6 +12,79 @@ constexpr int AGGR_COUNT = 2; // count of nonzero entries constexpr int AGGR_SQSUM = 4; // sum of squared values // sparse -> dense aggregate (CSR by cells), mask per cell, cats per cell + +// Grouped CSR aggregate: per (group, gene) sum / count / sum of squares via the +// tile sweep (see minor_tiles.cuh); rows are pre-sorted by group so a block +// accumulates for one group. Layout: one double plane per requested MASK bit. +template +struct AggrOp { + const T* data; + double* out_sum; + double* out_count; + double* out_sqsum; + size_t n_genes; + int tile_size; + static constexpr int n_planes = ((MASK & AGGR_SUM) ? 1 : 0) + + ((MASK & AGGR_COUNT) ? 1 : 0) + + ((MASK & AGGR_SQSUM) ? 1 : 0); + static constexpr int i_sum = 0; + static constexpr int i_count = (MASK & AGGR_SUM) ? 1 : 0; + static constexpr int i_sqsum = i_count + ((MASK & AGGR_COUNT) ? 1 : 0); + static constexpr size_t bytes_per_col = n_planes * sizeof(double); + static constexpr bool needs_rows = true; + __device__ double* plane(char* acc, int i) const { + return reinterpret_cast(acc) + (size_t)i * tile_size; + } + __device__ bool row_active(int) const { + return true; + } + __device__ void zero_col(char* acc, int g, int) const { + for (int i = 0; i < n_planes; ++i) plane(acc, i)[g] = 0.0; + } + __device__ void add(char* acc, long long q, int g) const { + const double v = static_cast(data[q]); + if constexpr (MASK & AGGR_SUM) atomicAdd(&plane(acc, i_sum)[g], v); + if constexpr (MASK & AGGR_COUNT) + atomicAdd(&plane(acc, i_count)[g], 1.0); + if constexpr (MASK & AGGR_SQSUM) + atomicAdd(&plane(acc, i_sqsum)[g], v * v); + } + __device__ void flush_col(const char* acc, int group, int col, + int g) const { + char* a = const_cast(acc); + bool any = false; + for (int i = 0; i < n_planes; ++i) any |= plane(a, i)[g] != 0.0; + if (!any) return; + const size_t idx = (size_t)group * n_genes + (size_t)col; + if constexpr (MASK & AGGR_SUM) + atomicAdd(&out_sum[idx], plane(a, i_sum)[g]); + if constexpr (MASK & AGGR_COUNT) + atomicAdd(&out_count[idx], plane(a, i_count)[g]); + if constexpr (MASK & AGGR_SQSUM) + atomicAdd(&out_sqsum[idx], plane(a, i_sqsum)[g]); + } + __device__ void add_global(long long q, int col, int group) const { + const size_t idx = (size_t)group * n_genes + (size_t)col; + const double v = static_cast(data[q]); + if constexpr (MASK & AGGR_SUM) atomicAdd(&out_sum[idx], v); + if constexpr (MASK & AGGR_COUNT) atomicAdd(&out_count[idx], 1.0); + if constexpr (MASK & AGGR_SQSUM) atomicAdd(&out_sqsum[idx], v * v); + } + void zero_outputs(int minor, int n_groups, cudaStream_t stream) const { + const size_t bytes = (size_t)minor * n_groups * sizeof(double); + if constexpr (MASK & AGGR_SUM) + cudaMemsetAsync(out_sum, 0, bytes, stream); + if constexpr (MASK & AGGR_COUNT) + cudaMemsetAsync(out_count, 0, bytes, stream); + if constexpr (MASK & AGGR_SQSUM) + cudaMemsetAsync(out_sqsum, 0, bytes, stream); + } +}; + +// sparse -> dense aggregate (CSR by cells), mask per cell, cats per cell. +// Kept for the host-streaming path (rank_stream.cu): it stages row blocks and +// overlaps copies with compute, so the per-block sync of the tile sweep would +// cost it more than the atomics do. template __global__ void csr_aggr_kernel( const IdxT* __restrict__ indptr, const IdxT* __restrict__ index, diff --git a/src/rapids_singlecell/_cuda/hvg/hvg.cu b/src/rapids_singlecell/_cuda/hvg/hvg.cu index e7457f690..dd8c65be1 100644 --- a/src/rapids_singlecell/_cuda/hvg/hvg.cu +++ b/src/rapids_singlecell/_cuda/hvg/hvg.cu @@ -1,4 +1,5 @@ #include +#include "../minor_tiles.cuh" #include "../nb_types.h" using namespace nb::literals; @@ -32,6 +33,73 @@ static void launch_expected_zeros(const T* scaled_means, const T* total_counts, CUDA_CHECK_LAST_ERROR(expected_zeros_kernel); } +/// Per-column sum and sum-of-squares of min(value, clip[col]) for seurat_v3 +/// (see minor_tiles.cuh). Layout: double sq-sums, double sums, double clips. +template +struct ClipSumOp { + const T* data; + const double* clip; + double* sq_sum; + double* sum; + int tile_size; + static constexpr size_t bytes_per_col = 3 * sizeof(double); + static constexpr bool needs_rows = false; + __device__ bool row_active(int) const { + return true; + } + __device__ void zero_col(char* acc, int g, int col) const { + double* s = reinterpret_cast(acc); + s[g] = 0.0; + s[tile_size + g] = 0.0; + s[2 * tile_size + g] = clip[col]; + } + __device__ void add(char* acc, long long q, int g) const { + double* s = reinterpret_cast(acc); + const double e = + fmin(static_cast(data[q]), s[2 * tile_size + g]); + atomicAdd(&s[g], e * e); + atomicAdd(&s[tile_size + g], e); + } + __device__ void flush_col(const char* acc, int, int col, int g) const { + const double* s = reinterpret_cast(acc); + if (s[g] != 0.0 || s[tile_size + g] != 0.0) { + atomicAdd(&sq_sum[col], s[g]); + atomicAdd(&sum[col], s[tile_size + g]); + } + } + __device__ void add_global(long long q, int col, int) const { + const double e = fmin(static_cast(data[q]), clip[col]); + atomicAdd(&sq_sum[col], e * e); + atomicAdd(&sum[col], e); + } + void zero_outputs(int minor, int, cudaStream_t stream) const { + cudaMemsetAsync(sq_sum, 0, (size_t)minor * sizeof(double), stream); + cudaMemsetAsync(sum, 0, (size_t)minor * sizeof(double), stream); + } +}; + +// Returns whether unsorted rows were detected. +template +void def_clip_square_sum(nb::module_& m) { + m.def( + "clip_square_sum", + [](gpu_array_c indptr, + gpu_array_c indices, + gpu_array_c data, + gpu_array_c clip_val, + gpu_array_c sq_sum, gpu_array_c sum, + bool assume_unsorted, std::uintptr_t stream) { + ClipSumOp op{data.data(), clip_val.data(), sq_sum.data(), + sum.data(), 0}; + return minor_reduce( + indptr.data(), indices.data(), op, (int)indptr.shape(0) - 1, + (int)sum.shape(0), (long long)data.shape(0), assume_unsorted, + (cudaStream_t)stream); + }, + "indptr"_a, "indices"_a, "data"_a, nb::kw_only(), "clip_val"_a, + "sq_sum"_a, "sum"_a, "assume_unsorted"_a = false, "stream"_a = 0); +} + template void def_expected_zeros(nb::module_& m) { m.def( @@ -52,8 +120,14 @@ template void register_bindings(nb::module_& m) { def_expected_zeros(m); def_expected_zeros(m); + + def_clip_square_sum(m); + def_clip_square_sum(m); + def_clip_square_sum(m); + def_clip_square_sum(m); } NB_MODULE(_hvg_cuda, m) { REGISTER_GPU_BINDINGS(register_bindings, m); + register_scratch_allocator(m); } diff --git a/src/rapids_singlecell/_cuda/ligrec/kernels_ligrec.cuh b/src/rapids_singlecell/_cuda/ligrec/kernels_ligrec.cuh index 0ee79200f..edbee1096 100644 --- a/src/rapids_singlecell/_cuda/ligrec/kernels_ligrec.cuh +++ b/src/rapids_singlecell/_cuda/ligrec/kernels_ligrec.cuh @@ -1,6 +1,7 @@ #pragma once #include +#include "../minor_tiles.cuh" template __global__ void sum_and_count_dense_kernel(const T* __restrict__ data, @@ -27,30 +28,6 @@ __global__ void sum_and_count_dense_kernel(const T* __restrict__ data, } } -template -__global__ void sum_and_count_sparse_kernel(const IdxT* __restrict__ indptr, - const IdxT* __restrict__ index, - const T* __restrict__ data, - const int* __restrict__ clusters, - T* __restrict__ sum_gt0, - int* __restrict__ count_gt0, - int nrows, int n_cls) { - int cell = blockDim.x * blockIdx.x + threadIdx.x; - if (cell >= nrows) return; - IdxT start_idx = indptr[cell]; - IdxT stop_idx = indptr[cell + 1]; - int cluster = clusters[cell]; - for (IdxT gene = start_idx; gene < stop_idx; gene++) { - T value = data[gene]; - IdxT gene_number = index[gene]; - if (value > (T)0) { - long long out_idx = (long long)gene_number * n_cls + cluster; - atomicAdd(&sum_gt0[out_idx], value); - atomicAdd(&count_gt0[out_idx], 1); - } - } -} - template __global__ void mean_dense_kernel(const T* __restrict__ data, const int* __restrict__ clusters, @@ -70,28 +47,6 @@ __global__ void mean_dense_kernel(const T* __restrict__ data, } } -template -__global__ void mean_sparse_kernel(const IdxT* __restrict__ indptr, - const IdxT* __restrict__ index, - const T* __restrict__ data, - const int* __restrict__ clusters, - T* __restrict__ sum_gt0, int nrows, - int n_cls) { - int cell = blockDim.x * blockIdx.x + threadIdx.x; - if (cell >= nrows) return; - IdxT start_idx = indptr[cell]; - IdxT stop_idx = indptr[cell + 1]; - int cluster = clusters[cell]; - for (IdxT gene = start_idx; gene < stop_idx; gene++) { - T value = data[gene]; - IdxT gene_number = index[gene]; - if (value > (T)0) { - long long out_idx = (long long)gene_number * n_cls + cluster; - atomicAdd(&sum_gt0[out_idx], value); - } - } -} - template __global__ void elementwise_diff_kernel(T* __restrict__ g_cluster, const T* __restrict__ total_counts, @@ -180,3 +135,63 @@ __global__ void res_mean_kernel(const int* __restrict__ interactions, } } } + +/// Per (gene, cluster) sum and optionally count of positive values via the +/// grouped tile sweep (see minor_tiles.cuh); output is gene-major with the +/// cluster fastest. Layout: double sums, then int counts. +template +struct LigrecOp { + const T* data; + T* sum; + int* count; // unused when !WITH_COUNT + int n_cls; + int tile_size; + static constexpr size_t bytes_per_col = + sizeof(double) + (WITH_COUNT ? sizeof(int) : 0); + static constexpr bool needs_rows = true; + __device__ double* s_sum(char* acc) const { + return reinterpret_cast(acc); + } + __device__ int* s_cnt(char* acc) const { + return reinterpret_cast(acc + (size_t)tile_size * sizeof(double)); + } + __device__ bool row_active(int) const { + return true; + } + __device__ void zero_col(char* acc, int g, int) const { + s_sum(acc)[g] = 0.0; + if constexpr (WITH_COUNT) s_cnt(acc)[g] = 0; + } + __device__ void add(char* acc, long long q, int g) const { + const T v = data[q]; + if (v > (T)0) { + atomicAdd(&s_sum(acc)[g], static_cast(v)); + if constexpr (WITH_COUNT) atomicAdd(&s_cnt(acc)[g], 1); + } + } + __device__ void flush_col(const char* acc, int group, int col, + int g) const { + char* a = const_cast(acc); + const double s = s_sum(a)[g]; + const long long idx = (long long)col * n_cls + group; + if (s != 0.0) atomicAdd(&sum[idx], static_cast(s)); + if constexpr (WITH_COUNT) { + const int c = s_cnt(a)[g]; + if (c != 0) atomicAdd(&count[idx], c); + } + } + __device__ void add_global(long long q, int col, int group) const { + const T v = data[q]; + if (v > (T)0) { + const long long idx = (long long)col * n_cls + group; + atomicAdd(&sum[idx], v); + if constexpr (WITH_COUNT) atomicAdd(&count[idx], 1); + } + } + void zero_outputs(int minor, int n_groups, cudaStream_t stream) const { + cudaMemsetAsync(sum, 0, (size_t)minor * n_groups * sizeof(T), stream); + if constexpr (WITH_COUNT) + cudaMemsetAsync(count, 0, (size_t)minor * n_groups * sizeof(int), + stream); + } +}; diff --git a/src/rapids_singlecell/_cuda/ligrec/ligrec.cu b/src/rapids_singlecell/_cuda/ligrec/ligrec.cu index fd79a16f1..c86f64e11 100644 --- a/src/rapids_singlecell/_cuda/ligrec/ligrec.cu +++ b/src/rapids_singlecell/_cuda/ligrec/ligrec.cu @@ -5,7 +5,6 @@ using namespace nb::literals; -constexpr int SPARSE_BLOCK_SIZE = 32; constexpr int DENSE_BLOCK_DIM = 32; template @@ -21,17 +20,19 @@ static inline void launch_sum_count_dense(const T* data, const int* clusters, CUDA_CHECK_LAST_ERROR(sum_and_count_dense_kernel); } +// Grouped tile sweep over cells sorted by cluster. Returns whether unsorted +// rows were detected. template -static inline void launch_sum_count_sparse(const IdxT* indptr, - const IdxT* index, const T* data, - const int* clusters, T* sum, - int* count, int rows, int ncls, - cudaStream_t stream) { - dim3 block(SPARSE_BLOCK_SIZE); - dim3 grid((rows + SPARSE_BLOCK_SIZE - 1) / SPARSE_BLOCK_SIZE); - sum_and_count_sparse_kernel<<>>( - indptr, index, data, clusters, sum, count, rows, ncls); - CUDA_CHECK_LAST_ERROR(sum_and_count_sparse_kernel); +static inline bool launch_sum_count_sparse( + const IdxT* indptr, const IdxT* index, const T* data, const int* clusters, + T* sum, int* count, int rows, int ncls, int n_genes, long long nnz, + bool assume_unsorted, cudaStream_t stream) { + RmmScratchPool pool; + const GroupedRows groups = + build_grouped_rows(pool, clusters, nullptr, rows, ncls, stream); + LigrecOp op{data, sum, count, ncls, 0}; + return minor_reduce(indptr, index, op, rows, n_genes, nnz, + assume_unsorted, stream, &groups); } template @@ -47,14 +48,17 @@ static inline void launch_mean_dense(const T* data, const int* clusters, T* g, } template -static inline void launch_mean_sparse(const IdxT* indptr, const IdxT* index, +static inline bool launch_mean_sparse(const IdxT* indptr, const IdxT* index, const T* data, const int* clusters, T* g, - int rows, int ncls, cudaStream_t stream) { - dim3 block(SPARSE_BLOCK_SIZE); - dim3 grid((rows + SPARSE_BLOCK_SIZE - 1) / SPARSE_BLOCK_SIZE); - mean_sparse_kernel<<>>( - indptr, index, data, clusters, g, rows, ncls); - CUDA_CHECK_LAST_ERROR(mean_sparse_kernel); + int rows, int ncls, int n_genes, + long long nnz, bool assume_unsorted, + cudaStream_t stream) { + RmmScratchPool pool; + const GroupedRows groups = + build_grouped_rows(pool, clusters, nullptr, rows, ncls, stream); + LigrecOp op{data, g, nullptr, ncls, 0}; + return minor_reduce(indptr, index, op, rows, n_genes, nnz, + assume_unsorted, stream, &groups); } template @@ -128,13 +132,18 @@ void def_sum_count_sparse(nb::module_& m) { gpu_array_c data, gpu_array_c clusters, gpu_array_c sum, gpu_array_c count, int rows, int ncls, - std::uintptr_t stream) { - launch_sum_count_sparse( + bool assume_unsorted, std::uintptr_t stream) { + // sum is (n_genes, ncls); derive n_genes from the element count so + // any 2-D shape with the same size works. + return launch_sum_count_sparse( indptr.data(), index.data(), data.data(), clusters.data(), - sum.data(), count.data(), rows, ncls, (cudaStream_t)stream); + sum.data(), count.data(), rows, ncls, (int)(sum.size() / ncls), + (long long)data.shape(0), assume_unsorted, + (cudaStream_t)stream); }, "indptr"_a, "index"_a, "data"_a, nb::kw_only(), "clusters"_a, "sum"_a, - "count"_a, "rows"_a, "ncls"_a, "stream"_a = 0); + "count"_a, "rows"_a, "ncls"_a, "assume_unsorted"_a = false, + "stream"_a = 0); } template @@ -159,13 +168,15 @@ void def_mean_sparse(nb::module_& m) { gpu_array_c index, gpu_array_c data, gpu_array_c clusters, gpu_array_c g, - int rows, int ncls, std::uintptr_t stream) { - launch_mean_sparse(indptr.data(), index.data(), - data.data(), clusters.data(), g.data(), - rows, ncls, (cudaStream_t)stream); + int rows, int ncls, bool assume_unsorted, std::uintptr_t stream) { + return launch_mean_sparse( + indptr.data(), index.data(), data.data(), clusters.data(), + g.data(), rows, ncls, (int)(g.size() / ncls), + (long long)data.shape(0), assume_unsorted, + (cudaStream_t)stream); }, "indptr"_a, "index"_a, "data"_a, nb::kw_only(), "clusters"_a, "g"_a, - "rows"_a, "ncls"_a, "stream"_a = 0); + "rows"_a, "ncls"_a, "assume_unsorted"_a = false, "stream"_a = 0); } template @@ -248,4 +259,5 @@ void register_bindings(nb::module_& m) { NB_MODULE(_ligrec_cuda, m) { REGISTER_GPU_BINDINGS(register_bindings, m); + register_scratch_allocator(m); } diff --git a/src/rapids_singlecell/_cuda/mean_var/kernels_mv.cuh b/src/rapids_singlecell/_cuda/mean_var/kernels_mv.cuh index f4883a94c..ee64374ae 100644 --- a/src/rapids_singlecell/_cuda/mean_var/kernels_mv.cuh +++ b/src/rapids_singlecell/_cuda/mean_var/kernels_mv.cuh @@ -1,6 +1,7 @@ #pragma once #include +#include "../minor_tiles.cuh" constexpr int BLOCK_SIZE_MAJOR = 64; @@ -45,18 +46,45 @@ __global__ void mean_var_major_kernel(const IdxT* __restrict__ indptr, } } -template -__global__ void mean_var_minor_kernel(const IdxT* __restrict__ indices, - const T* __restrict__ data, - double* __restrict__ means, - double* __restrict__ vars, - long long nnz) { - const long long stride = (long long)blockDim.x * gridDim.x; - for (long long idx = (long long)blockDim.x * blockIdx.x + threadIdx.x; - idx < nnz; idx += stride) { - double value = static_cast(data[idx]); - IdxT minor_pos = indices[idx]; - atomicAdd(&means[minor_pos], value); - atomicAdd(&vars[minor_pos], value * value); +/// Minor-axis sum / sum-of-squares per column (see minor_tiles.cuh). +template +struct MeanVarOp { + const T* data; + double* means; + double* vars; + int tile_size; + static constexpr size_t bytes_per_col = 2 * sizeof(double); + static constexpr bool needs_rows = false; + __device__ bool row_active(int) const { + return true; } -} + __device__ void zero_col(char* acc, int g, int) const { + double* s = reinterpret_cast(acc); + s[g] = 0.0; + s[tile_size + g] = 0.0; + } + __device__ void add(char* acc, long long q, int g) const { + double* s = reinterpret_cast(acc); + const double v = static_cast(data[q]); + atomicAdd(&s[g], v); + atomicAdd(&s[tile_size + g], v * v); + } + __device__ void flush_col(const char* acc, int, int col, int g) const { + const double* s = reinterpret_cast(acc); + const double sq = s[tile_size + g]; + // Zero only when no nonzero of this column landed in the block. + if (sq != 0.0) { + atomicAdd(&means[col], s[g]); + atomicAdd(&vars[col], sq); + } + } + __device__ void add_global(long long q, int col, int) const { + const double v = static_cast(data[q]); + atomicAdd(&means[col], v); + atomicAdd(&vars[col], v * v); + } + void zero_outputs(int minor, int, cudaStream_t stream) const { + cudaMemsetAsync(means, 0, (size_t)minor * sizeof(double), stream); + cudaMemsetAsync(vars, 0, (size_t)minor * sizeof(double), stream); + } +}; diff --git a/src/rapids_singlecell/_cuda/mean_var/mean_var.cu b/src/rapids_singlecell/_cuda/mean_var/mean_var.cu index 56f625158..cf3cd15cd 100644 --- a/src/rapids_singlecell/_cuda/mean_var/mean_var.cu +++ b/src/rapids_singlecell/_cuda/mean_var/mean_var.cu @@ -5,8 +5,6 @@ using namespace nb::literals; -constexpr int BLOCK_SIZE_MINOR = 256; - template static inline void launch_mean_var_major(const IdxT* indptr, const IdxT* indices, const T* data, @@ -19,17 +17,6 @@ static inline void launch_mean_var_major(const IdxT* indptr, CUDA_CHECK_LAST_ERROR(mean_var_major_kernel); } -template -static inline void launch_mean_var_minor(const IdxT* indices, const T* data, - double* means, double* vars, - long long nnz, cudaStream_t stream) { - int block = BLOCK_SIZE_MINOR; - unsigned int grid = strided_grid(nnz, block); - mean_var_minor_kernel - <<>>(indices, data, means, vars, nnz); - CUDA_CHECK_LAST_ERROR(mean_var_minor_kernel); -} - template void def_mean_var_major(nb::module_& m) { m.def( @@ -47,6 +34,7 @@ void def_mean_var_major(nb::module_& m) { "major"_a, "minor"_a, "stream"_a = 0); } +// Order-agnostic minor-axis sums (one atomic per nonzero); no indptr needed. template void def_mean_var_minor(nb::module_& m) { m.def( @@ -55,14 +43,56 @@ void def_mean_var_minor(nb::module_& m) { gpu_array_c data, gpu_array_c means, gpu_array_c vars, long long nnz, std::uintptr_t stream) { - launch_mean_var_minor(indices.data(), data.data(), - means.data(), vars.data(), nnz, - (cudaStream_t)stream); + MeanVarOp op{data.data(), means.data(), vars.data(), 0}; + minor_reduce_flat(indices.data(), op, nnz, + (cudaStream_t)stream); }, "indices"_a, "data"_a, "means"_a, "vars"_a, nb::kw_only(), "nnz"_a, "stream"_a = 0); } +// Shared-memory tile sweep with atomic fallback. Returns whether unsorted rows +// were detected, so the caller can skip the attempt next time. +template +void def_mean_var_minor_tiled(nb::module_& m) { + m.def( + "mean_var_minor_tiled", + [](gpu_array_c indptr, + gpu_array_c indices, + gpu_array_c data, gpu_array_c means, + gpu_array_c vars, int major, int minor, + long long nnz, bool assume_unsorted, std::uintptr_t stream) { + MeanVarOp op{data.data(), means.data(), vars.data(), 0}; + return minor_reduce(indptr.data(), indices.data(), op, major, + minor, nnz, assume_unsorted, + (cudaStream_t)stream); + }, + "indptr"_a, "indices"_a, "data"_a, "means"_a, "vars"_a, nb::kw_only(), + "major"_a, "minor"_a, "nnz"_a, "assume_unsorted"_a = false, + "stream"_a = 0); +} + +// Expose the planner so callers can ask which path a matrix would take: +// tile width, tile count, and whether unsorted rows would rescan or go atomic. +void def_tile_plan(nb::module_& m) { + m.def( + "tile_plan", + [](long long nnz, int n_rows, int n_cols, size_t bytes_per_col) { + const TilePlan p = plan_tiles(nnz, n_rows, n_cols, bytes_per_col); + nb::dict d; + d["use_tiled"] = p.use_tiled; + d["tile_size"] = p.tile_size; + d["n_tiles"] = p.n_tiles; + d["rows_per_block"] = p.rows_per_block; + d["smem_bytes"] = p.smem_bytes; + d["rescan_if_unsorted"] = + p.use_tiled && p.n_tiles <= SWEEP_MAX_RESCAN_TILES; + return d; + }, + "nnz"_a, "n_rows"_a, "n_cols"_a, + "bytes_per_col"_a = 2 * sizeof(double)); +} + template void register_bindings(nb::module_& m) { def_mean_var_major(m); @@ -74,8 +104,15 @@ void register_bindings(nb::module_& m) { def_mean_var_minor(m); def_mean_var_minor(m); def_mean_var_minor(m); + + def_mean_var_minor_tiled(m); + def_mean_var_minor_tiled(m); + def_mean_var_minor_tiled(m); + def_mean_var_minor_tiled(m); } NB_MODULE(_mean_var_cuda, m) { REGISTER_GPU_BINDINGS(register_bindings, m); + def_tile_plan(m); + register_scratch_allocator(m); } diff --git a/src/rapids_singlecell/_cuda/minor_tiles.cuh b/src/rapids_singlecell/_cuda/minor_tiles.cuh new file mode 100644 index 000000000..9fb6e2d9c --- /dev/null +++ b/src/rapids_singlecell/_cuda/minor_tiles.cuh @@ -0,0 +1,644 @@ +#pragma once + +#include +#include + +#include +#include + +#include "nb_types.h" +#include "rmm_scratch.h" + +// ============================================================================= +// Shared-memory tile sweep for minor-axis reductions over CSR/CSC matrices. +// ============================================================================= +// +// The job: for every column of a compressed matrix, accumulate something over +// the nonzeros in that column (a sum, a count, a sum of squares, ...). Rows are +// contiguous in memory but columns are scattered, so the naive kernel issues a +// global atomic per nonzero and is bound by the L2 atomic units. +// +// The design, in plain language: +// * A block owns a range of rows and a private scratch pad in shared memory. +// It adds into the pad with cheap on-chip atomics and writes each column's +// subtotal to global memory once, at the end. +// * The pad only holds a stripe of columns, a TILE. The block works tile by +// tile. +// * Each row is split into one SLICE per tile: the run of its nonzeros whose +// columns fall inside the tile. In a row with sorted columns the slices sit +// back to back, so a BOOKMARK per row (where this tile's slice ended) tells +// the next tile where its slice starts. Every nonzero is read exactly once. +// * Sorted rows are not required. On a later tile, a column below the tile's +// lower bound can only be one that an earlier tile skipped, which never +// happens in a sorted row, so the kernel raises a flag. The host then +// redoes the work order-agnostically: RESCAN (each tile reads the whole +// row) when there are few tiles, the per-nonzero atomic kernel when there +// are many. +// * The tile width comes from the device's shared-memory budget. +// +// Vocabulary used in the code: row, column, tile [tile_begin, tile_end), +// slice, bookmark, tile_col (a column's slot inside the tile), nnz_pos (an +// index into indices/data). +// +// An Op describes one reduction: a POD passed by value that holds the device +// pointers and defines +// static constexpr size_t bytes_per_col; // scratch-pad bytes per column +// static constexpr bool needs_rows; // fallback must know the row +// int tile_size; // filled in by minor_reduce() +// __device__ bool row_active(int row) const; +// __device__ void zero_col(char* pad, int tile_col, int col) const; +// __device__ void add(char* pad, long long nnz_pos, int tile_col) const; +// __device__ void flush_col(const char* pad, int group, int col, +// int tile_col) const; +// __device__ void add_global(long long nnz_pos, int col, int group) const; +// void zero_outputs(int n_cols, int n_groups, cudaStream_t stream) const; +// The pad is 16-byte aligned; lay sub-arrays out in decreasing alignment +// (doubles, then 4-byte types, then bytes). Grouped reductions (aggregate, +// ligrec) first sort rows by group so a block only ever accumulates for one +// group; see build_grouped_rows(). + +// ---- tunables -------------------------------------------------------------- + +constexpr int WARP = 32; +// 16 warps: one row slice per warp, coalesced 32-wide reads. +constexpr int SWEEP_BLOCK_THREADS = 512; +// Occupancy target; the per-block shared-memory budget follows from it. +constexpr int SWEEP_BLOCKS_PER_SM = 2; +// Rows per block are chosen for this many blocks in flight per SM: more blocks +// balance better, fewer blocks flush fewer subtotals. +constexpr int SWEEP_BLOCKS_IN_FLIGHT_PER_SM = 12; +constexpr int SWEEP_MIN_ROWS_PER_BLOCK = SWEEP_BLOCK_THREADS / WARP; +constexpr int SWEEP_MAX_ROWS_PER_BLOCK = 1024; // bookmark array size +// Below this many nonzeros per slice the per-tile bookkeeping dominates and the +// atomic kernel is faster (measured crossover between 10 and 20). +constexpr long long SWEEP_MIN_NNZ_PER_SLICE = 16; +// Rescan re-reads 4 B per nonzero per extra tile; it beats the atomic kernel up +// to this many tiles (measured: 0.43x to 0.75x its time at 1 to 4 tiles, +// break-even at 5). +constexpr int SWEEP_MAX_RESCAN_TILES = 4; +// Above this, dynamic shared memory needs an explicit per-kernel opt-in. +constexpr size_t SMEM_DEFAULT_LIMIT = 48 * 1024; +constexpr int ATOMIC_BLOCK_THREADS = 256; + +// ---- plan ------------------------------------------------------------------ + +struct TilePlan { + bool use_tiled; // false: run the atomic fallback instead + int tile_size; // columns per tile + int n_tiles; + int rows_per_block; + size_t smem_bytes; // scratch pad + bookmarks +}; + +/// Scratch-pad bytes for a tile, padded so the bookmarks behind it stay +/// aligned. +__host__ __device__ inline size_t pad_bytes(int tile_size, + size_t bytes_per_col) { + return ((size_t)tile_size * bytes_per_col + 7) / 8 * 8; +} + +/// Size the tile from the per-SM shared memory at the target occupancy, capped +/// by the per-block opt-in limit, so it adapts from a T4 (64 KB per SM) to +/// datacenter parts without hardcoding either. +inline TilePlan plan_tiles(long long nnz, int n_rows, int n_cols, + size_t bytes_per_col) { + const DeviceSmemLimits& lim = device_smem_limits(); + size_t budget = lim.per_sm / SWEEP_BLOCKS_PER_SM; + budget = + budget > lim.reserved_per_block ? budget - lim.reserved_per_block : 0; + budget = std::min(budget, lim.per_block_optin); + budget = budget > 8 ? budget - 8 : 0; // alignment padding + + int rows_per_block = n_rows / (lim.n_sms * SWEEP_BLOCKS_IN_FLIGHT_PER_SM); + rows_per_block = std::clamp(rows_per_block, SWEEP_MIN_ROWS_PER_BLOCK, + SWEEP_MAX_ROWS_PER_BLOCK); + + const size_t bookmark_bytes = (size_t)rows_per_block * sizeof(int); + TilePlan plan{false, 0, 0, rows_per_block, 0}; + if (budget < bookmark_bytes + bytes_per_col) return plan; + + const size_t max_tile_cols = (budget - bookmark_bytes) / bytes_per_col; + plan.tile_size = (int)std::min(max_tile_cols, (size_t)n_cols); + plan.n_tiles = (n_cols + plan.tile_size - 1) / plan.tile_size; + const long long nnz_per_slice = nnz / ((long long)n_rows * plan.n_tiles); + if (plan.n_tiles > 1 && nnz_per_slice < SWEEP_MIN_NNZ_PER_SLICE) + return plan; + + plan.smem_bytes = pad_bytes(plan.tile_size, bytes_per_col) + bookmark_bytes; + plan.use_tiled = true; + return plan; +} + +/// Opt a kernel into dynamic shared memory above the 48 KB default. +template +inline void opt_in_dynamic_smem(Kernel kernel, size_t bytes, const char* what) { + if (bytes > SMEM_DEFAULT_LIMIT) { + cuda_check(cudaFuncSetAttribute( + kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, + (int)bytes), + what); + } +} + +// ---- device: one block, one tile ------------------------------------------- + +/// Rows [first, last) of the row list that one block owns, all in `group`. +struct BlockRows { + int first; + int last; + int group; +}; + +__device__ inline BlockRows block_rows(const BlockRows* __restrict__ per_block, + int rows_per_block, int n_rows) { + if (per_block != nullptr) return per_block[blockIdx.x]; + const int first = blockIdx.x * rows_per_block; + return {first, min(first + rows_per_block, n_rows), 0}; +} + +/// Sweep one row's slice of the current tile, resuming at `bookmark` (an offset +/// into the row). Returns where the slice ended, which is the bookmark for the +/// next tile. Sets `out_of_order` when a column below `tile_begin` shows up: +/// that column was skipped by an earlier tile, so the row is not sorted. +template +__device__ inline int warp_sweep_slice(const IdxT* __restrict__ indices, + IdxT row_begin, IdxT row_end, + int bookmark, IdxT tile_begin, + IdxT tile_end, const Op& op, char* pad, + bool& out_of_order) { + const int lane = threadIdx.x & (WARP - 1); + IdxT pos = row_begin + bookmark; + while (pos < row_end) { + const IdxT nnz_pos = pos + lane; + // Lanes past the row read a sentinel that ends the slice exactly like a + // column past the tile would. + const IdxT col = (nnz_pos < row_end) ? indices[nnz_pos] : tile_end; + const bool in_tile = col < tile_end; + if (in_tile) { + if (col < tile_begin) { + out_of_order = true; + } else { + op.add(pad, (long long)nnz_pos, + static_cast(col - tile_begin)); + } + } + const unsigned past_tile = __ballot_sync(0xffffffffu, !in_tile); + if (past_tile) { + // The slice ends at the first lane whose column is past the tile. + return static_cast(pos + (__ffs(past_tile) - 1) - row_begin); + } + pos += WARP; + } + return static_cast(row_end - row_begin); +} + +/// Order-agnostic alternative: read the whole row, keep what falls in the tile. +template +__device__ inline void warp_rescan_row(const IdxT* __restrict__ indices, + IdxT row_begin, IdxT row_end, + IdxT tile_begin, IdxT tile_end, + const Op& op, char* pad) { + const int lane = threadIdx.x & (WARP - 1); + for (IdxT nnz_pos = row_begin + lane; nnz_pos < row_end; nnz_pos += WARP) { + const IdxT col = indices[nnz_pos]; + if (col >= tile_begin && col < tile_end) { + op.add(pad, (long long)nnz_pos, static_cast(col - tile_begin)); + } + } +} + +template +__device__ inline void zero_tile(const Op& op, char* pad, int tile_begin, + int width) { + for (int tile_col = threadIdx.x; tile_col < width; tile_col += blockDim.x) { + op.zero_col(pad, tile_col, tile_begin + tile_col); + } +} + +template +__device__ inline void flush_tile(const Op& op, const char* pad, int tile_begin, + int width, int group) { + for (int tile_col = threadIdx.x; tile_col < width; tile_col += blockDim.x) { + op.flush_col(pad, group, tile_begin + tile_col, tile_col); + } +} + +/// Every warp sweeps its rows' slices of the current tile. Rows are +/// `row_order[i]` when a reordering is given, else `i`. +template +__device__ inline void sweep_tile(const IdxT* __restrict__ indptr, + const IdxT* __restrict__ indices, + const int* __restrict__ row_order, + int* bookmarks, BlockRows rows, + int tile_begin, int tile_size, bool rescan, + int* __restrict__ out_of_order_flag, + const Op& op, char* pad) { + const int lane = threadIdx.x & (WARP - 1); + const int warp = threadIdx.x / WARP; + const int warps_per_block = blockDim.x / WARP; + const IdxT tile_lo = static_cast(tile_begin); + const IdxT tile_hi = tile_lo + tile_size; + bool out_of_order = false; + + for (int i = rows.first + warp; i < rows.last; i += warps_per_block) { + const int row = row_order != nullptr ? row_order[i] : i; + if (!op.row_active(row)) continue; + const IdxT row_begin = indptr[row]; + const IdxT row_end = indptr[row + 1]; + if (rescan) { + warp_rescan_row(indices, row_begin, row_end, tile_lo, tile_hi, op, + pad); + continue; + } + const int bookmark = warp_sweep_slice( + indices, row_begin, row_end, bookmarks[i - rows.first], tile_lo, + tile_hi, op, pad, out_of_order); + if (lane == 0) bookmarks[i - rows.first] = bookmark; + } + if (out_of_order && out_of_order_flag != nullptr) *out_of_order_flag = 1; +} + +/// The tile sweep: zero the pad, sweep every row's slice, flush the subtotals, +/// tile after tile. grid = number of row blocks. +template +__global__ void __launch_bounds__(SWEEP_BLOCK_THREADS, SWEEP_BLOCKS_PER_SM) + tile_sweep_kernel(const IdxT* __restrict__ indptr, + const IdxT* __restrict__ indices, + const int* __restrict__ row_order, + const BlockRows* __restrict__ per_block_rows, Op op, + int* __restrict__ out_of_order_flag, int n_rows, + int n_cols, int n_tiles, int rows_per_block, + bool rescan) { + extern __shared__ __align__(16) char smem[]; + char* pad = smem; // the Op's accumulators for the current tile + int* bookmarks = reinterpret_cast( + smem + pad_bytes(op.tile_size, Op::bytes_per_col)); + + const BlockRows rows = block_rows(per_block_rows, rows_per_block, n_rows); + for (int r = threadIdx.x; r < rows.last - rows.first; r += blockDim.x) { + bookmarks[r] = 0; // every row starts at its beginning + } + + for (int tile = 0; tile < n_tiles; ++tile) { + const int tile_begin = tile * op.tile_size; + const int width = min(op.tile_size, n_cols - tile_begin); + zero_tile(op, pad, tile_begin, width); + __syncthreads(); + sweep_tile(indptr, indices, row_order, bookmarks, rows, tile_begin, + op.tile_size, rescan, out_of_order_flag, op, pad); + __syncthreads(); + flush_tile(op, pad, tile_begin, width, rows.group); + __syncthreads(); + } +} + +// ---- device: order-agnostic fallbacks -------------------------------------- + +/// One global atomic per nonzero; for Ops without groups or row filters. +template +__global__ void atomic_per_nonzero_kernel(const IdxT* __restrict__ indices, + Op op, long long nnz) { + const long long stride = (long long)blockDim.x * gridDim.x; + for (long long nnz_pos = (long long)blockIdx.x * blockDim.x + threadIdx.x; + nnz_pos < nnz; nnz_pos += stride) { + op.add_global(nnz_pos, static_cast(indices[nnz_pos]), 0); + } +} + +/// Warp per row, lanes stride the nonzeros; knows the row, so it serves grouped +/// Ops and Ops with a row filter. +template +__global__ void atomic_per_row_kernel( + const IdxT* __restrict__ indptr, const IdxT* __restrict__ indices, + const int* __restrict__ row_order, + const BlockRows* __restrict__ per_block_rows, Op op, int n_rows, + int rows_per_block) { + const BlockRows rows = block_rows(per_block_rows, rows_per_block, n_rows); + const int lane = threadIdx.x & (WARP - 1); + const int warp = threadIdx.x / WARP; + const int warps_per_block = blockDim.x / WARP; + for (int i = rows.first + warp; i < rows.last; i += warps_per_block) { + const int row = row_order != nullptr ? row_order[i] : i; + if (!op.row_active(row)) continue; + const IdxT row_end = indptr[row + 1]; + for (IdxT nnz_pos = indptr[row] + lane; nnz_pos < row_end; + nnz_pos += WARP) { + op.add_global((long long)nnz_pos, + static_cast(indices[nnz_pos]), rows.group); + } + } +} + +// ---- grouped rows ----------------------------------------------------------- + +/// Rows sorted by group: `row_order` (device, owned by the caller's scratch +/// pool) plus host group offsets into it. +struct GroupedRows { + const int* row_order = nullptr; + std::vector offsets; // n_groups + 1 + int n_active = 0; +}; + +__global__ inline void group_keys_kernel(const int* __restrict__ cats, + const bool* __restrict__ mask, + int n_rows, int n_groups, + int* __restrict__ keys, + int* __restrict__ vals, + int* __restrict__ counts) { + const int stride = blockDim.x * gridDim.x; + for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < n_rows; + i += stride) { + int k = (mask == nullptr || mask[i]) ? cats[i] : n_groups; + if (k < 0 || k > n_groups) k = n_groups; // sorts last, excluded + keys[i] = k; + vals[i] = i; + if (k < n_groups) atomicAdd(&counts[k], 1); + } +} + +/// Sort rows by category (masked-out or invalid rows are dropped) so a grouped +/// reduction can hand each block rows of a single group. All device scratch, +/// including the returned row order, lives in `pool`. +inline GroupedRows build_grouped_rows(RmmScratchPool& pool, const int* cats, + const bool* mask, int n_rows, + int n_groups, cudaStream_t stream) { + GroupedRows g; + g.offsets.assign((size_t)n_groups + 1, 0); + if (n_rows <= 0 || n_groups <= 0) return g; + int* keys = pool.alloc((size_t)n_rows); + int* keys_sorted = pool.alloc((size_t)n_rows); + int* vals = pool.alloc((size_t)n_rows); + int* counts = pool.alloc((size_t)n_groups); + int* row_order = pool.alloc((size_t)n_rows); + g.row_order = row_order; + + cuda_check( + cudaMemsetAsync(counts, 0, (size_t)n_groups * sizeof(int), stream), + "cudaMemsetAsync(group counts)"); + group_keys_kernel<<>>( + cats, mask, n_rows, n_groups, keys, vals, counts); + CUDA_CHECK_LAST_ERROR(group_keys_kernel); + + int end_bit = 1; + while ((1 << end_bit) <= n_groups) ++end_bit; + size_t temp_bytes = 0; + cuda_check(cub::DeviceRadixSort::SortPairs(nullptr, temp_bytes, keys, + keys_sorted, vals, row_order, + n_rows, 0, end_bit, stream), + "cub::DeviceRadixSort::SortPairs(size)"); + void* temp = pool.alloc(temp_bytes); + cuda_check(cub::DeviceRadixSort::SortPairs(temp, temp_bytes, keys, + keys_sorted, vals, row_order, + n_rows, 0, end_bit, stream), + "cub::DeviceRadixSort::SortPairs"); + + std::vector h_counts((size_t)n_groups); + cuda_check( + cudaMemcpyAsync(h_counts.data(), counts, (size_t)n_groups * sizeof(int), + cudaMemcpyDeviceToHost, stream), + "cudaMemcpyAsync(group counts)"); + cuda_check(cudaStreamSynchronize(stream), "cudaStreamSynchronize"); + for (int k = 0; k < n_groups; ++k) { + g.offsets[(size_t)k + 1] = g.offsets[(size_t)k] + h_counts[(size_t)k]; + } + g.n_active = g.offsets[(size_t)n_groups]; + return g; +} + +/// Split each group's run of rows into blocks of at most `rows_per_block`. +inline std::vector block_rows_by_group( + const std::vector& offsets, int rows_per_block) { + std::vector blocks; + for (size_t k = 0; k + 1 < offsets.size(); ++k) { + for (int first = offsets[k]; first < offsets[k + 1]; + first += rows_per_block) { + blocks.push_back({first, + std::min(first + rows_per_block, offsets[k + 1]), + (int)k}); + } + } + return blocks; +} + +// ---- host entry points +// ------------------------------------------------------- + +/// Run an Op over the minor axis. Returns whether unsorted rows were detected, +/// so the caller can remember it and skip the bookmark sweep next time. +template +bool minor_reduce(const IdxT* indptr, const IdxT* indices, Op op, int n_rows, + int n_cols, long long nnz, bool assume_unsorted, + cudaStream_t stream, const GroupedRows* groups = nullptr) { + if (n_rows <= 0 || n_cols <= 0 || nnz <= 0) return false; + if (groups != nullptr && groups->n_active == 0) return false; + const int n_groups = + groups != nullptr ? (int)groups->offsets.size() - 1 : 1; + const TilePlan plan = + plan_tiles(nnz, groups != nullptr ? groups->n_active : n_rows, n_cols, + Op::bytes_per_col); + op.tile_size = plan.tile_size; + + // Which rows each block owns: contiguous chunks, or per-group chunks of the + // sorted row order when grouped. + RmmScratchPool pool; // flag + block table; released after the launches + const int* row_order = groups != nullptr ? groups->row_order : nullptr; + const BlockRows* per_block_rows = nullptr; + unsigned n_blocks = + (unsigned)((n_rows + plan.rows_per_block - 1) / plan.rows_per_block); + if (groups != nullptr) { + const std::vector h_blocks = + block_rows_by_group(groups->offsets, plan.rows_per_block); + BlockRows* d_blocks = pool.alloc(h_blocks.size()); + cuda_check(cudaMemcpyAsync(d_blocks, h_blocks.data(), + h_blocks.size() * sizeof(BlockRows), + cudaMemcpyHostToDevice, stream), + "cudaMemcpyAsync(block rows)"); + per_block_rows = d_blocks; + n_blocks = (unsigned)h_blocks.size(); + } + + auto launch_atomic = [&]() { + if (groups != nullptr || Op::needs_rows) { + atomic_per_row_kernel + <<>>( + indptr, indices, row_order, per_block_rows, op, n_rows, + plan.rows_per_block); + CUDA_CHECK_LAST_ERROR(atomic_per_row_kernel); + } else { + atomic_per_nonzero_kernel + <<>>(indices, op, nnz); + CUDA_CHECK_LAST_ERROR(atomic_per_nonzero_kernel); + } + }; + auto launch_tiled = [&](bool rescan, int* out_of_order_flag) { + opt_in_dynamic_smem(tile_sweep_kernel, plan.smem_bytes, + "cudaFuncSetAttribute(tile_sweep_kernel)"); + tile_sweep_kernel + <<>>( + indptr, indices, row_order, per_block_rows, op, + out_of_order_flag, n_rows, n_cols, plan.n_tiles, + plan.rows_per_block, rescan); + CUDA_CHECK_LAST_ERROR(tile_sweep_kernel); + }; + const bool few_tiles = plan.n_tiles <= SWEEP_MAX_RESCAN_TILES; + + // 1. The planner may have decided tiles do not pay off here. + if (!plan.use_tiled) { + launch_atomic(); + return false; + } + // 2. Rows already known to be unsorted: skip the bookmark sweep. + if (assume_unsorted) { + if (few_tiles) { + launch_tiled(true, nullptr); + } else { + launch_atomic(); + } + return false; + } + // 3. The bookmark sweep, with out-of-order detection. + int* flag = pool.alloc(1); + cuda_check(cudaMemsetAsync(flag, 0, sizeof(int), stream), + "cudaMemsetAsync(out-of-order flag)"); + launch_tiled(false, flag); + if (plan.n_tiles == 1) return false; // one tile cannot skip anything + // 4. Did a row turn out to be unsorted? Then redo it order-agnostically. + int h_flag = 0; + cuda_check(cudaMemcpyAsync(&h_flag, flag, sizeof(int), + cudaMemcpyDeviceToHost, stream), + "cudaMemcpyAsync(out-of-order flag)"); + cuda_check(cudaStreamSynchronize(stream), "cudaStreamSynchronize"); + if (h_flag == 0) return false; + op.zero_outputs(n_cols, n_groups, stream); + if (few_tiles) { + launch_tiled(true, nullptr); + } else { + launch_atomic(); + } + return true; +} + +/// Order-agnostic minor-axis reduction without indptr (one atomic per nonzero). +template +void minor_reduce_flat(const IdxT* indices, Op op, long long nnz, + cudaStream_t stream) { + if (nnz <= 0) return; + atomic_per_nonzero_kernel + <<>>(indices, op, nnz); + CUDA_CHECK_LAST_ERROR(atomic_per_nonzero_kernel); +} + +// ---- major-axis companion +// ----------------------------------------------------- + +/// Per compressed row: the sum (optionally over masked columns only) and the +/// number of stored entries. Warp per row, no atomics. +template +__global__ void row_reduce_kernel(const IdxT* __restrict__ indptr, + const IdxT* __restrict__ indices, + const T* __restrict__ data, + const bool* __restrict__ col_mask, + T* __restrict__ sums, + int* __restrict__ counts, int n_rows) { + const int lane = threadIdx.x & (WARP - 1); + const int warps_total = (gridDim.x * blockDim.x) / WARP; + for (int row = (blockIdx.x * blockDim.x + threadIdx.x) / WARP; row < n_rows; + row += warps_total) { + const IdxT row_end = indptr[row + 1]; + double acc = 0.0; + int cnt = 0; + for (IdxT nnz_pos = indptr[row] + lane; nnz_pos < row_end; + nnz_pos += WARP) { + if (col_mask != nullptr && !col_mask[indices[nnz_pos]]) continue; + acc += static_cast(data[nnz_pos]); + ++cnt; + } +#pragma unroll + for (int o = WARP / 2; o > 0; o >>= 1) { + acc += __shfl_down_sync(0xffffffffu, acc, o); + cnt += __shfl_down_sync(0xffffffffu, cnt, o); + } + if (lane == 0) { + if (sums != nullptr) sums[row] = static_cast(acc); + if (counts != nullptr) counts[row] = cnt; + } + } +} + +template +inline void row_reduce(const IdxT* indptr, const IdxT* indices, const T* data, + const bool* col_mask, T* sums, int* counts, int n_rows, + cudaStream_t stream) { + if (n_rows <= 0) return; + row_reduce_kernel + <<>>(indptr, indices, data, col_mask, + sums, counts, n_rows); + CUDA_CHECK_LAST_ERROR(row_reduce_kernel); +} + +// ---- Ops shared by several modules ----------------------------------------- + +/// Sum of values per column, optionally restricted to active rows. +template +struct MinorSumOp { + const T* data; + T* out; + const bool* row_mask; // nullable + int tile_size; + static constexpr size_t bytes_per_col = sizeof(double); + static constexpr bool needs_rows = true; + __device__ bool row_active(int row) const { + return row_mask == nullptr || row_mask[row]; + } + __device__ void zero_col(char* pad, int tile_col, int) const { + reinterpret_cast(pad)[tile_col] = 0.0; + } + __device__ void add(char* pad, long long nnz_pos, int tile_col) const { + atomicAdd(&reinterpret_cast(pad)[tile_col], + static_cast(data[nnz_pos])); + } + __device__ void flush_col(const char* pad, int, int col, + int tile_col) const { + const double s = reinterpret_cast(pad)[tile_col]; + if (s != 0.0) atomicAdd(&out[col], static_cast(s)); + } + __device__ void add_global(long long nnz_pos, int col, int) const { + atomicAdd(&out[col], data[nnz_pos]); + } + void zero_outputs(int n_cols, int, cudaStream_t stream) const { + cudaMemsetAsync(out, 0, (size_t)n_cols * sizeof(T), stream); + } +}; + +/// Number of stored entries per column. +struct MinorCountOp { + int* out; + int n_cols; + int tile_size; + static constexpr size_t bytes_per_col = sizeof(int); + static constexpr bool needs_rows = false; + __device__ bool row_active(int) const { + return true; + } + __device__ void zero_col(char* pad, int tile_col, int) const { + reinterpret_cast(pad)[tile_col] = 0; + } + __device__ void add(char* pad, long long, int tile_col) const { + atomicAdd(&reinterpret_cast(pad)[tile_col], 1); + } + __device__ void flush_col(const char* pad, int, int col, + int tile_col) const { + const int c = reinterpret_cast(pad)[tile_col]; + if (c != 0) atomicAdd(&out[col], c); + } + __device__ void add_global(long long, int col, int) const { + if (col >= 0 && col < n_cols) atomicAdd(&out[col], 1); + } + void zero_outputs(int n_cols_, int, cudaStream_t stream) const { + cudaMemsetAsync(out, 0, (size_t)n_cols_ * sizeof(int), stream); + } +}; diff --git a/src/rapids_singlecell/_cuda/nanmean/kernels_nanmean.cuh b/src/rapids_singlecell/_cuda/nanmean/kernels_nanmean.cuh index ae7087c90..35dedc8ff 100644 --- a/src/rapids_singlecell/_cuda/nanmean/kernels_nanmean.cuh +++ b/src/rapids_singlecell/_cuda/nanmean/kernels_nanmean.cuh @@ -1,29 +1,7 @@ #pragma once #include - -template -__global__ void nan_mean_minor_kernel(const IdxT* __restrict__ index, - const T* __restrict__ data, - double* __restrict__ means, - int* __restrict__ nans, - const bool* __restrict__ mask, - long long nnz) { - const long long stride = (long long)blockDim.x * gridDim.x; - for (long long idx = (long long)blockIdx.x * blockDim.x + threadIdx.x; - idx < nnz; idx += stride) { - IdxT minor_pos = index[idx]; - if (mask[minor_pos] == false) { - continue; - } - T v = data[idx]; - if (isnan((double)v)) { - atomicAdd(&nans[minor_pos], 1); - } else { - atomicAdd(&means[minor_pos], (double)v); - } - } -} +#include "../minor_tiles.cuh" template __global__ void nan_mean_major_kernel(const IdxT* __restrict__ indptr, @@ -73,3 +51,64 @@ __global__ void nan_mean_major_kernel(const IdxT* __restrict__ indptr, nans[major_idx] = nan_place[0]; } } + +/// Minor-axis NaN-aware sum and NaN count per masked column (see +/// minor_tiles.cuh). Layout: double sums, int NaN counts, bool mask. +template +struct NanMeanOp { + const T* data; + double* means; + int* nans; + const bool* mask; + int tile_size; + static constexpr size_t bytes_per_col = + sizeof(double) + sizeof(int) + sizeof(bool); + static constexpr bool needs_rows = false; + __device__ double* s_sum(char* acc) const { + return reinterpret_cast(acc); + } + __device__ int* s_nan(char* acc) const { + return reinterpret_cast(acc + (size_t)tile_size * sizeof(double)); + } + __device__ bool* s_mask(char* acc) const { + return reinterpret_cast( + acc + (size_t)tile_size * (sizeof(double) + sizeof(int))); + } + __device__ bool row_active(int) const { + return true; + } + __device__ void zero_col(char* acc, int g, int col) const { + s_sum(acc)[g] = 0.0; + s_nan(acc)[g] = 0; + s_mask(acc)[g] = mask[col]; + } + __device__ void add(char* acc, long long q, int g) const { + if (!s_mask(acc)[g]) return; + const double v = static_cast(data[q]); + if (isnan(v)) { + atomicAdd(&s_nan(acc)[g], 1); + } else { + atomicAdd(&s_sum(acc)[g], v); + } + } + __device__ void flush_col(const char* acc, int, int col, int g) const { + char* a = const_cast(acc); + const int n = s_nan(a)[g]; + const double s = s_sum(a)[g]; + if (n != 0) atomicAdd(&nans[col], n); + if (s != 0.0) atomicAdd(&means[col], s); + } + __device__ void add_global(long long q, int col, int) const { + if (!mask[col]) return; + const double v = static_cast(data[q]); + if (isnan(v)) { + atomicAdd(&nans[col], 1); + } else { + atomicAdd(&means[col], v); + } + } + void zero_outputs(int minor, int, cudaStream_t stream) const { + cudaMemsetAsync(means, 0, (size_t)minor * sizeof(double), stream); + cudaMemsetAsync(nans, 0, (size_t)minor * sizeof(int), stream); + } +}; diff --git a/src/rapids_singlecell/_cuda/nanmean/nanmean.cu b/src/rapids_singlecell/_cuda/nanmean/nanmean.cu index c2a98394f..d6b72af3b 100644 --- a/src/rapids_singlecell/_cuda/nanmean/nanmean.cu +++ b/src/rapids_singlecell/_cuda/nanmean/nanmean.cu @@ -5,21 +5,8 @@ using namespace nb::literals; -constexpr int BLOCK_SIZE_MINOR = 32; constexpr int BLOCK_SIZE_MAJOR = 64; -template -static inline void launch_nan_mean_minor(const IdxT* index, const T* data, - double* means, int* nans, - const bool* mask, long long nnz, - cudaStream_t stream) { - dim3 block(BLOCK_SIZE_MINOR); - dim3 grid(strided_grid(nnz, BLOCK_SIZE_MINOR)); - nan_mean_minor_kernel - <<>>(index, data, means, nans, mask, nnz); - CUDA_CHECK_LAST_ERROR(nan_mean_minor_kernel); -} - template static inline void launch_nan_mean_major(const IdxT* indptr, const IdxT* index, const T* data, double* means, @@ -32,6 +19,7 @@ static inline void launch_nan_mean_major(const IdxT* indptr, const IdxT* index, CUDA_CHECK_LAST_ERROR(nan_mean_major_kernel); } +// Order-agnostic minor-axis sums (one atomic per nonzero); no indptr needed. template void def_nan_mean_minor(nb::module_& m) { m.def( @@ -40,14 +28,38 @@ void def_nan_mean_minor(nb::module_& m) { gpu_array_c data, gpu_array_c means, gpu_array_c nans, gpu_array_c mask, long long nnz, std::uintptr_t stream) { - launch_nan_mean_minor( - index.data(), data.data(), means.data(), nans.data(), - mask.data(), nnz, (cudaStream_t)stream); + NanMeanOp op{data.data(), means.data(), nans.data(), mask.data(), + 0}; + minor_reduce_flat(index.data(), op, nnz, + (cudaStream_t)stream); }, "index"_a, "data"_a, nb::kw_only(), "means"_a, "nans"_a, "mask"_a, "nnz"_a, "stream"_a = 0); } +// Shared-memory tile sweep with atomic fallback. Returns whether unsorted rows +// were detected, so the caller can skip the attempt next time. +template +void def_nan_mean_minor_tiled(nb::module_& m) { + m.def( + "nan_mean_minor_tiled", + [](gpu_array_c indptr, + gpu_array_c index, + gpu_array_c data, gpu_array_c means, + gpu_array_c nans, gpu_array_c mask, + int major, int minor, long long nnz, bool assume_unsorted, + std::uintptr_t stream) { + NanMeanOp op{data.data(), means.data(), nans.data(), mask.data(), + 0}; + return minor_reduce(indptr.data(), index.data(), op, major, + minor, nnz, assume_unsorted, + (cudaStream_t)stream); + }, + "indptr"_a, "index"_a, "data"_a, nb::kw_only(), "means"_a, "nans"_a, + "mask"_a, "major"_a, "minor"_a, "nnz"_a, "assume_unsorted"_a = false, + "stream"_a = 0); +} + template void def_nan_mean_major(nb::module_& m) { m.def( @@ -72,6 +84,11 @@ void register_bindings(nb::module_& m) { def_nan_mean_minor(m); def_nan_mean_minor(m); + def_nan_mean_minor_tiled(m); + def_nan_mean_minor_tiled(m); + def_nan_mean_minor_tiled(m); + def_nan_mean_minor_tiled(m); + def_nan_mean_major(m); def_nan_mean_major(m); def_nan_mean_major(m); @@ -80,4 +97,5 @@ void register_bindings(nb::module_& m) { NB_MODULE(_nanmean_cuda, m) { REGISTER_GPU_BINDINGS(register_bindings, m); + register_scratch_allocator(m); } diff --git a/src/rapids_singlecell/_cuda/nb_types.h b/src/rapids_singlecell/_cuda/nb_types.h index dc27d4f1b..8e1370ee0 100644 --- a/src/rapids_singlecell/_cuda/nb_types.h +++ b/src/rapids_singlecell/_cuda/nb_types.h @@ -83,6 +83,43 @@ inline unsigned int strided_grid_y(long long nwork, int block_size) { return (unsigned int)(capped < 1 ? 1 : capped); } +/// Per-device shared-memory limits and SM count, cached like `max_grid_dims()`. +/// Kernels size dynamic smem from these instead of assuming the 48 KB default. +struct DeviceSmemLimits { + size_t per_block_optin; // cudaDevAttrMaxSharedMemoryPerBlockOptin + size_t per_sm; // cudaDevAttrMaxSharedMemoryPerMultiprocessor + size_t reserved_per_block; // cudaDevAttrReservedSharedMemoryPerBlock + int n_sms; +}; + +inline const DeviceSmemLimits& device_smem_limits() { + // Safe fallback: the universal 48 KB default and one block's reservation. + constexpr DeviceSmemLimits kFallback{48 * 1024, 48 * 1024, 1024, 1}; + static thread_local int cached_dev = -1; + static thread_local DeviceSmemLimits cached = kFallback; + int device = 0; + cudaGetDevice(&device); + if (device != cached_dev) { + cached = kFallback; + int v = 0; + if (cudaDeviceGetAttribute(&v, cudaDevAttrMaxSharedMemoryPerBlockOptin, + device) == cudaSuccess) + cached.per_block_optin = (size_t)v; + if (cudaDeviceGetAttribute(&v, + cudaDevAttrMaxSharedMemoryPerMultiprocessor, + device) == cudaSuccess) + cached.per_sm = (size_t)v; + if (cudaDeviceGetAttribute(&v, cudaDevAttrReservedSharedMemoryPerBlock, + device) == cudaSuccess) + cached.reserved_per_block = (size_t)v; + if (cudaDeviceGetAttribute(&v, cudaDevAttrMultiProcessorCount, + device) == cudaSuccess) + cached.n_sms = v; + cached_dev = device; + } + return cached; +} + // GPU array aliases for nanobind bindings, parameterized on device type. // CUDA and managed-memory variants both preserve CPU/GPU type safety. diff --git a/src/rapids_singlecell/_cuda/pr/kernels_pr_hvg.cuh b/src/rapids_singlecell/_cuda/pr/kernels_pr_hvg.cuh index 03b067748..ee1ccb39f 100644 --- a/src/rapids_singlecell/_cuda/pr/kernels_pr_hvg.cuh +++ b/src/rapids_singlecell/_cuda/pr/kernels_pr_hvg.cuh @@ -2,29 +2,6 @@ #include -// Compute column sums (sums_genes) and row sums (sums_cells) from CSC sparse -// matrix One thread per column (gene), atomicAdd for row sums -template -__global__ void sparse_sum_csc_kernel(const IdxT* __restrict__ indptr, - const IdxT* __restrict__ index, - const T* __restrict__ data, - T* __restrict__ sums_genes, - T* __restrict__ sums_cells, int n_genes) { - int gene = blockDim.x * blockIdx.x + threadIdx.x; - if (gene >= n_genes) { - return; - } - IdxT start = indptr[gene]; - IdxT stop = indptr[gene + 1]; - T col_sum = (T)0; - for (IdxT i = start; i < stop; ++i) { - T val = data[i]; - col_sum += val; - atomicAdd(&sums_cells[index[i]], val); - } - sums_genes[gene] = col_sum; -} - // Welford's single-pass algorithm for variance of clipped Pearson residuals // (CSC sparse) template diff --git a/src/rapids_singlecell/_cuda/pr/pr.cu b/src/rapids_singlecell/_cuda/pr/pr.cu index 74b68c8ad..5c96aa3de 100644 --- a/src/rapids_singlecell/_cuda/pr/pr.cu +++ b/src/rapids_singlecell/_cuda/pr/pr.cu @@ -1,4 +1,5 @@ #include +#include "../minor_tiles.cuh" #include "../nb_types.h" #include "kernels_pr.cuh" @@ -51,16 +52,20 @@ static inline void launch_dense_norm_res(const T* X, T* residuals, CUDA_CHECK_LAST_ERROR(dense_norm_res_kernel); } +// Gene sums along the compressed axis plus cell sums along the minor axis. +// Returns whether unsorted rows were detected. template -static inline void launch_sparse_sum_csc(const IdxT* indptr, const IdxT* index, +static inline bool launch_sparse_sum_csc(const IdxT* indptr, const IdxT* index, const T* data, T* sums_genes, T* sums_cells, int n_genes, + int n_cells, long long nnz, + bool assume_unsorted, cudaStream_t stream) { - dim3 block(SPARSE_BLOCK_SIZE); - dim3 grid((n_genes + SPARSE_BLOCK_SIZE - 1) / SPARSE_BLOCK_SIZE); - sparse_sum_csc_kernel<<>>( - indptr, index, data, sums_genes, sums_cells, n_genes); - CUDA_CHECK_LAST_ERROR(sparse_sum_csc_kernel); + row_reduce(indptr, index, data, nullptr, sums_genes, nullptr, + n_genes, stream); + MinorSumOp op{data, sums_cells, nullptr, 0}; + return minor_reduce(indptr, index, op, n_genes, n_cells, nnz, + assume_unsorted, stream); } template @@ -163,14 +168,17 @@ void def_sparse_sum_csc(nb::module_& m) { [](gpu_array_c indptr, gpu_array_c index, gpu_array_c data, gpu_array_c sums_genes, - gpu_array_c sums_cells, int n_genes, + gpu_array_c sums_cells, int n_genes, bool assume_unsorted, std::uintptr_t stream) { - launch_sparse_sum_csc( + return launch_sparse_sum_csc( indptr.data(), index.data(), data.data(), sums_genes.data(), - sums_cells.data(), n_genes, (cudaStream_t)stream); + sums_cells.data(), n_genes, (int)sums_cells.shape(0), + (long long)data.shape(0), assume_unsorted, + (cudaStream_t)stream); }, "indptr"_a, "index"_a, "data"_a, nb::kw_only(), "sums_genes"_a, - "sums_cells"_a, "n_genes"_a, "stream"_a = 0); + "sums_cells"_a, "n_genes"_a, "assume_unsorted"_a = false, + "stream"_a = 0); } // Helper to define csc_hvg_res for a given dtype and index type @@ -252,4 +260,5 @@ void register_bindings(nb::module_& m) { NB_MODULE(_pr_cuda, m) { REGISTER_GPU_BINDINGS(register_bindings, m); + register_scratch_allocator(m); } diff --git a/src/rapids_singlecell/_cuda/qc/kernels_qc.cuh b/src/rapids_singlecell/_cuda/qc/kernels_qc.cuh index c9b5ddb3c..f74905c4c 100644 --- a/src/rapids_singlecell/_cuda/qc/kernels_qc.cuh +++ b/src/rapids_singlecell/_cuda/qc/kernels_qc.cuh @@ -1,58 +1,7 @@ #pragma once #include - -template -__global__ void qc_csc_kernel(const IdxT* __restrict__ indptr, - const IdxT* __restrict__ index, - const T* __restrict__ data, - T* __restrict__ sums_cells, - T* __restrict__ sums_genes, - int* __restrict__ cell_ex, - int* __restrict__ gene_ex, int n_genes) { - int gene = blockDim.x * blockIdx.x + threadIdx.x; - if (gene >= n_genes) return; - IdxT start_idx = indptr[gene]; - IdxT stop_idx = indptr[gene + 1]; - T sums_genes_i = T(0); - int gene_ex_i = 0; - for (IdxT p = start_idx; p < stop_idx; ++p) { - T v = data[p]; - IdxT cell = index[p]; - sums_genes_i += v; - atomicAdd(&sums_cells[cell], v); - ++gene_ex_i; - atomicAdd(&cell_ex[cell], 1); - } - sums_genes[gene] = sums_genes_i; - gene_ex[gene] = gene_ex_i; -} - -template -__global__ void qc_csr_kernel(const IdxT* __restrict__ indptr, - const IdxT* __restrict__ index, - const T* __restrict__ data, - T* __restrict__ sums_cells, - T* __restrict__ sums_genes, - int* __restrict__ cell_ex, - int* __restrict__ gene_ex, int n_cells) { - int cell = blockDim.x * blockIdx.x + threadIdx.x; - if (cell >= n_cells) return; - IdxT start_idx = indptr[cell]; - IdxT stop_idx = indptr[cell + 1]; - T sums_cells_i = T(0); - int cell_ex_i = 0; - for (IdxT p = start_idx; p < stop_idx; ++p) { - T v = data[p]; - IdxT gene = index[p]; - atomicAdd(&sums_genes[gene], v); - sums_cells_i += v; - atomicAdd(&gene_ex[gene], 1); - ++cell_ex_i; - } - sums_cells[cell] = sums_cells_i; - cell_ex[cell] = cell_ex_i; -} +#include "../minor_tiles.cuh" template __global__ void qc_dense_kernel(const T* __restrict__ data, @@ -74,41 +23,6 @@ __global__ void qc_dense_kernel(const T* __restrict__ data, } } -template -__global__ void qc_csc_sub_kernel(const IdxT* __restrict__ indptr, - const IdxT* __restrict__ index, - const T* __restrict__ data, - T* __restrict__ sums_cells, - const bool* __restrict__ mask, int n_genes) { - int gene = blockDim.x * blockIdx.x + threadIdx.x; - if (gene >= n_genes) return; - if (!mask[gene]) return; - IdxT start_idx = indptr[gene]; - IdxT stop_idx = indptr[gene + 1]; - for (IdxT p = start_idx; p < stop_idx; ++p) { - IdxT cell = index[p]; - atomicAdd(&sums_cells[cell], data[p]); - } -} - -template -__global__ void qc_csr_sub_kernel(const IdxT* __restrict__ indptr, - const IdxT* __restrict__ index, - const T* __restrict__ data, - T* __restrict__ sums_cells, - const bool* __restrict__ mask, int n_cells) { - int cell = blockDim.x * blockIdx.x + threadIdx.x; - if (cell >= n_cells) return; - IdxT start_idx = indptr[cell]; - IdxT stop_idx = indptr[cell + 1]; - T sums_cells_i = T(0); - for (IdxT p = start_idx; p < stop_idx; ++p) { - IdxT gene = index[p]; - if (mask[gene]) sums_cells_i += data[p]; - } - sums_cells[cell] = sums_cells_i; -} - template __global__ void qc_dense_sub_kernel(const T* __restrict__ data, T* __restrict__ sums_cells, @@ -121,3 +35,48 @@ __global__ void qc_dense_sub_kernel(const T* __restrict__ data, long long idx = (long long)cell * n_genes + gene; atomicAdd(&sums_cells[cell], data[idx]); } + +/// Minor-axis sum (as T) and stored-entry count per column (see +/// minor_tiles.cuh). Layout: double sums, int counts. +template +struct QcOp { + const T* data; + T* sums; + int* counts; + int tile_size; + static constexpr size_t bytes_per_col = sizeof(double) + sizeof(int); + static constexpr bool needs_rows = false; + __device__ double* s_sum(char* acc) const { + return reinterpret_cast(acc); + } + __device__ int* s_cnt(char* acc) const { + return reinterpret_cast(acc + (size_t)tile_size * sizeof(double)); + } + __device__ bool row_active(int) const { + return true; + } + __device__ void zero_col(char* acc, int g, int) const { + s_sum(acc)[g] = 0.0; + s_cnt(acc)[g] = 0; + } + __device__ void add(char* acc, long long q, int g) const { + atomicAdd(&s_sum(acc)[g], static_cast(data[q])); + atomicAdd(&s_cnt(acc)[g], 1); + } + __device__ void flush_col(const char* acc, int, int col, int g) const { + char* a = const_cast(acc); + const int c = s_cnt(a)[g]; + if (c != 0) { + atomicAdd(&sums[col], static_cast(s_sum(a)[g])); + atomicAdd(&counts[col], c); + } + } + __device__ void add_global(long long q, int col, int) const { + atomicAdd(&sums[col], data[q]); + atomicAdd(&counts[col], 1); + } + void zero_outputs(int minor, int, cudaStream_t stream) const { + cudaMemsetAsync(sums, 0, (size_t)minor * sizeof(T), stream); + cudaMemsetAsync(counts, 0, (size_t)minor * sizeof(int), stream); + } +}; diff --git a/src/rapids_singlecell/_cuda/qc/qc.cu b/src/rapids_singlecell/_cuda/qc/qc.cu index 9c5d6dd50..ce1f7a96f 100644 --- a/src/rapids_singlecell/_cuda/qc/qc.cu +++ b/src/rapids_singlecell/_cuda/qc/qc.cu @@ -5,31 +5,21 @@ using namespace nb::literals; -constexpr int SPARSE_BLOCK_SIZE = 32; constexpr int DENSE_BLOCK_DIM = 16; +// Sparse QC = major-axis sums/counts (warp per row) + minor-axis QcOp sweep. +// Returns whether unsorted rows were detected. template -static inline void launch_qc_csc(const IdxT* indptr, const IdxT* index, - const T* data, T* sums_cells, T* sums_genes, - int* cell_ex, int* gene_ex, int n_genes, - cudaStream_t stream) { - dim3 block(SPARSE_BLOCK_SIZE); - dim3 grid((n_genes + SPARSE_BLOCK_SIZE - 1) / SPARSE_BLOCK_SIZE); - qc_csc_kernel<<>>( - indptr, index, data, sums_cells, sums_genes, cell_ex, gene_ex, n_genes); - CUDA_CHECK_LAST_ERROR(qc_csc_kernel); -} - -template -static inline void launch_qc_csr(const IdxT* indptr, const IdxT* index, - const T* data, T* sums_cells, T* sums_genes, - int* cell_ex, int* gene_ex, int n_cells, - cudaStream_t stream) { - dim3 block(SPARSE_BLOCK_SIZE); - dim3 grid((n_cells + SPARSE_BLOCK_SIZE - 1) / SPARSE_BLOCK_SIZE); - qc_csr_kernel<<>>( - indptr, index, data, sums_cells, sums_genes, cell_ex, gene_ex, n_cells); - CUDA_CHECK_LAST_ERROR(qc_csr_kernel); +static inline bool launch_qc_sparse(const IdxT* indptr, const IdxT* index, + const T* data, T* sums_major, int* ex_major, + T* sums_minor, int* ex_minor, int major, + int minor, long long nnz, + bool assume_unsorted, cudaStream_t stream) { + row_reduce(indptr, index, data, nullptr, sums_major, ex_major, + major, stream); + QcOp op{data, sums_minor, ex_minor, 0}; + return minor_reduce(indptr, index, op, major, minor, nnz, + assume_unsorted, stream); } template @@ -44,30 +34,6 @@ static inline void launch_qc_dense(const T* data, T* sums_cells, T* sums_genes, CUDA_CHECK_LAST_ERROR(qc_dense_kernel); } -template -static inline void launch_qc_csc_sub(const IdxT* indptr, const IdxT* index, - const T* data, T* sums_cells, - const bool* mask, int n_genes, - cudaStream_t stream) { - dim3 block(SPARSE_BLOCK_SIZE); - dim3 grid((n_genes + SPARSE_BLOCK_SIZE - 1) / SPARSE_BLOCK_SIZE); - qc_csc_sub_kernel<<>>( - indptr, index, data, sums_cells, mask, n_genes); - CUDA_CHECK_LAST_ERROR(qc_csc_sub_kernel); -} - -template -static inline void launch_qc_csr_sub(const IdxT* indptr, const IdxT* index, - const T* data, T* sums_cells, - const bool* mask, int n_cells, - cudaStream_t stream) { - dim3 block(SPARSE_BLOCK_SIZE); - dim3 grid((n_cells + SPARSE_BLOCK_SIZE - 1) / SPARSE_BLOCK_SIZE); - qc_csr_sub_kernel<<>>( - indptr, index, data, sums_cells, mask, n_cells); - CUDA_CHECK_LAST_ERROR(qc_csr_sub_kernel); -} - template static inline void launch_qc_dense_sub(const T* data, T* sums_cells, const bool* mask, int n_cells, @@ -88,15 +54,17 @@ void def_sparse_qc_csc(nb::module_& m) { gpu_array_c index, gpu_array_c data, gpu_array_c sums_cells, gpu_array_c sums_genes, gpu_array_c cell_ex, - gpu_array_c gene_ex, int n_genes, + gpu_array_c gene_ex, int n_genes, bool assume_unsorted, std::uintptr_t stream) { - launch_qc_csc(indptr.data(), index.data(), data.data(), - sums_cells.data(), sums_genes.data(), - cell_ex.data(), gene_ex.data(), n_genes, - (cudaStream_t)stream); + return launch_qc_sparse( + indptr.data(), index.data(), data.data(), sums_genes.data(), + gene_ex.data(), sums_cells.data(), cell_ex.data(), n_genes, + (int)sums_cells.shape(0), (long long)data.shape(0), + assume_unsorted, (cudaStream_t)stream); }, "indptr"_a, "index"_a, "data"_a, nb::kw_only(), "sums_cells"_a, - "sums_genes"_a, "cell_ex"_a, "gene_ex"_a, "n_genes"_a, "stream"_a = 0); + "sums_genes"_a, "cell_ex"_a, "gene_ex"_a, "n_genes"_a, + "assume_unsorted"_a = false, "stream"_a = 0); } template @@ -107,15 +75,38 @@ void def_sparse_qc_csr(nb::module_& m) { gpu_array_c index, gpu_array_c data, gpu_array_c sums_cells, gpu_array_c sums_genes, gpu_array_c cell_ex, - gpu_array_c gene_ex, int n_cells, + gpu_array_c gene_ex, int n_cells, bool assume_unsorted, std::uintptr_t stream) { - launch_qc_csr(indptr.data(), index.data(), data.data(), - sums_cells.data(), sums_genes.data(), - cell_ex.data(), gene_ex.data(), n_cells, - (cudaStream_t)stream); + return launch_qc_sparse( + indptr.data(), index.data(), data.data(), sums_cells.data(), + cell_ex.data(), sums_genes.data(), gene_ex.data(), n_cells, + (int)sums_genes.shape(0), (long long)data.shape(0), + assume_unsorted, (cudaStream_t)stream); }, "indptr"_a, "index"_a, "data"_a, nb::kw_only(), "sums_cells"_a, - "sums_genes"_a, "cell_ex"_a, "gene_ex"_a, "n_cells"_a, "stream"_a = 0); + "sums_genes"_a, "cell_ex"_a, "gene_ex"_a, "n_cells"_a, + "assume_unsorted"_a = false, "stream"_a = 0); +} + +// Minor-axis half of sparse_qc_csr alone (Dask chunks reduce the gene side +// separately from the cell side). +template +void def_sparse_qc_genes(nb::module_& m) { + m.def( + "sparse_qc_genes", + [](gpu_array_c indptr, + gpu_array_c index, + gpu_array_c data, gpu_array_c sums_genes, + gpu_array_c gene_ex, bool assume_unsorted, + std::uintptr_t stream) { + QcOp op{data.data(), sums_genes.data(), gene_ex.data(), 0}; + return minor_reduce( + indptr.data(), index.data(), op, (int)indptr.shape(0) - 1, + (int)sums_genes.shape(0), (long long)data.shape(0), + assume_unsorted, (cudaStream_t)stream); + }, + "indptr"_a, "index"_a, "data"_a, nb::kw_only(), "sums_genes"_a, + "gene_ex"_a, "assume_unsorted"_a = false, "stream"_a = 0); } template @@ -135,6 +126,7 @@ void def_sparse_qc_dense(nb::module_& m) { "gene_ex"_a, "n_cells"_a, "n_genes"_a, "stream"_a = 0); } +// Masked genes (rows of the CSC) summed per cell (minor axis). template void def_sparse_qc_csc_sub(nb::module_& m) { m.def( @@ -143,15 +135,18 @@ void def_sparse_qc_csc_sub(nb::module_& m) { gpu_array_c index, gpu_array_c data, gpu_array_c sums_cells, gpu_array_c mask, int n_genes, - std::uintptr_t stream) { - launch_qc_csc_sub(indptr.data(), index.data(), data.data(), - sums_cells.data(), mask.data(), n_genes, - (cudaStream_t)stream); + bool assume_unsorted, std::uintptr_t stream) { + MinorSumOp op{data.data(), sums_cells.data(), mask.data(), 0}; + return minor_reduce(indptr.data(), index.data(), op, n_genes, + (int)sums_cells.shape(0), + (long long)data.shape(0), assume_unsorted, + (cudaStream_t)stream); }, "indptr"_a, "index"_a, "data"_a, nb::kw_only(), "sums_cells"_a, - "mask"_a, "n_genes"_a, "stream"_a = 0); + "mask"_a, "n_genes"_a, "assume_unsorted"_a = false, "stream"_a = 0); } +// Masked genes summed per cell (major axis of the CSR); no atomics. template void def_sparse_qc_csr_sub(nb::module_& m) { m.def( @@ -161,9 +156,9 @@ void def_sparse_qc_csr_sub(nb::module_& m) { gpu_array_c data, gpu_array_c sums_cells, gpu_array_c mask, int n_cells, std::uintptr_t stream) { - launch_qc_csr_sub(indptr.data(), index.data(), data.data(), - sums_cells.data(), mask.data(), n_cells, - (cudaStream_t)stream); + row_reduce(indptr.data(), index.data(), data.data(), + mask.data(), sums_cells.data(), nullptr, + n_cells, (cudaStream_t)stream); }, "indptr"_a, "index"_a, "data"_a, nb::kw_only(), "sums_cells"_a, "mask"_a, "n_cells"_a, "stream"_a = 0); @@ -195,6 +190,11 @@ void register_bindings(nb::module_& m) { def_sparse_qc_csr(m); def_sparse_qc_csr(m); + def_sparse_qc_genes(m); + def_sparse_qc_genes(m); + def_sparse_qc_genes(m); + def_sparse_qc_genes(m); + def_sparse_qc_csc_sub(m); def_sparse_qc_csc_sub(m); def_sparse_qc_csc_sub(m); @@ -213,4 +213,5 @@ void register_bindings(nb::module_& m) { NB_MODULE(_qc_cuda, m) { REGISTER_GPU_BINDINGS(register_bindings, m); + register_scratch_allocator(m); } diff --git a/src/rapids_singlecell/_cuda/qc_dask/kernels_qcd.cuh b/src/rapids_singlecell/_cuda/qc_dask/kernels_qcd.cuh index 214f1a3f0..7089b4a43 100644 --- a/src/rapids_singlecell/_cuda/qc_dask/kernels_qcd.cuh +++ b/src/rapids_singlecell/_cuda/qc_dask/kernels_qcd.cuh @@ -1,41 +1,8 @@ #pragma once #include - -template -__global__ void qc_csr_cells_kernel(const IdxT* __restrict__ indptr, - const IdxT* __restrict__ index, - const T* __restrict__ data, - T* __restrict__ sums_cells, - int* __restrict__ cell_ex, int n_cells) { - int cell = blockDim.x * blockIdx.x + threadIdx.x; - if (cell >= n_cells) return; - IdxT start_idx = indptr[cell]; - IdxT stop_idx = indptr[cell + 1]; - T sums = T(0); - int ex = 0; - for (IdxT p = start_idx; p < stop_idx; ++p) { - sums += data[p]; - ++ex; - } - sums_cells[cell] = sums; - cell_ex[cell] = ex; -} - -template -__global__ void qc_csr_genes_kernel(const IdxT* __restrict__ index, - const T* __restrict__ data, - T* __restrict__ sums_genes, - int* __restrict__ gene_ex, long long nnz) { - const long long stride = (long long)blockDim.x * gridDim.x; - for (long long i = (long long)blockDim.x * blockIdx.x + threadIdx.x; - i < nnz; i += stride) { - IdxT g = index[i]; - T v = data[i]; - atomicAdd(&sums_genes[g], v); - atomicAdd(&gene_ex[g], 1); - } -} +#include "../minor_tiles.cuh" +#include "../qc/kernels_qc.cuh" template __global__ void qc_dense_cells_kernel(const T* __restrict__ data, diff --git a/src/rapids_singlecell/_cuda/qc_dask/qc_kernels_dask.cu b/src/rapids_singlecell/_cuda/qc_dask/qc_kernels_dask.cu index 4057742b6..e577871a6 100644 --- a/src/rapids_singlecell/_cuda/qc_dask/qc_kernels_dask.cu +++ b/src/rapids_singlecell/_cuda/qc_dask/qc_kernels_dask.cu @@ -5,33 +5,8 @@ using namespace nb::literals; -constexpr int SPARSE_BLOCK_SIZE = 32; -constexpr int GENES_BLOCK_SIZE = 256; constexpr int DENSE_BLOCK_DIM = 16; -template -static inline void launch_qc_csr_cells(const IdxT* indptr, const IdxT* index, - const T* data, T* sums_cells, - int* cell_ex, int n_cells, - cudaStream_t stream) { - dim3 block(SPARSE_BLOCK_SIZE); - dim3 grid((n_cells + SPARSE_BLOCK_SIZE - 1) / SPARSE_BLOCK_SIZE); - qc_csr_cells_kernel<<>>( - indptr, index, data, sums_cells, cell_ex, n_cells); - CUDA_CHECK_LAST_ERROR(qc_csr_cells_kernel); -} - -template -static inline void launch_qc_csr_genes(const IdxT* index, const T* data, - T* sums_genes, int* gene_ex, - long long nnz, cudaStream_t stream) { - int block = GENES_BLOCK_SIZE; - unsigned int grid = strided_grid(nnz, block); - qc_csr_genes_kernel - <<>>(index, data, sums_genes, gene_ex, nnz); - CUDA_CHECK_LAST_ERROR(qc_csr_genes_kernel); -} - template static inline void launch_qc_dense_cells(const T* data, T* sums_cells, int* cell_ex, int n_cells, int n_genes, @@ -65,9 +40,9 @@ void def_sparse_qc_csr_cells(nb::module_& m) { gpu_array_c data, gpu_array_c sums_cells, gpu_array_c cell_ex, int n_cells, std::uintptr_t stream) { - launch_qc_csr_cells( - indptr.data(), index.data(), data.data(), sums_cells.data(), - cell_ex.data(), n_cells, (cudaStream_t)stream); + row_reduce(indptr.data(), index.data(), data.data(), + nullptr, sums_cells.data(), cell_ex.data(), + n_cells, (cudaStream_t)stream); }, "indptr"_a, "index"_a, "data"_a, nb::kw_only(), "sums_cells"_a, "cell_ex"_a, "n_cells"_a, "stream"_a = 0); @@ -81,9 +56,9 @@ void def_sparse_qc_csr_genes(nb::module_& m) { gpu_array_c data, gpu_array_c sums_genes, gpu_array_c gene_ex, long long nnz, std::uintptr_t stream) { - launch_qc_csr_genes(index.data(), data.data(), - sums_genes.data(), gene_ex.data(), nnz, - (cudaStream_t)stream); + QcOp op{data.data(), sums_genes.data(), gene_ex.data(), 0}; + minor_reduce_flat(index.data(), op, nnz, + (cudaStream_t)stream); }, "index"_a, "data"_a, nb::kw_only(), "sums_genes"_a, "gene_ex"_a, "nnz"_a, "stream"_a = 0); @@ -140,4 +115,5 @@ void register_bindings(nb::module_& m) { NB_MODULE(_qc_dask_cuda, m) { REGISTER_GPU_BINDINGS(register_bindings, m); + register_scratch_allocator(m); } diff --git a/src/rapids_singlecell/_cuda/spca/kernels_spca.cuh b/src/rapids_singlecell/_cuda/spca/kernels_spca.cuh index 02b81a546..6e6b1ed3d 100644 --- a/src/rapids_singlecell/_cuda/spca/kernels_spca.cuh +++ b/src/rapids_singlecell/_cuda/spca/kernels_spca.cuh @@ -63,15 +63,3 @@ __global__ void cov_from_gram_kernel(T* cov_values, const T* gram_matrix, } } } - -template -__global__ void check_zero_genes_kernel(const IdxT* indices, int* genes, - long long nnz, int num_genes) { - const long long stride = (long long)blockDim.x * gridDim.x; - for (long long value = (long long)blockIdx.x * blockDim.x + threadIdx.x; - value < nnz; value += stride) { - long long gene_index = static_cast(indices[value]); - if (gene_index < 0 || gene_index >= num_genes) continue; - atomicAdd(&genes[gene_index], 1); - } -} diff --git a/src/rapids_singlecell/_cuda/spca/spca.cu b/src/rapids_singlecell/_cuda/spca/spca.cu index a85b86879..f14d8bb66 100644 --- a/src/rapids_singlecell/_cuda/spca/spca.cu +++ b/src/rapids_singlecell/_cuda/spca/spca.cu @@ -1,4 +1,9 @@ #include +#include + +#include + +#include "../minor_tiles.cuh" #include "../nb_types.h" #include "kernels_spca.cuh" @@ -7,7 +12,6 @@ using namespace nb::literals; constexpr int GRAM_BLOCK_SIZE = 128; constexpr int MATRIX_BLOCK_DIM = 32; -constexpr int ELEMENTWISE_BLOCK_SIZE = 32; template static inline void launch_gram_csr_upper(const IdxT* indptr, const IdxT* index, @@ -43,19 +47,6 @@ static inline void launch_cov_from_gram(T* cov, const T* gram, const T* meanx, CUDA_CHECK_LAST_ERROR(cov_from_gram_kernel); } -template -static inline void launch_check_zero_genes(const IdxT* indices, int* genes, - long long nnz, int num_genes, - cudaStream_t stream) { - if (nnz > 0) { - dim3 block(ELEMENTWISE_BLOCK_SIZE); - dim3 grid(strided_grid(nnz, ELEMENTWISE_BLOCK_SIZE)); - check_zero_genes_kernel - <<>>(indices, genes, nnz, num_genes); - CUDA_CHECK_LAST_ERROR(check_zero_genes_kernel); - } -} - template void def_gram_csr_upper(nb::module_& m) { m.def( @@ -98,18 +89,28 @@ void def_cov_from_gram(nb::module_& m) { "stream"_a = 0); } +// Stored entries per gene. With `indptr` the tile sweep is used; returns +// whether unsorted rows were detected. template void def_check_zero_genes(nb::module_& m) { m.def( "check_zero_genes", [](gpu_array_c indices, gpu_array_c out, long long nnz, int num_genes, - std::uintptr_t stream) { - launch_check_zero_genes(indices.data(), out.data(), nnz, - num_genes, (cudaStream_t)stream); + std::optional> indptr, + bool assume_unsorted, std::uintptr_t stream) { + MinorCountOp op{out.data(), num_genes, 0}; + if (!indptr) { + minor_reduce_flat(indices.data(), op, nnz, + (cudaStream_t)stream); + return false; + } + return minor_reduce(indptr->data(), indices.data(), op, + (int)indptr->shape(0) - 1, num_genes, nnz, + assume_unsorted, (cudaStream_t)stream); }, "indices"_a, nb::kw_only(), "out"_a, "nnz"_a, "num_genes"_a, - "stream"_a = 0); + "indptr"_a = nb::none(), "assume_unsorted"_a = false, "stream"_a = 0); } template @@ -131,4 +132,5 @@ void register_bindings(nb::module_& m) { NB_MODULE(_spca_cuda, m) { REGISTER_GPU_BINDINGS(register_bindings, m); + register_scratch_allocator(m); } diff --git a/src/rapids_singlecell/_utils/_sparse_rows.py b/src/rapids_singlecell/_utils/_sparse_rows.py new file mode 100644 index 000000000..7f64de9f1 --- /dev/null +++ b/src/rapids_singlecell/_utils/_sparse_rows.py @@ -0,0 +1,25 @@ +"""Bookkeeping for the shared-memory minor-axis reductions in ``_cuda/minor_tiles.cuh``.""" + +from __future__ import annotations + + +def _known_unsorted(X) -> bool: + """ + Whether ``X`` is already known to have unsorted indices within its rows. + + Reads cupyx's cached flag without triggering its check kernel. An unknown + flag returns ``False`` so the tiled kernel detects the order itself. + """ + return getattr(X, "_has_canonical_format", None) is False + + +def _minor_reduce(X, kernel, *args, **kwargs) -> None: + """ + Run a minor-axis reduction binding on ``X``. + + The binding tries the tile sweep unless ``X`` is known to be unsorted and + reports when it detected unsorted rows; that is remembered on ``X`` so the + next call goes straight to the atomic kernel. + """ + if kernel(*args, assume_unsorted=_known_unsorted(X), **kwargs): + X.has_canonical_format = False diff --git a/src/rapids_singlecell/get/_aggregated.py b/src/rapids_singlecell/get/_aggregated.py index eea8cab0a..3cabc84d9 100644 --- a/src/rapids_singlecell/get/_aggregated.py +++ b/src/rapids_singlecell/get/_aggregated.py @@ -16,6 +16,7 @@ from rapids_singlecell._compat import DaskArray, _meta_dense from rapids_singlecell._cuda import _aggr_cuda from rapids_singlecell._settings import Preset, settings +from rapids_singlecell._utils._sparse_rows import _minor_reduce from rapids_singlecell.preprocessing._utils import _check_gpu_X from ._utils import ( @@ -239,7 +240,9 @@ def __aggregate_dask(X_part, mask_part, groupby_part): out_sqsum = out[0, i_sqsum] if need_sqsum else None if is_sparse: - _aggr_cuda.sparse_aggr( + _minor_reduce( + X_part, + _aggr_cuda.sparse_aggr, X_part.indptr, X_part.indices, X_part.data, @@ -324,7 +327,9 @@ def count_mean_var_sparse(self, funcs=None, *, dof: int = 1): ) mask = self._get_mask() - _aggr_cuda.sparse_aggr( + _minor_reduce( + self.data, + _aggr_cuda.sparse_aggr, self.data.indptr, self.data.indices, self.data.data, diff --git a/src/rapids_singlecell/preprocessing/_hvg/_pearson_residuals.py b/src/rapids_singlecell/preprocessing/_hvg/_pearson_residuals.py index 4b882a71c..8a4239e66 100644 --- a/src/rapids_singlecell/preprocessing/_hvg/_pearson_residuals.py +++ b/src/rapids_singlecell/preprocessing/_hvg/_pearson_residuals.py @@ -9,6 +9,7 @@ from cupyx.scipy.sparse import issparse from rapids_singlecell._cuda import _pr_cuda +from rapids_singlecell._utils._sparse_rows import _minor_reduce from rapids_singlecell.get import _get_obs_rep from rapids_singlecell.preprocessing._utils import ( _check_gpu_X, @@ -90,7 +91,9 @@ def _highly_variable_pearson_residuals( # Compute sums using custom kernel (single pass, no CSC->CSR conversion) sums_genes = cp.zeros(n_genes, dtype=dtype) sums_cells = cp.zeros(n_cells, dtype=dtype) - _pr_cuda.sparse_sum_csc( + _minor_reduce( + X_batch, + _pr_cuda.sparse_sum_csc, X_batch.indptr, X_batch.indices, X_batch.data, diff --git a/src/rapids_singlecell/preprocessing/_hvg/_seurat_v3.py b/src/rapids_singlecell/preprocessing/_hvg/_seurat_v3.py index 910c55b9e..3990d6033 100644 --- a/src/rapids_singlecell/preprocessing/_hvg/_seurat_v3.py +++ b/src/rapids_singlecell/preprocessing/_hvg/_seurat_v3.py @@ -9,6 +9,7 @@ from cupyx.scipy.sparse import issparse, isspmatrix_csc from rapids_singlecell._compat import DaskArray +from rapids_singlecell._utils._sparse_rows import _minor_reduce from rapids_singlecell.get import _get_obs_rep from rapids_singlecell.preprocessing._utils import ( _check_gpu_X, @@ -20,31 +21,24 @@ from anndata import AnnData -_seurat_v3_elementwise_kernel = cp.ElementwiseKernel( - "T data, S idx, raw D clip_val", - "raw D sq_sum, raw D sum", - """ - D element = min((double)data, clip_val[idx]); - atomicAdd(&sq_sum[idx], element * element); - atomicAdd(&sum[idx], element); - """, - "seurat_v3_elementwise_kernel", - no_return=True, -) - - def _clip_square_sum_sparse(X, clip_val): """Compute clipped sum and sum-of-squares for a sparse CSR matrix.""" + from rapids_singlecell._cuda import _hvg_cuda + if isspmatrix_csc(X): X = X.tocsr() squared_batch_counts_sum = cp.zeros(clip_val.shape, dtype=cp.float64) batch_counts_sum = cp.zeros(clip_val.shape, dtype=cp.float64) - _seurat_v3_elementwise_kernel( - X.data, + _minor_reduce( + X, + _hvg_cuda.clip_square_sum, + X.indptr, X.indices, - clip_val, - squared_batch_counts_sum, - batch_counts_sum, + X.data, + clip_val=cp.ascontiguousarray(clip_val, dtype=cp.float64), + sq_sum=squared_batch_counts_sum, + sum=batch_counts_sum, + stream=cp.cuda.get_current_stream().ptr, ) return squared_batch_counts_sum, batch_counts_sum diff --git a/src/rapids_singlecell/preprocessing/_qc.py b/src/rapids_singlecell/preprocessing/_qc.py index 249ce3a04..5b90f8c40 100644 --- a/src/rapids_singlecell/preprocessing/_qc.py +++ b/src/rapids_singlecell/preprocessing/_qc.py @@ -6,6 +6,7 @@ from cupyx.scipy import sparse from rapids_singlecell._compat import DaskArray +from rapids_singlecell._utils._sparse_rows import _minor_reduce from rapids_singlecell.get import _get_obs_rep from ._utils import _check_gpu_X @@ -136,7 +137,9 @@ def _basic_qc( else: raise ValueError("Please use a csr or csc matrix") - sparse_qc( + _minor_reduce( + X, + sparse_qc, X.indptr, X.indices, X.data, @@ -188,14 +191,18 @@ def __qc_calc_1(X_part): return cp.stack([sums_cells, genes_per_cell.astype(X_part.dtype)], axis=1) def __qc_calc_2(X_part): + from rapids_singlecell._cuda import _qc_cuda as _qc + sums_genes = cp.zeros(X_part.shape[1], dtype=X_part.dtype) cells_per_gene = cp.zeros(X_part.shape[1], dtype=cp.int32) - _qcd.sparse_qc_csr_genes( + _minor_reduce( + X_part, + _qc.sparse_qc_genes, + X_part.indptr, X_part.indices, X_part.data, sums_genes=sums_genes, gene_ex=cells_per_gene, - nnz=X_part.nnz, stream=cp.cuda.get_current_stream().ptr, ) return cp.vstack([sums_genes, cells_per_gene.astype(X_part.dtype)])[ @@ -289,7 +296,9 @@ def _geneset_qc(X: ArrayTypesDask, mask: cp.ndarray) -> cp.ndarray: stream=cp.cuda.get_current_stream().ptr, ) elif sparse.isspmatrix_csc(X): - _qc.sparse_qc_csc_sub( + _minor_reduce( + X, + _qc.sparse_qc_csc_sub, X.indptr, X.indices, X.data, diff --git a/src/rapids_singlecell/preprocessing/_sparse_pca/_helper.py b/src/rapids_singlecell/preprocessing/_sparse_pca/_helper.py index 89423e960..2dbab2348 100644 --- a/src/rapids_singlecell/preprocessing/_sparse_pca/_helper.py +++ b/src/rapids_singlecell/preprocessing/_sparse_pca/_helper.py @@ -7,6 +7,7 @@ if TYPE_CHECKING: from cupyx.scipy.sparse import spmatrix from rapids_singlecell._cuda import _spca_cuda as _spca +from rapids_singlecell._utils._sparse_rows import _minor_reduce def _copy_gram(gram_matrix: cp.ndarray, n_cols: int) -> cp.ndarray: @@ -35,11 +36,14 @@ def _compute_cov( def _check_matrix_for_zero_genes(X: spmatrix) -> None: gene_ex = cp.zeros(X.shape[1], dtype=cp.int32) if X.nnz > 0: - _spca.check_zero_genes( + _minor_reduce( + X, + _spca.check_zero_genes, X.indices, out=gene_ex, nnz=X.nnz, num_genes=X.shape[1], + indptr=X.indptr, stream=cp.cuda.get_current_stream().ptr, ) if cp.any(gene_ex == 0): diff --git a/src/rapids_singlecell/preprocessing/_utils.py b/src/rapids_singlecell/preprocessing/_utils.py index fc87d0110..7cbace19d 100644 --- a/src/rapids_singlecell/preprocessing/_utils.py +++ b/src/rapids_singlecell/preprocessing/_utils.py @@ -10,6 +10,7 @@ from pandas.api.types import infer_dtype from rapids_singlecell._compat import DaskArray +from rapids_singlecell._utils._sparse_rows import _minor_reduce if TYPE_CHECKING: from anndata import AnnData @@ -83,19 +84,40 @@ def _mean_var_major(X, major, minor): return mean, var -def _mean_var_minor(X, major, minor): +def _mean_var_minor_sums(X, mean, var): + """ + Accumulate the minor-axis sums and squared sums of a compressed matrix. + + Parameters + ---------- + X + CSR or CSC matrix + mean + Zero-initialised float64 output for the sums + var + Zero-initialised float64 output for the squared sums + """ from rapids_singlecell._cuda import _mean_var_cuda as _mv - mean = cp.zeros(minor, dtype=cp.float64) - var = cp.zeros(minor, dtype=cp.float64) - _mv.mean_var_minor( + _minor_reduce( + X, + _mv.mean_var_minor_tiled, + X.indptr, X.indices, X.data, mean, var, + major=X.indptr.size - 1, + minor=mean.size, nnz=X.nnz, stream=cp.cuda.get_current_stream().ptr, ) + + +def _mean_var_minor(X, major, minor): + mean = cp.zeros(minor, dtype=cp.float64) + var = cp.zeros(minor, dtype=cp.float64) + _mean_var_minor_sums(X, mean, var) mean /= major var /= major var -= mean**2 @@ -108,19 +130,10 @@ def _mean_var_minor_dask(X, major, minor): Implements sum operation for dask array when the backend is cupy sparse csr matrix """ - from rapids_singlecell._cuda import _mean_var_cuda as _mv - def __mean_var(X_part): mean = cp.zeros(minor, dtype=cp.float64) var = cp.zeros(minor, dtype=cp.float64) - _mv.mean_var_minor( - X_part.indices, - X_part.data, - mean, - var, - nnz=X_part.nnz, - stream=cp.cuda.get_current_stream().ptr, - ) + _mean_var_minor_sums(X_part, mean, var) return cp.vstack([mean, var])[None, ...] # new axis for summing n_blocks = X.blocks.size diff --git a/src/rapids_singlecell/squidpy_gpu/_ligrec.py b/src/rapids_singlecell/squidpy_gpu/_ligrec.py index f1b730b81..0f7bea111 100644 --- a/src/rapids_singlecell/squidpy_gpu/_ligrec.py +++ b/src/rapids_singlecell/squidpy_gpu/_ligrec.py @@ -14,6 +14,8 @@ from cupyx.scipy.sparse import issparse as cpissparse from scipy.sparse import csc_matrix, issparse +from rapids_singlecell._utils._sparse_rows import _minor_reduce + from ._utils import _assert_categorical_obs, _create_sparse_df SOURCE = "source" @@ -474,7 +476,9 @@ def find_min_gene_in_complex(_complex: str | None) -> str | None: else: sum_gt0 = cp.zeros((data_cp.shape[1], n_clusters), dtype=cp.float32, order="C") count_gt0 = cp.zeros((data_cp.shape[1], n_clusters), dtype=cp.int32, order="C") - _lc.sum_count_sparse( + _minor_reduce( + data_cp, + _lc.sum_count_sparse, data_cp.indptr, data_cp.indices, data_cp.data, @@ -505,7 +509,9 @@ def find_min_gene_in_complex(_complex: str | None) -> str | None: for _i in range(n_perms): cp.random.shuffle(clustering_use) g = cp.zeros((data_cp.shape[1], n_cls), dtype=cp.float32, order="C") - _lc.mean_sparse( + _minor_reduce( + data_cp, + _lc.mean_sparse, data_cp.indptr, data_cp.indices, data_cp.data, diff --git a/src/rapids_singlecell/tools/_utils.py b/src/rapids_singlecell/tools/_utils.py index 6da7eccbc..62bfbd335 100644 --- a/src/rapids_singlecell/tools/_utils.py +++ b/src/rapids_singlecell/tools/_utils.py @@ -7,6 +7,7 @@ from rapids_singlecell._compat import DaskArray from rapids_singlecell._keys import _embedding_keys, _existing_preset_keys from rapids_singlecell._settings import settings +from rapids_singlecell._utils._sparse_rows import _minor_reduce from . import pca @@ -67,21 +68,44 @@ def _choose_representation(adata, use_rep=None, n_pcs=None): return X -def _nan_mean_minor_dask_sparse(X, major, minor, *, mask=None, n_features=None): +def _nan_mean_minor_sums(X, mean, nans, mask): + """ + Accumulate the minor-axis NaN-aware sums and NaN counts of a compressed matrix. + + Parameters + ---------- + X + CSR or CSC matrix + mean + Zero-initialised float64 output for the sums + nans + Zero-initialised int32 output for the NaN counts + mask + Boolean mask over the minor axis + """ from rapids_singlecell._cuda import _nanmean_cuda as _nm + _minor_reduce( + X, + _nm.nan_mean_minor_tiled, + X.indptr, + X.indices, + X.data, + means=mean, + nans=nans, + mask=mask, + major=X.indptr.size - 1, + minor=mean.size, + nnz=X.nnz, + stream=cp.cuda.get_current_stream().ptr, + ) + + +def _nan_mean_minor_dask_sparse(X, major, minor, *, mask=None, n_features=None): def __nan_mean_minor(X_part): mean = cp.zeros(minor, dtype=cp.float64) nans = cp.zeros(minor, dtype=cp.int32) - _nm.nan_mean_minor( - X_part.indices, - X_part.data, - means=mean, - nans=nans, - mask=mask, - nnz=X_part.nnz, - stream=cp.cuda.get_current_stream().ptr, - ) + _nan_mean_minor_sums(X_part, mean, nans, mask) return cp.vstack([mean, nans.astype(cp.float64)])[None, ...] n_blocks = X.blocks.size @@ -156,19 +180,9 @@ def __nan_mean_dense(X_part): def _nan_mean_minor(X, major, minor, *, mask=None, n_features=None): - from rapids_singlecell._cuda import _nanmean_cuda as _nm - mean = cp.zeros(minor, dtype=cp.float64) nans = cp.zeros(minor, dtype=cp.int32) - _nm.nan_mean_minor( - X.indices, - X.data, - means=mean, - nans=nans, - mask=mask, - nnz=X.nnz, - stream=cp.cuda.get_current_stream().ptr, - ) + _nan_mean_minor_sums(X, mean, nans, mask) mean /= n_features - nans return mean From f1a1397994015dfece54c568a79bd4a9e2da98e7 Mon Sep 17 00:00:00 2001 From: Intron7 Date: Mon, 7 Sep 2026 19:23:03 +0200 Subject: [PATCH 2/4] adress issues Signed-off-by: Intron7 --- src/rapids_singlecell/_cuda/aggr/aggr.cu | 21 +++++++++--- .../_cuda/aggr/kernels_aggr.cuh | 9 ++++-- src/rapids_singlecell/_cuda/hvg/hvg.cu | 23 ++++++++++--- .../_cuda/ligrec/kernels_ligrec.cuh | 10 ++++-- .../_cuda/mean_var/kernels_mv.cuh | 8 +++-- .../_cuda/mean_var/mean_var.cu | 4 +++ src/rapids_singlecell/_cuda/minor_tiles.cuh | 32 +++++++++++++++++-- .../_cuda/nanmean/kernels_nanmean.cuh | 8 +++-- src/rapids_singlecell/_cuda/qc/kernels_qc.cuh | 7 ++-- src/rapids_singlecell/_cuda/spca/spca.cu | 9 ++++++ 10 files changed, 109 insertions(+), 22 deletions(-) diff --git a/src/rapids_singlecell/_cuda/aggr/aggr.cu b/src/rapids_singlecell/_cuda/aggr/aggr.cu index c62c30096..be53015ed 100644 --- a/src/rapids_singlecell/_cuda/aggr/aggr.cu +++ b/src/rapids_singlecell/_cuda/aggr/aggr.cu @@ -132,10 +132,23 @@ void def_sparse_aggr(nb::module_& m) { double* pq = out_sqsum ? out_sqsum->data() : nullptr; int active = (ps ? AGGR_SUM : 0) | (pc ? AGGR_COUNT : 0) | (pq ? AGGR_SQSUM : 0); - const int n_groups = out_sum ? (int)out_sum->shape(0) - : out_count ? (int)out_count->shape(0) - : out_sqsum ? (int)out_sqsum->shape(0) - : 0; + // Every provided plane must be (n_groups, n_genes); n_groups is + // read off shape(0), so a flattened plane would over-size the + // memsets and the grouping. + int n_groups = 0; + for (const auto* plane : {&out_sum, &out_count, &out_sqsum}) { + if (!*plane) continue; + require_arg( + (*plane)->ndim() == 2 && (*plane)->shape(1) == n_genes, + "sparse_aggr: outputs must be (n_groups, n_genes)"); + const int g = (int)(*plane)->shape(0); + require_arg(n_groups == 0 || g == n_groups, + "sparse_aggr: outputs must have the same shape"); + n_groups = g; + } + require_csr_arrays("sparse_aggr", indptr, index, data); + require_arg(cats.shape(0) == n_cells && mask.shape(0) == n_cells, + "sparse_aggr: cats and mask must have n_cells entries"); bool unsorted = false; #define LAUNCH(M) \ unsorted = launch_sparse_aggr( \ diff --git a/src/rapids_singlecell/_cuda/aggr/kernels_aggr.cuh b/src/rapids_singlecell/_cuda/aggr/kernels_aggr.cuh index 94458d921..7ec260616 100644 --- a/src/rapids_singlecell/_cuda/aggr/kernels_aggr.cuh +++ b/src/rapids_singlecell/_cuda/aggr/kernels_aggr.cuh @@ -73,11 +73,14 @@ struct AggrOp { void zero_outputs(int minor, int n_groups, cudaStream_t stream) const { const size_t bytes = (size_t)minor * n_groups * sizeof(double); if constexpr (MASK & AGGR_SUM) - cudaMemsetAsync(out_sum, 0, bytes, stream); + cuda_check(cudaMemsetAsync(out_sum, 0, bytes, stream), + "cudaMemsetAsync(AggrOp outputs)"); if constexpr (MASK & AGGR_COUNT) - cudaMemsetAsync(out_count, 0, bytes, stream); + cuda_check(cudaMemsetAsync(out_count, 0, bytes, stream), + "cudaMemsetAsync(AggrOp outputs)"); if constexpr (MASK & AGGR_SQSUM) - cudaMemsetAsync(out_sqsum, 0, bytes, stream); + cuda_check(cudaMemsetAsync(out_sqsum, 0, bytes, stream), + "cudaMemsetAsync(AggrOp outputs)"); } }; diff --git a/src/rapids_singlecell/_cuda/hvg/hvg.cu b/src/rapids_singlecell/_cuda/hvg/hvg.cu index dd8c65be1..15c6c2863 100644 --- a/src/rapids_singlecell/_cuda/hvg/hvg.cu +++ b/src/rapids_singlecell/_cuda/hvg/hvg.cu @@ -33,6 +33,12 @@ static void launch_expected_zeros(const T* scaled_means, const T* total_counts, CUDA_CHECK_LAST_ERROR(expected_zeros_kernel); } +/// min(value, clip) that propagates NaN like cupy.minimum on the dense path +/// (fmin would return the finite operand). +__device__ inline double clip_min(double value, double clip) { + return (isnan(value) || isnan(clip)) ? NAN : fmin(value, clip); +} + /// Per-column sum and sum-of-squares of min(value, clip[col]) for seurat_v3 /// (see minor_tiles.cuh). Layout: double sq-sums, double sums, double clips. template @@ -56,7 +62,7 @@ struct ClipSumOp { __device__ void add(char* acc, long long q, int g) const { double* s = reinterpret_cast(acc); const double e = - fmin(static_cast(data[q]), s[2 * tile_size + g]); + clip_min(static_cast(data[q]), s[2 * tile_size + g]); atomicAdd(&s[g], e * e); atomicAdd(&s[tile_size + g], e); } @@ -68,13 +74,17 @@ struct ClipSumOp { } } __device__ void add_global(long long q, int col, int) const { - const double e = fmin(static_cast(data[q]), clip[col]); + const double e = clip_min(static_cast(data[q]), clip[col]); atomicAdd(&sq_sum[col], e * e); atomicAdd(&sum[col], e); } void zero_outputs(int minor, int, cudaStream_t stream) const { - cudaMemsetAsync(sq_sum, 0, (size_t)minor * sizeof(double), stream); - cudaMemsetAsync(sum, 0, (size_t)minor * sizeof(double), stream); + cuda_check( + cudaMemsetAsync(sq_sum, 0, (size_t)minor * sizeof(double), stream), + "cudaMemsetAsync(ClipSumOp outputs)"); + cuda_check( + cudaMemsetAsync(sum, 0, (size_t)minor * sizeof(double), stream), + "cudaMemsetAsync(ClipSumOp outputs)"); } }; @@ -89,6 +99,11 @@ void def_clip_square_sum(nb::module_& m) { gpu_array_c clip_val, gpu_array_c sq_sum, gpu_array_c sum, bool assume_unsorted, std::uintptr_t stream) { + require_csr_arrays("clip_square_sum", indptr, indices, data); + require_arg(clip_val.shape(0) == sum.shape(0) && + sq_sum.shape(0) == sum.shape(0), + "clip_square_sum: clip_val, sq_sum and sum must have " + "one entry per column"); ClipSumOp op{data.data(), clip_val.data(), sq_sum.data(), sum.data(), 0}; return minor_reduce( diff --git a/src/rapids_singlecell/_cuda/ligrec/kernels_ligrec.cuh b/src/rapids_singlecell/_cuda/ligrec/kernels_ligrec.cuh index edbee1096..e63f9d497 100644 --- a/src/rapids_singlecell/_cuda/ligrec/kernels_ligrec.cuh +++ b/src/rapids_singlecell/_cuda/ligrec/kernels_ligrec.cuh @@ -189,9 +189,13 @@ struct LigrecOp { } } void zero_outputs(int minor, int n_groups, cudaStream_t stream) const { - cudaMemsetAsync(sum, 0, (size_t)minor * n_groups * sizeof(T), stream); + cuda_check(cudaMemsetAsync(sum, 0, (size_t)minor * n_groups * sizeof(T), + stream), + "cudaMemsetAsync(LigrecOp outputs)"); if constexpr (WITH_COUNT) - cudaMemsetAsync(count, 0, (size_t)minor * n_groups * sizeof(int), - stream); + cuda_check( + cudaMemsetAsync(count, 0, + (size_t)minor * n_groups * sizeof(int), stream), + "cudaMemsetAsync(LigrecOp outputs)"); } }; diff --git a/src/rapids_singlecell/_cuda/mean_var/kernels_mv.cuh b/src/rapids_singlecell/_cuda/mean_var/kernels_mv.cuh index ee64374ae..9e94bf007 100644 --- a/src/rapids_singlecell/_cuda/mean_var/kernels_mv.cuh +++ b/src/rapids_singlecell/_cuda/mean_var/kernels_mv.cuh @@ -84,7 +84,11 @@ struct MeanVarOp { atomicAdd(&vars[col], v * v); } void zero_outputs(int minor, int, cudaStream_t stream) const { - cudaMemsetAsync(means, 0, (size_t)minor * sizeof(double), stream); - cudaMemsetAsync(vars, 0, (size_t)minor * sizeof(double), stream); + cuda_check( + cudaMemsetAsync(means, 0, (size_t)minor * sizeof(double), stream), + "cudaMemsetAsync(MeanVarOp outputs)"); + cuda_check( + cudaMemsetAsync(vars, 0, (size_t)minor * sizeof(double), stream), + "cudaMemsetAsync(MeanVarOp outputs)"); } }; diff --git a/src/rapids_singlecell/_cuda/mean_var/mean_var.cu b/src/rapids_singlecell/_cuda/mean_var/mean_var.cu index cf3cd15cd..ef29ba807 100644 --- a/src/rapids_singlecell/_cuda/mean_var/mean_var.cu +++ b/src/rapids_singlecell/_cuda/mean_var/mean_var.cu @@ -78,6 +78,10 @@ void def_tile_plan(nb::module_& m) { m.def( "tile_plan", [](long long nnz, int n_rows, int n_cols, size_t bytes_per_col) { + require_arg( + nnz >= 0 && n_rows > 0 && n_cols > 0 && bytes_per_col > 0, + "tile_plan: nnz must be >= 0 and n_rows, n_cols and " + "bytes_per_col must be positive"); const TilePlan p = plan_tiles(nnz, n_rows, n_cols, bytes_per_col); nb::dict d; d["use_tiled"] = p.use_tiled; diff --git a/src/rapids_singlecell/_cuda/minor_tiles.cuh b/src/rapids_singlecell/_cuda/minor_tiles.cuh index 9fb6e2d9c..3ae7d8490 100644 --- a/src/rapids_singlecell/_cuda/minor_tiles.cuh +++ b/src/rapids_singlecell/_cuda/minor_tiles.cuh @@ -4,6 +4,8 @@ #include #include +#include +#include #include #include "nb_types.h" @@ -57,6 +59,29 @@ // ligrec) first sort rows by group so a block only ever accumulates for one // group; see build_grouped_rows(). +// ---- argument validation --------------------------------------------------- + +/// Reject a malformed binding argument with a clear Python error instead of an +/// out-of-bounds kernel read. +inline void require_arg(bool cond, const char* what) { + if (!cond) throw std::invalid_argument(what); +} + +/// The host-side shape checks every compressed-matrix binding can afford: +/// a 1-D indptr and equally long indices/data. Offsets are not read back from +/// the device here; cupyx guarantees their structure. +template +inline void require_csr_arrays(const char* what, const IndptrArr& indptr, + const IdxArr& indices, const DataArr& data) { + require_arg( + indptr.ndim() == 1 && indptr.shape(0) >= 1, + (std::string(what) + ": indptr must be 1-D and non-empty").c_str()); + require_arg( + indices.shape(0) == data.shape(0), + (std::string(what) + ": indices and data must have equal length") + .c_str()); +} + // ---- tunables -------------------------------------------------------------- constexpr int WARP = 32; @@ -610,7 +635,8 @@ struct MinorSumOp { atomicAdd(&out[col], data[nnz_pos]); } void zero_outputs(int n_cols, int, cudaStream_t stream) const { - cudaMemsetAsync(out, 0, (size_t)n_cols * sizeof(T), stream); + cuda_check(cudaMemsetAsync(out, 0, (size_t)n_cols * sizeof(T), stream), + "cudaMemsetAsync(MinorSumOp outputs)"); } }; @@ -639,6 +665,8 @@ struct MinorCountOp { if (col >= 0 && col < n_cols) atomicAdd(&out[col], 1); } void zero_outputs(int n_cols_, int, cudaStream_t stream) const { - cudaMemsetAsync(out, 0, (size_t)n_cols_ * sizeof(int), stream); + cuda_check( + cudaMemsetAsync(out, 0, (size_t)n_cols_ * sizeof(int), stream), + "cudaMemsetAsync(MinorCountOp outputs)"); } }; diff --git a/src/rapids_singlecell/_cuda/nanmean/kernels_nanmean.cuh b/src/rapids_singlecell/_cuda/nanmean/kernels_nanmean.cuh index 35dedc8ff..3a266bf93 100644 --- a/src/rapids_singlecell/_cuda/nanmean/kernels_nanmean.cuh +++ b/src/rapids_singlecell/_cuda/nanmean/kernels_nanmean.cuh @@ -108,7 +108,11 @@ struct NanMeanOp { } } void zero_outputs(int minor, int, cudaStream_t stream) const { - cudaMemsetAsync(means, 0, (size_t)minor * sizeof(double), stream); - cudaMemsetAsync(nans, 0, (size_t)minor * sizeof(int), stream); + cuda_check( + cudaMemsetAsync(means, 0, (size_t)minor * sizeof(double), stream), + "cudaMemsetAsync(NanMeanOp outputs)"); + cuda_check( + cudaMemsetAsync(nans, 0, (size_t)minor * sizeof(int), stream), + "cudaMemsetAsync(NanMeanOp outputs)"); } }; diff --git a/src/rapids_singlecell/_cuda/qc/kernels_qc.cuh b/src/rapids_singlecell/_cuda/qc/kernels_qc.cuh index f74905c4c..99c48033a 100644 --- a/src/rapids_singlecell/_cuda/qc/kernels_qc.cuh +++ b/src/rapids_singlecell/_cuda/qc/kernels_qc.cuh @@ -76,7 +76,10 @@ struct QcOp { atomicAdd(&counts[col], 1); } void zero_outputs(int minor, int, cudaStream_t stream) const { - cudaMemsetAsync(sums, 0, (size_t)minor * sizeof(T), stream); - cudaMemsetAsync(counts, 0, (size_t)minor * sizeof(int), stream); + cuda_check(cudaMemsetAsync(sums, 0, (size_t)minor * sizeof(T), stream), + "cudaMemsetAsync(QcOp outputs)"); + cuda_check( + cudaMemsetAsync(counts, 0, (size_t)minor * sizeof(int), stream), + "cudaMemsetAsync(QcOp outputs)"); } }; diff --git a/src/rapids_singlecell/_cuda/spca/spca.cu b/src/rapids_singlecell/_cuda/spca/spca.cu index f14d8bb66..53de2ad36 100644 --- a/src/rapids_singlecell/_cuda/spca/spca.cu +++ b/src/rapids_singlecell/_cuda/spca/spca.cu @@ -99,6 +99,15 @@ void def_check_zero_genes(nb::module_& m) { gpu_array_c out, long long nnz, int num_genes, std::optional> indptr, bool assume_unsorted, std::uintptr_t stream) { + require_arg(nnz >= 0 && (size_t)nnz <= indices.shape(0), + "check_zero_genes: nnz exceeds the indices length"); + require_arg((int)out.shape(0) == num_genes, + "check_zero_genes: out must have num_genes entries"); + if (indptr) { + require_arg( + indptr->ndim() == 1 && indptr->shape(0) >= 1, + "check_zero_genes: indptr must be 1-D and non-empty"); + } MinorCountOp op{out.data(), num_genes, 0}; if (!indptr) { minor_reduce_flat(indices.data(), op, nnz, From a0cc1d00d1fc0fe71ad9d33271ac65d975e74ee5 Mon Sep 17 00:00:00 2001 From: Intron7 Date: Mon, 7 Sep 2026 20:22:11 +0200 Subject: [PATCH 3/4] fix hvg Signed-off-by: Intron7 --- src/rapids_singlecell/_cuda/hvg/hvg.cu | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/src/rapids_singlecell/_cuda/hvg/hvg.cu b/src/rapids_singlecell/_cuda/hvg/hvg.cu index 15c6c2863..6a33f175e 100644 --- a/src/rapids_singlecell/_cuda/hvg/hvg.cu +++ b/src/rapids_singlecell/_cuda/hvg/hvg.cu @@ -34,9 +34,11 @@ static void launch_expected_zeros(const T* scaled_means, const T* total_counts, } /// min(value, clip) that propagates NaN like cupy.minimum on the dense path -/// (fmin would return the finite operand). +/// (fmin would return the finite operand). Branch-free: a NaN value fails +/// `value < clip` but passes `value != value`; a NaN clip fails both and is +/// selected as is. __device__ inline double clip_min(double value, double clip) { - return (isnan(value) || isnan(clip)) ? NAN : fmin(value, clip); + return (value < clip || value != value) ? value : clip; } /// Per-column sum and sum-of-squares of min(value, clip[col]) for seurat_v3 From 9d14d4a93bcd3b34c5a1b29ad2e7718211af9fd8 Mon Sep 17 00:00:00 2001 From: Intron7 Date: Tue, 8 Sep 2026 13:06:56 +0200 Subject: [PATCH 4/4] small fixes Signed-off-by: Intron7 --- src/rapids_singlecell/_cuda/ligrec/ligrec.cu | 2 + .../_cuda/mean_var/kernels_mv.cuh | 9 ++- .../_cuda/mean_var/mean_var.cu | 8 +-- src/rapids_singlecell/_cuda/minor_tiles.cuh | 58 +++++++++++++------ 4 files changed, 51 insertions(+), 26 deletions(-) diff --git a/src/rapids_singlecell/_cuda/ligrec/ligrec.cu b/src/rapids_singlecell/_cuda/ligrec/ligrec.cu index c86f64e11..91838cbe5 100644 --- a/src/rapids_singlecell/_cuda/ligrec/ligrec.cu +++ b/src/rapids_singlecell/_cuda/ligrec/ligrec.cu @@ -133,6 +133,7 @@ void def_sum_count_sparse(nb::module_& m) { gpu_array_c clusters, gpu_array_c sum, gpu_array_c count, int rows, int ncls, bool assume_unsorted, std::uintptr_t stream) { + if (ncls <= 0 || rows <= 0) return false; // nothing to reduce // sum is (n_genes, ncls); derive n_genes from the element count so // any 2-D shape with the same size works. return launch_sum_count_sparse( @@ -169,6 +170,7 @@ void def_mean_sparse(nb::module_& m) { gpu_array_c data, gpu_array_c clusters, gpu_array_c g, int rows, int ncls, bool assume_unsorted, std::uintptr_t stream) { + if (ncls <= 0 || rows <= 0) return false; // nothing to reduce return launch_mean_sparse( indptr.data(), index.data(), data.data(), clusters.data(), g.data(), rows, ncls, (int)(g.size() / ncls), diff --git a/src/rapids_singlecell/_cuda/mean_var/kernels_mv.cuh b/src/rapids_singlecell/_cuda/mean_var/kernels_mv.cuh index 9e94bf007..689ad7c50 100644 --- a/src/rapids_singlecell/_cuda/mean_var/kernels_mv.cuh +++ b/src/rapids_singlecell/_cuda/mean_var/kernels_mv.cuh @@ -71,10 +71,13 @@ struct MeanVarOp { } __device__ void flush_col(const char* acc, int, int col, int g) const { const double* s = reinterpret_cast(acc); + const double sum = s[g]; const double sq = s[tile_size + g]; - // Zero only when no nonzero of this column landed in the block. - if (sq != 0.0) { - atomicAdd(&means[col], s[g]); + // Both zero only when nothing landed in this column (or it cancelled + // exactly, where adding zero is a no-op). Testing the squares alone + // would drop tiny values whose squares underflow. + if (sum != 0.0 || sq != 0.0) { + atomicAdd(&means[col], sum); atomicAdd(&vars[col], sq); } } diff --git a/src/rapids_singlecell/_cuda/mean_var/mean_var.cu b/src/rapids_singlecell/_cuda/mean_var/mean_var.cu index ef29ba807..41e25e862 100644 --- a/src/rapids_singlecell/_cuda/mean_var/mean_var.cu +++ b/src/rapids_singlecell/_cuda/mean_var/mean_var.cu @@ -78,10 +78,10 @@ void def_tile_plan(nb::module_& m) { m.def( "tile_plan", [](long long nnz, int n_rows, int n_cols, size_t bytes_per_col) { - require_arg( - nnz >= 0 && n_rows > 0 && n_cols > 0 && bytes_per_col > 0, - "tile_plan: nnz must be >= 0 and n_rows, n_cols and " - "bytes_per_col must be positive"); + require_arg(nnz >= 0 && n_rows > 0 && n_cols > 0 && + bytes_per_col > 0 && bytes_per_col <= 4096, + "tile_plan: nnz must be >= 0, n_rows and n_cols " + "positive, and bytes_per_col in [1, 4096]"); const TilePlan p = plan_tiles(nnz, n_rows, n_cols, bytes_per_col); nb::dict d; d["use_tiled"] = p.use_tiled; diff --git a/src/rapids_singlecell/_cuda/minor_tiles.cuh b/src/rapids_singlecell/_cuda/minor_tiles.cuh index 3ae7d8490..427b230ad 100644 --- a/src/rapids_singlecell/_cuda/minor_tiles.cuh +++ b/src/rapids_singlecell/_cuda/minor_tiles.cuh @@ -138,13 +138,18 @@ inline TilePlan plan_tiles(long long nnz, int n_rows, int n_cols, rows_per_block = std::clamp(rows_per_block, SWEEP_MIN_ROWS_PER_BLOCK, SWEEP_MAX_ROWS_PER_BLOCK); - const size_t bookmark_bytes = (size_t)rows_per_block * sizeof(int); + const size_t bookmark_bytes = (size_t)rows_per_block * sizeof(long long); TilePlan plan{false, 0, 0, rows_per_block, 0}; - if (budget < bookmark_bytes + bytes_per_col) return plan; + if (bytes_per_col == 0 || budget < bookmark_bytes || + budget - bookmark_bytes < bytes_per_col) { + return plan; + } const size_t max_tile_cols = (budget - bookmark_bytes) / bytes_per_col; plan.tile_size = (int)std::min(max_tile_cols, (size_t)n_cols); - plan.n_tiles = (n_cols + plan.tile_size - 1) / plan.tile_size; + if (plan.tile_size <= 0) return plan; + plan.n_tiles = + (int)(((long long)n_cols + plan.tile_size - 1) / plan.tile_size); const long long nnz_per_slice = nnz / ((long long)n_rows * plan.n_tiles); if (plan.n_tiles > 1 && nnz_per_slice < SWEEP_MIN_NNZ_PER_SLICE) return plan; @@ -186,18 +191,20 @@ __device__ inline BlockRows block_rows(const BlockRows* __restrict__ per_block, /// next tile. Sets `out_of_order` when a column below `tile_begin` shows up: /// that column was skipped by an earlier tile, so the row is not sorted. template -__device__ inline int warp_sweep_slice(const IdxT* __restrict__ indices, - IdxT row_begin, IdxT row_end, - int bookmark, IdxT tile_begin, - IdxT tile_end, const Op& op, char* pad, - bool& out_of_order) { +__device__ inline long long warp_sweep_slice(const IdxT* __restrict__ indices, + IdxT row_begin, IdxT row_end, + long long bookmark, + long long tile_begin, + long long tile_end, const Op& op, + char* pad, bool& out_of_order) { const int lane = threadIdx.x & (WARP - 1); IdxT pos = row_begin + bookmark; while (pos < row_end) { const IdxT nnz_pos = pos + lane; // Lanes past the row read a sentinel that ends the slice exactly like a // column past the tile would. - const IdxT col = (nnz_pos < row_end) ? indices[nnz_pos] : tile_end; + const long long col = + (nnz_pos < row_end) ? (long long)indices[nnz_pos] : tile_end; const bool in_tile = col < tile_end; if (in_tile) { if (col < tile_begin) { @@ -210,22 +217,22 @@ __device__ inline int warp_sweep_slice(const IdxT* __restrict__ indices, const unsigned past_tile = __ballot_sync(0xffffffffu, !in_tile); if (past_tile) { // The slice ends at the first lane whose column is past the tile. - return static_cast(pos + (__ffs(past_tile) - 1) - row_begin); + return (long long)(pos + (__ffs(past_tile) - 1) - row_begin); } pos += WARP; } - return static_cast(row_end - row_begin); + return (long long)(row_end - row_begin); } /// Order-agnostic alternative: read the whole row, keep what falls in the tile. template __device__ inline void warp_rescan_row(const IdxT* __restrict__ indices, IdxT row_begin, IdxT row_end, - IdxT tile_begin, IdxT tile_end, + long long tile_begin, long long tile_end, const Op& op, char* pad) { const int lane = threadIdx.x & (WARP - 1); for (IdxT nnz_pos = row_begin + lane; nnz_pos < row_end; nnz_pos += WARP) { - const IdxT col = indices[nnz_pos]; + const long long col = (long long)indices[nnz_pos]; if (col >= tile_begin && col < tile_end) { op.add(pad, (long long)nnz_pos, static_cast(col - tile_begin)); } @@ -254,15 +261,15 @@ template __device__ inline void sweep_tile(const IdxT* __restrict__ indptr, const IdxT* __restrict__ indices, const int* __restrict__ row_order, - int* bookmarks, BlockRows rows, + long long* bookmarks, BlockRows rows, int tile_begin, int tile_size, bool rescan, int* __restrict__ out_of_order_flag, const Op& op, char* pad) { const int lane = threadIdx.x & (WARP - 1); const int warp = threadIdx.x / WARP; const int warps_per_block = blockDim.x / WARP; - const IdxT tile_lo = static_cast(tile_begin); - const IdxT tile_hi = tile_lo + tile_size; + const long long tile_lo = tile_begin; + const long long tile_hi = tile_lo + tile_size; // 64-bit: cannot overflow bool out_of_order = false; for (int i = rows.first + warp; i < rows.last; i += warps_per_block) { @@ -275,7 +282,7 @@ __device__ inline void sweep_tile(const IdxT* __restrict__ indptr, pad); continue; } - const int bookmark = warp_sweep_slice( + const long long bookmark = warp_sweep_slice( indices, row_begin, row_end, bookmarks[i - rows.first], tile_lo, tile_hi, op, pad, out_of_order); if (lane == 0) bookmarks[i - rows.first] = bookmark; @@ -296,7 +303,7 @@ __global__ void __launch_bounds__(SWEEP_BLOCK_THREADS, SWEEP_BLOCKS_PER_SM) bool rescan) { extern __shared__ __align__(16) char smem[]; char* pad = smem; // the Op's accumulators for the current tile - int* bookmarks = reinterpret_cast( + long long* bookmarks = reinterpret_cast( smem + pad_bytes(op.tile_size, Op::bytes_per_col)); const BlockRows rows = block_rows(per_block_rows, rows_per_block, n_rows); @@ -466,7 +473,19 @@ bool minor_reduce(const IdxT* indptr, const IdxT* indices, Op op, int n_rows, // Which rows each block owns: contiguous chunks, or per-group chunks of the // sorted row order when grouped. - RmmScratchPool pool; // flag + block table; released after the launches + // Scratch is stream-ordered device memory: neither the pool nor the + // caller's grouped row order may be released before every launch issued + // here has finished, so this guard drains the stream before the pool + // destructs (declared after the pool, hence destroyed before it). + RmmScratchPool pool; + bool needs_sync = groups != nullptr; + struct SyncOnExit { + cudaStream_t stream; + const bool* armed; + ~SyncOnExit() { + if (*armed) cudaStreamSynchronize(stream); + } + } sync_on_exit{stream, &needs_sync}; const int* row_order = groups != nullptr ? groups->row_order : nullptr; const BlockRows* per_block_rows = nullptr; unsigned n_blocks = @@ -525,6 +544,7 @@ bool minor_reduce(const IdxT* indptr, const IdxT* indices, Op op, int n_rows, } // 3. The bookmark sweep, with out-of-order detection. int* flag = pool.alloc(1); + needs_sync = true; cuda_check(cudaMemsetAsync(flag, 0, sizeof(int), stream), "cudaMemsetAsync(out-of-order flag)"); launch_tiled(false, flag);