Gpu cusparse build acceleration

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

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

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
+)

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 <benchmark/benchmark.h>
+
+#include <cstdint>
+#include <cstdlib>
+#include <fstream>
+#include <iostream>
+#include <memory>
+#include <stdexcept>
+#include <string>
+#include <vector>
+
+#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<nsparse::idx_t> indptr;
+    std::vector<nsparse::term_t> indices;
+    std::vector<float> 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<char*>(sizes), sizeof(sizes));
+    m.nrow = sizes[0];
+    m.ncol = sizes[1];
+    m.nnz = sizes[2];
+
+    std::vector<int64_t> indptr64(m.nrow + 1);
+    f.read(reinterpret_cast<char*>(indptr64.data()),
+           static_cast<std::streamsize>((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<nsparse::idx_t>(indptr64[i]);
+    }
+
+    std::vector<int32_t> indices32(m.nnz);
+    f.read(reinterpret_cast<char*>(indices32.data()),
+           static_cast<std::streamsize>(m.nnz * sizeof(int32_t)));
+    m.indices.resize(m.nnz);
+    for (int64_t i = 0; i < m.nnz; ++i) {
+        m.indices[i] = static_cast<nsparse::term_t>(indices32[i]);
+    }
+
+    m.data.resize(m.nnz);
+    f.read(reinterpret_cast<char*>(m.data.data()),
+           static_cast<std::streamsize>(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<nsparse::SeismicIndex>(
+            static_cast<int>(data.ncol), kParams);
+        index->add(static_cast<nsparse::idx_t>(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<double>(data.nrow);
+    state.counters["nnz"] = static_cast<double>(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();

nsparse/CMakeLists.txt

@@ -138,6 +138,25 @@ if(NSPARSE_IS_X86)
   target_include_directories(nsparse_avx512 PUBLIC $<BUILD_INTERFACE:${PROJECT_SOURCE_DIR}>)
 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)