From bf4df0285e87131390863aa36d26f52cd0c2f830 Mon Sep 17 00:00:00 2001 From: Liyun Xiu Date: Tue, 7 Jul 2026 01:53:20 +0000 Subject: [PATCH 1/8] Add GPU-accelerated SeismicIndex build via cuSPARSE Accelerate the index-building (build()) path with NVIDIA cuSPARSE while keeping search() entirely on CPU. The k-means document->centroid assignment inside build() is a sparse(docs CSR) x dense(centroids) matrix product, which maps onto cusparseSpMM plus a per-row argmax. Implementation (nsparse/gpu/gpu_cluster_assigner.{h,cu}): - Full corpus CSR uploaded to the GPU once, cached by (pointer, n_vectors, nnz) so repeated per-list calls reuse it. - Doc sub-matrix A and dense centroid matrix B are built on-device via gather/scatter kernels reading the resident corpus, so only small id arrays cross PCIe per inverted list. - Per-thread cuSPARSE handle + CUDA stream + grow-on-demand scratch buffers (no global mutex), preserving the 32-way OpenMP concurrency of build_inverted_lists_clusters(). Assignments are bit-identical to the scalar CPU path (skips centroids, strict-greater argmax with ties to the lowest cluster index); float (U32) weights only, with automatic CPU fallback otherwise or on any GPU error. The hook in map_docs_to_clusters() is gated on NSPARSE_WITH_CUDA && !__AVX512F__. Build wiring: NSPARSE_ENABLE_CUDA CMake option (default OFF) plus cmake/cuda.cmake; CPU-only builds and CI are unaffected. Tests: tests/gpu_cluster_assigner_test.cpp asserts GPU==CPU cluster membership (built only with CUDA). All 363 tests pass on an L4. Benchmark (benchmarks/index_build_benchmark.cpp, Big-ANN SPLADE, L4 vs 32-core, Release): - base_small (100K docs, 12.7M nnz): CPU 33.5s -> GPU 6.58s (5.1x) - base_full (8.84M docs, 1.12B nnz): CPU 483s -> GPU 81.1s (5.96x) Signed-off-by: Liyun Xiu --- CMakeLists.txt | 9 + benchmarks/CMakeLists.txt | 13 + benchmarks/index_build_benchmark.cpp | 193 ++++++++++++ cmake/cuda.cmake | 48 +++ nsparse/CMakeLists.txt | 20 ++ nsparse/cluster/kmeans_utils.cpp | 19 ++ nsparse/gpu/gpu_cluster_assigner.cu | 452 +++++++++++++++++++++++++++ nsparse/gpu/gpu_cluster_assigner.h | 101 ++++++ tests/CMakeLists.txt | 5 + tests/gpu_cluster_assigner_test.cpp | 160 ++++++++++ 10 files changed, 1020 insertions(+) create mode 100644 benchmarks/index_build_benchmark.cpp create mode 100644 cmake/cuda.cmake create mode 100644 nsparse/gpu/gpu_cluster_assigner.cu create mode 100644 nsparse/gpu/gpu_cluster_assigner.h create mode 100644 tests/gpu_cluster_assigner_test.cpp diff --git a/CMakeLists.txt b/CMakeLists.txt index fd17b0b..90719bc 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -22,6 +22,15 @@ include(cmake/third_party.cmake) option(NSPARSE_OPT_LEVEL "" "generic") option(NSPARSE_ENABLE_PYTHON "Build Python bindings" OFF) +# GPU acceleration for the index-building (build()) path via NVIDIA cuSPARSE. +# The search() path always stays on CPU. Default OFF so CPU-only builds and CI +# are unaffected when no CUDA toolkit is present. +option(NSPARSE_ENABLE_CUDA "Build GPU-accelerated index building via cuSPARSE" OFF) + +if(NSPARSE_ENABLE_CUDA) + include(cmake/cuda.cmake) +endif() + add_subdirectory(nsparse) # Python bindings diff --git a/benchmarks/CMakeLists.txt b/benchmarks/CMakeLists.txt index 0e7ae4e..c059942 100644 --- a/benchmarks/CMakeLists.txt +++ b/benchmarks/CMakeLists.txt @@ -41,3 +41,16 @@ if(NSPARSE_BENCHMARK_SRC) absl::flat_hash_map ) endif() + +# Build-time benchmark: SeismicIndex build() on CPU vs GPU (cuSPARSE). The GPU +# case is compiled in only when the nsparse target carries CUDA support (the +# NSPARSE_WITH_CUDA define propagates from the nsparse target). +add_executable(nsparse_build_benchmark index_build_benchmark.cpp) +target_include_directories(nsparse_build_benchmark PRIVATE ${PROJECT_SOURCE_DIR}) +target_link_libraries(nsparse_build_benchmark PRIVATE + nsparse + OpenMP::OpenMP_CXX + benchmark::benchmark + absl::flat_hash_set + absl::flat_hash_map +) diff --git a/benchmarks/index_build_benchmark.cpp b/benchmarks/index_build_benchmark.cpp new file mode 100644 index 0000000..dbd51ce --- /dev/null +++ b/benchmarks/index_build_benchmark.cpp @@ -0,0 +1,193 @@ +/** + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + * + * Benchmarks SeismicIndex build() time on CPU vs GPU (cuSPARSE). + * + * The GPU path accelerates the k-means document->centroid assignment step of + * build() via cuSPARSE SpMM; the rest of build() (inverted-list construction, + * summarization) stays on CPU. GPU vs CPU is selected at run time through the + * NSPARSE_GPU_MIN_DOCS gate so both paths run from the same binary: + * - CPU: set the gate above the largest inverted list (no list offloads). + * - GPU: set the gate to 1 (every eligible list offloads). + * + * Data is the Big-ANN sparse-vectors (SPLADE MS MARCO) CSR file, e.g. + * data/base_small.csr from cpp-sparse-ann's dataset.py. + * + * Usage: + * NSPARSE_DATA_CSR=/path/to/base_small.csr \ + * ./nsparse_build_benchmark --benchmark_min_time=1x + */ + +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +#include "nsparse/seismic_common.h" +#include "nsparse/seismic_index.h" +#include "nsparse/types.h" + +#ifdef NSPARSE_WITH_CUDA +#include "nsparse/gpu/gpu_cluster_assigner.h" +#endif + +namespace { + +struct CSRMatrix { + int64_t nrow; + int64_t ncol; + int64_t nnz; + std::vector indptr; + std::vector indices; + std::vector data; +}; + +CSRMatrix read_csr(const std::string& path) { + std::ifstream f(path, std::ios::binary); + if (!f.is_open()) { + throw std::runtime_error("Cannot open CSR file: " + path); + } + + CSRMatrix m; + int64_t sizes[3]; + f.read(reinterpret_cast(sizes), sizeof(sizes)); + m.nrow = sizes[0]; + m.ncol = sizes[1]; + m.nnz = sizes[2]; + + std::vector indptr64(m.nrow + 1); + f.read(reinterpret_cast(indptr64.data()), + static_cast((m.nrow + 1) * sizeof(int64_t))); + m.indptr.resize(m.nrow + 1); + for (int64_t i = 0; i <= m.nrow; ++i) { + m.indptr[i] = static_cast(indptr64[i]); + } + + std::vector indices32(m.nnz); + f.read(reinterpret_cast(indices32.data()), + static_cast(m.nnz * sizeof(int32_t))); + m.indices.resize(m.nnz); + for (int64_t i = 0; i < m.nnz; ++i) { + m.indices[i] = static_cast(indices32[i]); + } + + m.data.resize(m.nnz); + f.read(reinterpret_cast(m.data.data()), + static_cast(m.nnz * sizeof(float))); + + return m; +} + +std::string get_env_or_die(const char* name) { + const char* val = std::getenv(name); + if (val == nullptr || val[0] == '\0') { + throw std::runtime_error(std::string("Environment variable not set: ") + + name); + } + return val; +} + +// Loads the data CSR exactly once and shares it across benchmark cases. +const CSRMatrix& shared_data() { + static CSRMatrix data = [] { + std::string path = get_env_or_die("NSPARSE_DATA_CSR"); + std::cout << "Loading data CSR: " << path << "\n"; + CSRMatrix m = read_csr(path); + std::cout << " rows=" << m.nrow << " cols=" << m.ncol + << " nnz=" << m.nnz << "\n"; + return m; + }(); + return data; +} + +// Seismic cluster parameters comparable to the search benchmark's index. +constexpr nsparse::SeismicClusterParameters kParams = { + .lambda = 6000, .beta = 400, .alpha = 0.4F}; + +// Force the GPU-offload gate one way or the other for this process. Setting the +// minimum-docs threshold to 1 offloads every eligible list; setting it beyond +// the largest inverted list keeps the whole build on the CPU. +void set_gpu_offload(bool enabled) { + if (enabled) { + ::setenv("NSPARSE_GPU_MIN_DOCS", "1", /*overwrite=*/1); + } else { + // Larger than lambda (max pruned posting length), so no list offloads. + ::setenv("NSPARSE_GPU_MIN_DOCS", "1000000000", /*overwrite=*/1); + } +} + +// Builds a fresh SeismicIndex from the shared corpus and times only build(). +void run_build(benchmark::State& state, bool use_gpu) { + const CSRMatrix& data = shared_data(); + for (auto _ : state) { + state.PauseTiming(); + auto index = std::make_unique( + static_cast(data.ncol), kParams); + index->add(static_cast(data.nrow), data.indptr.data(), + data.indices.data(), data.data.data()); + set_gpu_offload(use_gpu); + state.ResumeTiming(); + + index->build(); + + benchmark::DoNotOptimize(index.get()); + benchmark::ClobberMemory(); + } + state.counters["docs"] = static_cast(data.nrow); + state.counters["nnz"] = static_cast(data.nnz); +} + +// Iterations / repetitions are configurable so the same binary can run a +// robust multi-sample sweep on small corpora and a single build on very large +// ones (where one full build already takes many minutes). +int env_int(const char* name, int fallback) { + if (const char* v = std::getenv(name); v != nullptr && v[0] != '\0') { + return std::atoi(v); + } + return fallback; +} + +} // namespace + +// --------------------------------------------------------------------------- +// CPU build (assignment on CPU) +// --------------------------------------------------------------------------- +static void BM_Seismic_Build_CPU(benchmark::State& state) { + run_build(state, /*use_gpu=*/false); +} +BENCHMARK(BM_Seismic_Build_CPU) + ->Unit(benchmark::kMillisecond) + ->UseRealTime() + ->Iterations(env_int("NSPARSE_BENCH_ITERS", 3)) + ->Repetitions(env_int("NSPARSE_BENCH_REPS", 3)); + +// --------------------------------------------------------------------------- +// GPU build (assignment via cuSPARSE) - only when compiled with CUDA +// --------------------------------------------------------------------------- +#ifdef NSPARSE_WITH_CUDA +static void BM_Seismic_Build_GPU(benchmark::State& state) { + if (!nsparse::detail::GpuClusterAssigner::available()) { + state.SkipWithError("No CUDA-capable GPU available"); + return; + } + run_build(state, /*use_gpu=*/true); +} +BENCHMARK(BM_Seismic_Build_GPU) + ->Unit(benchmark::kMillisecond) + ->UseRealTime() + ->Iterations(env_int("NSPARSE_BENCH_ITERS", 3)) + ->Repetitions(env_int("NSPARSE_BENCH_REPS", 3)); +#endif + +BENCHMARK_MAIN(); diff --git a/cmake/cuda.cmake b/cmake/cuda.cmake new file mode 100644 index 0000000..d71b058 --- /dev/null +++ b/cmake/cuda.cmake @@ -0,0 +1,48 @@ +# Copyright OpenSearch Contributors +# SPDX-License-Identifier: Apache-2.0 +# +# The OpenSearch Contributors require contributions made to +# this file be licensed under the Apache-2.0 license or a +# compatible open source license. +# +# CUDA / cuSPARSE toolchain setup for GPU-accelerated index building. +# +# Enable with -DNSPARSE_ENABLE_CUDA=ON. When the CUDA toolkit is installed in a +# non-standard location (for example, shipped inside a pip `nvidia-cu13` wheel), +# point CMake at it with: +# -DNSPARSE_CUDA_TOOLKIT_ROOT=/path/to/nvidia/cu13 +# and CMake will pick up nvcc, the headers and the runtime/cusparse libraries. + +# Allow the caller to hint the toolkit location (nvcc lives in /bin). +if(NSPARSE_CUDA_TOOLKIT_ROOT) + set(CUDAToolkit_ROOT "${NSPARSE_CUDA_TOOLKIT_ROOT}" CACHE PATH "" FORCE) + if(NOT CMAKE_CUDA_COMPILER) + set(CMAKE_CUDA_COMPILER "${NSPARSE_CUDA_TOOLKIT_ROOT}/bin/nvcc" CACHE FILEPATH "" FORCE) + endif() + # Wheel-packaged toolkits keep shared libs under lib/ rather than lib64/. + list(APPEND CMAKE_LIBRARY_PATH + "${NSPARSE_CUDA_TOOLKIT_ROOT}/lib" + "${NSPARSE_CUDA_TOOLKIT_ROOT}/lib64") +endif() + +# Turn on the CUDA language now that the compiler is (optionally) hinted. +enable_language(CUDA) + +# CUDAToolkit provides the CUDA::cusparse and CUDA::cudart imported targets. +find_package(CUDAToolkit REQUIRED) + +# Default to the common data-center / Ada architectures. L4 (this project's +# reference GPU) is compute capability 8.9. CMake picks a conservative default +# during language enablement, so override it unless the caller pinned one +# explicitly via -DNSPARSE_CUDA_ARCHITECTURES. +if(NSPARSE_CUDA_ARCHITECTURES) + set(CMAKE_CUDA_ARCHITECTURES "${NSPARSE_CUDA_ARCHITECTURES}" CACHE STRING "" FORCE) +else() + set(CMAKE_CUDA_ARCHITECTURES 80 89 CACHE STRING "CUDA architectures to build for" FORCE) +endif() + +set(CMAKE_CUDA_STANDARD 17) +set(CMAKE_CUDA_STANDARD_REQUIRED ON) + +message(STATUS "nsparse: CUDA enabled (nvcc=${CMAKE_CUDA_COMPILER}, " + "archs=${CMAKE_CUDA_ARCHITECTURES})") diff --git a/nsparse/CMakeLists.txt b/nsparse/CMakeLists.txt index 8786d60..0e29371 100644 --- a/nsparse/CMakeLists.txt +++ b/nsparse/CMakeLists.txt @@ -138,6 +138,26 @@ if(NSPARSE_IS_X86) target_include_directories(nsparse_avx512 PUBLIC $) endif() +# GPU acceleration for the build() path via cuSPARSE (search() stays on CPU). +# Compiled into every nsparse variant only when configured with +# -DNSPARSE_ENABLE_CUDA=ON. Guarded by the NSPARSE_WITH_CUDA compile define so +# the CPU sources pick up the GPU assignment hook. +if(NSPARSE_ENABLE_CUDA) + set(NSPARSE_CUDA_SRC gpu/gpu_cluster_assigner.cu) + set(NSPARSE_CUDA_TARGETS nsparse) + if(NSPARSE_IS_X86) + list(APPEND NSPARSE_CUDA_TARGETS nsparse_avx2 nsparse_avx512) + elseif(NSPARSE_IS_ARM AND NOT APPLE) + list(APPEND NSPARSE_CUDA_TARGETS nsparse_sve) + endif() + foreach(_tgt IN LISTS NSPARSE_CUDA_TARGETS) + target_sources(${_tgt} PRIVATE ${NSPARSE_CUDA_SRC}) + target_compile_definitions(${_tgt} PUBLIC NSPARSE_WITH_CUDA) + target_link_libraries(${_tgt} PRIVATE CUDA::cusparse CUDA::cudart) + set_target_properties(${_tgt} PROPERTIES CUDA_SEPARABLE_COMPILATION ON) + endforeach() +endif() + find_package(OpenMP REQUIRED) # Link base library (always) target_link_libraries(nsparse PUBLIC OpenMP::OpenMP_CXX absl::flat_hash_set absl::flat_hash_map) diff --git a/nsparse/cluster/kmeans_utils.cpp b/nsparse/cluster/kmeans_utils.cpp index 7f8ebeb..8af08a7 100644 --- a/nsparse/cluster/kmeans_utils.cpp +++ b/nsparse/cluster/kmeans_utils.cpp @@ -18,6 +18,9 @@ #include "nsparse/sparse_vectors.h" #include "nsparse/types.h" #include "nsparse/utils/distance_simd.h" +#ifdef NSPARSE_WITH_CUDA +#include "nsparse/gpu/gpu_cluster_assigner.h" +#endif #if defined(__AVX512F__) #include #include @@ -170,6 +173,22 @@ void map_docs_to_clusters(const SparseVectors* vectors, if (n_clusters == 0 || n_docs == 0) { return; } +#if defined(NSPARSE_WITH_CUDA) && !defined(__AVX512F__) + // GPU-accelerated assignment via cuSPARSE for float (U32) weights on large + // problems. Produces assignments identical to the scalar CPU reference + // below (skips centroids, ties break to the lowest cluster index); on any + // failure we fall through to the CPU path. Disabled for AVX-512 builds + // whose assignment semantics differ. + if (vectors->get_element_size() == U32 && + should_offload_assignment_to_gpu(n_docs, n_clusters)) { + try { + GpuClusterAssigner::instance().assign(vectors, docs, clusters); + return; + } catch (const std::exception&) { + // Fall back to CPU on any GPU error (OOM, driver, etc.). + } + } +#endif #if defined(__AVX512F__) map_docs_to_clusters_avx512(vectors, docs, clusters); return; diff --git a/nsparse/gpu/gpu_cluster_assigner.cu b/nsparse/gpu/gpu_cluster_assigner.cu new file mode 100644 index 0000000..288a611 --- /dev/null +++ b/nsparse/gpu/gpu_cluster_assigner.cu @@ -0,0 +1,452 @@ +/** + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +#include "nsparse/gpu/gpu_cluster_assigner.h" + +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +#include "absl/container/flat_hash_set.h" +#include "cuda_runtime.h" +#include "nsparse/sparse_vectors.h" +#include "nsparse/types.h" + +namespace nsparse::detail { +namespace { + +void check_cuda(cudaError_t status, const char* what) { + if (status != cudaSuccess) { + std::ostringstream oss; + oss << "CUDA error in " << what << ": " << cudaGetErrorString(status); + throw std::runtime_error(oss.str()); + } +} + +void check_cusparse(cusparseStatus_t status, const char* what) { + if (status != CUSPARSE_STATUS_SUCCESS) { + std::ostringstream oss; + oss << "cuSPARSE error in " << what << ": " + << cusparseGetErrorString(status); + throw std::runtime_error(oss.str()); + } +} + +#define NSPARSE_CUDA_CHECK(expr) check_cuda((expr), #expr) +#define NSPARSE_CUSPARSE_CHECK(expr) check_cusparse((expr), #expr) + +// One thread per document row. Scans the n_clusters scores for that row and +// selects the argmax, breaking ties toward the lowest cluster index (strict +// greater-than update) to match the CPU reference in map_docs_to_clusters(). +__global__ void row_argmax_kernel(const float* __restrict__ scores, + int n_rows, int n_clusters, + int32_t* __restrict__ best_cluster) { + int row = blockIdx.x * blockDim.x + threadIdx.x; + if (row >= n_rows) { + return; + } + const float* row_scores = scores + static_cast(row) * n_clusters; + float best = row_scores[0]; + int best_j = 0; + for (int j = 1; j < n_clusters; ++j) { + if (row_scores[j] > best) { + best = row_scores[j]; + best_j = j; + } + } + best_cluster[row] = best_j; +} + +// Builds the compact CSR values/columns of the document sub-matrix A by +// gathering rows straight from the resident corpus, so the large doc data never +// crosses the PCIe bus per list. One thread per output row. +__global__ void gather_csr_kernel(const int32_t* __restrict__ corpus_indptr, + const int32_t* __restrict__ corpus_indices, + const float* __restrict__ corpus_values, + const int32_t* __restrict__ docs, int n_docs, + const int32_t* __restrict__ a_row_ptr, + int32_t* __restrict__ a_col, + float* __restrict__ a_val) { + int row = blockIdx.x * blockDim.x + threadIdx.x; + if (row >= n_docs) { + return; + } + const int32_t d = docs[row]; + const int32_t src = corpus_indptr[d]; + const int32_t len = corpus_indptr[d + 1] - src; + const int32_t dst = a_row_ptr[row]; + for (int32_t t = 0; t < len; ++t) { + a_col[dst + t] = corpus_indices[src + t]; + a_val[dst + t] = corpus_values[src + t]; + } +} + +// Scatters centroid rows from the resident corpus into the dense B matrix +// (dim x n_clusters, row-major, ldb = n_clusters). Column j of B is centroid j. +// One thread per centroid; each writes a disjoint column, so there are no +// races. B must be zeroed first. +__global__ void scatter_dense_kernel(const int32_t* __restrict__ corpus_indptr, + const int32_t* __restrict__ corpus_indices, + const float* __restrict__ corpus_values, + const int32_t* __restrict__ centroids, + int n_clusters, int ldb, + float* __restrict__ b) { + int j = blockIdx.x * blockDim.x + threadIdx.x; + if (j >= n_clusters) { + return; + } + const int32_t c = centroids[j]; + const int32_t src = corpus_indptr[c]; + const int32_t len = corpus_indptr[c + 1] - src; + for (int32_t t = 0; t < len; ++t) { + const int32_t col = corpus_indices[src + t]; + b[static_cast(col) * ldb + j] = corpus_values[src + t]; + } +} + +// The full corpus, uploaded to the GPU once and shared read-only across all +// threads/lists. Indices are widened to int32 on upload to feed both the gather +// kernel and cuSPARSE's 32-bit CSR indices. +struct DeviceCorpus { + int32_t* indptr = nullptr; // n_vectors + 1 + int32_t* indices = nullptr; // nnz + float* values = nullptr; // nnz + int64_t nnz = 0; + size_t dim = 0; + size_t n_vectors = 0; + // Identity of the resident corpus. Pointer alone is unsafe: a new + // SparseVectors can reuse a freed one's address, so the residency check + // also compares n_vectors and nnz. + const SparseVectors* owner = nullptr; + + bool matches(const SparseVectors* v, size_t nv, int64_t nz) const { + return owner == v && n_vectors == nv && nnz == nz; + } + + void free() { + cudaFree(indptr); + cudaFree(indices); + cudaFree(values); + indptr = nullptr; + indices = nullptr; + values = nullptr; + owner = nullptr; + n_vectors = 0; + nnz = 0; + } +}; + +// Per-thread GPU resources: an independent cuSPARSE handle + stream (handles +// are not shareable across threads) and scratch buffers reused across lists, +// grown on demand so steady-state builds do no per-list cudaMalloc. +struct ThreadCtx { + cusparseHandle_t handle{}; + cudaStream_t stream{}; + bool init = false; + + int32_t* d_docs = nullptr; size_t docs_cap = 0; + int32_t* d_centroids = nullptr; size_t cent_cap = 0; + int32_t* d_a_row_ptr = nullptr; size_t rowptr_cap = 0; + int32_t* d_a_col = nullptr; size_t col_cap = 0; + float* d_a_val = nullptr; size_t val_cap = 0; + float* d_b = nullptr; size_t b_cap = 0; + float* d_c = nullptr; size_t c_cap = 0; + int32_t* d_best = nullptr; size_t best_cap = 0; + void* d_spmm = nullptr; size_t spmm_cap = 0; + + ~ThreadCtx() { + cudaFree(d_docs); + cudaFree(d_centroids); + cudaFree(d_a_row_ptr); + cudaFree(d_a_col); + cudaFree(d_a_val); + cudaFree(d_b); + cudaFree(d_c); + cudaFree(d_best); + cudaFree(d_spmm); + if (init) { + cusparseDestroy(handle); + cudaStreamDestroy(stream); + } + } +}; + +// Reallocate *ptr only when it must grow. Capacity tracked in bytes. +template +void ensure_capacity(T** ptr, size_t& cap_bytes, size_t need_bytes) { + if (cap_bytes < need_bytes) { + cudaFree(*ptr); + void* p = nullptr; + check_cuda(cudaMalloc(&p, need_bytes), "cudaMalloc(ensure_capacity)"); + *ptr = static_cast(p); + cap_bytes = need_bytes; + } +} + +// Lazily created once per worker thread; destroyed at thread exit. +thread_local std::unique_ptr t_ctx; + +ThreadCtx& thread_ctx() { + if (t_ctx == nullptr) { + t_ctx = std::make_unique(); + check_cuda(cudaStreamCreate(&t_ctx->stream), "cudaStreamCreate"); + check_cusparse(cusparseCreate(&t_ctx->handle), "cusparseCreate"); + check_cusparse(cusparseSetStream(t_ctx->handle, t_ctx->stream), + "cusparseSetStream"); + t_ctx->init = true; + } + return *t_ctx; +} + +} // namespace + +struct GpuClusterAssigner::Impl { + std::mutex corpus_mutex; // guards residency check/upload only + DeviceCorpus corpus; + bool usable = false; +}; + +GpuClusterAssigner::GpuClusterAssigner() : impl_(new Impl()) { + int device_count = 0; + cudaError_t status = cudaGetDeviceCount(&device_count); + if (status != cudaSuccess || device_count == 0) { + return; // impl_->usable stays false; callers fall back to CPU + } + impl_->usable = true; +} + +GpuClusterAssigner::~GpuClusterAssigner() { + if (impl_ != nullptr) { + impl_->corpus.free(); + delete impl_; + } +} + +GpuClusterAssigner& GpuClusterAssigner::instance() { + static GpuClusterAssigner singleton; + return singleton; +} + +bool GpuClusterAssigner::available() { + return instance().impl_ != nullptr && instance().impl_->usable; +} + +// Uploads the corpus to the GPU once. Safe to call from many threads; the first +// caller uploads while the rest wait, and subsequent calls are a cheap pointer +// comparison. Returns a stable view (device pointers are not mutated once set). +// Declared as a member so it can touch the private Impl; the header exposes it +// as an opaque const void* to avoid leaking CUDA types. +const void* GpuClusterAssigner::ensure_corpus_resident( + const SparseVectors* vectors) { + Impl* impl = impl_; + const size_t n_vectors = vectors->num_vectors(); + const idx_t* indptr = vectors->indptr_data(); + const term_t* indices = vectors->indices_data(); + const float* values = vectors->values_data_float(); + const int64_t nnz = indptr[n_vectors]; + + std::lock_guard lock(impl->corpus_mutex); + if (impl->corpus.matches(vectors, n_vectors, nnz)) { + return &impl->corpus; + } + impl->corpus.free(); + + // Widen indices (term_t == uint16) to int32 for the gather kernel and + // cuSPARSE. Done once per corpus. + std::vector indices32(static_cast(nnz)); + for (int64_t i = 0; i < nnz; ++i) { + indices32[i] = static_cast(indices[i]); + } + + check_cuda(cudaMalloc(&impl->corpus.indptr, + (n_vectors + 1) * sizeof(int32_t)), + "cudaMalloc(corpus.indptr)"); + check_cuda(cudaMalloc(&impl->corpus.indices, nnz * sizeof(int32_t)), + "cudaMalloc(corpus.indices)"); + check_cuda(cudaMalloc(&impl->corpus.values, nnz * sizeof(float)), + "cudaMalloc(corpus.values)"); + check_cuda(cudaMemcpy(impl->corpus.indptr, indptr, + (n_vectors + 1) * sizeof(int32_t), + cudaMemcpyHostToDevice), + "cudaMemcpy(corpus.indptr)"); + check_cuda(cudaMemcpy(impl->corpus.indices, indices32.data(), + nnz * sizeof(int32_t), cudaMemcpyHostToDevice), + "cudaMemcpy(corpus.indices)"); + check_cuda(cudaMemcpy(impl->corpus.values, values, nnz * sizeof(float), + cudaMemcpyHostToDevice), + "cudaMemcpy(corpus.values)"); + + impl->corpus.nnz = nnz; + impl->corpus.dim = vectors->get_dimension(); + impl->corpus.n_vectors = n_vectors; + impl->corpus.owner = vectors; + return &impl->corpus; +} + +void GpuClusterAssigner::assign(const SparseVectors* vectors, + const std::vector& docs, + std::vector>& clusters) { + const size_t n_docs = docs.size(); + const size_t n_clusters = clusters.size(); + if (n_docs == 0 || n_clusters == 0) { + return; + } + + const idx_t* indptr = vectors->indptr_data(); + const size_t dim = vectors->get_dimension(); + + // Centroids are the first element of each cluster. Collect them (host) and + // mark which input docs are centroids so they are not re-added, as on CPU. + std::vector centroid_docs(n_clusters); + absl::flat_hash_set centroid_set; + centroid_set.reserve(n_clusters); + for (size_t j = 0; j < n_clusters; ++j) { + centroid_docs[j] = clusters[j].front(); + centroid_set.insert(clusters[j].front()); + } + + // Row pointers of the document sub-matrix A (host-side; the actual column + // indices and values are gathered on-device from the resident corpus). + std::vector h_a_row_ptr(n_docs + 1, 0); + for (size_t i = 0; i < n_docs; ++i) { + const idx_t d = docs[i]; + h_a_row_ptr[i + 1] = + h_a_row_ptr[i] + static_cast(indptr[d + 1] - indptr[d]); + } + const int64_t nnz_a = h_a_row_ptr[n_docs]; + + const DeviceCorpus& corpus = + *static_cast(ensure_corpus_resident(vectors)); + ThreadCtx& ctx = thread_ctx(); + + // Size (grow-on-demand) all per-thread scratch buffers for this list. + ensure_capacity(&ctx.d_docs, ctx.docs_cap, n_docs * sizeof(int32_t)); + ensure_capacity(&ctx.d_centroids, ctx.cent_cap, + n_clusters * sizeof(int32_t)); + ensure_capacity(&ctx.d_a_row_ptr, ctx.rowptr_cap, + (n_docs + 1) * sizeof(int32_t)); + ensure_capacity(&ctx.d_a_col, ctx.col_cap, + static_cast(nnz_a) * sizeof(int32_t)); + ensure_capacity(&ctx.d_a_val, ctx.val_cap, + static_cast(nnz_a) * sizeof(float)); + ensure_capacity(&ctx.d_b, ctx.b_cap, dim * n_clusters * sizeof(float)); + ensure_capacity(&ctx.d_c, ctx.c_cap, n_docs * n_clusters * sizeof(float)); + ensure_capacity(&ctx.d_best, ctx.best_cap, n_docs * sizeof(int32_t)); + + cudaStream_t stream = ctx.stream; + + // Upload only the small per-list metadata; doc/centroid payloads stay on + // the GPU (gathered/scattered from the resident corpus below). + NSPARSE_CUDA_CHECK(cudaMemcpyAsync(ctx.d_docs, docs.data(), + n_docs * sizeof(int32_t), + cudaMemcpyHostToDevice, stream)); + NSPARSE_CUDA_CHECK(cudaMemcpyAsync(ctx.d_centroids, centroid_docs.data(), + n_clusters * sizeof(int32_t), + cudaMemcpyHostToDevice, stream)); + NSPARSE_CUDA_CHECK(cudaMemcpyAsync(ctx.d_a_row_ptr, h_a_row_ptr.data(), + (n_docs + 1) * sizeof(int32_t), + cudaMemcpyHostToDevice, stream)); + + // Gather A's CSR payload from the corpus. + constexpr int kBlock = 256; + const int docs_grid = static_cast((n_docs + kBlock - 1) / kBlock); + gather_csr_kernel<<>>( + corpus.indptr, corpus.indices, corpus.values, ctx.d_docs, + static_cast(n_docs), ctx.d_a_row_ptr, ctx.d_a_col, ctx.d_a_val); + NSPARSE_CUDA_CHECK(cudaGetLastError()); + + // Build dense B (dim x n_clusters, row-major) from centroid rows. + NSPARSE_CUDA_CHECK(cudaMemsetAsync( + ctx.d_b, 0, dim * n_clusters * sizeof(float), stream)); + const int cent_grid = + static_cast((n_clusters + kBlock - 1) / kBlock); + scatter_dense_kernel<<>>( + corpus.indptr, corpus.indices, corpus.values, ctx.d_centroids, + static_cast(n_clusters), static_cast(n_clusters), ctx.d_b); + NSPARSE_CUDA_CHECK(cudaGetLastError()); + + // C = A * B, with A sparse (CSR), B/C dense row-major. + cusparseSpMatDescr_t mat_a; + cusparseDnMatDescr_t mat_b; + cusparseDnMatDescr_t mat_c; + NSPARSE_CUSPARSE_CHECK(cusparseCreateCsr( + &mat_a, static_cast(n_docs), static_cast(dim), nnz_a, + ctx.d_a_row_ptr, ctx.d_a_col, ctx.d_a_val, CUSPARSE_INDEX_32I, + CUSPARSE_INDEX_32I, CUSPARSE_INDEX_BASE_ZERO, CUDA_R_32F)); + NSPARSE_CUSPARSE_CHECK(cusparseCreateDnMat( + &mat_b, static_cast(dim), static_cast(n_clusters), + static_cast(n_clusters), ctx.d_b, CUDA_R_32F, + CUSPARSE_ORDER_ROW)); + NSPARSE_CUSPARSE_CHECK(cusparseCreateDnMat( + &mat_c, static_cast(n_docs), static_cast(n_clusters), + static_cast(n_clusters), ctx.d_c, CUDA_R_32F, + CUSPARSE_ORDER_ROW)); + + const float alpha = 1.0F; + const float beta = 0.0F; + size_t buffer_size = 0; + NSPARSE_CUSPARSE_CHECK(cusparseSpMM_bufferSize( + ctx.handle, CUSPARSE_OPERATION_NON_TRANSPOSE, + CUSPARSE_OPERATION_NON_TRANSPOSE, &alpha, mat_a, mat_b, &beta, mat_c, + CUDA_R_32F, CUSPARSE_SPMM_ALG_DEFAULT, &buffer_size)); + ensure_capacity(reinterpret_cast(&ctx.d_spmm), ctx.spmm_cap, + buffer_size == 0 ? 1 : buffer_size); + NSPARSE_CUSPARSE_CHECK(cusparseSpMM( + ctx.handle, CUSPARSE_OPERATION_NON_TRANSPOSE, + CUSPARSE_OPERATION_NON_TRANSPOSE, &alpha, mat_a, mat_b, &beta, mat_c, + CUDA_R_32F, CUSPARSE_SPMM_ALG_DEFAULT, ctx.d_spmm)); + + row_argmax_kernel<<>>( + ctx.d_c, static_cast(n_docs), static_cast(n_clusters), + ctx.d_best); + NSPARSE_CUDA_CHECK(cudaGetLastError()); + + std::vector h_best(n_docs); + NSPARSE_CUDA_CHECK(cudaMemcpyAsync(h_best.data(), ctx.d_best, + n_docs * sizeof(int32_t), + cudaMemcpyDeviceToHost, stream)); + NSPARSE_CUDA_CHECK(cudaStreamSynchronize(stream)); + + cusparseDestroySpMat(mat_a); + cusparseDestroyDnMat(mat_b); + cusparseDestroyDnMat(mat_c); + + // Append assignments on the host, skipping documents that are themselves a + // centroid (matching the CPU reference exactly). + for (size_t i = 0; i < n_docs; ++i) { + const idx_t d = docs[i]; + if (centroid_set.contains(d)) { + continue; + } + clusters[h_best[i]].push_back(d); + } +} + +bool should_offload_assignment_to_gpu(size_t n_docs, size_t n_clusters) { + if (!GpuClusterAssigner::available()) { + return false; + } + // With on-device gather/scatter and per-thread streams the offload is + // efficient, but very small lists still favor the CPU. Overridable for + // tuning/benchmarking. + size_t min_docs = 1024; + if (const char* env = std::getenv("NSPARSE_GPU_MIN_DOCS")) { + min_docs = static_cast(std::strtoull(env, nullptr, 10)); + } + return n_docs >= min_docs && n_clusters >= 2; +} + +} // namespace nsparse::detail diff --git a/nsparse/gpu/gpu_cluster_assigner.h b/nsparse/gpu/gpu_cluster_assigner.h new file mode 100644 index 0000000..b7268ce --- /dev/null +++ b/nsparse/gpu/gpu_cluster_assigner.h @@ -0,0 +1,101 @@ +/** + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +#ifndef GPU_CLUSTER_ASSIGNER_H +#define GPU_CLUSTER_ASSIGNER_H + +#include + +#include "nsparse/sparse_vectors.h" +#include "nsparse/types.h" + +namespace nsparse::detail { + +/** + * @brief GPU-accelerated k-means document-to-cluster assignment used by the + * index-building (build()) path. + * + * The assignment step of Seismic clustering computes, for every document, the + * dot product against each cluster centroid and picks the highest-scoring + * cluster. Across a whole inverted list this is a sparse-times-dense matrix + * product (documents as a CSR matrix, centroids as a dense matrix), which maps + * directly onto NVIDIA cuSPARSE SpMM followed by a per-row argmax. + * + * This class mirrors the CPU map_docs_to_clusters() semantics exactly for + * float (U32) weights: + * - documents equal to a cluster centroid are left untouched (not re-added), + * - ties are broken toward the lowest cluster index (strict-greater update). + * + * The search() path is unaffected and always runs on the CPU. + * + * The implementation lives in gpu_cluster_assigner.cu and is only compiled when + * the project is configured with -DNSPARSE_ENABLE_CUDA=ON (which defines + * NSPARSE_WITH_CUDA). The declarations below are always visible so that callers + * can guard usage with a single runtime available() check. + */ +class GpuClusterAssigner { +public: + /// Process-wide singleton. Owns the cuSPARSE handle and serializes GPU + /// access so it is safe to call from OpenMP worker threads. + static GpuClusterAssigner& instance(); + + /// True when the library was built with CUDA support and a usable GPU is + /// present. When false, callers must use the CPU path. + static bool available(); + + /** + * @brief Assign each non-centroid document to its best cluster on the GPU. + * + * For every document id in @p docs (interpreted as a row of @p vectors), + * computes the dot product against each cluster centroid + * (clusters[j].front()) and appends the document to the highest-scoring + * cluster. Documents that are themselves a centroid are skipped, matching + * the CPU reference. + * + * Only float (U32) weights are supported on the GPU. Callers must check + * vectors->get_element_size() == U32 and fall back to the CPU path + * otherwise. + * + * @param vectors full corpus of sparse vectors (CSR-backed). + * @param docs document ids belonging to this inverted list. + * @param clusters in/out clusters; clusters[j].front() is the centroid of + * cluster j and assigned documents are appended. + */ + void assign(const SparseVectors* vectors, const std::vector& docs, + std::vector>& clusters); + + GpuClusterAssigner(const GpuClusterAssigner&) = delete; + GpuClusterAssigner& operator=(const GpuClusterAssigner&) = delete; + +private: + GpuClusterAssigner(); + ~GpuClusterAssigner(); + + struct Impl; + // Uploads the corpus to the GPU on first use (or when it changes) and + // returns an opaque handle to the resident device buffers. Thread-safe. + const void* ensure_corpus_resident(const SparseVectors* vectors); + + Impl* impl_; +}; + +/** + * @brief Heuristic gate for the automatic GPU path in map_docs_to_clusters(). + * + * The per-inverted-list problems in a Seismic build are individually small and + * the CPU already parallelizes across lists, so offloading tiny problems to the + * GPU loses to transfer/launch overhead. This returns true only once the + * problem is large enough to benefit. The threshold on the number of documents + * can be overridden with the NSPARSE_GPU_MIN_DOCS environment variable. + */ +bool should_offload_assignment_to_gpu(size_t n_docs, size_t n_clusters); + +} // namespace nsparse::detail + +#endif // GPU_CLUSTER_ASSIGNER_H diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index 2d59fa7..e8fa816 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -31,6 +31,11 @@ set(NSPARSE_TEST_SRC vector_process_test.cpp ) +# GPU parity test is only meaningful when CUDA support is compiled in. +if(NSPARSE_ENABLE_CUDA) + list(APPEND NSPARSE_TEST_SRC gpu_cluster_assigner_test.cpp) +endif() + # Fetch GoogleTest include(FetchContent) FetchContent_Declare( diff --git a/tests/gpu_cluster_assigner_test.cpp b/tests/gpu_cluster_assigner_test.cpp new file mode 100644 index 0000000..95ec047 --- /dev/null +++ b/tests/gpu_cluster_assigner_test.cpp @@ -0,0 +1,160 @@ +/** + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +#include "nsparse/gpu/gpu_cluster_assigner.h" + +#include + +#include +#include +#include +#include + +#include "nsparse/cluster/kmeans_utils.h" +#include "nsparse/sparse_vectors.h" +#include "nsparse/types.h" + +namespace nsparse::detail { +namespace { + +// CPU reference assignment, copied to mirror the scalar map_docs_to_clusters() +// semantics: strict-greater argmax (ties -> lowest cluster index) and skipping +// documents that are themselves a centroid. +std::vector> cpu_reference_assign( + const SparseVectors* vectors, const std::vector& docs, + std::vector> clusters) { + const idx_t* indptr = vectors->indptr_data(); + const term_t* indices = vectors->indices_data(); + const float* values = vectors->values_data_float(); + const size_t n_clusters = clusters.size(); + for (size_t i = 0; i < docs.size(); ++i) { + const auto dense = vectors->get_dense_vector_float(docs[i]); + float best = std::numeric_limits::lowest(); + size_t best_j = 0; + bool is_centroid = false; + for (size_t j = 0; j < n_clusters; ++j) { + const idx_t c = clusters[j].front(); + if (docs[i] == c) { + is_centroid = true; + break; + } + const idx_t start = indptr[c]; + const size_t len = indptr[c + 1] - start; + float score = 0.0F; + for (size_t t = 0; t < len; ++t) { + score += dense[indices[start + t]] * values[start + t]; + } + if (score > best) { + best = score; + best_j = j; + } + } + if (!is_centroid) { + clusters[best_j].push_back(docs[i]); + } + } + return clusters; +} + +SparseVectors make_random_corpus(size_t n_docs, size_t dim, size_t nnz_per_doc, + unsigned seed) { + SparseVectors vectors({.element_size = U32, .dimension = dim}); + std::mt19937 gen(seed); + std::uniform_int_distribution term_dist( + 0, static_cast(dim - 1)); + std::uniform_real_distribution val_dist(0.1F, 1.0F); + for (size_t d = 0; d < n_docs; ++d) { + std::vector terms; + std::vector vals; + // Deduplicate terms so each column appears once per row. + std::vector chosen; + for (size_t k = 0; k < nnz_per_doc; ++k) { + chosen.push_back(term_dist(gen)); + } + std::ranges::sort(chosen); + chosen.erase(std::unique(chosen.begin(), chosen.end()), chosen.end()); + for (term_t t : chosen) { + terms.push_back(t); + vals.push_back(val_dist(gen)); + } + vectors.add_vector( + terms.data(), terms.size(), + reinterpret_cast(vals.data()), + vals.size() * sizeof(float)); + } + return vectors; +} + +// Build initial clusters: the first n_clusters docs are the centroids. +std::vector> seed_clusters(const std::vector& docs, + size_t n_clusters) { + std::vector> clusters(n_clusters); + for (size_t j = 0; j < n_clusters; ++j) { + clusters[j].push_back(docs[j]); + } + return clusters; +} + +TEST(GpuClusterAssignerTest, MatchesCpuReference) { + if (!GpuClusterAssigner::available()) { + GTEST_SKIP() << "No CUDA-capable GPU available"; + } + constexpr size_t kNumDocs = 6000; + constexpr size_t kDim = 2000; + constexpr size_t kNnz = 40; + constexpr size_t kNumClusters = 32; + + SparseVectors vectors = make_random_corpus(kNumDocs, kDim, kNnz, 1234); + + std::vector docs(kNumDocs); + std::iota(docs.begin(), docs.end(), 0); + + auto expected = + cpu_reference_assign(&vectors, docs, seed_clusters(docs, kNumClusters)); + + auto actual = seed_clusters(docs, kNumClusters); + GpuClusterAssigner::instance().assign(&vectors, docs, actual); + + ASSERT_EQ(actual.size(), expected.size()); + for (size_t j = 0; j < expected.size(); ++j) { + // Assignment order within a cluster is preserved (docs iterated in + // order on both paths), so compare directly. + EXPECT_EQ(actual[j], expected[j]) << "cluster " << j << " differs"; + } +} + +TEST(GpuClusterAssignerTest, AutoPathViaMapDocsMatchesReference) { + if (!GpuClusterAssigner::available()) { + GTEST_SKIP() << "No CUDA-capable GPU available"; + } + // Force the auto-offload gate to fire for this size. + ::setenv("NSPARSE_GPU_MIN_DOCS", "1", /*overwrite=*/1); + constexpr size_t kNumDocs = 5000; + constexpr size_t kDim = 1500; + constexpr size_t kNnz = 30; + constexpr size_t kNumClusters = 16; + + SparseVectors vectors = make_random_corpus(kNumDocs, kDim, kNnz, 99); + std::vector docs(kNumDocs); + std::iota(docs.begin(), docs.end(), 0); + + auto expected = + cpu_reference_assign(&vectors, docs, seed_clusters(docs, kNumClusters)); + + auto actual = seed_clusters(docs, kNumClusters); + map_docs_to_clusters(&vectors, docs, actual); // routes to GPU + + ASSERT_EQ(actual.size(), expected.size()); + for (size_t j = 0; j < expected.size(); ++j) { + EXPECT_EQ(actual[j], expected[j]) << "cluster " << j << " differs"; + } +} + +} // namespace +} // namespace nsparse::detail From ce5b75af133dbd33c4e38104be1117e29f2afe40 Mon Sep 17 00:00:00 2001 From: Liyun Xiu Date: Tue, 7 Jul 2026 02:08:32 +0000 Subject: [PATCH 2/8] Add opt-in GPU profiling and env-configurable benchmark iters - NSPARSE_GPU_PROFILE=1 env var enables lightweight phase profiling that prints host_prep / gpu_exec / host_assign breakdown at exit. Negligible overhead when disabled (a conditional branch per list). - NSPARSE_BENCH_ITERS / NSPARSE_BENCH_REPS env vars control benchmark iterations/repetitions (default 3/3), so base_full runs can use 1/1 without rebuilding. Profiling findings (base_full, 8.84M docs, L4): gpu_exec dominates at 97% of wall-time (141s aggregate / 84s wall). host_prep and host_assign together are ~5%. GPU SM utilization ~60% during the build; cuSPARSE SpMM throughput is the bottleneck. A custom fused warp-per-row kernel was tested but proved slightly slower than cuSPARSE on base_full's 4400x440 problems (cuSPARSE's tiled algorithm has better B-matrix cache behavior). Kept cuSPARSE. Signed-off-by: Liyun Xiu --- nsparse/gpu/gpu_cluster_assigner.cu | 134 ++++++++++++++++++++-------- 1 file changed, 95 insertions(+), 39 deletions(-) diff --git a/nsparse/gpu/gpu_cluster_assigner.cu b/nsparse/gpu/gpu_cluster_assigner.cu index 288a611..42f0fce 100644 --- a/nsparse/gpu/gpu_cluster_assigner.cu +++ b/nsparse/gpu/gpu_cluster_assigner.cu @@ -11,7 +11,10 @@ #include +#include +#include #include +#include #include #include #include @@ -48,6 +51,39 @@ void check_cusparse(cusparseStatus_t status, const char* what) { #define NSPARSE_CUDA_CHECK(expr) check_cuda((expr), #expr) #define NSPARSE_CUSPARSE_CHECK(expr) check_cusparse((expr), #expr) +// Optional phase profiling, enabled by NSPARSE_GPU_PROFILE=1. Accumulates +// wall-time (ns) across all list calls into a few buckets so the per-phase +// split of the GPU path can be inspected without an external profiler. Cheap +// (a handful of atomics per list) and compiled to near-nothing when disabled. +struct Profile { + std::atomic host_prep{0}; // row_ptr/centroid host setup + std::atomic gpu_exec{0}; // upload+kernels+SpMM+sync + std::atomic host_assign{0}; // scatter results into clusters + std::atomic calls{0}; + bool enabled = false; + + Profile() { + const char* v = std::getenv("NSPARSE_GPU_PROFILE"); + enabled = (v != nullptr && v[0] == '1'); + } + ~Profile() { + if (!enabled) return; + std::fprintf(stderr, + "[nsparse gpu] calls=%lld host_prep=%.3fs gpu_exec=%.3fs " + "host_assign=%.3fs\n", + static_cast(calls.load()), + host_prep.load() / 1e9, gpu_exec.load() / 1e9, + host_assign.load() / 1e9); + } +}; +Profile g_profile; + +inline int64_t now_ns() { + return std::chrono::duration_cast( + std::chrono::steady_clock::now().time_since_epoch()) + .count(); +} + // One thread per document row. Scans the n_clusters scores for that row and // selects the argmax, breaking ties toward the lowest cluster index (strict // greater-than update) to match the CPU reference in map_docs_to_clusters(). @@ -94,6 +130,7 @@ __global__ void gather_csr_kernel(const int32_t* __restrict__ corpus_indptr, } } + // Scatters centroid rows from the resident corpus into the dense B matrix // (dim x n_clusters, row-major, ldb = n_clusters). Column j of B is centroid j. // One thread per centroid; each writes a disjoint column, so there are no @@ -305,6 +342,9 @@ void GpuClusterAssigner::assign(const SparseVectors* vectors, return; } + const bool prof = g_profile.enabled; + const int64_t t0 = prof ? now_ns() : 0; + const idx_t* indptr = vectors->indptr_data(); const size_t dim = vectors->get_dimension(); @@ -328,47 +368,32 @@ void GpuClusterAssigner::assign(const SparseVectors* vectors, } const int64_t nnz_a = h_a_row_ptr[n_docs]; + const int64_t t1 = prof ? now_ns() : 0; + const DeviceCorpus& corpus = *static_cast(ensure_corpus_resident(vectors)); ThreadCtx& ctx = thread_ctx(); + cudaStream_t stream = ctx.stream; - // Size (grow-on-demand) all per-thread scratch buffers for this list. + // Upload only the small per-list metadata (doc/centroid ids). The actual + // doc/centroid payloads live on the GPU (resident corpus). ensure_capacity(&ctx.d_docs, ctx.docs_cap, n_docs * sizeof(int32_t)); ensure_capacity(&ctx.d_centroids, ctx.cent_cap, n_clusters * sizeof(int32_t)); - ensure_capacity(&ctx.d_a_row_ptr, ctx.rowptr_cap, - (n_docs + 1) * sizeof(int32_t)); - ensure_capacity(&ctx.d_a_col, ctx.col_cap, - static_cast(nnz_a) * sizeof(int32_t)); - ensure_capacity(&ctx.d_a_val, ctx.val_cap, - static_cast(nnz_a) * sizeof(float)); - ensure_capacity(&ctx.d_b, ctx.b_cap, dim * n_clusters * sizeof(float)); - ensure_capacity(&ctx.d_c, ctx.c_cap, n_docs * n_clusters * sizeof(float)); ensure_capacity(&ctx.d_best, ctx.best_cap, n_docs * sizeof(int32_t)); + ensure_capacity(&ctx.d_b, ctx.b_cap, dim * n_clusters * sizeof(float)); - cudaStream_t stream = ctx.stream; - - // Upload only the small per-list metadata; doc/centroid payloads stay on - // the GPU (gathered/scattered from the resident corpus below). NSPARSE_CUDA_CHECK(cudaMemcpyAsync(ctx.d_docs, docs.data(), n_docs * sizeof(int32_t), cudaMemcpyHostToDevice, stream)); NSPARSE_CUDA_CHECK(cudaMemcpyAsync(ctx.d_centroids, centroid_docs.data(), n_clusters * sizeof(int32_t), cudaMemcpyHostToDevice, stream)); - NSPARSE_CUDA_CHECK(cudaMemcpyAsync(ctx.d_a_row_ptr, h_a_row_ptr.data(), - (n_docs + 1) * sizeof(int32_t), - cudaMemcpyHostToDevice, stream)); - // Gather A's CSR payload from the corpus. + // Build dense B (dim x n_clusters, row-major) from centroid rows. The + // fused kernel below reads B but not A's CSR separately; docs are gathered + // directly from the corpus inside the fused kernel. constexpr int kBlock = 256; - const int docs_grid = static_cast((n_docs + kBlock - 1) / kBlock); - gather_csr_kernel<<>>( - corpus.indptr, corpus.indices, corpus.values, ctx.d_docs, - static_cast(n_docs), ctx.d_a_row_ptr, ctx.d_a_col, ctx.d_a_val); - NSPARSE_CUDA_CHECK(cudaGetLastError()); - - // Build dense B (dim x n_clusters, row-major) from centroid rows. NSPARSE_CUDA_CHECK(cudaMemsetAsync( ctx.d_b, 0, dim * n_clusters * sizeof(float), stream)); const int cent_grid = @@ -378,51 +403,74 @@ void GpuClusterAssigner::assign(const SparseVectors* vectors, static_cast(n_clusters), static_cast(n_clusters), ctx.d_b); NSPARSE_CUDA_CHECK(cudaGetLastError()); - // C = A * B, with A sparse (CSR), B/C dense row-major. + // cuSPARSE SpMM path: gathers A from corpus, then runs SpMM + argmax. + ensure_capacity(&ctx.d_a_row_ptr, ctx.rowptr_cap, + (n_docs + 1) * sizeof(int32_t)); + ensure_capacity(&ctx.d_a_col, ctx.col_cap, + static_cast(nnz_a) * sizeof(int32_t)); + ensure_capacity(&ctx.d_a_val, ctx.val_cap, + static_cast(nnz_a) * sizeof(float)); + ensure_capacity(&ctx.d_c, ctx.c_cap, + n_docs * n_clusters * sizeof(float)); + + NSPARSE_CUDA_CHECK(cudaMemcpyAsync(ctx.d_a_row_ptr, h_a_row_ptr.data(), + (n_docs + 1) * sizeof(int32_t), + cudaMemcpyHostToDevice, stream)); + const int docs_grid = static_cast((n_docs + kBlock - 1) / kBlock); + gather_csr_kernel<<>>( + corpus.indptr, corpus.indices, corpus.values, ctx.d_docs, + static_cast(n_docs), ctx.d_a_row_ptr, ctx.d_a_col, ctx.d_a_val); + NSPARSE_CUDA_CHECK(cudaGetLastError()); + cusparseSpMatDescr_t mat_a; cusparseDnMatDescr_t mat_b; cusparseDnMatDescr_t mat_c; NSPARSE_CUSPARSE_CHECK(cusparseCreateCsr( - &mat_a, static_cast(n_docs), static_cast(dim), nnz_a, - ctx.d_a_row_ptr, ctx.d_a_col, ctx.d_a_val, CUSPARSE_INDEX_32I, - CUSPARSE_INDEX_32I, CUSPARSE_INDEX_BASE_ZERO, CUDA_R_32F)); + &mat_a, static_cast(n_docs), static_cast(dim), + nnz_a, ctx.d_a_row_ptr, ctx.d_a_col, ctx.d_a_val, + CUSPARSE_INDEX_32I, CUSPARSE_INDEX_32I, CUSPARSE_INDEX_BASE_ZERO, + CUDA_R_32F)); NSPARSE_CUSPARSE_CHECK(cusparseCreateDnMat( - &mat_b, static_cast(dim), static_cast(n_clusters), + &mat_b, static_cast(dim), + static_cast(n_clusters), static_cast(n_clusters), ctx.d_b, CUDA_R_32F, CUSPARSE_ORDER_ROW)); NSPARSE_CUSPARSE_CHECK(cusparseCreateDnMat( - &mat_c, static_cast(n_docs), static_cast(n_clusters), + &mat_c, static_cast(n_docs), + static_cast(n_clusters), static_cast(n_clusters), ctx.d_c, CUDA_R_32F, CUSPARSE_ORDER_ROW)); - const float alpha = 1.0F; - const float beta = 0.0F; + const float alpha_v = 1.0F; + const float beta_v = 0.0F; size_t buffer_size = 0; NSPARSE_CUSPARSE_CHECK(cusparseSpMM_bufferSize( ctx.handle, CUSPARSE_OPERATION_NON_TRANSPOSE, - CUSPARSE_OPERATION_NON_TRANSPOSE, &alpha, mat_a, mat_b, &beta, mat_c, - CUDA_R_32F, CUSPARSE_SPMM_ALG_DEFAULT, &buffer_size)); + CUSPARSE_OPERATION_NON_TRANSPOSE, &alpha_v, mat_a, mat_b, &beta_v, + mat_c, CUDA_R_32F, CUSPARSE_SPMM_ALG_DEFAULT, &buffer_size)); ensure_capacity(reinterpret_cast(&ctx.d_spmm), ctx.spmm_cap, buffer_size == 0 ? 1 : buffer_size); NSPARSE_CUSPARSE_CHECK(cusparseSpMM( ctx.handle, CUSPARSE_OPERATION_NON_TRANSPOSE, - CUSPARSE_OPERATION_NON_TRANSPOSE, &alpha, mat_a, mat_b, &beta, mat_c, - CUDA_R_32F, CUSPARSE_SPMM_ALG_DEFAULT, ctx.d_spmm)); + CUSPARSE_OPERATION_NON_TRANSPOSE, &alpha_v, mat_a, mat_b, &beta_v, + mat_c, CUDA_R_32F, CUSPARSE_SPMM_ALG_DEFAULT, ctx.d_spmm)); row_argmax_kernel<<>>( ctx.d_c, static_cast(n_docs), static_cast(n_clusters), ctx.d_best); NSPARSE_CUDA_CHECK(cudaGetLastError()); + cusparseDestroySpMat(mat_a); + cusparseDestroyDnMat(mat_b); + cusparseDestroyDnMat(mat_c); + std::vector h_best(n_docs); NSPARSE_CUDA_CHECK(cudaMemcpyAsync(h_best.data(), ctx.d_best, n_docs * sizeof(int32_t), cudaMemcpyDeviceToHost, stream)); NSPARSE_CUDA_CHECK(cudaStreamSynchronize(stream)); - cusparseDestroySpMat(mat_a); - cusparseDestroyDnMat(mat_b); - cusparseDestroyDnMat(mat_c); + const int64_t t2 = prof ? now_ns() : 0; // Append assignments on the host, skipping documents that are themselves a // centroid (matching the CPU reference exactly). @@ -433,6 +481,14 @@ void GpuClusterAssigner::assign(const SparseVectors* vectors, } clusters[h_best[i]].push_back(d); } + + if (prof) { + const int64_t t3 = now_ns(); + g_profile.host_prep.fetch_add(t1 - t0, std::memory_order_relaxed); + g_profile.gpu_exec.fetch_add(t2 - t1, std::memory_order_relaxed); + g_profile.host_assign.fetch_add(t3 - t2, std::memory_order_relaxed); + g_profile.calls.fetch_add(1, std::memory_order_relaxed); + } } bool should_offload_assignment_to_gpu(size_t n_docs, size_t n_clusters) { From 47087d71723b9c079399af30ab8d8269026383ab Mon Sep 17 00:00:00 2001 From: Liyun Xiu Date: Tue, 7 Jul 2026 04:12:28 +0000 Subject: [PATCH 3/8] Offload summarize() max-pool to GPU, batched per inverted list MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The per-term max-pool is ~67% of summarize() and, after the assignment offload, summarize became the dominant remaining CPU cost in build() (profiling: assignment dropped 84x, leaving summarize as the bottleneck that capped GPU utilization at ~45%). Offload the max-pool to the GPU as ONE kernel launch per inverted list (one thread block per cluster), not per cluster. Clusters average only ~10 docs, so a per-cluster launch/sync is dominated by GPU round-trip overhead — an initial per-cluster version was ~5x slower than CPU. The batched per-list kernel collapses ~27M blocking syncs down to ~30K. The tie-order-sensitive sort/truncate stays on the CPU, so output ordering is preserved; results are bit-identical to the CPU max-pool on tie-free data (verified by SummarizeListMaxpool* tests). float (U32) only; default on when a GPU is present, opt out with NSPARSE_GPU_SUMMARIZE=0. Also adds NSPARSE_BUILD_PROFILE / NSPARSE_SUMMARIZE_PROFILE phase timers. Benchmark (base_full, 8.84M docs, L4): build 83.6s -> 74.9s; summarize thread-seconds 1436 -> 760. Overall vs 32-core CPU baseline (483s): 5.96x -> 6.45x. (base_small's ~10-doc clusters are below the break-even point, so GPU summarize is a slight net loss there; the win is at scale.) All 365 tests pass. Signed-off-by: Liyun Xiu --- nsparse/cluster/inverted_list_clusters.cpp | 117 ++++++++-- nsparse/gpu/gpu_cluster_assigner.cu | 250 +++++++++++++++++++++ nsparse/gpu/gpu_cluster_assigner.h | 41 ++++ nsparse/seismic_common.h | 53 ++++- tests/gpu_cluster_assigner_test.cpp | 118 ++++++++++ 5 files changed, 557 insertions(+), 22 deletions(-) diff --git a/nsparse/cluster/inverted_list_clusters.cpp b/nsparse/cluster/inverted_list_clusters.cpp index 24ace98..2a7f023 100644 --- a/nsparse/cluster/inverted_list_clusters.cpp +++ b/nsparse/cluster/inverted_list_clusters.cpp @@ -10,17 +10,50 @@ #include "nsparse/cluster/inverted_list_clusters.h" #include +#include +#include #include +#include +#include #include +#include #include #include #include "nsparse/sparse_vectors.h" #include "nsparse/types.h" +#ifdef NSPARSE_WITH_CUDA +#include "nsparse/gpu/gpu_cluster_assigner.h" +#endif namespace nsparse { namespace { +// Opt-in split timing for summarize (NSPARSE_SUMMARIZE_PROFILE=1): how much of +// summarize is the memory-bound max-pool vs the per-cluster sort/truncate. +struct SummarizeProfile { + std::atomic maxpool_ns{0}; + std::atomic sort_ns{0}; + bool enabled = false; + SummarizeProfile() { + const char* v = std::getenv("NSPARSE_SUMMARIZE_PROFILE"); + enabled = (v != nullptr && v[0] == '1'); + } + ~SummarizeProfile() { + if (!enabled) return; + std::fprintf(stderr, + "[nsparse summarize] maxpool=%.1fs sort_truncate=%.1fs\n", + maxpool_ns.load() / 1e9, sort_ns.load() / 1e9); + } +}; +SummarizeProfile g_summ_profile; + +inline int64_t summ_now_ns() { + return std::chrono::duration_cast( + std::chrono::steady_clock::now().time_since_epoch()) + .count(); +} + /** * @brief Generage summary sparse vector for posting lists * @@ -43,28 +76,69 @@ SparseVectors summarize_(const SparseVectors* vectors, const auto& indptr_data = vectors->indptr_data(); const auto& indices_data = vectors->indices_data(); const auto& values_data = vectors->values_data(); + const bool prof = g_summ_profile.enabled; + int64_t maxpool_acc = 0; + int64_t sort_acc = 0; + // The per-term max-pool (the memory-bound majority of summarize) can be + // offloaded to the GPU for float (U32) weights, batched as ONE launch for + // the whole list's clusters (per-cluster launches are dominated by GPU + // round-trip overhead). The tie-order-sensitive sort/truncate below stays + // on the CPU so output ordering is preserved. + bool gpu_ok = false; +#ifdef NSPARSE_WITH_CUDA + std::vector gpu_clusters; + if constexpr (std::is_same_v) { + if (detail::should_offload_summarize_to_gpu()) { + gpu_ok = detail::GpuClusterAssigner::instance() + .summarize_list_maxpool(vectors, + group_of_doc_ids.data(), + offsets.data(), + offsets.size() - 1, + gpu_clusters); + } + } +#endif + for (size_t i = 0; i < offsets.size() - 1; ++i) { + const int64_t ts0 = prof ? summ_now_ns() : 0; size_t n_docs = offsets[i + 1] - offsets[i]; - std::unordered_map summary_map; + std::vector> summary_vec; float sum = 0.0F; - auto doc_ids = std::span( - group_of_doc_ids.data() + offsets[i], n_docs); - for (const auto& doc_id : doc_ids) { - int start = indptr_data[doc_id]; - int end = indptr_data[doc_id + 1]; - for (size_t j = start; j < end; ++j) { - const auto old = summary_map[indices_data[j]]; - auto& value = summary_map[indices_data[j]]; - // j is element index, need byte offset for T access - value = std::max(value, *reinterpret_cast( - values_data + j * sizeof(T))); - sum += value - old; + + bool used_gpu = false; +#ifdef NSPARSE_WITH_CUDA + if constexpr (std::is_same_v) { + if (gpu_ok) { + const auto& gc = gpu_clusters[i]; + summary_vec.reserve(gc.terms.size()); + for (size_t t = 0; t < gc.terms.size(); ++t) { + summary_vec.emplace_back(gc.terms[t], gc.values[t]); + } + sum = gc.sum; + used_gpu = true; + } + } +#endif + if (!used_gpu) { + std::unordered_map summary_map; + auto doc_ids = std::span( + group_of_doc_ids.data() + offsets[i], n_docs); + for (const auto& doc_id : doc_ids) { + int start = indptr_data[doc_id]; + int end = indptr_data[doc_id + 1]; + for (size_t j = start; j < end; ++j) { + const auto old = summary_map[indices_data[j]]; + auto& value = summary_map[indices_data[j]]; + // j is element index, need byte offset for T access + value = std::max(value, *reinterpret_cast( + values_data + j * sizeof(T))); + sum += value - old; + } } + summary_vec.assign(summary_map.begin(), summary_map.end()); } - // Convert summary_map to vector of pairs - std::vector> summary_vec(summary_map.begin(), - summary_map.end()); + const int64_t ts1 = prof ? summ_now_ns() : 0; // Sort by value in descending order std::ranges::sort(summary_vec, [](const auto& a, const auto& b) { @@ -100,6 +174,17 @@ SparseVectors summarize_(const SparseVectors* vectors, terms.data(), terms.size(), reinterpret_cast(values.data()), values.size() * sizeof(T)); + + if (prof) { + const int64_t ts2 = summ_now_ns(); + maxpool_acc += ts1 - ts0; + sort_acc += ts2 - ts1; + } + } + if (prof) { + g_summ_profile.maxpool_ns.fetch_add(maxpool_acc, + std::memory_order_relaxed); + g_summ_profile.sort_ns.fetch_add(sort_acc, std::memory_order_relaxed); } return summarized_vectors; } diff --git a/nsparse/gpu/gpu_cluster_assigner.cu b/nsparse/gpu/gpu_cluster_assigner.cu index 42f0fce..ea461d1 100644 --- a/nsparse/gpu/gpu_cluster_assigner.cu +++ b/nsparse/gpu/gpu_cluster_assigner.cu @@ -154,6 +154,98 @@ __global__ void scatter_dense_kernel(const int32_t* __restrict__ corpus_indptr, } } +// ========================================================================= +// GPU max-pool for summarize(). For one cluster (posting list) this computes, +// per term, the maximum weight across all documents in the cluster, then +// compacts the touched terms. SPLADE weights are >= 0, so a non-negative float +// compares monotonically as a reinterpreted int32 — allowing atomicMax on the +// integer view of the accumulator. +// +// Buffers (all sized to corpus dim, reused across clusters, left clean on exit +// so no per-cluster dim-wide memset is needed): +// acc[dim] : int32 view of the running max weight per term (0 == absent) +// seen[dim] : 0/1 first-touch flag used to append each term once to touched +// touched[] : compact list of term ids that appeared in this cluster +// n_touched : length of touched[] (atomic counter) +// ========================================================================= + +// Batched per-list max-pool: ONE block per cluster, so an entire inverted +// list's clusters are summarized in a single kernel launch with a single +// stream sync (instead of two blocking syncs per cluster, which serialized ~27M +// GPU round-trips across the build). Clusters are independent and each is owned +// by one block, so only intra-block synchronization is needed. +// +// Buffers (per thread/list, reused across lists, kept clean via touched-reset): +// acc[n_clusters*dim], seen[n_clusters*dim] : per-cluster per-term max / flag, +// partitioned so block b owns region [b*dim, (b+1)*dim). Start zeroed and +// are restored to zero for touched slots before the kernel returns. +// out_term / out_val : compacted results, laid out per cluster using +// out_base[b] (the input nnz prefix, an upper bound on distinct terms). +// out_count[b] : number of distinct terms produced for cluster b. +// out_sum[b] : sum of the cluster's max values. +__global__ void summarize_list_kernel( + const int32_t* __restrict__ corpus_indptr, + const int32_t* __restrict__ corpus_indices, + const float* __restrict__ corpus_values, + const int32_t* __restrict__ docs, const int32_t* __restrict__ offsets, + const int32_t* __restrict__ out_base, int dim, + int32_t* __restrict__ acc, int32_t* __restrict__ seen, + int32_t* __restrict__ out_term, float* __restrict__ out_val, + int32_t* __restrict__ out_count, float* __restrict__ out_sum) { + const int b = blockIdx.x; // one block per cluster + const int start = offsets[b]; + const int end = offsets[b + 1]; + const int64_t base = static_cast(b) * dim; + const int32_t obase = out_base[b]; + + __shared__ int s_count; + if (threadIdx.x == 0) { + s_count = 0; + } + __syncthreads(); + + // Scatter-max over the cluster's documents. Each distinct term is appended + // once (via seen[]) to this cluster's output segment at out_term[obase..]. + for (int i = start + threadIdx.x; i < end; i += blockDim.x) { + const int32_t d = docs[i]; + const int32_t src = corpus_indptr[d]; + const int32_t len = corpus_indptr[d + 1] - src; + for (int32_t t = 0; t < len; ++t) { + const int32_t term = corpus_indices[src + t]; + const float v = corpus_values[src + t]; + // Non-negative float -> monotonic int bits; atomicMax on int view. + atomicMax(&acc[base + term], __float_as_int(v)); + if (atomicCAS(&seen[base + term], 0, 1) == 0) { + const int slot = atomicAdd(&s_count, 1); + out_term[obase + slot] = term; + } + } + } + __syncthreads(); + + // Compact: read each touched term's max, accumulate the sum, and reset + // acc/seen for reuse. Sum via a block-wide atomic into a shared scalar. + __shared__ float s_sum; + if (threadIdx.x == 0) { + s_sum = 0.0F; + out_count[b] = s_count; + } + __syncthreads(); + + for (int k = threadIdx.x; k < s_count; k += blockDim.x) { + const int32_t term = out_term[obase + k]; + const float v = __int_as_float(acc[base + term]); + out_val[obase + k] = v; + atomicAdd(&s_sum, v); + acc[base + term] = 0; // reset for next list reusing this buffer + seen[base + term] = 0; + } + __syncthreads(); + if (threadIdx.x == 0) { + out_sum[b] = s_sum; + } +} + // The full corpus, uploaded to the GPU once and shared read-only across all // threads/lists. Indices are widened to int32 on upload to feed both the gather // kernel and cuSPARSE's 32-bit CSR indices. @@ -204,6 +296,19 @@ struct ThreadCtx { int32_t* d_best = nullptr; size_t best_cap = 0; void* d_spmm = nullptr; size_t spmm_cap = 0; + // summarize() batched per-list max-pool scratch. acc/seen are partitioned + // as n_clusters x dim and kept clean between lists (each cluster resets the + // slots it touched). Sized to the largest list seen so far. + int32_t* d_sum_docs = nullptr; size_t sum_docs_cap = 0; + int32_t* d_offsets = nullptr; size_t offsets_cap = 0; + int32_t* d_out_base = nullptr; size_t out_base_cap = 0; + int32_t* d_acc = nullptr; size_t acc_cap = 0; + int32_t* d_seen = nullptr; size_t seen_cap = 0; + int32_t* d_out_term = nullptr; size_t out_term_cap = 0; + float* d_out_val = nullptr; size_t out_val_cap = 0; + int32_t* d_out_count = nullptr; size_t out_count_cap = 0; + float* d_out_sum = nullptr; size_t out_sum_cap = 0; + ~ThreadCtx() { cudaFree(d_docs); cudaFree(d_centroids); @@ -214,6 +319,15 @@ struct ThreadCtx { cudaFree(d_c); cudaFree(d_best); cudaFree(d_spmm); + cudaFree(d_sum_docs); + cudaFree(d_offsets); + cudaFree(d_out_base); + cudaFree(d_acc); + cudaFree(d_seen); + cudaFree(d_out_term); + cudaFree(d_out_val); + cudaFree(d_out_count); + cudaFree(d_out_sum); if (init) { cusparseDestroy(handle); cudaStreamDestroy(stream); @@ -491,6 +605,128 @@ void GpuClusterAssigner::assign(const SparseVectors* vectors, } } +bool GpuClusterAssigner::summarize_list_maxpool( + const SparseVectors* vectors, const idx_t* docs, const idx_t* offsets, + size_t n_clusters, std::vector& out) { + if (!impl_->usable || n_clusters == 0) { + return false; + } + const size_t dim = vectors->get_dimension(); + const idx_t* indptr = vectors->indptr_data(); + + // The flattened doc array spans offsets[0]..offsets[n_clusters]. Compute, + // per cluster, the nnz prefix (out_base) used to place each cluster's + // compacted output in a private segment; the total nnz is the output cap. + const size_t n_docs = static_cast(offsets[n_clusters] - offsets[0]); + if (n_docs == 0) { + return false; + } + std::vector h_out_base(n_clusters + 1, 0); + for (size_t b = 0; b < n_clusters; ++b) { + int64_t nnz_b = 0; + for (idx_t i = offsets[b]; i < offsets[b + 1]; ++i) { + const idx_t d = docs[i]; + nnz_b += indptr[d + 1] - indptr[d]; + } + h_out_base[b + 1] = h_out_base[b] + static_cast(nnz_b); + } + const int64_t total_nnz = h_out_base[n_clusters]; + if (total_nnz == 0) { + return false; + } + + const DeviceCorpus& corpus = + *static_cast(ensure_corpus_resident(vectors)); + ThreadCtx& ctx = thread_ctx(); + cudaStream_t stream = ctx.stream; + + // acc/seen are partitioned as n_clusters x dim; must start (and remain) + // zeroed. Zero only when (re)grown — the kernel restores touched slots. + const size_t acc_bytes = n_clusters * dim * sizeof(int32_t); + const bool acc_grew = ctx.acc_cap < acc_bytes; + ensure_capacity(&ctx.d_acc, ctx.acc_cap, acc_bytes); + ensure_capacity(&ctx.d_seen, ctx.seen_cap, acc_bytes); + if (acc_grew) { + NSPARSE_CUDA_CHECK(cudaMemsetAsync(ctx.d_acc, 0, ctx.acc_cap, stream)); + NSPARSE_CUDA_CHECK( + cudaMemsetAsync(ctx.d_seen, 0, ctx.seen_cap, stream)); + } + ensure_capacity(&ctx.d_sum_docs, ctx.sum_docs_cap, + n_docs * sizeof(int32_t)); + ensure_capacity(&ctx.d_offsets, ctx.offsets_cap, + (n_clusters + 1) * sizeof(int32_t)); + ensure_capacity(&ctx.d_out_base, ctx.out_base_cap, + (n_clusters + 1) * sizeof(int32_t)); + ensure_capacity(&ctx.d_out_term, ctx.out_term_cap, + static_cast(total_nnz) * sizeof(int32_t)); + ensure_capacity(&ctx.d_out_val, ctx.out_val_cap, + static_cast(total_nnz) * sizeof(float)); + ensure_capacity(&ctx.d_out_count, ctx.out_count_cap, + n_clusters * sizeof(int32_t)); + ensure_capacity(&ctx.d_out_sum, ctx.out_sum_cap, + n_clusters * sizeof(float)); + + // Offsets are relative to offsets[0]; rebase to 0 for the flat doc upload. + std::vector h_offsets(n_clusters + 1); + const idx_t base0 = offsets[0]; + for (size_t b = 0; b <= n_clusters; ++b) { + h_offsets[b] = static_cast(offsets[b] - base0); + } + + NSPARSE_CUDA_CHECK(cudaMemcpyAsync(ctx.d_sum_docs, docs + base0, + n_docs * sizeof(int32_t), + cudaMemcpyHostToDevice, stream)); + NSPARSE_CUDA_CHECK(cudaMemcpyAsync(ctx.d_offsets, h_offsets.data(), + (n_clusters + 1) * sizeof(int32_t), + cudaMemcpyHostToDevice, stream)); + NSPARSE_CUDA_CHECK(cudaMemcpyAsync(ctx.d_out_base, h_out_base.data(), + (n_clusters + 1) * sizeof(int32_t), + cudaMemcpyHostToDevice, stream)); + + // One block per cluster. + constexpr int kBlock = 128; + summarize_list_kernel<<(n_clusters), kBlock, 0, stream>>>( + corpus.indptr, corpus.indices, corpus.values, ctx.d_sum_docs, + ctx.d_offsets, ctx.d_out_base, static_cast(dim), ctx.d_acc, + ctx.d_seen, ctx.d_out_term, ctx.d_out_val, ctx.d_out_count, + ctx.d_out_sum); + NSPARSE_CUDA_CHECK(cudaGetLastError()); + + // Copy back the compact outputs (single sync for the whole list). + std::vector h_term(total_nnz); + std::vector h_val(total_nnz); + std::vector h_count(n_clusters); + std::vector h_sum(n_clusters); + NSPARSE_CUDA_CHECK(cudaMemcpyAsync(h_term.data(), ctx.d_out_term, + total_nnz * sizeof(int32_t), + cudaMemcpyDeviceToHost, stream)); + NSPARSE_CUDA_CHECK(cudaMemcpyAsync(h_val.data(), ctx.d_out_val, + total_nnz * sizeof(float), + cudaMemcpyDeviceToHost, stream)); + NSPARSE_CUDA_CHECK(cudaMemcpyAsync(h_count.data(), ctx.d_out_count, + n_clusters * sizeof(int32_t), + cudaMemcpyDeviceToHost, stream)); + NSPARSE_CUDA_CHECK(cudaMemcpyAsync(h_sum.data(), ctx.d_out_sum, + n_clusters * sizeof(float), + cudaMemcpyDeviceToHost, stream)); + NSPARSE_CUDA_CHECK(cudaStreamSynchronize(stream)); + + out.resize(n_clusters); + for (size_t b = 0; b < n_clusters; ++b) { + const int32_t base = h_out_base[b]; + const int32_t cnt = h_count[b]; + ClusterSummary& cs = out[b]; + cs.terms.resize(cnt); + cs.values.resize(cnt); + for (int32_t k = 0; k < cnt; ++k) { + cs.terms[k] = static_cast(h_term[base + k]); + cs.values[k] = h_val[base + k]; + } + cs.sum = h_sum[b]; + } + return true; +} + bool should_offload_assignment_to_gpu(size_t n_docs, size_t n_clusters) { if (!GpuClusterAssigner::available()) { return false; @@ -505,4 +741,18 @@ bool should_offload_assignment_to_gpu(size_t n_docs, size_t n_clusters) { return n_docs >= min_docs && n_clusters >= 2; } +bool should_offload_summarize_to_gpu() { + if (!GpuClusterAssigner::available()) { + return false; + } + // On by default when a GPU is present. The max-pool is batched as one + // kernel launch per inverted list (all clusters at once), so it avoids the + // per-cluster sync storm that made a naive offload far slower. Opt out with + // NSPARSE_GPU_SUMMARIZE=0. + if (const char* env = std::getenv("NSPARSE_GPU_SUMMARIZE")) { + return env[0] != '0'; + } + return true; +} + } // namespace nsparse::detail diff --git a/nsparse/gpu/gpu_cluster_assigner.h b/nsparse/gpu/gpu_cluster_assigner.h index b7268ce..b0631e4 100644 --- a/nsparse/gpu/gpu_cluster_assigner.h +++ b/nsparse/gpu/gpu_cluster_assigner.h @@ -70,6 +70,39 @@ class GpuClusterAssigner { void assign(const SparseVectors* vectors, const std::vector& docs, std::vector>& clusters); + /// Per-cluster max-pool result from summarize_list_maxpool(). Terms are in + /// GPU touched-append order (not sorted); the CPU sort/truncate handles + /// ordering, preserving the existing output semantics. + struct ClusterSummary { + std::vector terms; + std::vector values; + float sum = 0.0F; + }; + + /** + * @brief GPU max-pool for a whole inverted list's clusters in one launch. + * + * Given the flattened cluster layout of one inverted list (@p docs with + * @p offsets, where cluster b owns docs[offsets[b]..offsets[b+1])), computes + * on the GPU, per cluster and per term, the maximum weight across the + * cluster's documents — the per-term max-pool that dominates summarize(). + * + * Processing the entire list in a single kernel launch + single stream sync + * (one thread block per cluster) is essential: clusters average only a + * handful of documents, so a per-cluster launch/sync is dominated by GPU + * round-trip overhead. The tie-order-sensitive sort/truncate stays on the + * CPU. Requires the corpus resident (shares assign()'s cache); float (U32) + * only. Returns false (caller falls back to CPU) if the GPU is unavailable + * or the list is empty. + * + * @param n_clusters offsets.size() - 1. + * @param out resized to n_clusters; out[b] holds cluster b's result. + */ + bool summarize_list_maxpool(const SparseVectors* vectors, + const idx_t* docs, const idx_t* offsets, + size_t n_clusters, + std::vector& out); + GpuClusterAssigner(const GpuClusterAssigner&) = delete; GpuClusterAssigner& operator=(const GpuClusterAssigner&) = delete; @@ -96,6 +129,14 @@ class GpuClusterAssigner { */ bool should_offload_assignment_to_gpu(size_t n_docs, size_t n_clusters); +/** + * @brief Gate for the GPU summarize() max-pool offload. + * + * Returns true when a usable GPU is present and the offload is enabled (on by + * default when built with CUDA; can be disabled with NSPARSE_GPU_SUMMARIZE=0). + */ +bool should_offload_summarize_to_gpu(); + } // namespace nsparse::detail #endif // GPU_CLUSTER_ASSIGNER_H diff --git a/nsparse/seismic_common.h b/nsparse/seismic_common.h index b73e09b..cf1948c 100644 --- a/nsparse/seismic_common.h +++ b/nsparse/seismic_common.h @@ -10,10 +10,16 @@ #ifndef SEISMIC_COMMON_H #define SEISMIC_COMMON_H +#include +#include #include #include #include +#ifdef _OPENMP +#include +#endif + #include "nsparse/cluster/inverted_list_clusters.h" #include "nsparse/cluster/random_kmeans.h" #include "nsparse/id_selector.h" @@ -139,17 +145,52 @@ inline std::vector build_inverted_lists_clusters( size_t inverted_lists_size = inverted_lists->size(); std::vector clustered_inverted_lists( inverted_lists_size); -#pragma omp parallel for schedule(dynamic, 64) + + // Optional coarse phase timing, enabled by NSPARSE_BUILD_PROFILE=1, to see + // how build CPU time splits across prune / train (k-means assignment) / + // summarize. Accumulates thread-seconds (sum across threads). + const bool build_prof = + std::getenv("NSPARSE_BUILD_PROFILE") != nullptr && + std::getenv("NSPARSE_BUILD_PROFILE")[0] == '1'; + double t_prune = 0.0; + double t_train = 0.0; + double t_summ = 0.0; + +#pragma omp parallel for schedule(dynamic, 64) \ + reduction(+ : t_prune, t_train, t_summ) for (int64_t idx = 0; idx < static_cast(inverted_lists_size); ++idx) { auto& invlist = (*inverted_lists)[idx]; - const auto& doc_ids = invlist.prune_and_keep_doc_ids(lambda); - InvertedListClusters inverted_list_clusters( - detail::RandomKMeans::train(vectors, doc_ids, beta)); - inverted_list_clusters.summarize(vectors, seismic_cluster_params.alpha); - clustered_inverted_lists[idx] = std::move(inverted_list_clusters); + if (build_prof) { + double a = omp_get_wtime(); + const auto doc_ids = invlist.prune_and_keep_doc_ids(lambda); + double b = omp_get_wtime(); + InvertedListClusters inverted_list_clusters( + detail::RandomKMeans::train(vectors, doc_ids, beta)); + double c = omp_get_wtime(); + inverted_list_clusters.summarize(vectors, + seismic_cluster_params.alpha); + double d = omp_get_wtime(); + t_prune += b - a; + t_train += c - b; + t_summ += d - c; + clustered_inverted_lists[idx] = std::move(inverted_list_clusters); + } else { + const auto& doc_ids = invlist.prune_and_keep_doc_ids(lambda); + InvertedListClusters inverted_list_clusters( + detail::RandomKMeans::train(vectors, doc_ids, beta)); + inverted_list_clusters.summarize(vectors, + seismic_cluster_params.alpha); + clustered_inverted_lists[idx] = std::move(inverted_list_clusters); + } invlist.clear(); } + if (build_prof) { + std::fprintf(stderr, + "[nsparse build] thread-seconds: prune=%.1f train=%.1f " + "summarize=%.1f\n", + t_prune, t_train, t_summ); + } return clustered_inverted_lists; } diff --git a/tests/gpu_cluster_assigner_test.cpp b/tests/gpu_cluster_assigner_test.cpp index 95ec047..704ee15 100644 --- a/tests/gpu_cluster_assigner_test.cpp +++ b/tests/gpu_cluster_assigner_test.cpp @@ -12,6 +12,7 @@ #include #include +#include #include #include #include @@ -156,5 +157,122 @@ TEST(GpuClusterAssignerTest, AutoPathViaMapDocsMatchesReference) { } } +// CPU reference for the summarize max-pool: per-term max weight over the docs +// of a cluster, plus the sum of maxes. Returns terms in ascending order. +void cpu_reference_maxpool(const SparseVectors* vectors, + const std::vector& doc_ids, + std::vector& terms, + std::vector& values, float& sum) { + const idx_t* indptr = vectors->indptr_data(); + const term_t* indices = vectors->indices_data(); + const float* vals = vectors->values_data_float(); + std::map m; + for (idx_t d : doc_ids) { + for (idx_t j = indptr[d]; j < indptr[d + 1]; ++j) { + auto& v = m[indices[j]]; + v = std::max(v, vals[j]); + } + } + sum = 0.0F; + terms.clear(); + values.clear(); + for (auto& [t, v] : m) { + terms.push_back(t); + values.push_back(v); + sum += v; + } +} + +// Compares GPU per-cluster max-pool against the CPU reference for a whole list +// of several clusters processed in one batched call. +TEST(GpuClusterAssignerTest, SummarizeListMaxpoolMatchesCpu) { + if (!GpuClusterAssigner::available()) { + GTEST_SKIP() << "No CUDA-capable GPU available"; + } + constexpr size_t kNumDocs = 8000; + constexpr size_t kDim = 3000; + constexpr size_t kNnz = 50; + SparseVectors vectors = make_random_corpus(kNumDocs, kDim, kNnz, 7); + + // Build a list of clusters: partition docs into varied-size groups. + std::vector flat_docs; + std::vector offsets = {0}; + const std::vector sizes = {5, 137, 1, 900, 42, 2000}; + idx_t d = 0; + for (size_t sz : sizes) { + for (size_t k = 0; k < sz && d < static_cast(kNumDocs); ++k) { + flat_docs.push_back(d++); + } + offsets.push_back(static_cast(flat_docs.size())); + } + const size_t n_clusters = offsets.size() - 1; + + std::vector gpu_out; + ASSERT_TRUE(GpuClusterAssigner::instance().summarize_list_maxpool( + &vectors, flat_docs.data(), offsets.data(), n_clusters, gpu_out)); + ASSERT_EQ(gpu_out.size(), n_clusters); + + for (size_t b = 0; b < n_clusters; ++b) { + std::vector cluster(flat_docs.begin() + offsets[b], + flat_docs.begin() + offsets[b + 1]); + std::vector exp_terms; + std::vector exp_vals; + float exp_sum = 0.0F; + cpu_reference_maxpool(&vectors, cluster, exp_terms, exp_vals, exp_sum); + + // GPU returns terms in touched-append order; sort by term to compare. + const auto& gc = gpu_out[b]; + std::vector order(gc.terms.size()); + std::iota(order.begin(), order.end(), 0); + std::ranges::sort(order, [&](size_t a, size_t c) { + return gc.terms[a] < gc.terms[c]; + }); + ASSERT_EQ(gc.terms.size(), exp_terms.size()) << "cluster " << b; + for (size_t i = 0; i < order.size(); ++i) { + EXPECT_EQ(gc.terms[order[i]], exp_terms[i]) + << "cluster " << b << " term " << i; + EXPECT_FLOAT_EQ(gc.values[order[i]], exp_vals[i]) + << "cluster " << b << " value " << i; + } + EXPECT_NEAR(gc.sum, exp_sum, exp_sum * 1e-5F + 1e-5F) << "cluster " << b; + } +} + +// A second list reusing the buffers must not leak accumulator state from the +// first — verifies the touched-slot reset path. +TEST(GpuClusterAssignerTest, SummarizeListMaxpoolReuseIsClean) { + if (!GpuClusterAssigner::available()) { + GTEST_SKIP() << "No CUDA-capable GPU available"; + } + SparseVectors vectors = make_random_corpus(5000, 2500, 40, 21); + auto& gpu = GpuClusterAssigner::instance(); + + // First list. + std::vector docs1; + for (idx_t d = 0; d < 2000; ++d) docs1.push_back(d); + std::vector off1 = {0, 800, 2000}; + std::vector out1; + gpu.summarize_list_maxpool(&vectors, docs1.data(), off1.data(), 2, out1); + + // Second, different list must match its own CPU reference. + std::vector docs2; + for (idx_t d = 2000; d < 5000; ++d) docs2.push_back(d); + std::vector off2 = {0, 1500, 3000}; + std::vector out2; + ASSERT_TRUE( + gpu.summarize_list_maxpool(&vectors, docs2.data(), off2.data(), 2, + out2)); + for (size_t b = 0; b < 2; ++b) { + std::vector cluster(docs2.begin() + off2[b], + docs2.begin() + off2[b + 1]); + std::vector et; + std::vector ev; + float es = 0.0F; + cpu_reference_maxpool(&vectors, cluster, et, ev, es); + EXPECT_EQ(out2[b].terms.size(), et.size()) << "cluster " << b; + EXPECT_NEAR(out2[b].sum, es, es * 1e-5F + 1e-5F) << "cluster " << b; + } +} + } // namespace } // namespace nsparse::detail From ca97235589e9eccbf7cab37dfe29b5955e34df0a Mon Sep 17 00:00:00 2001 From: Liyun Xiu Date: Thu, 9 Jul 2026 09:50:30 +0000 Subject: [PATCH 4/8] Rename to NSPARSE_ENABLE_GPU, flat-array CPU summarize, opt-in GPU summarize Interface cleanup (align with FAISS's GPU model): - Rename CMake option NSPARSE_ENABLE_CUDA -> NSPARSE_ENABLE_GPU and the compile define NSPARSE_WITH_CUDA -> NSPARSE_WITH_GPU. - Remove the NSPARSE_GPU_MIN_DOCS assignment threshold: offload is now unconditional when built with GPU support and a device is present. - Consolidate NSPARSE_BUILD_PROFILE / NSPARSE_SUMMARIZE_PROFILE into the single NSPARSE_GPU_PROFILE env var. - The build benchmark now selects CPU vs GPU at compile time (build with or without NSPARSE_ENABLE_GPU) rather than via a runtime gate. summarize() max-pool: default to a flat dim-sized accumulator on the CPU, replacing the per-cluster std::unordered_map. A per-cluster epoch marks live slots so the dim-sized buffers are reused across clusters with no reset; first-touch is detected via the epoch (not acc==0), robust to zero weights. ~2.8x less CPU work than the map; bit-identical output. GPU summarize becomes opt-in via NSPARSE_GPU_SUMMARIZE=1 (default: CPU flat-array). Measured base_full (8.84M docs, L4): - 32-core: CPU flat-array 63.0s vs GPU summarize 74.9s -> CPU wins - 8-core: CPU flat-array 113.3s vs GPU summarize 101.4s -> GPU wins So the default (CPU) is best on high-core hosts; enable the GPU offload on GPU-rich / low-core hosts where CPU summarize dominates wall-time. All 365 tests pass in both modes (default and NSPARSE_GPU_SUMMARIZE=1). Signed-off-by: Liyun Xiu --- CMakeLists.txt | 4 +- benchmarks/CMakeLists.txt | 2 +- benchmarks/index_build_benchmark.cpp | 61 ++++++--------------- cmake/cuda.cmake | 2 +- nsparse/CMakeLists.txt | 18 +++---- nsparse/cluster/inverted_list_clusters.cpp | 63 +++++++++++++++------- nsparse/cluster/kmeans_utils.cpp | 14 ++--- nsparse/gpu/gpu_cluster_assigner.cu | 29 ++++------ nsparse/gpu/gpu_cluster_assigner.h | 21 ++++---- nsparse/seismic_common.h | 6 +-- tests/CMakeLists.txt | 4 +- tests/gpu_cluster_assigner_test.cpp | 4 +- 12 files changed, 110 insertions(+), 118 deletions(-) diff --git a/CMakeLists.txt b/CMakeLists.txt index 90719bc..597f7a0 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -25,9 +25,9 @@ option(NSPARSE_ENABLE_PYTHON "Build Python bindings" OFF) # GPU acceleration for the index-building (build()) path via NVIDIA cuSPARSE. # The search() path always stays on CPU. Default OFF so CPU-only builds and CI # are unaffected when no CUDA toolkit is present. -option(NSPARSE_ENABLE_CUDA "Build GPU-accelerated index building via cuSPARSE" OFF) +option(NSPARSE_ENABLE_GPU "Build GPU-accelerated index building via cuSPARSE" OFF) -if(NSPARSE_ENABLE_CUDA) +if(NSPARSE_ENABLE_GPU) include(cmake/cuda.cmake) endif() diff --git a/benchmarks/CMakeLists.txt b/benchmarks/CMakeLists.txt index c059942..20630d5 100644 --- a/benchmarks/CMakeLists.txt +++ b/benchmarks/CMakeLists.txt @@ -44,7 +44,7 @@ endif() # Build-time benchmark: SeismicIndex build() on CPU vs GPU (cuSPARSE). The GPU # case is compiled in only when the nsparse target carries CUDA support (the -# NSPARSE_WITH_CUDA define propagates from the nsparse target). +# NSPARSE_WITH_GPU define propagates from the nsparse target). add_executable(nsparse_build_benchmark index_build_benchmark.cpp) target_include_directories(nsparse_build_benchmark PRIVATE ${PROJECT_SOURCE_DIR}) target_link_libraries(nsparse_build_benchmark PRIVATE diff --git a/benchmarks/index_build_benchmark.cpp b/benchmarks/index_build_benchmark.cpp index dbd51ce..d0aee4f 100644 --- a/benchmarks/index_build_benchmark.cpp +++ b/benchmarks/index_build_benchmark.cpp @@ -6,14 +6,14 @@ * this file be licensed under the Apache-2.0 license or a * compatible open source license. * - * Benchmarks SeismicIndex build() time on CPU vs GPU (cuSPARSE). + * Benchmarks SeismicIndex build() time. * - * The GPU path accelerates the k-means document->centroid assignment step of - * build() via cuSPARSE SpMM; the rest of build() (inverted-list construction, - * summarization) stays on CPU. GPU vs CPU is selected at run time through the - * NSPARSE_GPU_MIN_DOCS gate so both paths run from the same binary: - * - CPU: set the gate above the largest inverted list (no list offloads). - * - GPU: set the gate to 1 (every eligible list offloads). + * The GPU path accelerates the k-means document->centroid assignment and the + * summarize() max-pool of build() via cuSPARSE and custom kernels; the rest of + * build() (inverted-list construction, pruning) stays on CPU. Whether build() + * uses the GPU is decided at compile time (NSPARSE_ENABLE_GPU) plus device + * availability — there is no runtime gate. To compare CPU vs GPU, build this + * binary once without NSPARSE_ENABLE_GPU (CPU) and once with it (GPU). * * Data is the Big-ANN sparse-vectors (SPLADE MS MARCO) CSR file, e.g. * data/base_small.csr from cpp-sparse-ann's dataset.py. @@ -38,7 +38,7 @@ #include "nsparse/seismic_index.h" #include "nsparse/types.h" -#ifdef NSPARSE_WITH_CUDA +#ifdef NSPARSE_WITH_GPU #include "nsparse/gpu/gpu_cluster_assigner.h" #endif @@ -115,20 +115,11 @@ const CSRMatrix& shared_data() { constexpr nsparse::SeismicClusterParameters kParams = { .lambda = 6000, .beta = 400, .alpha = 0.4F}; -// Force the GPU-offload gate one way or the other for this process. Setting the -// minimum-docs threshold to 1 offloads every eligible list; setting it beyond -// the largest inverted list keeps the whole build on the CPU. -void set_gpu_offload(bool enabled) { - if (enabled) { - ::setenv("NSPARSE_GPU_MIN_DOCS", "1", /*overwrite=*/1); - } else { - // Larger than lambda (max pruned posting length), so no list offloads. - ::setenv("NSPARSE_GPU_MIN_DOCS", "1000000000", /*overwrite=*/1); - } -} - // Builds a fresh SeismicIndex from the shared corpus and times only build(). -void run_build(benchmark::State& state, bool use_gpu) { +// Whether build() runs on CPU or offloads to the GPU is decided at compile time +// (NSPARSE_ENABLE_GPU) plus device availability — there is no runtime gate. To +// compare CPU vs GPU, build this binary once without and once with GPU support. +void run_build(benchmark::State& state) { const CSRMatrix& data = shared_data(); for (auto _ : state) { state.PauseTiming(); @@ -136,7 +127,6 @@ void run_build(benchmark::State& state, bool use_gpu) { static_cast(data.ncol), kParams); index->add(static_cast(data.nrow), data.indptr.data(), data.indices.data(), data.data.data()); - set_gpu_offload(use_gpu); state.ResumeTiming(); index->build(); @@ -161,33 +151,16 @@ int env_int(const char* name, int fallback) { } // namespace // --------------------------------------------------------------------------- -// CPU build (assignment on CPU) +// SeismicIndex build(). Runs on GPU when the library was built with +// NSPARSE_ENABLE_GPU and a device is available; otherwise on CPU. // --------------------------------------------------------------------------- -static void BM_Seismic_Build_CPU(benchmark::State& state) { - run_build(state, /*use_gpu=*/false); +static void BM_Seismic_Build(benchmark::State& state) { + run_build(state); } -BENCHMARK(BM_Seismic_Build_CPU) +BENCHMARK(BM_Seismic_Build) ->Unit(benchmark::kMillisecond) ->UseRealTime() ->Iterations(env_int("NSPARSE_BENCH_ITERS", 3)) ->Repetitions(env_int("NSPARSE_BENCH_REPS", 3)); -// --------------------------------------------------------------------------- -// GPU build (assignment via cuSPARSE) - only when compiled with CUDA -// --------------------------------------------------------------------------- -#ifdef NSPARSE_WITH_CUDA -static void BM_Seismic_Build_GPU(benchmark::State& state) { - if (!nsparse::detail::GpuClusterAssigner::available()) { - state.SkipWithError("No CUDA-capable GPU available"); - return; - } - run_build(state, /*use_gpu=*/true); -} -BENCHMARK(BM_Seismic_Build_GPU) - ->Unit(benchmark::kMillisecond) - ->UseRealTime() - ->Iterations(env_int("NSPARSE_BENCH_ITERS", 3)) - ->Repetitions(env_int("NSPARSE_BENCH_REPS", 3)); -#endif - BENCHMARK_MAIN(); diff --git a/cmake/cuda.cmake b/cmake/cuda.cmake index d71b058..c97e7e0 100644 --- a/cmake/cuda.cmake +++ b/cmake/cuda.cmake @@ -7,7 +7,7 @@ # # CUDA / cuSPARSE toolchain setup for GPU-accelerated index building. # -# Enable with -DNSPARSE_ENABLE_CUDA=ON. When the CUDA toolkit is installed in a +# Enable with -DNSPARSE_ENABLE_GPU=ON. When the CUDA toolkit is installed in a # non-standard location (for example, shipped inside a pip `nvidia-cu13` wheel), # point CMake at it with: # -DNSPARSE_CUDA_TOOLKIT_ROOT=/path/to/nvidia/cu13 diff --git a/nsparse/CMakeLists.txt b/nsparse/CMakeLists.txt index 0e29371..afe2102 100644 --- a/nsparse/CMakeLists.txt +++ b/nsparse/CMakeLists.txt @@ -140,19 +140,19 @@ endif() # GPU acceleration for the build() path via cuSPARSE (search() stays on CPU). # Compiled into every nsparse variant only when configured with -# -DNSPARSE_ENABLE_CUDA=ON. Guarded by the NSPARSE_WITH_CUDA compile define so +# -DNSPARSE_ENABLE_GPU=ON. Guarded by the NSPARSE_WITH_GPU compile define so # the CPU sources pick up the GPU assignment hook. -if(NSPARSE_ENABLE_CUDA) - set(NSPARSE_CUDA_SRC gpu/gpu_cluster_assigner.cu) - set(NSPARSE_CUDA_TARGETS nsparse) +if(NSPARSE_ENABLE_GPU) + set(NSPARSE_GPU_SRC gpu/gpu_cluster_assigner.cu) + set(NSPARSE_GPU_TARGETS nsparse) if(NSPARSE_IS_X86) - list(APPEND NSPARSE_CUDA_TARGETS nsparse_avx2 nsparse_avx512) + list(APPEND NSPARSE_GPU_TARGETS nsparse_avx2 nsparse_avx512) elseif(NSPARSE_IS_ARM AND NOT APPLE) - list(APPEND NSPARSE_CUDA_TARGETS nsparse_sve) + list(APPEND NSPARSE_GPU_TARGETS nsparse_sve) endif() - foreach(_tgt IN LISTS NSPARSE_CUDA_TARGETS) - target_sources(${_tgt} PRIVATE ${NSPARSE_CUDA_SRC}) - target_compile_definitions(${_tgt} PUBLIC NSPARSE_WITH_CUDA) + foreach(_tgt IN LISTS NSPARSE_GPU_TARGETS) + target_sources(${_tgt} PRIVATE ${NSPARSE_GPU_SRC}) + target_compile_definitions(${_tgt} PUBLIC NSPARSE_WITH_GPU) target_link_libraries(${_tgt} PRIVATE CUDA::cusparse CUDA::cudart) set_target_properties(${_tgt} PROPERTIES CUDA_SEPARABLE_COMPILATION ON) endforeach() diff --git a/nsparse/cluster/inverted_list_clusters.cpp b/nsparse/cluster/inverted_list_clusters.cpp index 2a7f023..50b7e6d 100644 --- a/nsparse/cluster/inverted_list_clusters.cpp +++ b/nsparse/cluster/inverted_list_clusters.cpp @@ -22,21 +22,21 @@ #include "nsparse/sparse_vectors.h" #include "nsparse/types.h" -#ifdef NSPARSE_WITH_CUDA +#ifdef NSPARSE_WITH_GPU #include "nsparse/gpu/gpu_cluster_assigner.h" #endif namespace nsparse { namespace { -// Opt-in split timing for summarize (NSPARSE_SUMMARIZE_PROFILE=1): how much of +// Opt-in split timing for summarize (NSPARSE_GPU_PROFILE=1): how much of // summarize is the memory-bound max-pool vs the per-cluster sort/truncate. struct SummarizeProfile { std::atomic maxpool_ns{0}; std::atomic sort_ns{0}; bool enabled = false; SummarizeProfile() { - const char* v = std::getenv("NSPARSE_SUMMARIZE_PROFILE"); + const char* v = std::getenv("NSPARSE_GPU_PROFILE"); enabled = (v != nullptr && v[0] == '1'); } ~SummarizeProfile() { @@ -79,13 +79,26 @@ SparseVectors summarize_(const SparseVectors* vectors, const bool prof = g_summ_profile.enabled; int64_t maxpool_acc = 0; int64_t sort_acc = 0; - // The per-term max-pool (the memory-bound majority of summarize) can be - // offloaded to the GPU for float (U32) weights, batched as ONE launch for - // the whole list's clusters (per-cluster launches are dominated by GPU - // round-trip overhead). The tie-order-sensitive sort/truncate below stays - // on the CPU so output ordering is preserved. + + // Per-term max-pool over a flat, dim-sized accumulator reused across all + // clusters in this list (replaces a per-cluster std::unordered_map, which is + // slow for the many small clusters). A monotonically increasing per-cluster + // epoch marks which accumulator slots belong to the current cluster, so no + // per-cluster reset of the dim-sized buffers is needed. First touch of a + // term is detected via the epoch (not acc[t]==0), which is robust to + // legitimately-zero stored weights. + const size_t dim = vectors->get_dimension(); + std::vector acc(dim, T(0)); + std::vector epoch(dim, 0); + std::vector touched; + uint32_t cur_epoch = 0; + + // The per-term max-pool can optionally be offloaded to the GPU (batched as + // one kernel launch per list) instead of the CPU flat-array path above. + // Opt-in via NSPARSE_GPU_SUMMARIZE=1 (see should_offload_summarize_to_gpu); + // a net win on GPU-rich / low-core hosts. Everything downstream is identical. +#ifdef NSPARSE_WITH_GPU bool gpu_ok = false; -#ifdef NSPARSE_WITH_CUDA std::vector gpu_clusters; if constexpr (std::is_same_v) { if (detail::should_offload_summarize_to_gpu()) { @@ -102,11 +115,11 @@ SparseVectors summarize_(const SparseVectors* vectors, for (size_t i = 0; i < offsets.size() - 1; ++i) { const int64_t ts0 = prof ? summ_now_ns() : 0; size_t n_docs = offsets[i + 1] - offsets[i]; - std::vector> summary_vec; float sum = 0.0F; + std::vector> summary_vec; bool used_gpu = false; -#ifdef NSPARSE_WITH_CUDA +#ifdef NSPARSE_WITH_GPU if constexpr (std::is_same_v) { if (gpu_ok) { const auto& gc = gpu_clusters[i]; @@ -120,22 +133,36 @@ SparseVectors summarize_(const SparseVectors* vectors, } #endif if (!used_gpu) { - std::unordered_map summary_map; + ++cur_epoch; + touched.clear(); auto doc_ids = std::span( group_of_doc_ids.data() + offsets[i], n_docs); for (const auto& doc_id : doc_ids) { int start = indptr_data[doc_id]; int end = indptr_data[doc_id + 1]; for (size_t j = start; j < end; ++j) { - const auto old = summary_map[indices_data[j]]; - auto& value = summary_map[indices_data[j]]; + const term_t term = indices_data[j]; // j is element index, need byte offset for T access - value = std::max(value, *reinterpret_cast( - values_data + j * sizeof(T))); - sum += value - old; + const T v = *reinterpret_cast( + values_data + j * sizeof(T)); + if (epoch[term] != cur_epoch) { + // First occurrence in this cluster. Matches the old map + // semantics: value = max(0, v). + epoch[term] = cur_epoch; + const T value = std::max(T(0), v); + acc[term] = value; + touched.push_back(term); + sum += value; + } else if (v > acc[term]) { + sum += v - acc[term]; + acc[term] = v; + } } } - summary_vec.assign(summary_map.begin(), summary_map.end()); + summary_vec.reserve(touched.size()); + for (const term_t term : touched) { + summary_vec.emplace_back(term, acc[term]); + } } const int64_t ts1 = prof ? summ_now_ns() : 0; diff --git a/nsparse/cluster/kmeans_utils.cpp b/nsparse/cluster/kmeans_utils.cpp index 8af08a7..89e8626 100644 --- a/nsparse/cluster/kmeans_utils.cpp +++ b/nsparse/cluster/kmeans_utils.cpp @@ -18,7 +18,7 @@ #include "nsparse/sparse_vectors.h" #include "nsparse/types.h" #include "nsparse/utils/distance_simd.h" -#ifdef NSPARSE_WITH_CUDA +#ifdef NSPARSE_WITH_GPU #include "nsparse/gpu/gpu_cluster_assigner.h" #endif #if defined(__AVX512F__) @@ -173,12 +173,12 @@ void map_docs_to_clusters(const SparseVectors* vectors, if (n_clusters == 0 || n_docs == 0) { return; } -#if defined(NSPARSE_WITH_CUDA) && !defined(__AVX512F__) - // GPU-accelerated assignment via cuSPARSE for float (U32) weights on large - // problems. Produces assignments identical to the scalar CPU reference - // below (skips centroids, ties break to the lowest cluster index); on any - // failure we fall through to the CPU path. Disabled for AVX-512 builds - // whose assignment semantics differ. +#if defined(NSPARSE_WITH_GPU) && !defined(__AVX512F__) + // GPU-accelerated assignment via cuSPARSE for float (U32) weights. Produces + // assignments identical to the scalar CPU reference below (skips centroids, + // ties break to the lowest cluster index); on any failure we fall through + // to the CPU path. Disabled for AVX-512 builds whose assignment semantics + // differ. if (vectors->get_element_size() == U32 && should_offload_assignment_to_gpu(n_docs, n_clusters)) { try { diff --git a/nsparse/gpu/gpu_cluster_assigner.cu b/nsparse/gpu/gpu_cluster_assigner.cu index ea461d1..469581c 100644 --- a/nsparse/gpu/gpu_cluster_assigner.cu +++ b/nsparse/gpu/gpu_cluster_assigner.cu @@ -728,31 +728,22 @@ bool GpuClusterAssigner::summarize_list_maxpool( } bool should_offload_assignment_to_gpu(size_t n_docs, size_t n_clusters) { - if (!GpuClusterAssigner::available()) { - return false; - } - // With on-device gather/scatter and per-thread streams the offload is - // efficient, but very small lists still favor the CPU. Overridable for - // tuning/benchmarking. - size_t min_docs = 1024; - if (const char* env = std::getenv("NSPARSE_GPU_MIN_DOCS")) { - min_docs = static_cast(std::strtoull(env, nullptr, 10)); - } - return n_docs >= min_docs && n_clusters >= 2; + // Offload whenever GPU support is compiled in and a device is available. + // k-means assignment needs at least two clusters to be meaningful. + return GpuClusterAssigner::available() && n_docs > 0 && n_clusters >= 2; } bool should_offload_summarize_to_gpu() { if (!GpuClusterAssigner::available()) { return false; } - // On by default when a GPU is present. The max-pool is batched as one - // kernel launch per inverted list (all clusters at once), so it avoids the - // per-cluster sync storm that made a naive offload far slower. Opt out with - // NSPARSE_GPU_SUMMARIZE=0. - if (const char* env = std::getenv("NSPARSE_GPU_SUMMARIZE")) { - return env[0] != '0'; - } - return true; + // Opt-in. The CPU flat-array max-pool is faster on high-core hosts, so the + // default keeps summarize on the CPU (bit-identical, no GPU dependency). + // On GPU-rich / low-core hosts (e.g. an 8-vCPU single-GPU instance), the + // CPU summarize dominates wall-time and offloading its max-pool to the + // otherwise-idle GPU is a net win — enable it with NSPARSE_GPU_SUMMARIZE=1. + const char* env = std::getenv("NSPARSE_GPU_SUMMARIZE"); + return env != nullptr && env[0] == '1'; } } // namespace nsparse::detail diff --git a/nsparse/gpu/gpu_cluster_assigner.h b/nsparse/gpu/gpu_cluster_assigner.h index b0631e4..b17f20b 100644 --- a/nsparse/gpu/gpu_cluster_assigner.h +++ b/nsparse/gpu/gpu_cluster_assigner.h @@ -35,8 +35,8 @@ namespace nsparse::detail { * The search() path is unaffected and always runs on the CPU. * * The implementation lives in gpu_cluster_assigner.cu and is only compiled when - * the project is configured with -DNSPARSE_ENABLE_CUDA=ON (which defines - * NSPARSE_WITH_CUDA). The declarations below are always visible so that callers + * the project is configured with -DNSPARSE_ENABLE_GPU=ON (which defines + * NSPARSE_WITH_GPU). The declarations below are always visible so that callers * can guard usage with a single runtime available() check. */ class GpuClusterAssigner { @@ -119,21 +119,22 @@ class GpuClusterAssigner { }; /** - * @brief Heuristic gate for the automatic GPU path in map_docs_to_clusters(). + * @brief Gate for the GPU assignment path in map_docs_to_clusters(). * - * The per-inverted-list problems in a Seismic build are individually small and - * the CPU already parallelizes across lists, so offloading tiny problems to the - * GPU loses to transfer/launch overhead. This returns true only once the - * problem is large enough to benefit. The threshold on the number of documents - * can be overridden with the NSPARSE_GPU_MIN_DOCS environment variable. + * Returns true when the library was built with GPU support and a usable device + * is present (and the assignment has at least two clusters). Offload is + * unconditional beyond that — there is no runtime size threshold. */ bool should_offload_assignment_to_gpu(size_t n_docs, size_t n_clusters); /** * @brief Gate for the GPU summarize() max-pool offload. * - * Returns true when a usable GPU is present and the offload is enabled (on by - * default when built with CUDA; can be disabled with NSPARSE_GPU_SUMMARIZE=0). + * Opt-in: returns true only when the library was built with GPU support, a + * usable device is present, and NSPARSE_GPU_SUMMARIZE=1 is set. The default + * (unset) keeps summarize on the CPU flat-array path, which is faster on + * high-core hosts and bit-identical. Enabling the GPU offload is a net win on + * GPU-rich / low-core hosts where CPU summarize dominates build wall-time. */ bool should_offload_summarize_to_gpu(); diff --git a/nsparse/seismic_common.h b/nsparse/seismic_common.h index cf1948c..de3d9e6 100644 --- a/nsparse/seismic_common.h +++ b/nsparse/seismic_common.h @@ -146,12 +146,12 @@ inline std::vector build_inverted_lists_clusters( std::vector clustered_inverted_lists( inverted_lists_size); - // Optional coarse phase timing, enabled by NSPARSE_BUILD_PROFILE=1, to see + // Optional coarse phase timing, enabled by NSPARSE_GPU_PROFILE=1, to see // how build CPU time splits across prune / train (k-means assignment) / // summarize. Accumulates thread-seconds (sum across threads). const bool build_prof = - std::getenv("NSPARSE_BUILD_PROFILE") != nullptr && - std::getenv("NSPARSE_BUILD_PROFILE")[0] == '1'; + std::getenv("NSPARSE_GPU_PROFILE") != nullptr && + std::getenv("NSPARSE_GPU_PROFILE")[0] == '1'; double t_prune = 0.0; double t_train = 0.0; double t_summ = 0.0; diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index e8fa816..69eb4e0 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -31,8 +31,8 @@ set(NSPARSE_TEST_SRC vector_process_test.cpp ) -# GPU parity test is only meaningful when CUDA support is compiled in. -if(NSPARSE_ENABLE_CUDA) +# GPU parity test is only meaningful when GPU support is compiled in. +if(NSPARSE_ENABLE_GPU) list(APPEND NSPARSE_TEST_SRC gpu_cluster_assigner_test.cpp) endif() diff --git a/tests/gpu_cluster_assigner_test.cpp b/tests/gpu_cluster_assigner_test.cpp index 704ee15..8b67227 100644 --- a/tests/gpu_cluster_assigner_test.cpp +++ b/tests/gpu_cluster_assigner_test.cpp @@ -134,8 +134,8 @@ TEST(GpuClusterAssignerTest, AutoPathViaMapDocsMatchesReference) { if (!GpuClusterAssigner::available()) { GTEST_SKIP() << "No CUDA-capable GPU available"; } - // Force the auto-offload gate to fire for this size. - ::setenv("NSPARSE_GPU_MIN_DOCS", "1", /*overwrite=*/1); + // With GPU support compiled in and a device present, map_docs_to_clusters + // routes through the GPU unconditionally (no size gate). constexpr size_t kNumDocs = 5000; constexpr size_t kDim = 1500; constexpr size_t kNnz = 30; From 6ab574a07a60d3ee95a65893644dd3e4bd9039e1 Mon Sep 17 00:00:00 2001 From: Liyun Xiu Date: Thu, 9 Jul 2026 10:09:20 +0000 Subject: [PATCH 5/8] Move CUDA build logic into nsparse/gpu/CMakeLists.txt Co-locate the GPU build config with the GPU sources: replace the top-level cmake/cuda.cmake with nsparse/gpu/CMakeLists.txt, which does the CUDA toolchain setup (enable_language, find CUDAToolkit, architectures) and defines a nsparse_attach_gpu() helper. nsparse/CMakeLists.txt include()s it (when NSPARSE_ENABLE_GPU=ON) after the library targets exist and calls the helper to compile gpu_cluster_assigner.cu into each variant. Also set CUDA_RESOLVE_DEVICE_SYMBOLS ON so device symbols are resolved into the nsparse library itself; pure-C++ consumers (tests, benchmarks) link it without a separate CUDA device-link step. (Previously enable_language ran at top-level scope, which handled this implicitly.) Verified on an L4: GPU build + all 365 tests pass; CPU-only build unaffected. Signed-off-by: Liyun Xiu --- CMakeLists.txt | 7 ++-- cmake/cuda.cmake | 48 -------------------------- nsparse/CMakeLists.txt | 13 ++++--- nsparse/gpu/CMakeLists.txt | 71 ++++++++++++++++++++++++++++++++++++++ 4 files changed, 79 insertions(+), 60 deletions(-) delete mode 100644 cmake/cuda.cmake create mode 100644 nsparse/gpu/CMakeLists.txt diff --git a/CMakeLists.txt b/CMakeLists.txt index 597f7a0..955eba3 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -24,13 +24,10 @@ option(NSPARSE_ENABLE_PYTHON "Build Python bindings" OFF) # GPU acceleration for the index-building (build()) path via NVIDIA cuSPARSE. # The search() path always stays on CPU. Default OFF so CPU-only builds and CI -# are unaffected when no CUDA toolkit is present. +# are unaffected when no CUDA toolkit is present. Toolchain setup and per-target +# wiring live in nsparse/gpu/CMakeLists.txt (included from nsparse/CMakeLists.txt). option(NSPARSE_ENABLE_GPU "Build GPU-accelerated index building via cuSPARSE" OFF) -if(NSPARSE_ENABLE_GPU) - include(cmake/cuda.cmake) -endif() - add_subdirectory(nsparse) # Python bindings diff --git a/cmake/cuda.cmake b/cmake/cuda.cmake deleted file mode 100644 index c97e7e0..0000000 --- a/cmake/cuda.cmake +++ /dev/null @@ -1,48 +0,0 @@ -# Copyright OpenSearch Contributors -# SPDX-License-Identifier: Apache-2.0 -# -# The OpenSearch Contributors require contributions made to -# this file be licensed under the Apache-2.0 license or a -# compatible open source license. -# -# CUDA / cuSPARSE toolchain setup for GPU-accelerated index building. -# -# Enable with -DNSPARSE_ENABLE_GPU=ON. When the CUDA toolkit is installed in a -# non-standard location (for example, shipped inside a pip `nvidia-cu13` wheel), -# point CMake at it with: -# -DNSPARSE_CUDA_TOOLKIT_ROOT=/path/to/nvidia/cu13 -# and CMake will pick up nvcc, the headers and the runtime/cusparse libraries. - -# Allow the caller to hint the toolkit location (nvcc lives in /bin). -if(NSPARSE_CUDA_TOOLKIT_ROOT) - set(CUDAToolkit_ROOT "${NSPARSE_CUDA_TOOLKIT_ROOT}" CACHE PATH "" FORCE) - if(NOT CMAKE_CUDA_COMPILER) - set(CMAKE_CUDA_COMPILER "${NSPARSE_CUDA_TOOLKIT_ROOT}/bin/nvcc" CACHE FILEPATH "" FORCE) - endif() - # Wheel-packaged toolkits keep shared libs under lib/ rather than lib64/. - list(APPEND CMAKE_LIBRARY_PATH - "${NSPARSE_CUDA_TOOLKIT_ROOT}/lib" - "${NSPARSE_CUDA_TOOLKIT_ROOT}/lib64") -endif() - -# Turn on the CUDA language now that the compiler is (optionally) hinted. -enable_language(CUDA) - -# CUDAToolkit provides the CUDA::cusparse and CUDA::cudart imported targets. -find_package(CUDAToolkit REQUIRED) - -# Default to the common data-center / Ada architectures. L4 (this project's -# reference GPU) is compute capability 8.9. CMake picks a conservative default -# during language enablement, so override it unless the caller pinned one -# explicitly via -DNSPARSE_CUDA_ARCHITECTURES. -if(NSPARSE_CUDA_ARCHITECTURES) - set(CMAKE_CUDA_ARCHITECTURES "${NSPARSE_CUDA_ARCHITECTURES}" CACHE STRING "" FORCE) -else() - set(CMAKE_CUDA_ARCHITECTURES 80 89 CACHE STRING "CUDA architectures to build for" FORCE) -endif() - -set(CMAKE_CUDA_STANDARD 17) -set(CMAKE_CUDA_STANDARD_REQUIRED ON) - -message(STATUS "nsparse: CUDA enabled (nvcc=${CMAKE_CUDA_COMPILER}, " - "archs=${CMAKE_CUDA_ARCHITECTURES})") diff --git a/nsparse/CMakeLists.txt b/nsparse/CMakeLists.txt index afe2102..70c2613 100644 --- a/nsparse/CMakeLists.txt +++ b/nsparse/CMakeLists.txt @@ -140,10 +140,12 @@ endif() # GPU acceleration for the build() path via cuSPARSE (search() stays on CPU). # Compiled into every nsparse variant only when configured with -# -DNSPARSE_ENABLE_GPU=ON. Guarded by the NSPARSE_WITH_GPU compile define so -# the CPU sources pick up the GPU assignment hook. +# -DNSPARSE_ENABLE_GPU=ON. gpu/CMakeLists.txt does the CUDA toolchain setup and +# defines nsparse_attach_gpu(); it is include()d (not add_subdirectory'd) here, +# after the library targets exist, so the function is visible in this scope and +# can wire each target. if(NSPARSE_ENABLE_GPU) - set(NSPARSE_GPU_SRC gpu/gpu_cluster_assigner.cu) + include(gpu/CMakeLists.txt) set(NSPARSE_GPU_TARGETS nsparse) if(NSPARSE_IS_X86) list(APPEND NSPARSE_GPU_TARGETS nsparse_avx2 nsparse_avx512) @@ -151,10 +153,7 @@ if(NSPARSE_ENABLE_GPU) list(APPEND NSPARSE_GPU_TARGETS nsparse_sve) endif() foreach(_tgt IN LISTS NSPARSE_GPU_TARGETS) - target_sources(${_tgt} PRIVATE ${NSPARSE_GPU_SRC}) - target_compile_definitions(${_tgt} PUBLIC NSPARSE_WITH_GPU) - target_link_libraries(${_tgt} PRIVATE CUDA::cusparse CUDA::cudart) - set_target_properties(${_tgt} PROPERTIES CUDA_SEPARABLE_COMPILATION ON) + nsparse_attach_gpu(${_tgt}) endforeach() endif() diff --git a/nsparse/gpu/CMakeLists.txt b/nsparse/gpu/CMakeLists.txt new file mode 100644 index 0000000..c57e826 --- /dev/null +++ b/nsparse/gpu/CMakeLists.txt @@ -0,0 +1,71 @@ +# Copyright OpenSearch Contributors +# SPDX-License-Identifier: Apache-2.0 +# +# The OpenSearch Contributors require contributions made to +# this file be licensed under the Apache-2.0 license or a +# compatible open source license. +# +# CUDA / cuSPARSE toolchain setup and target wiring for GPU-accelerated index +# building (the build() path; search() always stays on CPU). +# +# Included from nsparse/CMakeLists.txt (before the nsparse targets are defined) +# only when -DNSPARSE_ENABLE_GPU=ON. It enables the CUDA language, finds the +# toolkit, and defines nsparse_attach_gpu() which the parent uses to +# compile gpu_cluster_assigner.cu into each nsparse variant and link cuSPARSE. +# +# When the CUDA toolkit lives in a non-standard location (for example, shipped +# inside a pip `nvidia-cu13` wheel), point CMake at it with: +# -DNSPARSE_CUDA_TOOLKIT_ROOT=/path/to/nvidia/cu13 + +# Allow the caller to hint the toolkit location (nvcc lives in /bin). +if(NSPARSE_CUDA_TOOLKIT_ROOT) + set(CUDAToolkit_ROOT "${NSPARSE_CUDA_TOOLKIT_ROOT}" CACHE PATH "" FORCE) + if(NOT CMAKE_CUDA_COMPILER) + set(CMAKE_CUDA_COMPILER "${NSPARSE_CUDA_TOOLKIT_ROOT}/bin/nvcc" CACHE FILEPATH "" FORCE) + endif() + # Wheel-packaged toolkits keep shared libs under lib/ rather than lib64/. + list(APPEND CMAKE_LIBRARY_PATH + "${NSPARSE_CUDA_TOOLKIT_ROOT}/lib" + "${NSPARSE_CUDA_TOOLKIT_ROOT}/lib64") +endif() + +# Turn on the CUDA language now that the compiler is (optionally) hinted. +enable_language(CUDA) + +# CUDAToolkit provides the CUDA::cusparse and CUDA::cudart imported targets. +find_package(CUDAToolkit REQUIRED) + +# Default to the common data-center / Ada architectures. L4 (this project's +# reference GPU) is compute capability 8.9. Override with +# -DNSPARSE_CUDA_ARCHITECTURES=... if targeting different GPUs. +if(NSPARSE_CUDA_ARCHITECTURES) + set(CMAKE_CUDA_ARCHITECTURES "${NSPARSE_CUDA_ARCHITECTURES}" CACHE STRING "" FORCE) +else() + set(CMAKE_CUDA_ARCHITECTURES 80 89 CACHE STRING "CUDA architectures to build for" FORCE) +endif() + +set(CMAKE_CUDA_STANDARD 17) +set(CMAKE_CUDA_STANDARD_REQUIRED ON) + +# Absolute path to the CUDA translation unit, resolved once here so the helper +# can be called from the parent directory scope. This file is include()d from +# nsparse/CMakeLists.txt, so CMAKE_CURRENT_LIST_DIR is this gpu/ directory. +set(NSPARSE_GPU_CU "${CMAKE_CURRENT_LIST_DIR}/gpu_cluster_assigner.cu" + CACHE INTERNAL "nsparse GPU source") + +message(STATUS "nsparse: GPU enabled (nvcc=${CMAKE_CUDA_COMPILER}, " + "archs=${CMAKE_CUDA_ARCHITECTURES})") + +# Compile the GPU source into and link cuSPARSE/cudart. Guarded by the +# NSPARSE_WITH_GPU compile define so the CPU sources pick up the GPU hooks. +function(nsparse_attach_gpu target) + target_sources(${target} PRIVATE "${NSPARSE_GPU_CU}") + target_compile_definitions(${target} PUBLIC NSPARSE_WITH_GPU) + target_link_libraries(${target} PRIVATE CUDA::cusparse CUDA::cudart) + # Resolve device symbols into the library itself so pure-C++ consumers + # (tests, benchmarks) can link it without a separate CUDA device-link step. + set_target_properties(${target} PROPERTIES + CUDA_SEPARABLE_COMPILATION ON + CUDA_RESOLVE_DEVICE_SYMBOLS ON + CUDA_ARCHITECTURES "${CMAKE_CUDA_ARCHITECTURES}") +endfunction() From cb03500823928b2414d299452bf05dca5f0ddaa2 Mon Sep 17 00:00:00 2001 From: Liyun Xiu Date: Fri, 10 Jul 2026 03:51:07 +0000 Subject: [PATCH 6/8] Refactor GPU module: remove profiling, split out GpuSummarizer Address review feedback: - Remove NSPARSE_GPU_PROFILE and all phase-profiling code (Profile struct, SummarizeProfile, build_prof timers). Profiling will be reintroduced later in a better form. - Extract the GPU summarize max-pool into its own GpuSummarizer utility class (gpu_summarizer.{h,cu}); GpuClusterAssigner now only does assignment. Shared CUDA helpers and the single resident corpus (GpuCorpus) move to gpu_common.cuh so both helpers reuse one corpus upload (avoids duplicating the ~corpus-sized allocation). NSPARSE_GPU_SUMMARIZE=1 still selects GPU vs CPU summarize; default stays CPU flat-array. - Tighten comments for concision. - Document why GPU assignment is excluded from AVX-512 builds: the AVX-512 CPU path does not skip centroids while the scalar path (which the GPU path mirrors) does, so mixing them in one build would be inconsistent. All 365 tests pass in both summarize modes; CPU-only build unaffected. Signed-off-by: Liyun Xiu --- nsparse/cluster/inverted_list_clusters.cpp | 78 +--- nsparse/cluster/kmeans_utils.cpp | 10 +- nsparse/gpu/CMakeLists.txt | 14 +- nsparse/gpu/gpu_cluster_assigner.cu | 518 ++------------------- nsparse/gpu/gpu_cluster_assigner.h | 111 +---- nsparse/gpu/gpu_common.cuh | 149 ++++++ nsparse/gpu/gpu_summarizer.cu | 283 +++++++++++ nsparse/gpu/gpu_summarizer.h | 79 ++++ nsparse/seismic_common.h | 52 +-- tests/gpu_cluster_assigner_test.cpp | 19 +- 10 files changed, 602 insertions(+), 711 deletions(-) create mode 100644 nsparse/gpu/gpu_common.cuh create mode 100644 nsparse/gpu/gpu_summarizer.cu create mode 100644 nsparse/gpu/gpu_summarizer.h diff --git a/nsparse/cluster/inverted_list_clusters.cpp b/nsparse/cluster/inverted_list_clusters.cpp index 50b7e6d..e8d7214 100644 --- a/nsparse/cluster/inverted_list_clusters.cpp +++ b/nsparse/cluster/inverted_list_clusters.cpp @@ -10,11 +10,7 @@ #include "nsparse/cluster/inverted_list_clusters.h" #include -#include -#include #include -#include -#include #include #include #include @@ -23,37 +19,12 @@ #include "nsparse/sparse_vectors.h" #include "nsparse/types.h" #ifdef NSPARSE_WITH_GPU -#include "nsparse/gpu/gpu_cluster_assigner.h" +#include "nsparse/gpu/gpu_summarizer.h" #endif namespace nsparse { namespace { -// Opt-in split timing for summarize (NSPARSE_GPU_PROFILE=1): how much of -// summarize is the memory-bound max-pool vs the per-cluster sort/truncate. -struct SummarizeProfile { - std::atomic maxpool_ns{0}; - std::atomic sort_ns{0}; - bool enabled = false; - SummarizeProfile() { - const char* v = std::getenv("NSPARSE_GPU_PROFILE"); - enabled = (v != nullptr && v[0] == '1'); - } - ~SummarizeProfile() { - if (!enabled) return; - std::fprintf(stderr, - "[nsparse summarize] maxpool=%.1fs sort_truncate=%.1fs\n", - maxpool_ns.load() / 1e9, sort_ns.load() / 1e9); - } -}; -SummarizeProfile g_summ_profile; - -inline int64_t summ_now_ns() { - return std::chrono::duration_cast( - std::chrono::steady_clock::now().time_since_epoch()) - .count(); -} - /** * @brief Generage summary sparse vector for posting lists * @@ -76,44 +47,33 @@ SparseVectors summarize_(const SparseVectors* vectors, const auto& indptr_data = vectors->indptr_data(); const auto& indices_data = vectors->indices_data(); const auto& values_data = vectors->values_data(); - const bool prof = g_summ_profile.enabled; - int64_t maxpool_acc = 0; - int64_t sort_acc = 0; // Per-term max-pool over a flat, dim-sized accumulator reused across all - // clusters in this list (replaces a per-cluster std::unordered_map, which is - // slow for the many small clusters). A monotonically increasing per-cluster - // epoch marks which accumulator slots belong to the current cluster, so no - // per-cluster reset of the dim-sized buffers is needed. First touch of a - // term is detected via the epoch (not acc[t]==0), which is robust to - // legitimately-zero stored weights. + // clusters (replaces a per-cluster std::unordered_map, ~2.8x cheaper for + // the many small clusters). A per-cluster epoch marks which slots are live, + // so the dim-sized buffers need no per-cluster reset; first touch is + // detected via the epoch (not acc==0), robust to zero-valued weights. const size_t dim = vectors->get_dimension(); std::vector acc(dim, T(0)); std::vector epoch(dim, 0); std::vector touched; uint32_t cur_epoch = 0; - // The per-term max-pool can optionally be offloaded to the GPU (batched as - // one kernel launch per list) instead of the CPU flat-array path above. - // Opt-in via NSPARSE_GPU_SUMMARIZE=1 (see should_offload_summarize_to_gpu); - // a net win on GPU-rich / low-core hosts. Everything downstream is identical. + // The max-pool may instead be offloaded to the GPU (one launch per list) + // when NSPARSE_GPU_SUMMARIZE=1; downstream sort/truncate is identical. #ifdef NSPARSE_WITH_GPU bool gpu_ok = false; - std::vector gpu_clusters; + std::vector gpu_clusters; if constexpr (std::is_same_v) { if (detail::should_offload_summarize_to_gpu()) { - gpu_ok = detail::GpuClusterAssigner::instance() - .summarize_list_maxpool(vectors, - group_of_doc_ids.data(), - offsets.data(), - offsets.size() - 1, - gpu_clusters); + gpu_ok = detail::GpuSummarizer::instance().summarize_list( + vectors, group_of_doc_ids.data(), offsets.data(), + offsets.size() - 1, gpu_clusters); } } #endif for (size_t i = 0; i < offsets.size() - 1; ++i) { - const int64_t ts0 = prof ? summ_now_ns() : 0; size_t n_docs = offsets[i + 1] - offsets[i]; float sum = 0.0F; std::vector> summary_vec; @@ -146,8 +106,7 @@ SparseVectors summarize_(const SparseVectors* vectors, const T v = *reinterpret_cast( values_data + j * sizeof(T)); if (epoch[term] != cur_epoch) { - // First occurrence in this cluster. Matches the old map - // semantics: value = max(0, v). + // First occurrence in this cluster; value = max(0, v). epoch[term] = cur_epoch; const T value = std::max(T(0), v); acc[term] = value; @@ -165,8 +124,6 @@ SparseVectors summarize_(const SparseVectors* vectors, } } - const int64_t ts1 = prof ? summ_now_ns() : 0; - // Sort by value in descending order std::ranges::sort(summary_vec, [](const auto& a, const auto& b) { return a.second > b.second; @@ -201,17 +158,6 @@ SparseVectors summarize_(const SparseVectors* vectors, terms.data(), terms.size(), reinterpret_cast(values.data()), values.size() * sizeof(T)); - - if (prof) { - const int64_t ts2 = summ_now_ns(); - maxpool_acc += ts1 - ts0; - sort_acc += ts2 - ts1; - } - } - if (prof) { - g_summ_profile.maxpool_ns.fetch_add(maxpool_acc, - std::memory_order_relaxed); - g_summ_profile.sort_ns.fetch_add(sort_acc, std::memory_order_relaxed); } return summarized_vectors; } diff --git a/nsparse/cluster/kmeans_utils.cpp b/nsparse/cluster/kmeans_utils.cpp index 89e8626..aee9397 100644 --- a/nsparse/cluster/kmeans_utils.cpp +++ b/nsparse/cluster/kmeans_utils.cpp @@ -174,11 +174,11 @@ void map_docs_to_clusters(const SparseVectors* vectors, return; } #if defined(NSPARSE_WITH_GPU) && !defined(__AVX512F__) - // GPU-accelerated assignment via cuSPARSE for float (U32) weights. Produces - // assignments identical to the scalar CPU reference below (skips centroids, - // ties break to the lowest cluster index); on any failure we fall through - // to the CPU path. Disabled for AVX-512 builds whose assignment semantics - // differ. + // GPU assignment (cuSPARSE) for float (U32) weights. It matches the scalar + // CPU path below (skips centroids); on any failure we fall through to CPU. + // Excluded from AVX-512 builds because that path (unlike the scalar one) + // does not skip centroids, so mixing GPU- and AVX-512-assigned lists in one + // build would be inconsistent. if (vectors->get_element_size() == U32 && should_offload_assignment_to_gpu(n_docs, n_clusters)) { try { diff --git a/nsparse/gpu/CMakeLists.txt b/nsparse/gpu/CMakeLists.txt index c57e826..a45df3b 100644 --- a/nsparse/gpu/CMakeLists.txt +++ b/nsparse/gpu/CMakeLists.txt @@ -47,19 +47,21 @@ endif() set(CMAKE_CUDA_STANDARD 17) set(CMAKE_CUDA_STANDARD_REQUIRED ON) -# Absolute path to the CUDA translation unit, resolved once here so the helper -# can be called from the parent directory scope. This file is include()d from +# Absolute paths to the CUDA translation units, resolved here so the helper can +# be called from the parent directory scope. This file is include()d from # nsparse/CMakeLists.txt, so CMAKE_CURRENT_LIST_DIR is this gpu/ directory. -set(NSPARSE_GPU_CU "${CMAKE_CURRENT_LIST_DIR}/gpu_cluster_assigner.cu" - CACHE INTERNAL "nsparse GPU source") +set(NSPARSE_GPU_SRC + "${CMAKE_CURRENT_LIST_DIR}/gpu_cluster_assigner.cu" + "${CMAKE_CURRENT_LIST_DIR}/gpu_summarizer.cu" + CACHE INTERNAL "nsparse GPU sources") message(STATUS "nsparse: GPU enabled (nvcc=${CMAKE_CUDA_COMPILER}, " "archs=${CMAKE_CUDA_ARCHITECTURES})") -# Compile the GPU source into and link cuSPARSE/cudart. Guarded by the +# Compile the GPU sources into and link cuSPARSE/cudart. Guarded by the # NSPARSE_WITH_GPU compile define so the CPU sources pick up the GPU hooks. function(nsparse_attach_gpu target) - target_sources(${target} PRIVATE "${NSPARSE_GPU_CU}") + target_sources(${target} PRIVATE ${NSPARSE_GPU_SRC}) target_compile_definitions(${target} PUBLIC NSPARSE_WITH_GPU) target_link_libraries(${target} PRIVATE CUDA::cusparse CUDA::cudart) # Resolve device symbols into the library itself so pure-C++ consumers diff --git a/nsparse/gpu/gpu_cluster_assigner.cu b/nsparse/gpu/gpu_cluster_assigner.cu index 469581c..a4f58b2 100644 --- a/nsparse/gpu/gpu_cluster_assigner.cu +++ b/nsparse/gpu/gpu_cluster_assigner.cu @@ -11,84 +11,24 @@ #include -#include -#include #include -#include #include #include -#include -#include -#include -#include #include #include "absl/container/flat_hash_set.h" #include "cuda_runtime.h" +#include "nsparse/gpu/gpu_common.cuh" #include "nsparse/sparse_vectors.h" #include "nsparse/types.h" namespace nsparse::detail { namespace { -void check_cuda(cudaError_t status, const char* what) { - if (status != cudaSuccess) { - std::ostringstream oss; - oss << "CUDA error in " << what << ": " << cudaGetErrorString(status); - throw std::runtime_error(oss.str()); - } -} - -void check_cusparse(cusparseStatus_t status, const char* what) { - if (status != CUSPARSE_STATUS_SUCCESS) { - std::ostringstream oss; - oss << "cuSPARSE error in " << what << ": " - << cusparseGetErrorString(status); - throw std::runtime_error(oss.str()); - } -} - -#define NSPARSE_CUDA_CHECK(expr) check_cuda((expr), #expr) -#define NSPARSE_CUSPARSE_CHECK(expr) check_cusparse((expr), #expr) - -// Optional phase profiling, enabled by NSPARSE_GPU_PROFILE=1. Accumulates -// wall-time (ns) across all list calls into a few buckets so the per-phase -// split of the GPU path can be inspected without an external profiler. Cheap -// (a handful of atomics per list) and compiled to near-nothing when disabled. -struct Profile { - std::atomic host_prep{0}; // row_ptr/centroid host setup - std::atomic gpu_exec{0}; // upload+kernels+SpMM+sync - std::atomic host_assign{0}; // scatter results into clusters - std::atomic calls{0}; - bool enabled = false; - - Profile() { - const char* v = std::getenv("NSPARSE_GPU_PROFILE"); - enabled = (v != nullptr && v[0] == '1'); - } - ~Profile() { - if (!enabled) return; - std::fprintf(stderr, - "[nsparse gpu] calls=%lld host_prep=%.3fs gpu_exec=%.3fs " - "host_assign=%.3fs\n", - static_cast(calls.load()), - host_prep.load() / 1e9, gpu_exec.load() / 1e9, - host_assign.load() / 1e9); - } -}; -Profile g_profile; - -inline int64_t now_ns() { - return std::chrono::duration_cast( - std::chrono::steady_clock::now().time_since_epoch()) - .count(); -} - -// One thread per document row. Scans the n_clusters scores for that row and -// selects the argmax, breaking ties toward the lowest cluster index (strict -// greater-than update) to match the CPU reference in map_docs_to_clusters(). -__global__ void row_argmax_kernel(const float* __restrict__ scores, - int n_rows, int n_clusters, +// One thread per document row: argmax over the row's n_clusters scores, ties +// broken to the lowest cluster index (strict-greater) to match the CPU path. +__global__ void row_argmax_kernel(const float* __restrict__ scores, int n_rows, + int n_clusters, int32_t* __restrict__ best_cluster) { int row = blockIdx.x * blockDim.x + threadIdx.x; if (row >= n_rows) { @@ -106,9 +46,8 @@ __global__ void row_argmax_kernel(const float* __restrict__ scores, best_cluster[row] = best_j; } -// Builds the compact CSR values/columns of the document sub-matrix A by -// gathering rows straight from the resident corpus, so the large doc data never -// crosses the PCIe bus per list. One thread per output row. +// Gather the CSR of the document sub-matrix A from the resident corpus, so the +// bulk doc data never re-crosses PCIe per list. One thread per output row. __global__ void gather_csr_kernel(const int32_t* __restrict__ corpus_indptr, const int32_t* __restrict__ corpus_indices, const float* __restrict__ corpus_values, @@ -130,11 +69,9 @@ __global__ void gather_csr_kernel(const int32_t* __restrict__ corpus_indptr, } } - -// Scatters centroid rows from the resident corpus into the dense B matrix -// (dim x n_clusters, row-major, ldb = n_clusters). Column j of B is centroid j. -// One thread per centroid; each writes a disjoint column, so there are no -// races. B must be zeroed first. +// Scatter centroid rows into the dense B matrix (dim x n_clusters, row-major, +// ldb = n_clusters). One thread per centroid; each owns a disjoint column, so +// no races. B must be zeroed first. __global__ void scatter_dense_kernel(const int32_t* __restrict__ corpus_indptr, const int32_t* __restrict__ corpus_indices, const float* __restrict__ corpus_values, @@ -154,130 +91,6 @@ __global__ void scatter_dense_kernel(const int32_t* __restrict__ corpus_indptr, } } -// ========================================================================= -// GPU max-pool for summarize(). For one cluster (posting list) this computes, -// per term, the maximum weight across all documents in the cluster, then -// compacts the touched terms. SPLADE weights are >= 0, so a non-negative float -// compares monotonically as a reinterpreted int32 — allowing atomicMax on the -// integer view of the accumulator. -// -// Buffers (all sized to corpus dim, reused across clusters, left clean on exit -// so no per-cluster dim-wide memset is needed): -// acc[dim] : int32 view of the running max weight per term (0 == absent) -// seen[dim] : 0/1 first-touch flag used to append each term once to touched -// touched[] : compact list of term ids that appeared in this cluster -// n_touched : length of touched[] (atomic counter) -// ========================================================================= - -// Batched per-list max-pool: ONE block per cluster, so an entire inverted -// list's clusters are summarized in a single kernel launch with a single -// stream sync (instead of two blocking syncs per cluster, which serialized ~27M -// GPU round-trips across the build). Clusters are independent and each is owned -// by one block, so only intra-block synchronization is needed. -// -// Buffers (per thread/list, reused across lists, kept clean via touched-reset): -// acc[n_clusters*dim], seen[n_clusters*dim] : per-cluster per-term max / flag, -// partitioned so block b owns region [b*dim, (b+1)*dim). Start zeroed and -// are restored to zero for touched slots before the kernel returns. -// out_term / out_val : compacted results, laid out per cluster using -// out_base[b] (the input nnz prefix, an upper bound on distinct terms). -// out_count[b] : number of distinct terms produced for cluster b. -// out_sum[b] : sum of the cluster's max values. -__global__ void summarize_list_kernel( - const int32_t* __restrict__ corpus_indptr, - const int32_t* __restrict__ corpus_indices, - const float* __restrict__ corpus_values, - const int32_t* __restrict__ docs, const int32_t* __restrict__ offsets, - const int32_t* __restrict__ out_base, int dim, - int32_t* __restrict__ acc, int32_t* __restrict__ seen, - int32_t* __restrict__ out_term, float* __restrict__ out_val, - int32_t* __restrict__ out_count, float* __restrict__ out_sum) { - const int b = blockIdx.x; // one block per cluster - const int start = offsets[b]; - const int end = offsets[b + 1]; - const int64_t base = static_cast(b) * dim; - const int32_t obase = out_base[b]; - - __shared__ int s_count; - if (threadIdx.x == 0) { - s_count = 0; - } - __syncthreads(); - - // Scatter-max over the cluster's documents. Each distinct term is appended - // once (via seen[]) to this cluster's output segment at out_term[obase..]. - for (int i = start + threadIdx.x; i < end; i += blockDim.x) { - const int32_t d = docs[i]; - const int32_t src = corpus_indptr[d]; - const int32_t len = corpus_indptr[d + 1] - src; - for (int32_t t = 0; t < len; ++t) { - const int32_t term = corpus_indices[src + t]; - const float v = corpus_values[src + t]; - // Non-negative float -> monotonic int bits; atomicMax on int view. - atomicMax(&acc[base + term], __float_as_int(v)); - if (atomicCAS(&seen[base + term], 0, 1) == 0) { - const int slot = atomicAdd(&s_count, 1); - out_term[obase + slot] = term; - } - } - } - __syncthreads(); - - // Compact: read each touched term's max, accumulate the sum, and reset - // acc/seen for reuse. Sum via a block-wide atomic into a shared scalar. - __shared__ float s_sum; - if (threadIdx.x == 0) { - s_sum = 0.0F; - out_count[b] = s_count; - } - __syncthreads(); - - for (int k = threadIdx.x; k < s_count; k += blockDim.x) { - const int32_t term = out_term[obase + k]; - const float v = __int_as_float(acc[base + term]); - out_val[obase + k] = v; - atomicAdd(&s_sum, v); - acc[base + term] = 0; // reset for next list reusing this buffer - seen[base + term] = 0; - } - __syncthreads(); - if (threadIdx.x == 0) { - out_sum[b] = s_sum; - } -} - -// The full corpus, uploaded to the GPU once and shared read-only across all -// threads/lists. Indices are widened to int32 on upload to feed both the gather -// kernel and cuSPARSE's 32-bit CSR indices. -struct DeviceCorpus { - int32_t* indptr = nullptr; // n_vectors + 1 - int32_t* indices = nullptr; // nnz - float* values = nullptr; // nnz - int64_t nnz = 0; - size_t dim = 0; - size_t n_vectors = 0; - // Identity of the resident corpus. Pointer alone is unsafe: a new - // SparseVectors can reuse a freed one's address, so the residency check - // also compares n_vectors and nnz. - const SparseVectors* owner = nullptr; - - bool matches(const SparseVectors* v, size_t nv, int64_t nz) const { - return owner == v && n_vectors == nv && nnz == nz; - } - - void free() { - cudaFree(indptr); - cudaFree(indices); - cudaFree(values); - indptr = nullptr; - indices = nullptr; - values = nullptr; - owner = nullptr; - n_vectors = 0; - nnz = 0; - } -}; - // Per-thread GPU resources: an independent cuSPARSE handle + stream (handles // are not shareable across threads) and scratch buffers reused across lists, // grown on demand so steady-state builds do no per-list cudaMalloc. @@ -296,19 +109,6 @@ struct ThreadCtx { int32_t* d_best = nullptr; size_t best_cap = 0; void* d_spmm = nullptr; size_t spmm_cap = 0; - // summarize() batched per-list max-pool scratch. acc/seen are partitioned - // as n_clusters x dim and kept clean between lists (each cluster resets the - // slots it touched). Sized to the largest list seen so far. - int32_t* d_sum_docs = nullptr; size_t sum_docs_cap = 0; - int32_t* d_offsets = nullptr; size_t offsets_cap = 0; - int32_t* d_out_base = nullptr; size_t out_base_cap = 0; - int32_t* d_acc = nullptr; size_t acc_cap = 0; - int32_t* d_seen = nullptr; size_t seen_cap = 0; - int32_t* d_out_term = nullptr; size_t out_term_cap = 0; - float* d_out_val = nullptr; size_t out_val_cap = 0; - int32_t* d_out_count = nullptr; size_t out_count_cap = 0; - float* d_out_sum = nullptr; size_t out_sum_cap = 0; - ~ThreadCtx() { cudaFree(d_docs); cudaFree(d_centroids); @@ -319,15 +119,6 @@ struct ThreadCtx { cudaFree(d_c); cudaFree(d_best); cudaFree(d_spmm); - cudaFree(d_sum_docs); - cudaFree(d_offsets); - cudaFree(d_out_base); - cudaFree(d_acc); - cudaFree(d_seen); - cudaFree(d_out_term); - cudaFree(d_out_val); - cudaFree(d_out_count); - cudaFree(d_out_sum); if (init) { cusparseDestroy(handle); cudaStreamDestroy(stream); @@ -335,18 +126,6 @@ struct ThreadCtx { } }; -// Reallocate *ptr only when it must grow. Capacity tracked in bytes. -template -void ensure_capacity(T** ptr, size_t& cap_bytes, size_t need_bytes) { - if (cap_bytes < need_bytes) { - cudaFree(*ptr); - void* p = nullptr; - check_cuda(cudaMalloc(&p, need_bytes), "cudaMalloc(ensure_capacity)"); - *ptr = static_cast(p); - cap_bytes = need_bytes; - } -} - // Lazily created once per worker thread; destroyed at thread exit. thread_local std::unique_ptr t_ctx; @@ -362,29 +141,13 @@ ThreadCtx& thread_ctx() { return *t_ctx; } -} // namespace - -struct GpuClusterAssigner::Impl { - std::mutex corpus_mutex; // guards residency check/upload only - DeviceCorpus corpus; - bool usable = false; -}; - -GpuClusterAssigner::GpuClusterAssigner() : impl_(new Impl()) { +bool gpu_present() { int device_count = 0; cudaError_t status = cudaGetDeviceCount(&device_count); - if (status != cudaSuccess || device_count == 0) { - return; // impl_->usable stays false; callers fall back to CPU - } - impl_->usable = true; + return status == cudaSuccess && device_count > 0; } -GpuClusterAssigner::~GpuClusterAssigner() { - if (impl_ != nullptr) { - impl_->corpus.free(); - delete impl_; - } -} +} // namespace GpuClusterAssigner& GpuClusterAssigner::instance() { static GpuClusterAssigner singleton; @@ -392,59 +155,8 @@ GpuClusterAssigner& GpuClusterAssigner::instance() { } bool GpuClusterAssigner::available() { - return instance().impl_ != nullptr && instance().impl_->usable; -} - -// Uploads the corpus to the GPU once. Safe to call from many threads; the first -// caller uploads while the rest wait, and subsequent calls are a cheap pointer -// comparison. Returns a stable view (device pointers are not mutated once set). -// Declared as a member so it can touch the private Impl; the header exposes it -// as an opaque const void* to avoid leaking CUDA types. -const void* GpuClusterAssigner::ensure_corpus_resident( - const SparseVectors* vectors) { - Impl* impl = impl_; - const size_t n_vectors = vectors->num_vectors(); - const idx_t* indptr = vectors->indptr_data(); - const term_t* indices = vectors->indices_data(); - const float* values = vectors->values_data_float(); - const int64_t nnz = indptr[n_vectors]; - - std::lock_guard lock(impl->corpus_mutex); - if (impl->corpus.matches(vectors, n_vectors, nnz)) { - return &impl->corpus; - } - impl->corpus.free(); - - // Widen indices (term_t == uint16) to int32 for the gather kernel and - // cuSPARSE. Done once per corpus. - std::vector indices32(static_cast(nnz)); - for (int64_t i = 0; i < nnz; ++i) { - indices32[i] = static_cast(indices[i]); - } - - check_cuda(cudaMalloc(&impl->corpus.indptr, - (n_vectors + 1) * sizeof(int32_t)), - "cudaMalloc(corpus.indptr)"); - check_cuda(cudaMalloc(&impl->corpus.indices, nnz * sizeof(int32_t)), - "cudaMalloc(corpus.indices)"); - check_cuda(cudaMalloc(&impl->corpus.values, nnz * sizeof(float)), - "cudaMalloc(corpus.values)"); - check_cuda(cudaMemcpy(impl->corpus.indptr, indptr, - (n_vectors + 1) * sizeof(int32_t), - cudaMemcpyHostToDevice), - "cudaMemcpy(corpus.indptr)"); - check_cuda(cudaMemcpy(impl->corpus.indices, indices32.data(), - nnz * sizeof(int32_t), cudaMemcpyHostToDevice), - "cudaMemcpy(corpus.indices)"); - check_cuda(cudaMemcpy(impl->corpus.values, values, nnz * sizeof(float), - cudaMemcpyHostToDevice), - "cudaMemcpy(corpus.values)"); - - impl->corpus.nnz = nnz; - impl->corpus.dim = vectors->get_dimension(); - impl->corpus.n_vectors = n_vectors; - impl->corpus.owner = vectors; - return &impl->corpus; + static const bool present = gpu_present(); + return present; } void GpuClusterAssigner::assign(const SparseVectors* vectors, @@ -456,14 +168,11 @@ void GpuClusterAssigner::assign(const SparseVectors* vectors, return; } - const bool prof = g_profile.enabled; - const int64_t t0 = prof ? now_ns() : 0; - const idx_t* indptr = vectors->indptr_data(); const size_t dim = vectors->get_dimension(); - // Centroids are the first element of each cluster. Collect them (host) and - // mark which input docs are centroids so they are not re-added, as on CPU. + // Centroids are clusters[j].front(); collect them and record which input + // docs are centroids so they are not re-added (matches the CPU path). std::vector centroid_docs(n_clusters); absl::flat_hash_set centroid_set; centroid_set.reserve(n_clusters); @@ -472,8 +181,7 @@ void GpuClusterAssigner::assign(const SparseVectors* vectors, centroid_set.insert(clusters[j].front()); } - // Row pointers of the document sub-matrix A (host-side; the actual column - // indices and values are gathered on-device from the resident corpus). + // Row pointers of A (host); column indices/values are gathered on-device. std::vector h_a_row_ptr(n_docs + 1, 0); for (size_t i = 0; i < n_docs; ++i) { const idx_t d = docs[i]; @@ -482,18 +190,13 @@ void GpuClusterAssigner::assign(const SparseVectors* vectors, } const int64_t nnz_a = h_a_row_ptr[n_docs]; - const int64_t t1 = prof ? now_ns() : 0; - - const DeviceCorpus& corpus = - *static_cast(ensure_corpus_resident(vectors)); + const DeviceCorpus& corpus = GpuCorpus::instance().ensure_resident(vectors); ThreadCtx& ctx = thread_ctx(); cudaStream_t stream = ctx.stream; - // Upload only the small per-list metadata (doc/centroid ids). The actual - // doc/centroid payloads live on the GPU (resident corpus). + // Upload only the small per-list metadata (doc/centroid ids). ensure_capacity(&ctx.d_docs, ctx.docs_cap, n_docs * sizeof(int32_t)); - ensure_capacity(&ctx.d_centroids, ctx.cent_cap, - n_clusters * sizeof(int32_t)); + ensure_capacity(&ctx.d_centroids, ctx.cent_cap, n_clusters * sizeof(int32_t)); ensure_capacity(&ctx.d_best, ctx.best_cap, n_docs * sizeof(int32_t)); ensure_capacity(&ctx.d_b, ctx.b_cap, dim * n_clusters * sizeof(float)); @@ -504,28 +207,24 @@ void GpuClusterAssigner::assign(const SparseVectors* vectors, n_clusters * sizeof(int32_t), cudaMemcpyHostToDevice, stream)); - // Build dense B (dim x n_clusters, row-major) from centroid rows. The - // fused kernel below reads B but not A's CSR separately; docs are gathered - // directly from the corpus inside the fused kernel. + // Build dense B (dim x n_clusters, row-major) from centroid rows. constexpr int kBlock = 256; - NSPARSE_CUDA_CHECK(cudaMemsetAsync( - ctx.d_b, 0, dim * n_clusters * sizeof(float), stream)); - const int cent_grid = - static_cast((n_clusters + kBlock - 1) / kBlock); + NSPARSE_CUDA_CHECK( + cudaMemsetAsync(ctx.d_b, 0, dim * n_clusters * sizeof(float), stream)); + const int cent_grid = static_cast((n_clusters + kBlock - 1) / kBlock); scatter_dense_kernel<<>>( corpus.indptr, corpus.indices, corpus.values, ctx.d_centroids, static_cast(n_clusters), static_cast(n_clusters), ctx.d_b); NSPARSE_CUDA_CHECK(cudaGetLastError()); - // cuSPARSE SpMM path: gathers A from corpus, then runs SpMM + argmax. + // Gather A's CSR, then C = A * B (SpMM) and per-row argmax. ensure_capacity(&ctx.d_a_row_ptr, ctx.rowptr_cap, (n_docs + 1) * sizeof(int32_t)); ensure_capacity(&ctx.d_a_col, ctx.col_cap, static_cast(nnz_a) * sizeof(int32_t)); ensure_capacity(&ctx.d_a_val, ctx.val_cap, static_cast(nnz_a) * sizeof(float)); - ensure_capacity(&ctx.d_c, ctx.c_cap, - n_docs * n_clusters * sizeof(float)); + ensure_capacity(&ctx.d_c, ctx.c_cap, n_docs * n_clusters * sizeof(float)); NSPARSE_CUDA_CHECK(cudaMemcpyAsync(ctx.d_a_row_ptr, h_a_row_ptr.data(), (n_docs + 1) * sizeof(int32_t), @@ -540,18 +239,15 @@ void GpuClusterAssigner::assign(const SparseVectors* vectors, cusparseDnMatDescr_t mat_b; cusparseDnMatDescr_t mat_c; NSPARSE_CUSPARSE_CHECK(cusparseCreateCsr( - &mat_a, static_cast(n_docs), static_cast(dim), - nnz_a, ctx.d_a_row_ptr, ctx.d_a_col, ctx.d_a_val, - CUSPARSE_INDEX_32I, CUSPARSE_INDEX_32I, CUSPARSE_INDEX_BASE_ZERO, - CUDA_R_32F)); + &mat_a, static_cast(n_docs), static_cast(dim), nnz_a, + ctx.d_a_row_ptr, ctx.d_a_col, ctx.d_a_val, CUSPARSE_INDEX_32I, + CUSPARSE_INDEX_32I, CUSPARSE_INDEX_BASE_ZERO, CUDA_R_32F)); NSPARSE_CUSPARSE_CHECK(cusparseCreateDnMat( - &mat_b, static_cast(dim), - static_cast(n_clusters), + &mat_b, static_cast(dim), static_cast(n_clusters), static_cast(n_clusters), ctx.d_b, CUDA_R_32F, CUSPARSE_ORDER_ROW)); NSPARSE_CUSPARSE_CHECK(cusparseCreateDnMat( - &mat_c, static_cast(n_docs), - static_cast(n_clusters), + &mat_c, static_cast(n_docs), static_cast(n_clusters), static_cast(n_clusters), ctx.d_c, CUDA_R_32F, CUSPARSE_ORDER_ROW)); @@ -584,10 +280,7 @@ void GpuClusterAssigner::assign(const SparseVectors* vectors, cudaMemcpyDeviceToHost, stream)); NSPARSE_CUDA_CHECK(cudaStreamSynchronize(stream)); - const int64_t t2 = prof ? now_ns() : 0; - - // Append assignments on the host, skipping documents that are themselves a - // centroid (matching the CPU reference exactly). + // Append assignments, skipping docs that are themselves a centroid. for (size_t i = 0; i < n_docs; ++i) { const idx_t d = docs[i]; if (centroid_set.contains(d)) { @@ -595,155 +288,10 @@ void GpuClusterAssigner::assign(const SparseVectors* vectors, } clusters[h_best[i]].push_back(d); } - - if (prof) { - const int64_t t3 = now_ns(); - g_profile.host_prep.fetch_add(t1 - t0, std::memory_order_relaxed); - g_profile.gpu_exec.fetch_add(t2 - t1, std::memory_order_relaxed); - g_profile.host_assign.fetch_add(t3 - t2, std::memory_order_relaxed); - g_profile.calls.fetch_add(1, std::memory_order_relaxed); - } -} - -bool GpuClusterAssigner::summarize_list_maxpool( - const SparseVectors* vectors, const idx_t* docs, const idx_t* offsets, - size_t n_clusters, std::vector& out) { - if (!impl_->usable || n_clusters == 0) { - return false; - } - const size_t dim = vectors->get_dimension(); - const idx_t* indptr = vectors->indptr_data(); - - // The flattened doc array spans offsets[0]..offsets[n_clusters]. Compute, - // per cluster, the nnz prefix (out_base) used to place each cluster's - // compacted output in a private segment; the total nnz is the output cap. - const size_t n_docs = static_cast(offsets[n_clusters] - offsets[0]); - if (n_docs == 0) { - return false; - } - std::vector h_out_base(n_clusters + 1, 0); - for (size_t b = 0; b < n_clusters; ++b) { - int64_t nnz_b = 0; - for (idx_t i = offsets[b]; i < offsets[b + 1]; ++i) { - const idx_t d = docs[i]; - nnz_b += indptr[d + 1] - indptr[d]; - } - h_out_base[b + 1] = h_out_base[b] + static_cast(nnz_b); - } - const int64_t total_nnz = h_out_base[n_clusters]; - if (total_nnz == 0) { - return false; - } - - const DeviceCorpus& corpus = - *static_cast(ensure_corpus_resident(vectors)); - ThreadCtx& ctx = thread_ctx(); - cudaStream_t stream = ctx.stream; - - // acc/seen are partitioned as n_clusters x dim; must start (and remain) - // zeroed. Zero only when (re)grown — the kernel restores touched slots. - const size_t acc_bytes = n_clusters * dim * sizeof(int32_t); - const bool acc_grew = ctx.acc_cap < acc_bytes; - ensure_capacity(&ctx.d_acc, ctx.acc_cap, acc_bytes); - ensure_capacity(&ctx.d_seen, ctx.seen_cap, acc_bytes); - if (acc_grew) { - NSPARSE_CUDA_CHECK(cudaMemsetAsync(ctx.d_acc, 0, ctx.acc_cap, stream)); - NSPARSE_CUDA_CHECK( - cudaMemsetAsync(ctx.d_seen, 0, ctx.seen_cap, stream)); - } - ensure_capacity(&ctx.d_sum_docs, ctx.sum_docs_cap, - n_docs * sizeof(int32_t)); - ensure_capacity(&ctx.d_offsets, ctx.offsets_cap, - (n_clusters + 1) * sizeof(int32_t)); - ensure_capacity(&ctx.d_out_base, ctx.out_base_cap, - (n_clusters + 1) * sizeof(int32_t)); - ensure_capacity(&ctx.d_out_term, ctx.out_term_cap, - static_cast(total_nnz) * sizeof(int32_t)); - ensure_capacity(&ctx.d_out_val, ctx.out_val_cap, - static_cast(total_nnz) * sizeof(float)); - ensure_capacity(&ctx.d_out_count, ctx.out_count_cap, - n_clusters * sizeof(int32_t)); - ensure_capacity(&ctx.d_out_sum, ctx.out_sum_cap, - n_clusters * sizeof(float)); - - // Offsets are relative to offsets[0]; rebase to 0 for the flat doc upload. - std::vector h_offsets(n_clusters + 1); - const idx_t base0 = offsets[0]; - for (size_t b = 0; b <= n_clusters; ++b) { - h_offsets[b] = static_cast(offsets[b] - base0); - } - - NSPARSE_CUDA_CHECK(cudaMemcpyAsync(ctx.d_sum_docs, docs + base0, - n_docs * sizeof(int32_t), - cudaMemcpyHostToDevice, stream)); - NSPARSE_CUDA_CHECK(cudaMemcpyAsync(ctx.d_offsets, h_offsets.data(), - (n_clusters + 1) * sizeof(int32_t), - cudaMemcpyHostToDevice, stream)); - NSPARSE_CUDA_CHECK(cudaMemcpyAsync(ctx.d_out_base, h_out_base.data(), - (n_clusters + 1) * sizeof(int32_t), - cudaMemcpyHostToDevice, stream)); - - // One block per cluster. - constexpr int kBlock = 128; - summarize_list_kernel<<(n_clusters), kBlock, 0, stream>>>( - corpus.indptr, corpus.indices, corpus.values, ctx.d_sum_docs, - ctx.d_offsets, ctx.d_out_base, static_cast(dim), ctx.d_acc, - ctx.d_seen, ctx.d_out_term, ctx.d_out_val, ctx.d_out_count, - ctx.d_out_sum); - NSPARSE_CUDA_CHECK(cudaGetLastError()); - - // Copy back the compact outputs (single sync for the whole list). - std::vector h_term(total_nnz); - std::vector h_val(total_nnz); - std::vector h_count(n_clusters); - std::vector h_sum(n_clusters); - NSPARSE_CUDA_CHECK(cudaMemcpyAsync(h_term.data(), ctx.d_out_term, - total_nnz * sizeof(int32_t), - cudaMemcpyDeviceToHost, stream)); - NSPARSE_CUDA_CHECK(cudaMemcpyAsync(h_val.data(), ctx.d_out_val, - total_nnz * sizeof(float), - cudaMemcpyDeviceToHost, stream)); - NSPARSE_CUDA_CHECK(cudaMemcpyAsync(h_count.data(), ctx.d_out_count, - n_clusters * sizeof(int32_t), - cudaMemcpyDeviceToHost, stream)); - NSPARSE_CUDA_CHECK(cudaMemcpyAsync(h_sum.data(), ctx.d_out_sum, - n_clusters * sizeof(float), - cudaMemcpyDeviceToHost, stream)); - NSPARSE_CUDA_CHECK(cudaStreamSynchronize(stream)); - - out.resize(n_clusters); - for (size_t b = 0; b < n_clusters; ++b) { - const int32_t base = h_out_base[b]; - const int32_t cnt = h_count[b]; - ClusterSummary& cs = out[b]; - cs.terms.resize(cnt); - cs.values.resize(cnt); - for (int32_t k = 0; k < cnt; ++k) { - cs.terms[k] = static_cast(h_term[base + k]); - cs.values[k] = h_val[base + k]; - } - cs.sum = h_sum[b]; - } - return true; } bool should_offload_assignment_to_gpu(size_t n_docs, size_t n_clusters) { - // Offload whenever GPU support is compiled in and a device is available. - // k-means assignment needs at least two clusters to be meaningful. return GpuClusterAssigner::available() && n_docs > 0 && n_clusters >= 2; } -bool should_offload_summarize_to_gpu() { - if (!GpuClusterAssigner::available()) { - return false; - } - // Opt-in. The CPU flat-array max-pool is faster on high-core hosts, so the - // default keeps summarize on the CPU (bit-identical, no GPU dependency). - // On GPU-rich / low-core hosts (e.g. an 8-vCPU single-GPU instance), the - // CPU summarize dominates wall-time and offloading its max-pool to the - // otherwise-idle GPU is a net win — enable it with NSPARSE_GPU_SUMMARIZE=1. - const char* env = std::getenv("NSPARSE_GPU_SUMMARIZE"); - return env != nullptr && env[0] == '1'; -} - } // namespace nsparse::detail diff --git a/nsparse/gpu/gpu_cluster_assigner.h b/nsparse/gpu/gpu_cluster_assigner.h index b17f20b..9a67663 100644 --- a/nsparse/gpu/gpu_cluster_assigner.h +++ b/nsparse/gpu/gpu_cluster_assigner.h @@ -18,126 +18,49 @@ namespace nsparse::detail { /** - * @brief GPU-accelerated k-means document-to-cluster assignment used by the - * index-building (build()) path. + * @brief GPU k-means document-to-cluster assignment for the build() path. * - * The assignment step of Seismic clustering computes, for every document, the - * dot product against each cluster centroid and picks the highest-scoring - * cluster. Across a whole inverted list this is a sparse-times-dense matrix - * product (documents as a CSR matrix, centroids as a dense matrix), which maps - * directly onto NVIDIA cuSPARSE SpMM followed by a per-row argmax. + * Per inverted list, scoring every document against every centroid is a + * sparse-times-dense product (documents as CSR, centroids as dense), computed + * with cuSPARSE SpMM plus a per-row argmax. Output matches the scalar CPU + * map_docs_to_clusters(): centroids are not re-added, ties break to the lowest + * cluster index. float (U32) weights only; search() is unaffected (CPU). * - * This class mirrors the CPU map_docs_to_clusters() semantics exactly for - * float (U32) weights: - * - documents equal to a cluster centroid are left untouched (not re-added), - * - ties are broken toward the lowest cluster index (strict-greater update). - * - * The search() path is unaffected and always runs on the CPU. - * - * The implementation lives in gpu_cluster_assigner.cu and is only compiled when - * the project is configured with -DNSPARSE_ENABLE_GPU=ON (which defines - * NSPARSE_WITH_GPU). The declarations below are always visible so that callers - * can guard usage with a single runtime available() check. + * Compiled only with -DNSPARSE_ENABLE_GPU=ON (defines NSPARSE_WITH_GPU); the + * declarations stay visible so callers guard on a single available() check. */ class GpuClusterAssigner { public: - /// Process-wide singleton. Owns the cuSPARSE handle and serializes GPU - /// access so it is safe to call from OpenMP worker threads. + /// Process-wide singleton owning the per-thread cuSPARSE handles/streams. static GpuClusterAssigner& instance(); - /// True when the library was built with CUDA support and a usable GPU is - /// present. When false, callers must use the CPU path. + /// True when built with GPU support and a usable device is present. static bool available(); /** * @brief Assign each non-centroid document to its best cluster on the GPU. * - * For every document id in @p docs (interpreted as a row of @p vectors), - * computes the dot product against each cluster centroid - * (clusters[j].front()) and appends the document to the highest-scoring - * cluster. Documents that are themselves a centroid are skipped, matching - * the CPU reference. - * - * Only float (U32) weights are supported on the GPU. Callers must check - * vectors->get_element_size() == U32 and fall back to the CPU path - * otherwise. - * - * @param vectors full corpus of sparse vectors (CSR-backed). - * @param docs document ids belonging to this inverted list. - * @param clusters in/out clusters; clusters[j].front() is the centroid of - * cluster j and assigned documents are appended. + * @param vectors corpus of sparse vectors (CSR-backed). + * @param docs document ids in this inverted list. + * @param clusters in/out; clusters[j].front() is centroid j, assigned docs + * are appended. Documents that are a centroid are skipped. */ void assign(const SparseVectors* vectors, const std::vector& docs, std::vector>& clusters); - /// Per-cluster max-pool result from summarize_list_maxpool(). Terms are in - /// GPU touched-append order (not sorted); the CPU sort/truncate handles - /// ordering, preserving the existing output semantics. - struct ClusterSummary { - std::vector terms; - std::vector values; - float sum = 0.0F; - }; - - /** - * @brief GPU max-pool for a whole inverted list's clusters in one launch. - * - * Given the flattened cluster layout of one inverted list (@p docs with - * @p offsets, where cluster b owns docs[offsets[b]..offsets[b+1])), computes - * on the GPU, per cluster and per term, the maximum weight across the - * cluster's documents — the per-term max-pool that dominates summarize(). - * - * Processing the entire list in a single kernel launch + single stream sync - * (one thread block per cluster) is essential: clusters average only a - * handful of documents, so a per-cluster launch/sync is dominated by GPU - * round-trip overhead. The tie-order-sensitive sort/truncate stays on the - * CPU. Requires the corpus resident (shares assign()'s cache); float (U32) - * only. Returns false (caller falls back to CPU) if the GPU is unavailable - * or the list is empty. - * - * @param n_clusters offsets.size() - 1. - * @param out resized to n_clusters; out[b] holds cluster b's result. - */ - bool summarize_list_maxpool(const SparseVectors* vectors, - const idx_t* docs, const idx_t* offsets, - size_t n_clusters, - std::vector& out); - GpuClusterAssigner(const GpuClusterAssigner&) = delete; GpuClusterAssigner& operator=(const GpuClusterAssigner&) = delete; private: - GpuClusterAssigner(); - ~GpuClusterAssigner(); - - struct Impl; - // Uploads the corpus to the GPU on first use (or when it changes) and - // returns an opaque handle to the resident device buffers. Thread-safe. - const void* ensure_corpus_resident(const SparseVectors* vectors); - - Impl* impl_; + GpuClusterAssigner() = default; }; /** - * @brief Gate for the GPU assignment path in map_docs_to_clusters(). - * - * Returns true when the library was built with GPU support and a usable device - * is present (and the assignment has at least two clusters). Offload is - * unconditional beyond that — there is no runtime size threshold. + * @brief Whether to run assignment on the GPU: built with GPU support, a device + * is present, and the list has >= 2 clusters. No runtime size threshold. */ bool should_offload_assignment_to_gpu(size_t n_docs, size_t n_clusters); -/** - * @brief Gate for the GPU summarize() max-pool offload. - * - * Opt-in: returns true only when the library was built with GPU support, a - * usable device is present, and NSPARSE_GPU_SUMMARIZE=1 is set. The default - * (unset) keeps summarize on the CPU flat-array path, which is faster on - * high-core hosts and bit-identical. Enabling the GPU offload is a net win on - * GPU-rich / low-core hosts where CPU summarize dominates build wall-time. - */ -bool should_offload_summarize_to_gpu(); - } // namespace nsparse::detail #endif // GPU_CLUSTER_ASSIGNER_H diff --git a/nsparse/gpu/gpu_common.cuh b/nsparse/gpu/gpu_common.cuh new file mode 100644 index 0000000..0a9f18d --- /dev/null +++ b/nsparse/gpu/gpu_common.cuh @@ -0,0 +1,149 @@ +/** + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +#ifndef NSPARSE_GPU_COMMON_CUH +#define NSPARSE_GPU_COMMON_CUH + +#include + +#include +#include +#include +#include + +#include "cuda_runtime.h" +#include "nsparse/sparse_vectors.h" +#include "nsparse/types.h" + +namespace nsparse::detail { + +inline void check_cuda(cudaError_t status, const char* what) { + if (status != cudaSuccess) { + std::ostringstream oss; + oss << "CUDA error in " << what << ": " << cudaGetErrorString(status); + throw std::runtime_error(oss.str()); + } +} + +inline void check_cusparse(cusparseStatus_t status, const char* what) { + if (status != CUSPARSE_STATUS_SUCCESS) { + std::ostringstream oss; + oss << "cuSPARSE error in " << what << ": " + << cusparseGetErrorString(status); + throw std::runtime_error(oss.str()); + } +} + +#define NSPARSE_CUDA_CHECK(expr) check_cuda((expr), #expr) +#define NSPARSE_CUSPARSE_CHECK(expr) check_cusparse((expr), #expr) + +// Grow *ptr only when it must; capacity tracked in bytes. +template +void ensure_capacity(T** ptr, size_t& cap_bytes, size_t need_bytes) { + if (cap_bytes < need_bytes) { + cudaFree(*ptr); + void* p = nullptr; + check_cuda(cudaMalloc(&p, need_bytes), "cudaMalloc(ensure_capacity)"); + *ptr = static_cast(p); + cap_bytes = need_bytes; + } +} + +// The full corpus uploaded to the GPU once and shared read-only by the +// assigner and the summarizer (a single ~corpus-sized allocation, never +// duplicated). Indices are widened to int32 to feed both the gather kernel and +// cuSPARSE's 32-bit CSR indices. +struct DeviceCorpus { + int32_t* indptr = nullptr; // n_vectors + 1 + int32_t* indices = nullptr; // nnz + float* values = nullptr; // nnz + int64_t nnz = 0; + size_t dim = 0; + size_t n_vectors = 0; + // Residency identity. A pointer alone is unsafe (a freed SparseVectors + // address can be reused), so also compare n_vectors and nnz. + const SparseVectors* owner = nullptr; + + bool matches(const SparseVectors* v, size_t nv, int64_t nz) const { + return owner == v && n_vectors == nv && nnz == nz; + } +}; + +// Process-wide resident corpus, shared by all GPU build helpers. +class GpuCorpus { +public: + static GpuCorpus& instance() { + static GpuCorpus c; + return c; + } + + // Upload the corpus on first use (or when it changes) and return it. + // Thread-safe: the first caller uploads while the rest wait; later calls + // are a cheap identity check. + const DeviceCorpus& ensure_resident(const SparseVectors* vectors) { + const size_t n_vectors = vectors->num_vectors(); + const idx_t* indptr = vectors->indptr_data(); + const term_t* indices = vectors->indices_data(); + const float* values = vectors->values_data_float(); + const int64_t nnz = indptr[n_vectors]; + + std::lock_guard lock(mutex_); + if (corpus_.matches(vectors, n_vectors, nnz)) { + return corpus_; + } + free_locked(); + + std::vector indices32(static_cast(nnz)); + for (int64_t i = 0; i < nnz; ++i) { + indices32[i] = static_cast(indices[i]); + } + check_cuda(cudaMalloc(&corpus_.indptr, (n_vectors + 1) * sizeof(int32_t)), + "cudaMalloc(corpus.indptr)"); + check_cuda(cudaMalloc(&corpus_.indices, nnz * sizeof(int32_t)), + "cudaMalloc(corpus.indices)"); + check_cuda(cudaMalloc(&corpus_.values, nnz * sizeof(float)), + "cudaMalloc(corpus.values)"); + check_cuda(cudaMemcpy(corpus_.indptr, indptr, + (n_vectors + 1) * sizeof(int32_t), + cudaMemcpyHostToDevice), + "cudaMemcpy(corpus.indptr)"); + check_cuda(cudaMemcpy(corpus_.indices, indices32.data(), + nnz * sizeof(int32_t), cudaMemcpyHostToDevice), + "cudaMemcpy(corpus.indices)"); + check_cuda(cudaMemcpy(corpus_.values, values, nnz * sizeof(float), + cudaMemcpyHostToDevice), + "cudaMemcpy(corpus.values)"); + corpus_.nnz = nnz; + corpus_.dim = vectors->get_dimension(); + corpus_.n_vectors = n_vectors; + corpus_.owner = vectors; + return corpus_; + } + + GpuCorpus(const GpuCorpus&) = delete; + GpuCorpus& operator=(const GpuCorpus&) = delete; + +private: + GpuCorpus() = default; + ~GpuCorpus() { free_locked(); } + + void free_locked() { + cudaFree(corpus_.indptr); + cudaFree(corpus_.indices); + cudaFree(corpus_.values); + corpus_ = DeviceCorpus{}; + } + + std::mutex mutex_; + DeviceCorpus corpus_; +}; + +} // namespace nsparse::detail + +#endif // NSPARSE_GPU_COMMON_CUH diff --git a/nsparse/gpu/gpu_summarizer.cu b/nsparse/gpu/gpu_summarizer.cu new file mode 100644 index 0000000..862000a --- /dev/null +++ b/nsparse/gpu/gpu_summarizer.cu @@ -0,0 +1,283 @@ +/** + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +#include "nsparse/gpu/gpu_summarizer.h" + +#include +#include +#include +#include + +#include "cuda_runtime.h" +#include "nsparse/gpu/gpu_common.cuh" +#include "nsparse/sparse_vectors.h" +#include "nsparse/types.h" + +namespace nsparse::detail { +namespace { + +// Batched per-list max-pool: one block per cluster, so a whole inverted list is +// summarized in one launch + one sync (a per-cluster launch is dominated by GPU +// round-trip overhead). Clusters are independent, each owned by one block, so +// only intra-block sync is needed. SPLADE weights are >= 0, so a non-negative +// float orders monotonically as its reinterpreted int32 — allowing atomicMax on +// the accumulator's int view. +// +// Buffers (per thread, reused across lists, left clean via touched-reset): +// acc[n_clusters*dim] / seen[n_clusters*dim] : per-cluster per-term max / flag, +// block b owning region [b*dim, (b+1)*dim). Start zeroed; touched slots +// are restored to zero before return. +// out_term / out_val : compacted results, per cluster at out_base[b]. +// out_count[b] / out_sum[b] : distinct terms and sum of maxes for cluster b. +__global__ void summarize_list_kernel( + const int32_t* __restrict__ corpus_indptr, + const int32_t* __restrict__ corpus_indices, + const float* __restrict__ corpus_values, + const int32_t* __restrict__ docs, const int32_t* __restrict__ offsets, + const int32_t* __restrict__ out_base, int dim, + int32_t* __restrict__ acc, int32_t* __restrict__ seen, + int32_t* __restrict__ out_term, float* __restrict__ out_val, + int32_t* __restrict__ out_count, float* __restrict__ out_sum) { + const int b = blockIdx.x; // one block per cluster + const int start = offsets[b]; + const int end = offsets[b + 1]; + const int64_t base = static_cast(b) * dim; + const int32_t obase = out_base[b]; + + __shared__ int s_count; + if (threadIdx.x == 0) { + s_count = 0; + } + __syncthreads(); + + // Scatter-max over the cluster's documents; each distinct term is appended + // once (via seen[]) to this cluster's output segment. + for (int i = start + threadIdx.x; i < end; i += blockDim.x) { + const int32_t d = docs[i]; + const int32_t src = corpus_indptr[d]; + const int32_t len = corpus_indptr[d + 1] - src; + for (int32_t t = 0; t < len; ++t) { + const int32_t term = corpus_indices[src + t]; + const float v = corpus_values[src + t]; + atomicMax(&acc[base + term], __float_as_int(v)); + if (atomicCAS(&seen[base + term], 0, 1) == 0) { + out_term[obase + atomicAdd(&s_count, 1)] = term; + } + } + } + __syncthreads(); + + // Read each touched term's max, accumulate the sum, and reset acc/seen for + // reuse by the next list. + __shared__ float s_sum; + if (threadIdx.x == 0) { + s_sum = 0.0F; + out_count[b] = s_count; + } + __syncthreads(); + + for (int k = threadIdx.x; k < s_count; k += blockDim.x) { + const int32_t term = out_term[obase + k]; + const float v = __int_as_float(acc[base + term]); + out_val[obase + k] = v; + atomicAdd(&s_sum, v); + acc[base + term] = 0; + seen[base + term] = 0; + } + __syncthreads(); + if (threadIdx.x == 0) { + out_sum[b] = s_sum; + } +} + +// Per-thread scratch (reused across lists, grown on demand). acc/seen are +// n_clusters x dim and kept zeroed between lists via the kernel's touched-reset. +struct ThreadCtx { + cudaStream_t stream{}; + bool init = false; + + int32_t* d_docs = nullptr; size_t docs_cap = 0; + int32_t* d_offsets = nullptr; size_t offsets_cap = 0; + int32_t* d_out_base = nullptr; size_t out_base_cap = 0; + int32_t* d_acc = nullptr; size_t acc_cap = 0; + int32_t* d_seen = nullptr; size_t seen_cap = 0; + int32_t* d_out_term = nullptr; size_t out_term_cap = 0; + float* d_out_val = nullptr; size_t out_val_cap = 0; + int32_t* d_out_count = nullptr; size_t out_count_cap = 0; + float* d_out_sum = nullptr; size_t out_sum_cap = 0; + + ~ThreadCtx() { + cudaFree(d_docs); + cudaFree(d_offsets); + cudaFree(d_out_base); + cudaFree(d_acc); + cudaFree(d_seen); + cudaFree(d_out_term); + cudaFree(d_out_val); + cudaFree(d_out_count); + cudaFree(d_out_sum); + if (init) { + cudaStreamDestroy(stream); + } + } +}; + +thread_local std::unique_ptr t_ctx; + +ThreadCtx& thread_ctx() { + if (t_ctx == nullptr) { + t_ctx = std::make_unique(); + check_cuda(cudaStreamCreate(&t_ctx->stream), "cudaStreamCreate"); + t_ctx->init = true; + } + return *t_ctx; +} + +bool gpu_present() { + int device_count = 0; + cudaError_t status = cudaGetDeviceCount(&device_count); + return status == cudaSuccess && device_count > 0; +} + +} // namespace + +GpuSummarizer& GpuSummarizer::instance() { + static GpuSummarizer singleton; + return singleton; +} + +bool GpuSummarizer::available() { + static const bool present = gpu_present(); + return present; +} + +bool GpuSummarizer::summarize_list(const SparseVectors* vectors, + const idx_t* docs, const idx_t* offsets, + size_t n_clusters, + std::vector& out) { + if (!available() || n_clusters == 0) { + return false; + } + const size_t dim = vectors->get_dimension(); + const idx_t* indptr = vectors->indptr_data(); + + const size_t n_docs = static_cast(offsets[n_clusters] - offsets[0]); + if (n_docs == 0) { + return false; + } + // Per-cluster nnz prefix: places each cluster's compact output in a private + // segment; the total is the output capacity. + std::vector h_out_base(n_clusters + 1, 0); + for (size_t b = 0; b < n_clusters; ++b) { + int64_t nnz_b = 0; + for (idx_t i = offsets[b]; i < offsets[b + 1]; ++i) { + const idx_t d = docs[i]; + nnz_b += indptr[d + 1] - indptr[d]; + } + h_out_base[b + 1] = h_out_base[b] + static_cast(nnz_b); + } + const int64_t total_nnz = h_out_base[n_clusters]; + if (total_nnz == 0) { + return false; + } + + const DeviceCorpus& corpus = GpuCorpus::instance().ensure_resident(vectors); + ThreadCtx& ctx = thread_ctx(); + cudaStream_t stream = ctx.stream; + + // acc/seen are n_clusters x dim; must start (and remain) zeroed. Zero only + // when (re)grown — the kernel restores touched slots. + const size_t acc_bytes = n_clusters * dim * sizeof(int32_t); + const bool acc_grew = ctx.acc_cap < acc_bytes; + ensure_capacity(&ctx.d_acc, ctx.acc_cap, acc_bytes); + ensure_capacity(&ctx.d_seen, ctx.seen_cap, acc_bytes); + if (acc_grew) { + NSPARSE_CUDA_CHECK(cudaMemsetAsync(ctx.d_acc, 0, ctx.acc_cap, stream)); + NSPARSE_CUDA_CHECK(cudaMemsetAsync(ctx.d_seen, 0, ctx.seen_cap, stream)); + } + ensure_capacity(&ctx.d_docs, ctx.docs_cap, n_docs * sizeof(int32_t)); + ensure_capacity(&ctx.d_offsets, ctx.offsets_cap, + (n_clusters + 1) * sizeof(int32_t)); + ensure_capacity(&ctx.d_out_base, ctx.out_base_cap, + (n_clusters + 1) * sizeof(int32_t)); + ensure_capacity(&ctx.d_out_term, ctx.out_term_cap, + static_cast(total_nnz) * sizeof(int32_t)); + ensure_capacity(&ctx.d_out_val, ctx.out_val_cap, + static_cast(total_nnz) * sizeof(float)); + ensure_capacity(&ctx.d_out_count, ctx.out_count_cap, + n_clusters * sizeof(int32_t)); + ensure_capacity(&ctx.d_out_sum, ctx.out_sum_cap, n_clusters * sizeof(float)); + + // Rebase offsets to 0 for the flat doc upload. + std::vector h_offsets(n_clusters + 1); + const idx_t base0 = offsets[0]; + for (size_t b = 0; b <= n_clusters; ++b) { + h_offsets[b] = static_cast(offsets[b] - base0); + } + + NSPARSE_CUDA_CHECK(cudaMemcpyAsync(ctx.d_docs, docs + base0, + n_docs * sizeof(int32_t), + cudaMemcpyHostToDevice, stream)); + NSPARSE_CUDA_CHECK(cudaMemcpyAsync(ctx.d_offsets, h_offsets.data(), + (n_clusters + 1) * sizeof(int32_t), + cudaMemcpyHostToDevice, stream)); + NSPARSE_CUDA_CHECK(cudaMemcpyAsync(ctx.d_out_base, h_out_base.data(), + (n_clusters + 1) * sizeof(int32_t), + cudaMemcpyHostToDevice, stream)); + + constexpr int kBlock = 128; + summarize_list_kernel<<(n_clusters), kBlock, 0, stream>>>( + corpus.indptr, corpus.indices, corpus.values, ctx.d_docs, ctx.d_offsets, + ctx.d_out_base, static_cast(dim), ctx.d_acc, ctx.d_seen, + ctx.d_out_term, ctx.d_out_val, ctx.d_out_count, ctx.d_out_sum); + NSPARSE_CUDA_CHECK(cudaGetLastError()); + + std::vector h_term(total_nnz); + std::vector h_val(total_nnz); + std::vector h_count(n_clusters); + std::vector h_sum(n_clusters); + NSPARSE_CUDA_CHECK(cudaMemcpyAsync(h_term.data(), ctx.d_out_term, + total_nnz * sizeof(int32_t), + cudaMemcpyDeviceToHost, stream)); + NSPARSE_CUDA_CHECK(cudaMemcpyAsync(h_val.data(), ctx.d_out_val, + total_nnz * sizeof(float), + cudaMemcpyDeviceToHost, stream)); + NSPARSE_CUDA_CHECK(cudaMemcpyAsync(h_count.data(), ctx.d_out_count, + n_clusters * sizeof(int32_t), + cudaMemcpyDeviceToHost, stream)); + NSPARSE_CUDA_CHECK(cudaMemcpyAsync(h_sum.data(), ctx.d_out_sum, + n_clusters * sizeof(float), + cudaMemcpyDeviceToHost, stream)); + NSPARSE_CUDA_CHECK(cudaStreamSynchronize(stream)); + + out.resize(n_clusters); + for (size_t b = 0; b < n_clusters; ++b) { + const int32_t base = h_out_base[b]; + const int32_t cnt = h_count[b]; + ClusterSummary& cs = out[b]; + cs.terms.resize(cnt); + cs.values.resize(cnt); + for (int32_t k = 0; k < cnt; ++k) { + cs.terms[k] = static_cast(h_term[base + k]); + cs.values[k] = h_val[base + k]; + } + cs.sum = h_sum[b]; + } + return true; +} + +bool should_offload_summarize_to_gpu() { + if (!GpuSummarizer::available()) { + return false; + } + const char* env = std::getenv("NSPARSE_GPU_SUMMARIZE"); + return env != nullptr && env[0] == '1'; +} + +} // namespace nsparse::detail diff --git a/nsparse/gpu/gpu_summarizer.h b/nsparse/gpu/gpu_summarizer.h new file mode 100644 index 0000000..48f7977 --- /dev/null +++ b/nsparse/gpu/gpu_summarizer.h @@ -0,0 +1,79 @@ +/** + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +#ifndef GPU_SUMMARIZER_H +#define GPU_SUMMARIZER_H + +#include + +#include "nsparse/sparse_vectors.h" +#include "nsparse/types.h" + +namespace nsparse::detail { + +/** + * @brief GPU per-term max-pool for the summarize() step of the build() path. + * + * Computes, per cluster and per term, the maximum weight over the cluster's + * documents. The tie-order-sensitive sort/truncate stays on the CPU, so this + * offloads only the max-pool. A whole inverted list is processed in one kernel + * launch (one block per cluster): clusters average a handful of documents, so a + * per-cluster launch is dominated by GPU round-trip overhead. float (U32) only. + * + * Compiled only with -DNSPARSE_ENABLE_GPU=ON (defines NSPARSE_WITH_GPU). Shares + * the resident corpus with GpuClusterAssigner via GpuCorpus. + */ +class GpuSummarizer { +public: + /// Process-wide singleton owning the per-thread CUDA streams. + static GpuSummarizer& instance(); + + /// True when built with GPU support and a usable device is present. + static bool available(); + + /// Per-cluster max-pool result. Terms are in touched-append order (not + /// sorted); the CPU sort/truncate handles ordering. + struct ClusterSummary { + std::vector terms; + std::vector values; + float sum = 0.0F; + }; + + /** + * @brief Max-pool a whole inverted list's clusters in one launch. + * + * @param vectors corpus of sparse vectors (CSR-backed). + * @param docs flattened doc ids for the list. + * @param offsets cluster b owns docs[offsets[b]..offsets[b+1]). + * @param n_clusters offsets.size() - 1. + * @param out resized to n_clusters; out[b] is cluster b's result. + * @return false (caller falls back to CPU) if unavailable or list empty. + */ + bool summarize_list(const SparseVectors* vectors, const idx_t* docs, + const idx_t* offsets, size_t n_clusters, + std::vector& out); + + GpuSummarizer(const GpuSummarizer&) = delete; + GpuSummarizer& operator=(const GpuSummarizer&) = delete; + +private: + GpuSummarizer() = default; +}; + +/** + * @brief Whether to run summarize()'s max-pool on the GPU. Opt-in via + * NSPARSE_GPU_SUMMARIZE=1. Default (unset) uses the CPU flat-array path, which + * is faster on high-core hosts and bit-identical. The GPU offload is a net win + * on GPU-rich / low-core hosts where CPU summarize dominates wall-time. + */ +bool should_offload_summarize_to_gpu(); + +} // namespace nsparse::detail + +#endif // GPU_SUMMARIZER_H diff --git a/nsparse/seismic_common.h b/nsparse/seismic_common.h index de3d9e6..a07659a 100644 --- a/nsparse/seismic_common.h +++ b/nsparse/seismic_common.h @@ -10,16 +10,10 @@ #ifndef SEISMIC_COMMON_H #define SEISMIC_COMMON_H -#include -#include #include #include #include -#ifdef _OPENMP -#include -#endif - #include "nsparse/cluster/inverted_list_clusters.h" #include "nsparse/cluster/random_kmeans.h" #include "nsparse/id_selector.h" @@ -146,51 +140,17 @@ inline std::vector build_inverted_lists_clusters( std::vector clustered_inverted_lists( inverted_lists_size); - // Optional coarse phase timing, enabled by NSPARSE_GPU_PROFILE=1, to see - // how build CPU time splits across prune / train (k-means assignment) / - // summarize. Accumulates thread-seconds (sum across threads). - const bool build_prof = - std::getenv("NSPARSE_GPU_PROFILE") != nullptr && - std::getenv("NSPARSE_GPU_PROFILE")[0] == '1'; - double t_prune = 0.0; - double t_train = 0.0; - double t_summ = 0.0; - -#pragma omp parallel for schedule(dynamic, 64) \ - reduction(+ : t_prune, t_train, t_summ) +#pragma omp parallel for schedule(dynamic, 64) for (int64_t idx = 0; idx < static_cast(inverted_lists_size); ++idx) { auto& invlist = (*inverted_lists)[idx]; - if (build_prof) { - double a = omp_get_wtime(); - const auto doc_ids = invlist.prune_and_keep_doc_ids(lambda); - double b = omp_get_wtime(); - InvertedListClusters inverted_list_clusters( - detail::RandomKMeans::train(vectors, doc_ids, beta)); - double c = omp_get_wtime(); - inverted_list_clusters.summarize(vectors, - seismic_cluster_params.alpha); - double d = omp_get_wtime(); - t_prune += b - a; - t_train += c - b; - t_summ += d - c; - clustered_inverted_lists[idx] = std::move(inverted_list_clusters); - } else { - const auto& doc_ids = invlist.prune_and_keep_doc_ids(lambda); - InvertedListClusters inverted_list_clusters( - detail::RandomKMeans::train(vectors, doc_ids, beta)); - inverted_list_clusters.summarize(vectors, - seismic_cluster_params.alpha); - clustered_inverted_lists[idx] = std::move(inverted_list_clusters); - } + const auto& doc_ids = invlist.prune_and_keep_doc_ids(lambda); + InvertedListClusters inverted_list_clusters( + detail::RandomKMeans::train(vectors, doc_ids, beta)); + inverted_list_clusters.summarize(vectors, seismic_cluster_params.alpha); + clustered_inverted_lists[idx] = std::move(inverted_list_clusters); invlist.clear(); } - if (build_prof) { - std::fprintf(stderr, - "[nsparse build] thread-seconds: prune=%.1f train=%.1f " - "summarize=%.1f\n", - t_prune, t_train, t_summ); - } return clustered_inverted_lists; } diff --git a/tests/gpu_cluster_assigner_test.cpp b/tests/gpu_cluster_assigner_test.cpp index 8b67227..c4e9ab6 100644 --- a/tests/gpu_cluster_assigner_test.cpp +++ b/tests/gpu_cluster_assigner_test.cpp @@ -18,6 +18,7 @@ #include #include "nsparse/cluster/kmeans_utils.h" +#include "nsparse/gpu/gpu_summarizer.h" #include "nsparse/sparse_vectors.h" #include "nsparse/types.h" @@ -186,7 +187,7 @@ void cpu_reference_maxpool(const SparseVectors* vectors, // Compares GPU per-cluster max-pool against the CPU reference for a whole list // of several clusters processed in one batched call. TEST(GpuClusterAssignerTest, SummarizeListMaxpoolMatchesCpu) { - if (!GpuClusterAssigner::available()) { + if (!GpuSummarizer::available()) { GTEST_SKIP() << "No CUDA-capable GPU available"; } constexpr size_t kNumDocs = 8000; @@ -207,8 +208,8 @@ TEST(GpuClusterAssignerTest, SummarizeListMaxpoolMatchesCpu) { } const size_t n_clusters = offsets.size() - 1; - std::vector gpu_out; - ASSERT_TRUE(GpuClusterAssigner::instance().summarize_list_maxpool( + std::vector gpu_out; + ASSERT_TRUE(GpuSummarizer::instance().summarize_list( &vectors, flat_docs.data(), offsets.data(), n_clusters, gpu_out)); ASSERT_EQ(gpu_out.size(), n_clusters); @@ -241,26 +242,26 @@ TEST(GpuClusterAssignerTest, SummarizeListMaxpoolMatchesCpu) { // A second list reusing the buffers must not leak accumulator state from the // first — verifies the touched-slot reset path. TEST(GpuClusterAssignerTest, SummarizeListMaxpoolReuseIsClean) { - if (!GpuClusterAssigner::available()) { + if (!GpuSummarizer::available()) { GTEST_SKIP() << "No CUDA-capable GPU available"; } SparseVectors vectors = make_random_corpus(5000, 2500, 40, 21); - auto& gpu = GpuClusterAssigner::instance(); + auto& gpu = GpuSummarizer::instance(); // First list. std::vector docs1; for (idx_t d = 0; d < 2000; ++d) docs1.push_back(d); std::vector off1 = {0, 800, 2000}; - std::vector out1; - gpu.summarize_list_maxpool(&vectors, docs1.data(), off1.data(), 2, out1); + std::vector out1; + gpu.summarize_list(&vectors, docs1.data(), off1.data(), 2, out1); // Second, different list must match its own CPU reference. std::vector docs2; for (idx_t d = 2000; d < 5000; ++d) docs2.push_back(d); std::vector off2 = {0, 1500, 3000}; - std::vector out2; + std::vector out2; ASSERT_TRUE( - gpu.summarize_list_maxpool(&vectors, docs2.data(), off2.data(), 2, + gpu.summarize_list(&vectors, docs2.data(), off2.data(), 2, out2)); for (size_t b = 0; b < 2; ++b) { std::vector cluster(docs2.begin() + off2[b], From 33d45ad09610ff293bfa81904cd92eeae76237cf Mon Sep 17 00:00:00 2001 From: Liyun Xiu Date: Fri, 10 Jul 2026 05:33:46 +0000 Subject: [PATCH 7/8] Refactor summarize_ into a single CPU/GPU dispatch Replace the interleaved per-cluster GPU checks with a thin dispatcher that picks one path: summarize_with_gpu_ (opt-in, float-only, returns nullopt to fall back) or summarize_with_cpu_. Both feed the shared summarize_emit_cluster_ (sort / alpha-truncate / emit), so output is identical. No behavior change. All 365 tests pass in both summarize modes. Signed-off-by: Liyun Xiu --- nsparse/cluster/inverted_list_clusters.cpp | 248 +++++++++++---------- 1 file changed, 134 insertions(+), 114 deletions(-) diff --git a/nsparse/cluster/inverted_list_clusters.cpp b/nsparse/cluster/inverted_list_clusters.cpp index e8d7214..0901a8b 100644 --- a/nsparse/cluster/inverted_list_clusters.cpp +++ b/nsparse/cluster/inverted_list_clusters.cpp @@ -11,9 +11,10 @@ #include #include +#include #include #include -#include +#include #include #include "nsparse/sparse_vectors.h" @@ -25,141 +26,160 @@ namespace nsparse { namespace { -/** - * @brief Generage summary sparse vector for posting lists - * - * @param vectors inverted index - * @param group_of_doc_ids a list of posting list - * @param alpha prune ratio - * @return SparseVectors - */ +// Prune a cluster's (term, max-value) pairs at the alpha mass cutoff and append +// the survivors (in ascending term order) as a summary vector. Shared by the +// CPU and GPU max-pool paths so their output is identical. template -SparseVectors summarize_(const SparseVectors* vectors, - const std::vector& group_of_doc_ids, - const std::vector& offsets, float alpha) { - SparseVectors summarized_vectors( - {.element_size = vectors->get_element_size(), - .dimension = vectors->get_dimension()}); - if (offsets.size() <= 1) { - return summarized_vectors; +void summarize_emit_cluster_(std::vector>& pairs, + float sum, float alpha, SparseVectors& out) { + std::ranges::sort(pairs, [](const auto& a, const auto& b) { + return a.second > b.second; + }); + float addup = 0.0F; + for (size_t j = 0; j < pairs.size(); ++j) { + addup += pairs[j].second; + if (addup / sum >= alpha) { + pairs.erase(pairs.begin() + j + 1, pairs.end()); + break; + } } - const auto element_size = vectors->get_element_size(); + std::ranges::sort(pairs, [](const auto& a, const auto& b) { + return a.first < b.first; + }); + std::vector terms; + std::vector values; + terms.reserve(pairs.size()); + values.reserve(pairs.size()); + for (const auto& [term, value] : pairs) { + terms.push_back(term); + values.push_back(value); + } + out.add_vector(terms.data(), terms.size(), + reinterpret_cast(values.data()), + values.size() * sizeof(T)); +} + +// CPU per-term max-pool over a flat, dim-sized accumulator reused across all +// clusters (replaces a per-cluster std::unordered_map, ~2.8x cheaper for the +// many small clusters). A per-cluster epoch marks live slots, so the buffers +// need no per-cluster reset; first touch is detected via the epoch (not +// acc==0), robust to zero-valued weights. +template +SparseVectors summarize_with_cpu_(const SparseVectors* vectors, + const std::vector& group_of_doc_ids, + const std::vector& offsets, + float alpha) { + SparseVectors out({.element_size = vectors->get_element_size(), + .dimension = vectors->get_dimension()}); + const size_t n_clusters = offsets.size() - 1; const auto& indptr_data = vectors->indptr_data(); const auto& indices_data = vectors->indices_data(); const auto& values_data = vectors->values_data(); - - // Per-term max-pool over a flat, dim-sized accumulator reused across all - // clusters (replaces a per-cluster std::unordered_map, ~2.8x cheaper for - // the many small clusters). A per-cluster epoch marks which slots are live, - // so the dim-sized buffers need no per-cluster reset; first touch is - // detected via the epoch (not acc==0), robust to zero-valued weights. const size_t dim = vectors->get_dimension(); + std::vector acc(dim, T(0)); std::vector epoch(dim, 0); std::vector touched; uint32_t cur_epoch = 0; - // The max-pool may instead be offloaded to the GPU (one launch per list) - // when NSPARSE_GPU_SUMMARIZE=1; downstream sort/truncate is identical. -#ifdef NSPARSE_WITH_GPU - bool gpu_ok = false; - std::vector gpu_clusters; - if constexpr (std::is_same_v) { - if (detail::should_offload_summarize_to_gpu()) { - gpu_ok = detail::GpuSummarizer::instance().summarize_list( - vectors, group_of_doc_ids.data(), offsets.data(), - offsets.size() - 1, gpu_clusters); - } - } -#endif - - for (size_t i = 0; i < offsets.size() - 1; ++i) { - size_t n_docs = offsets[i + 1] - offsets[i]; + for (size_t i = 0; i < n_clusters; ++i) { + ++cur_epoch; + touched.clear(); float sum = 0.0F; - std::vector> summary_vec; - - bool used_gpu = false; -#ifdef NSPARSE_WITH_GPU - if constexpr (std::is_same_v) { - if (gpu_ok) { - const auto& gc = gpu_clusters[i]; - summary_vec.reserve(gc.terms.size()); - for (size_t t = 0; t < gc.terms.size(); ++t) { - summary_vec.emplace_back(gc.terms[t], gc.values[t]); + auto doc_ids = std::span( + group_of_doc_ids.data() + offsets[i], offsets[i + 1] - offsets[i]); + for (const auto& doc_id : doc_ids) { + int start = indptr_data[doc_id]; + int end = indptr_data[doc_id + 1]; + for (size_t j = start; j < end; ++j) { + const term_t term = indices_data[j]; + // j is element index, need byte offset for T access + const T v = + *reinterpret_cast(values_data + j * sizeof(T)); + if (epoch[term] != cur_epoch) { + // First occurrence in this cluster; value = max(0, v). + epoch[term] = cur_epoch; + const T value = std::max(T(0), v); + acc[term] = value; + touched.push_back(term); + sum += value; + } else if (v > acc[term]) { + sum += v - acc[term]; + acc[term] = v; } - sum = gc.sum; - used_gpu = true; } } -#endif - if (!used_gpu) { - ++cur_epoch; - touched.clear(); - auto doc_ids = std::span( - group_of_doc_ids.data() + offsets[i], n_docs); - for (const auto& doc_id : doc_ids) { - int start = indptr_data[doc_id]; - int end = indptr_data[doc_id + 1]; - for (size_t j = start; j < end; ++j) { - const term_t term = indices_data[j]; - // j is element index, need byte offset for T access - const T v = *reinterpret_cast( - values_data + j * sizeof(T)); - if (epoch[term] != cur_epoch) { - // First occurrence in this cluster; value = max(0, v). - epoch[term] = cur_epoch; - const T value = std::max(T(0), v); - acc[term] = value; - touched.push_back(term); - sum += value; - } else if (v > acc[term]) { - sum += v - acc[term]; - acc[term] = v; - } - } - } - summary_vec.reserve(touched.size()); - for (const term_t term : touched) { - summary_vec.emplace_back(term, acc[term]); - } + std::vector> pairs; + pairs.reserve(touched.size()); + for (const term_t term : touched) { + pairs.emplace_back(term, acc[term]); } + summarize_emit_cluster_(pairs, sum, alpha, out); + } + return out; +} - // Sort by value in descending order - std::ranges::sort(summary_vec, [](const auto& a, const auto& b) { - return a.second > b.second; - }); - - float addup = 0.0F; - for (int j = 0; j < summary_vec.size(); ++j) { - addup += summary_vec[j].second; - if (addup / sum >= alpha) { - summary_vec.erase(summary_vec.begin() + j + 1, - summary_vec.end()); - break; +#ifdef NSPARSE_WITH_GPU +// GPU per-term max-pool: one kernel launch for the whole list, then the shared +// CPU sort/truncate. float (U32) only. Returns std::nullopt when the GPU +// declines (unavailable / empty list) so the caller falls back to the CPU path. +template +std::optional summarize_with_gpu_( + const SparseVectors* vectors, const std::vector& group_of_doc_ids, + const std::vector& offsets, float alpha) { + if constexpr (!std::is_same_v) { + return std::nullopt; // GPU max-pool is float-only + } else { + const size_t n_clusters = offsets.size() - 1; + std::vector gpu_clusters; + if (!detail::GpuSummarizer::instance().summarize_list( + vectors, group_of_doc_ids.data(), offsets.data(), n_clusters, + gpu_clusters)) { + return std::nullopt; + } + SparseVectors out({.element_size = vectors->get_element_size(), + .dimension = vectors->get_dimension()}); + for (size_t i = 0; i < n_clusters; ++i) { + const auto& gc = gpu_clusters[i]; + std::vector> pairs; + pairs.reserve(gc.terms.size()); + for (size_t t = 0; t < gc.terms.size(); ++t) { + pairs.emplace_back(gc.terms[t], gc.values[t]); } + summarize_emit_cluster_(pairs, gc.sum, alpha, out); } + return out; + } +} +#endif - // Sort by term_t order - std::ranges::sort(summary_vec, [](const auto& a, const auto& b) { - return a.first < b.first; - }); - - // Break into separate terms and values vectors - std::vector terms; - std::vector values; - terms.reserve(summary_vec.size()); - values.reserve(summary_vec.size()); - for (const auto& [term, value] : summary_vec) { - terms.push_back(term); - values.push_back(value); +/** + * @brief Generate the summary sparse vectors for a list's posting-list clusters. + * + * @param vectors inverted index + * @param group_of_doc_ids flattened doc ids of all clusters + * @param offsets cluster boundaries into group_of_doc_ids + * @param alpha prune ratio + */ +template +SparseVectors summarize_(const SparseVectors* vectors, + const std::vector& group_of_doc_ids, + const std::vector& offsets, float alpha) { + if (offsets.size() <= 1) { + return SparseVectors({.element_size = vectors->get_element_size(), + .dimension = vectors->get_dimension()}); + } +#ifdef NSPARSE_WITH_GPU + // GPU max-pool is opt-in (NSPARSE_GPU_SUMMARIZE=1); fall back to CPU when + // disabled, unsupported for T, or the GPU declines. + if (detail::should_offload_summarize_to_gpu()) { + if (auto gpu_result = summarize_with_gpu_(vectors, group_of_doc_ids, + offsets, alpha)) { + return std::move(*gpu_result); } - - summarized_vectors.add_vector( - terms.data(), terms.size(), - reinterpret_cast(values.data()), - values.size() * sizeof(T)); } - return summarized_vectors; +#endif + return summarize_with_cpu_(vectors, group_of_doc_ids, offsets, alpha); } } // namespace From cb6da4a7e5e6aa21edf7edd4282b20ddc80a67fc Mon Sep 17 00:00:00 2001 From: Liyun Xiu Date: Fri, 10 Jul 2026 05:55:58 +0000 Subject: [PATCH 8/8] Document GPU acceleration in DEVELOPER_GUIDE Add GPU build info: the NSPARSE_ENABLE_GPU option, a GPU Acceleration section (what it accelerates, requirements, build commands, the NSPARSE_CUDA_TOOLKIT_ROOT / NSPARSE_CUDA_ARCHITECTURES knobs, and the opt-in NSPARSE_GPU_SUMMARIZE env var), the CUDA/cuSPARSE dependency, and the nsparse/gpu/ directory entry. Signed-off-by: Liyun Xiu --- DEVELOPER_GUIDE.md | 41 +++++++++++++++++++++++++++++++++++++++++ 1 file changed, 41 insertions(+) diff --git a/DEVELOPER_GUIDE.md b/DEVELOPER_GUIDE.md index 38638fc..f358b8a 100644 --- a/DEVELOPER_GUIDE.md +++ b/DEVELOPER_GUIDE.md @@ -7,6 +7,7 @@ - [Build](#build) - [Build Options](#build-options) - [SIMD Optimization Levels](#simd-optimization-levels) + - [GPU Acceleration](#gpu-acceleration) - [Run Tests](#run-tests) - [Run Benchmarks](#run-benchmarks) - [Python Bindings](#python-bindings) @@ -96,6 +97,7 @@ cmake --build build -j | `NSPARSE_ENABLE_PYTHON` | `OFF` | Build Python bindings | | `NSPARSE_ENABLE_TESTS` | `OFF` | Build unit tests | | `NSPARSE_ENABLE_BENCHMARKS` | `OFF` | Build benchmarks | +| `NSPARSE_ENABLE_GPU` | `OFF` | GPU-accelerate index building via cuSPARSE (see [GPU Acceleration](#gpu-acceleration)) | Example with multiple options: ```bash @@ -116,6 +118,43 @@ The `NSPARSE_OPT_LEVEL` option controls which SIMD instruction sets are compiled > Note: ARM NEON is used automatically on ARM platforms. SVE is not supported on Apple Silicon. +### GPU Acceleration + +`NSPARSE_ENABLE_GPU=ON` compiles optional NVIDIA GPU acceleration for the +**index-building** (`build()`) path of `SeismicIndex`, using cuSPARSE plus a few +custom CUDA kernels. This offloads the two heaviest build phases — the k-means +document→centroid assignment and the `summarize()` per-term max-pool. Search is +unaffected and always runs on the CPU. The option is `OFF` by default, so +CPU-only builds and platforms without a CUDA toolkit are unaffected. + +Requirements: + +- NVIDIA CUDA Toolkit (nvcc, cuSPARSE, cudart) and a CUDA-capable GPU +- `float` (`U32`) weights — quantized (`seismic_sq`) indices fall back to the CPU + +Build: + +```bash +cmake -S . -B build -DNSPARSE_ENABLE_GPU=ON +cmake --build build -j +``` + +If the CUDA toolkit is in a non-standard location (for example, a pip-packaged +`nvidia-cu*` wheel), point CMake at it: + +```bash +cmake -S . -B build -DNSPARSE_ENABLE_GPU=ON \ + -DNSPARSE_CUDA_TOOLKIT_ROOT=/path/to/nvidia/cuXX +``` + +Target GPU architectures default to `80;89` (Ampere / Ada); override with +`-DNSPARSE_CUDA_ARCHITECTURES=...`. When built with GPU support and a device is +present, assignment is offloaded unconditionally. The `summarize()` max-pool +offload is opt-in via the `NSPARSE_GPU_SUMMARIZE=1` environment variable; the +default keeps summarize on the CPU, which is faster on high-core hosts and +produces identical output. The GPU build wiring lives in +[`nsparse/gpu/CMakeLists.txt`](nsparse/gpu/CMakeLists.txt). + ## Run Tests Build with tests enabled and run: @@ -212,6 +251,7 @@ In VS Code, you can set breakpoints and debug directly from the IDE using the te | [Google Benchmark](https://github.com/google/benchmark) | Benchmarking framework | Auto-fetched via CMake | | OpenMP | Parallelism | System package | | SWIG | Python bindings generation | System package | +| CUDA Toolkit / cuSPARSE | GPU-accelerated index building (only with `NSPARSE_ENABLE_GPU=ON`) | System or CUDA pip wheel | ## Submitting Changes @@ -232,6 +272,7 @@ Try to put new classes into existing directories if the directory name abstracts - `nsparse/invlists/` — Inverted list storage - `nsparse/io/` — Serialization and I/O - `nsparse/utils/` — Utilities (distance functions, SIMD, quantization, ranker) +- `nsparse/gpu/` — Optional CUDA/cuSPARSE build-time acceleration (`NSPARSE_ENABLE_GPU`) - `nsparse/python/` — Python bindings (SWIG) ### Modular code