diff --git a/CMakeLists.txt b/CMakeLists.txt index fd17b0b..955eba3 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -22,6 +22,12 @@ 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. 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) + add_subdirectory(nsparse) # Python bindings 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 diff --git a/benchmarks/CMakeLists.txt b/benchmarks/CMakeLists.txt index 0e7ae4e..20630d5 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_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 + 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..d0aee4f --- /dev/null +++ b/benchmarks/index_build_benchmark.cpp @@ -0,0 +1,166 @@ +/** + * 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. + * + * 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. + * + * 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_GPU +#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}; + +// Builds a fresh SeismicIndex from the shared corpus and times only build(). +// 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(); + 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()); + 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 + +// --------------------------------------------------------------------------- +// 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(benchmark::State& state) { + run_build(state); +} +BENCHMARK(BM_Seismic_Build) + ->Unit(benchmark::kMillisecond) + ->UseRealTime() + ->Iterations(env_int("NSPARSE_BENCH_ITERS", 3)) + ->Repetitions(env_int("NSPARSE_BENCH_REPS", 3)); + +BENCHMARK_MAIN(); diff --git a/nsparse/CMakeLists.txt b/nsparse/CMakeLists.txt index 8786d60..70c2613 100644 --- a/nsparse/CMakeLists.txt +++ b/nsparse/CMakeLists.txt @@ -138,6 +138,25 @@ 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_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) + include(gpu/CMakeLists.txt) + set(NSPARSE_GPU_TARGETS nsparse) + if(NSPARSE_IS_X86) + list(APPEND NSPARSE_GPU_TARGETS nsparse_avx2 nsparse_avx512) + elseif(NSPARSE_IS_ARM AND NOT APPLE) + list(APPEND NSPARSE_GPU_TARGETS nsparse_sve) + endif() + foreach(_tgt IN LISTS NSPARSE_GPU_TARGETS) + nsparse_attach_gpu(${_tgt}) + 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/inverted_list_clusters.cpp b/nsparse/cluster/inverted_list_clusters.cpp index 24ace98..0901a8b 100644 --- a/nsparse/cluster/inverted_list_clusters.cpp +++ b/nsparse/cluster/inverted_list_clusters.cpp @@ -11,97 +11,175 @@ #include #include +#include #include -#include +#include +#include #include #include "nsparse/sparse_vectors.h" #include "nsparse/types.h" +#ifdef NSPARSE_WITH_GPU +#include "nsparse/gpu/gpu_summarizer.h" +#endif 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(); - for (size_t i = 0; i < offsets.size() - 1; ++i) { - size_t n_docs = offsets[i + 1] - offsets[i]; - std::unordered_map summary_map; + 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; + + for (size_t i = 0; i < n_clusters; ++i) { + ++cur_epoch; + touched.clear(); float sum = 0.0F; auto doc_ids = std::span( - group_of_doc_ids.data() + offsets[i], n_docs); + 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 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; 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; + } } } + 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; +} - // Convert summary_map to vector of pairs - std::vector> summary_vec(summary_map.begin(), - summary_map.end()); - - // 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 diff --git a/nsparse/cluster/kmeans_utils.cpp b/nsparse/cluster/kmeans_utils.cpp index 7f8ebeb..aee9397 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_GPU +#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_GPU) && !defined(__AVX512F__) + // 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 { + 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/CMakeLists.txt b/nsparse/gpu/CMakeLists.txt new file mode 100644 index 0000000..a45df3b --- /dev/null +++ b/nsparse/gpu/CMakeLists.txt @@ -0,0 +1,73 @@ +# 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 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_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 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_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 + # (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() diff --git a/nsparse/gpu/gpu_cluster_assigner.cu b/nsparse/gpu/gpu_cluster_assigner.cu new file mode 100644 index 0000000..a4f58b2 --- /dev/null +++ b/nsparse/gpu/gpu_cluster_assigner.cu @@ -0,0 +1,297 @@ +/** + * 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 "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 { + +// 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) { + 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; +} + +// 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, + 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]; + } +} + +// 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, + 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]; + } +} + +// 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); + } + } +}; + +// 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; +} + +bool gpu_present() { + int device_count = 0; + cudaError_t status = cudaGetDeviceCount(&device_count); + return status == cudaSuccess && device_count > 0; +} + +} // namespace + +GpuClusterAssigner& GpuClusterAssigner::instance() { + static GpuClusterAssigner singleton; + return singleton; +} + +bool GpuClusterAssigner::available() { + static const bool present = gpu_present(); + return present; +} + +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 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); + for (size_t j = 0; j < n_clusters; ++j) { + centroid_docs[j] = clusters[j].front(); + centroid_set.insert(clusters[j].front()); + } + + // 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]; + 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 = GpuCorpus::instance().ensure_resident(vectors); + ThreadCtx& ctx = thread_ctx(); + cudaStream_t stream = ctx.stream; + + // 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_best, ctx.best_cap, n_docs * sizeof(int32_t)); + ensure_capacity(&ctx.d_b, ctx.b_cap, dim * n_clusters * sizeof(float)); + + 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)); + + // 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); + 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()); + + // 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)); + + 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)); + 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_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_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_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)); + + // 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)) { + continue; + } + clusters[h_best[i]].push_back(d); + } +} + +bool should_offload_assignment_to_gpu(size_t n_docs, size_t n_clusters) { + return GpuClusterAssigner::available() && n_docs > 0 && 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..9a67663 --- /dev/null +++ b/nsparse/gpu/gpu_cluster_assigner.h @@ -0,0 +1,66 @@ +/** + * 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 k-means document-to-cluster assignment for the build() path. + * + * 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). + * + * 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 owning the per-thread cuSPARSE handles/streams. + static GpuClusterAssigner& instance(); + + /// 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. + * + * @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); + + GpuClusterAssigner(const GpuClusterAssigner&) = delete; + GpuClusterAssigner& operator=(const GpuClusterAssigner&) = delete; + +private: + GpuClusterAssigner() = default; +}; + +/** + * @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); + +} // 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 b73e09b..a07659a 100644 --- a/nsparse/seismic_common.h +++ b/nsparse/seismic_common.h @@ -139,6 +139,7 @@ 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) for (int64_t idx = 0; idx < static_cast(inverted_lists_size); ++idx) { diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index 2d59fa7..69eb4e0 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 GPU support is compiled in. +if(NSPARSE_ENABLE_GPU) + 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..c4e9ab6 --- /dev/null +++ b/tests/gpu_cluster_assigner_test.cpp @@ -0,0 +1,279 @@ +/** + * 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 "nsparse/cluster/kmeans_utils.h" +#include "nsparse/gpu/gpu_summarizer.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"; + } + // 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; + 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"; + } +} + +// 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 (!GpuSummarizer::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(GpuSummarizer::instance().summarize_list( + &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 (!GpuSummarizer::available()) { + GTEST_SKIP() << "No CUDA-capable GPU available"; + } + SparseVectors vectors = make_random_corpus(5000, 2500, 40, 21); + 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(&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(&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