From 6e1f2e48dd66dbd083769a111ebdaff017cddefb Mon Sep 17 00:00:00 2001 From: James Lamb Date: Thu, 10 Sep 2026 17:44:13 -0500 Subject: [PATCH 1/7] wheels: get 'tileiras' from system install, not 'cuda-toolkit' wheels (#2564) Follow-up to https://github.com/NVIDIA/cuvs/pull/2548 Contributes to https://github.com/rapidsai/build-planning/issues/324 * switches to wheel builds and `pip` devcontainers getting `tileiras` from the system-installed CTK, not `cuda-toolkit` wheels. ## Notes for Reviewers ### Benefits of the `tileiras` change * one less version to need to remember to update when we bump the CTK version we build against * consistent with how we get NVCC from the system install, not wheels * removes a source of patching in DLFW builds (where a system CTK is always preferred to one provided by wheels) ### How I tested this relied on CI --- dependencies.yaml | 19 ++----------------- python/libcuvs/pyproject.toml | 1 - 2 files changed, 2 insertions(+), 18 deletions(-) diff --git a/dependencies.yaml b/dependencies.yaml index 08b1209411..374949ff64 100644 --- a/dependencies.yaml +++ b/dependencies.yaml @@ -429,6 +429,8 @@ dependencies: - libcusolver-dev - libcusparse-dev - libnvjitlink-dev + # NOTE: 'tileiras' is needed too, but in CI that's provided by a system-installed CTK, so + # we do not use the [tileiras] extra recommended at https://pypi.org/project/cuda-tile/ cutile_python: specific: - output_types: conda @@ -436,16 +438,6 @@ dependencies: - matrix: cuda: "12.*" packages: - - matrix: - cuda: "13.3" - packages: - - cutile-python - - cuda-tileiras - - matrix: - cuda: "13.*" - packages: - - cutile-python - - cuda-tileiras - matrix: packages: - cutile-python @@ -455,22 +447,15 @@ dependencies: - matrix: cuda: "12.*" packages: - - matrix: - cuda: "13.3" - packages: - - cuda-tile - - cuda-toolkit[tileiras]==13.3.* - matrix: cuda: "13.*" packages: - &cutile_python_cu13 cuda-tile - - &cutile_toolkit_cu13 cuda-toolkit[tileiras]==13.* # if no matching matrix selectors passed, list the CUDA 13 requirement # (as a source of documentation in the generated pyproject.toml) - matrix: packages: - *cutile_python_cu13 - - *cutile_toolkit_cu13 cuda_wheels: specific: # cuVS needs 'nvJitLink>={whatever-cuvs-was-built-against}' at runtime, and mixing diff --git a/python/libcuvs/pyproject.toml b/python/libcuvs/pyproject.toml index 42b5e8e213..37d4533ef3 100644 --- a/python/libcuvs/pyproject.toml +++ b/python/libcuvs/pyproject.toml @@ -84,7 +84,6 @@ build-backend = "scikit_build_core.build" requires = [ "cmake>=4.0", "cuda-tile", - "cuda-toolkit[tileiras]==13.*", "libkvikio==26.10.*,>=0.0.0a0", "libraft==26.10.*,>=0.0.0a0", "librmm==26.10.*,>=0.0.0a0", From c4e23e4febfe0be11991fb0ace2808324c8d9963 Mon Sep 17 00:00:00 2001 From: Tarang Jain <40517122+tarang-jain@users.noreply.github.com> Date: Thu, 10 Sep 2026 17:32:13 -0700 Subject: [PATCH 2/7] [ANN_BENCH] Remove Stream Pool Creation from Individual Wrappers (#2526) - remove stream pool creation from individual wrappers Authors: - Tarang Jain (https://github.com/tarang-jain) - Mike Sarahan (https://github.com/msarahan) Approvers: - Artem M. Chirkin (https://github.com/achirkin) - Bradley Dice (https://github.com/bdice) URL: https://github.com/NVIDIA/cuvs/pull/2526 --- cpp/bench/ann/CMakeLists.txt | 7 +++-- cpp/bench/ann/src/common/util.hpp | 29 ++++++------------- cpp/bench/ann/src/cuvs/cuvs_ann_bench_utils.h | 5 +++- .../ann/src/cuvs/cuvs_ivf_flat_wrapper.h | 6 +--- cpp/bench/ann/src/cuvs/cuvs_ivf_pq_wrapper.h | 6 +--- .../ann/src/cuvs/cuvs_ivf_rabitq_wrapper.h | 6 +--- cpp/bench/ann/src/cuvs/cuvs_ivf_sq_wrapper.h | 6 +--- 7 files changed, 21 insertions(+), 44 deletions(-) diff --git a/cpp/bench/ann/CMakeLists.txt b/cpp/bench/ann/CMakeLists.txt index 90d23d9aef..9e3b69fdbe 100644 --- a/cpp/bench/ann/CMakeLists.txt +++ b/cpp/bench/ann/CMakeLists.txt @@ -1,6 +1,6 @@ # ============================================================================= # cmake-format: off -# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # cmake-format: on # ============================================================================= @@ -181,6 +181,7 @@ function(ConfigureAnnBench) sqlite3 Threads::Threads $<$:CUDA::cudart_static> + $<$:rmm::rmm> $ $ ) @@ -426,8 +427,8 @@ if(CUVS_ANN_BENCH_SINGLE_EXE) target_include_directories(ANN_BENCH PRIVATE ${CMAKE_CUDA_TOOLKIT_INCLUDE_DIRECTORIES}) target_link_libraries( - ANN_BENCH PRIVATE raft::raft nlohmann_json::nlohmann_json sqlite3 benchmark::benchmark dl - $<$:CUDA::nvtx3> + ANN_BENCH PRIVATE $ nlohmann_json::nlohmann_json sqlite3 + benchmark::benchmark dl $<$:CUDA::nvtx3> ) set_target_properties( ANN_BENCH diff --git a/cpp/bench/ann/src/common/util.hpp b/cpp/bench/ann/src/common/util.hpp index 6e4c57b35e..069d257717 100644 --- a/cpp/bench/ann/src/common/util.hpp +++ b/cpp/bench/ann/src/common/util.hpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2023-2025, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ #pragma once @@ -7,6 +7,10 @@ #include "ann_types.hpp" #include "cuda_stub.hpp" // cuda-related utils +#ifndef BUILD_CPU_ONLY +#include +#endif + #if __has_include() #define ANN_BENCH_NVTX3_HEADERS_FOUND #include @@ -145,25 +149,8 @@ struct cuda_timer { }; #ifndef BUILD_CPU_ONLY -// ATM, rmm::stream does not support passing in flags; hence this helper type. -struct non_blocking_stream { - non_blocking_stream() { cudaStreamCreateWithFlags(&stream_, cudaStreamNonBlocking); } - ~non_blocking_stream() noexcept - { - if (stream_ != nullptr) { cudaStreamDestroy(stream_); } - } - non_blocking_stream(non_blocking_stream const&) = delete; - non_blocking_stream(non_blocking_stream&& other) noexcept { std::swap(stream_, other.stream_); } - auto operator=(non_blocking_stream const&) -> non_blocking_stream& = delete; - auto operator=(non_blocking_stream&&) -> non_blocking_stream& = delete; - [[nodiscard]] auto view() const noexcept -> cudaStream_t { return stream_; } - - private: - cudaStream_t stream_{nullptr}; -}; - namespace detail { -inline std::vector global_stream_pool(0); +inline std::vector global_stream_pool(0); inline std::mutex gsp_mutex; } // namespace detail #endif @@ -180,7 +167,9 @@ inline auto get_stream_from_global_pool() -> cudaStream_t #ifndef BUILD_CPU_ONLY std::lock_guard guard(detail::gsp_mutex); if (static_cast(detail::global_stream_pool.size()) < benchmark_n_threads) { - detail::global_stream_pool.resize(benchmark_n_threads); + while (static_cast(detail::global_stream_pool.size()) < benchmark_n_threads) { + detail::global_stream_pool.emplace_back(rmm::cuda_stream::flags::non_blocking); + } } return detail::global_stream_pool[benchmark_thread_id].view(); #else diff --git a/cpp/bench/ann/src/cuvs/cuvs_ann_bench_utils.h b/cpp/bench/ann/src/cuvs/cuvs_ann_bench_utils.h index 1a276e8cc8..f699bdad1a 100644 --- a/cpp/bench/ann/src/cuvs/cuvs_ann_bench_utils.h +++ b/cpp/bench/ann/src/cuvs/cuvs_ann_bench_utils.h @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ #pragma once @@ -15,9 +15,11 @@ #include #include #include +#include #include #include +#include #include #include #include @@ -127,6 +129,7 @@ class configured_raft_resources { { raft::resource::set_large_workspace_resource( *res_, raft::mr::device_resource{shared_res_->get_large_memory_resource()}); + raft::resource::set_cuda_stream_pool(*res_, std::make_shared(1)); } /** Default constructor creates all resources anew. */ diff --git a/cpp/bench/ann/src/cuvs/cuvs_ivf_flat_wrapper.h b/cpp/bench/ann/src/cuvs/cuvs_ivf_flat_wrapper.h index ae334f70d7..8ed5adee26 100644 --- a/cpp/bench/ann/src/cuvs/cuvs_ivf_flat_wrapper.h +++ b/cpp/bench/ann/src/cuvs/cuvs_ivf_flat_wrapper.h @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2023-2024, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ #pragma once @@ -15,7 +15,6 @@ #include #include #include -#include #include #include @@ -88,9 +87,6 @@ class cuvs_ivf_flat : public algo, public algo_gpu { template void cuvs_ivf_flat::build(const T* dataset, size_t nrow) { - // Create a CUDA stream pool with 1 stream (besides main stream) for kernel/copy overlapping. - size_t n_streams = 1; - raft::resource::set_cuda_stream_pool(handle_, std::make_shared(n_streams)); index_ = std::make_shared>( std::move(cuvs::neighbors::ivf_flat::build( handle_, diff --git a/cpp/bench/ann/src/cuvs/cuvs_ivf_pq_wrapper.h b/cpp/bench/ann/src/cuvs/cuvs_ivf_pq_wrapper.h index 161563d8fe..363d065e91 100644 --- a/cpp/bench/ann/src/cuvs/cuvs_ivf_pq_wrapper.h +++ b/cpp/bench/ann/src/cuvs/cuvs_ivf_pq_wrapper.h @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2023-2025, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ #pragma once @@ -20,7 +20,6 @@ #include #include #include -#include #include @@ -118,9 +117,6 @@ void cuvs_ivf_pq::load(const std::string& file) template void cuvs_ivf_pq::build(const T* dataset, size_t nrow) { - // Create a CUDA stream pool with 1 stream (besides main stream) for kernel/copy overlapping. - size_t n_streams = 1; - raft::resource::set_cuda_stream_pool(handle_, std::make_shared(n_streams)); auto dataset_v = raft::make_device_matrix_view(dataset, IdxT(nrow), dim_); std::make_shared>( std::move(cuvs::neighbors::ivf_pq::build(handle_, index_params_, dataset_v))) diff --git a/cpp/bench/ann/src/cuvs/cuvs_ivf_rabitq_wrapper.h b/cpp/bench/ann/src/cuvs/cuvs_ivf_rabitq_wrapper.h index ca8f77b808..0c41ee1c96 100644 --- a/cpp/bench/ann/src/cuvs/cuvs_ivf_rabitq_wrapper.h +++ b/cpp/bench/ann/src/cuvs/cuvs_ivf_rabitq_wrapper.h @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ #pragma once @@ -18,7 +18,6 @@ #include #include #include -#include #include @@ -98,9 +97,6 @@ void cuvs_ivf_rabitq::load(const std::string& file) template void cuvs_ivf_rabitq::build(const T* dataset, size_t nrow) { - // Create a CUDA stream pool with 1 stream (besides main stream) for kernel/copy overlapping. - size_t n_streams = 1; - raft::resource::set_cuda_stream_pool(handle_, std::make_shared(n_streams)); auto dataset_v = raft::make_device_matrix_view(dataset, IdxT(nrow), dim_); std::make_shared>( std::move(cuvs::neighbors::ivf_rabitq::build(handle_, index_params_, dataset_v))) diff --git a/cpp/bench/ann/src/cuvs/cuvs_ivf_sq_wrapper.h b/cpp/bench/ann/src/cuvs/cuvs_ivf_sq_wrapper.h index 1503e6bb84..5bf0098eaa 100644 --- a/cpp/bench/ann/src/cuvs/cuvs_ivf_sq_wrapper.h +++ b/cpp/bench/ann/src/cuvs/cuvs_ivf_sq_wrapper.h @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ #pragma once @@ -12,9 +12,7 @@ #include #include #include -#include #include -#include #include #include @@ -83,8 +81,6 @@ class cuvs_ivf_sq : public algo, public algo_gpu { template void cuvs_ivf_sq::build(const T* dataset, size_t nrow) { - size_t n_streams = 1; - raft::resource::set_cuda_stream_pool(handle_, std::make_shared(n_streams)); index_ = std::make_shared>( std::move(cuvs::neighbors::ivf_sq::build( handle_, From 8c3cac929df806a410f07f0c793302aab7422efb Mon Sep 17 00:00:00 2001 From: Igor Motov Date: Thu, 10 Sep 2026 17:08:21 -1000 Subject: [PATCH 3/7] Match merged vectors by document id, not by ordinal position (#2556) `testMergeTwoSegsWithASingleDocPerSeg` and `testTwoVectorFieldsPerDoc` asserted that ordinal i of the merged segment holds the i-th document's vector. Lucene does not offer that: `MockRandomMergePolicy` shuffles the segments of a forced merge on purpose, so the document committed second can land at ordinal 0. Stock `Lucene99HnswVectorsFormat` fails the same tests on the same seeds, so no cuVS writer is involved. The assertions now resolve each ordinal to its document and check that the document kept its own vector, which leaves the randomized merge policy in play. ordToDoc is checked to be increasing so the id lookup cannot agree with a mapping that is itself garbled. Applied to the quantized and GPU-search formats too, which carried the same assumption unreported. Closes #2550 Authors: - Igor Motov (https://github.com/imotov) Approvers: - Corey J. Nolet (https://github.com/cjnolet) URL: https://github.com/NVIDIA/cuvs/pull/2556 --- .../cuvs/lucene/TestCuVSVectorsFormat.java | 29 +++------- ...tLucene99AcceleratedHNSWVectorsFormat.java | 29 +++------- .../lucene/TestQuantizedVectorsFormats.java | 27 +++------- .../com/nvidia/cuvs/lucene/TestUtils.java | 54 +++++++++++++++++++ 4 files changed, 76 insertions(+), 63 deletions(-) diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestCuVSVectorsFormat.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestCuVSVectorsFormat.java index 2b960bf3cd..6db380a3ce 100644 --- a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestCuVSVectorsFormat.java +++ b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestCuVSVectorsFormat.java @@ -4,6 +4,7 @@ */ package com.nvidia.cuvs.lucene; +import static com.nvidia.cuvs.lucene.TestUtils.assertVectorsKeepTheirDocuments; import static com.nvidia.cuvs.lucene.ThreadLocalCuVSResourcesProvider.isSupported; import static org.apache.lucene.index.VectorSimilarityFunction.EUCLIDEAN; @@ -14,7 +15,6 @@ import org.apache.lucene.document.KnnFloatVectorField; import org.apache.lucene.document.StringField; import org.apache.lucene.index.DirectoryReader; -import org.apache.lucene.index.FloatVectorValues; import org.apache.lucene.index.IndexWriter; import org.apache.lucene.index.LeafReader; import org.apache.lucene.index.LeafReaderContext; @@ -44,12 +44,12 @@ public void testMergeTwoSegsWithASingleDocPerSeg() throws Exception { try (Directory dir = newDirectory(); IndexWriter w = new IndexWriter(dir, newIndexWriterConfig())) { Document doc1 = new Document(); - doc1.add(new StringField("id", "0", Field.Store.NO)); + doc1.add(new StringField("id", "0", Field.Store.YES)); doc1.add(new KnnFloatVectorField("f", f[0], EUCLIDEAN)); w.addDocument(doc1); w.commit(); Document doc2 = new Document(); - doc2.add(new StringField("id", "1", Field.Store.NO)); + doc2.add(new StringField("id", "1", Field.Store.YES)); doc2.add(new KnnFloatVectorField("f", f[1], EUCLIDEAN)); w.addDocument(doc2); w.flush(); @@ -69,11 +69,7 @@ public void testMergeTwoSegsWithASingleDocPerSeg() throws Exception { // verify merged content try (DirectoryReader reader = DirectoryReader.open(w)) { LeafReader r = getOnlyLeafReader(reader); - FloatVectorValues values = r.getFloatVectorValues("f"); - assertNotNull(values); - assertEquals(2, values.size()); - assertArrayEquals(f[0], values.vectorValue(0), 0.0f); - assertArrayEquals(f[1], values.vectorValue(1), 0.0f); + assertVectorsKeepTheirDocuments(r, "f", f); } } } @@ -85,12 +81,12 @@ public void testTwoVectorFieldsPerDoc() throws Exception { try (Directory dir = newDirectory(); IndexWriter w = new IndexWriter(dir, newIndexWriterConfig())) { Document doc1 = new Document(); - doc1.add(new StringField("id", "0", Field.Store.NO)); + doc1.add(new StringField("id", "0", Field.Store.YES)); doc1.add(new KnnFloatVectorField("f1", f1[0], EUCLIDEAN)); doc1.add(new KnnFloatVectorField("f2", f2[0], EUCLIDEAN)); w.addDocument(doc1); Document doc2 = new Document(); - doc2.add(new StringField("id", "1", Field.Store.NO)); + doc2.add(new StringField("id", "1", Field.Store.YES)); doc2.add(new KnnFloatVectorField("f1", f1[1], EUCLIDEAN)); doc2.add(new KnnFloatVectorField("f2", f2[1], EUCLIDEAN)); w.addDocument(doc2); @@ -98,17 +94,8 @@ public void testTwoVectorFieldsPerDoc() throws Exception { try (DirectoryReader reader = DirectoryReader.open(w)) { LeafReader r = getOnlyLeafReader(reader); - FloatVectorValues values = r.getFloatVectorValues("f1"); - assertNotNull(values); - assertEquals(2, values.size()); - assertArrayEquals(f1[0], values.vectorValue(0), 0.0f); - assertArrayEquals(f1[1], values.vectorValue(1), 0.0f); - - values = r.getFloatVectorValues("f2"); - assertNotNull(values); - assertEquals(2, values.size()); - assertArrayEquals(f2[0], values.vectorValue(0), 0.0f); - assertArrayEquals(f2[1], values.vectorValue(1), 0.0f); + assertVectorsKeepTheirDocuments(r, "f1", f1); + assertVectorsKeepTheirDocuments(r, "f2", f2); // opportunistically check boundary condition - search with a 0 topK var topDocs = r.searchNearestVectors("f1", randomVector(384), 0, null, 10); diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestLucene99AcceleratedHNSWVectorsFormat.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestLucene99AcceleratedHNSWVectorsFormat.java index 9bb140228b..707ee35709 100644 --- a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestLucene99AcceleratedHNSWVectorsFormat.java +++ b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestLucene99AcceleratedHNSWVectorsFormat.java @@ -4,6 +4,7 @@ */ package com.nvidia.cuvs.lucene; +import static com.nvidia.cuvs.lucene.TestUtils.assertVectorsKeepTheirDocuments; import static com.nvidia.cuvs.lucene.ThreadLocalCuVSResourcesProvider.isSupported; import static org.apache.lucene.index.VectorSimilarityFunction.EUCLIDEAN; @@ -14,7 +15,6 @@ import org.apache.lucene.document.KnnFloatVectorField; import org.apache.lucene.document.StringField; import org.apache.lucene.index.DirectoryReader; -import org.apache.lucene.index.FloatVectorValues; import org.apache.lucene.index.IndexWriter; import org.apache.lucene.index.LeafReader; import org.apache.lucene.index.LeafReaderContext; @@ -44,12 +44,12 @@ public void testMergeTwoSegsWithASingleDocPerSeg() throws Exception { try (Directory dir = newDirectory(); IndexWriter w = new IndexWriter(dir, newIndexWriterConfig())) { Document doc1 = new Document(); - doc1.add(new StringField("id", "0", Field.Store.NO)); + doc1.add(new StringField("id", "0", Field.Store.YES)); doc1.add(new KnnFloatVectorField("f", f[0], EUCLIDEAN)); w.addDocument(doc1); w.commit(); Document doc2 = new Document(); - doc2.add(new StringField("id", "1", Field.Store.NO)); + doc2.add(new StringField("id", "1", Field.Store.YES)); doc2.add(new KnnFloatVectorField("f", f[1], EUCLIDEAN)); w.addDocument(doc2); w.flush(); @@ -69,11 +69,7 @@ public void testMergeTwoSegsWithASingleDocPerSeg() throws Exception { // verify merged content try (DirectoryReader reader = DirectoryReader.open(w)) { LeafReader r = getOnlyLeafReader(reader); - FloatVectorValues values = r.getFloatVectorValues("f"); - assertNotNull(values); - assertEquals(2, values.size()); - assertArrayEquals(f[0], values.vectorValue(0), 0.0f); - assertArrayEquals(f[1], values.vectorValue(1), 0.0f); + assertVectorsKeepTheirDocuments(r, "f", f); } } } @@ -85,12 +81,12 @@ public void testTwoVectorFieldsPerDoc() throws Exception { try (Directory dir = newDirectory(); IndexWriter w = new IndexWriter(dir, newIndexWriterConfig())) { Document doc1 = new Document(); - doc1.add(new StringField("id", "0", Field.Store.NO)); + doc1.add(new StringField("id", "0", Field.Store.YES)); doc1.add(new KnnFloatVectorField("f1", f1[0], EUCLIDEAN)); doc1.add(new KnnFloatVectorField("f2", f2[0], EUCLIDEAN)); w.addDocument(doc1); Document doc2 = new Document(); - doc2.add(new StringField("id", "1", Field.Store.NO)); + doc2.add(new StringField("id", "1", Field.Store.YES)); doc2.add(new KnnFloatVectorField("f1", f1[1], EUCLIDEAN)); doc2.add(new KnnFloatVectorField("f2", f2[1], EUCLIDEAN)); w.addDocument(doc2); @@ -98,17 +94,8 @@ public void testTwoVectorFieldsPerDoc() throws Exception { try (DirectoryReader reader = DirectoryReader.open(w)) { LeafReader r = getOnlyLeafReader(reader); - FloatVectorValues values = r.getFloatVectorValues("f1"); - assertNotNull(values); - assertEquals(2, values.size()); - assertArrayEquals(f1[0], values.vectorValue(0), 0.0f); - assertArrayEquals(f1[1], values.vectorValue(1), 0.0f); - - values = r.getFloatVectorValues("f2"); - assertNotNull(values); - assertEquals(2, values.size()); - assertArrayEquals(f2[0], values.vectorValue(0), 0.0f); - assertArrayEquals(f2[1], values.vectorValue(1), 0.0f); + assertVectorsKeepTheirDocuments(r, "f1", f1); + assertVectorsKeepTheirDocuments(r, "f2", f2); // opportunistically check boundary condition - search with a 0 topK var topDocs = r.searchNearestVectors("f1", randomVector(384), 0, null, 10); diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestQuantizedVectorsFormats.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestQuantizedVectorsFormats.java index 5140782fc7..3b05beaacf 100644 --- a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestQuantizedVectorsFormats.java +++ b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestQuantizedVectorsFormats.java @@ -4,6 +4,7 @@ */ package com.nvidia.cuvs.lucene; +import static com.nvidia.cuvs.lucene.TestUtils.assertVectorsKeepTheirDocuments; import static com.nvidia.cuvs.lucene.ThreadLocalCuVSResourcesProvider.isSupported; import static org.apache.lucene.index.VectorSimilarityFunction.COSINE; import static org.apache.lucene.index.VectorSimilarityFunction.EUCLIDEAN; @@ -77,7 +78,7 @@ public void testMergeTwoSegsWithASingleDocPerSeg() throws Exception { IndexWriter w = new IndexWriter(dir, newIndexWriterConfig())) { for (int i = 0; i < R; i++) { Document doc = new Document(); - doc.add(new StringField("id", String.valueOf(i), Field.Store.NO)); + doc.add(new StringField("id", String.valueOf(i), Field.Store.YES)); doc.add(new KnnFloatVectorField(F, f[i], EUCLIDEAN)); w.addDocument(doc); w.commit(); @@ -95,12 +96,7 @@ public void testMergeTwoSegsWithASingleDocPerSeg() throws Exception { try (DirectoryReader reader = DirectoryReader.open(w)) { LeafReader r = getOnlyLeafReader(reader); - FloatVectorValues values = r.getFloatVectorValues(F); - assertNotNull(values); - assertEquals(R, values.size()); - for (int i = 0; i < R; i++) { - assertArrayEquals(f[i], values.vectorValue(i), 0.0f); - } + assertVectorsKeepTheirDocuments(r, F, f); } } } @@ -121,7 +117,7 @@ public void testTwoVectorFieldsPerDoc() throws Exception { for (int i = 0; i < R; i++) { Document doc = new Document(); - doc.add(new StringField("id", String.valueOf(i), Field.Store.NO)); + doc.add(new StringField("id", String.valueOf(i), Field.Store.YES)); doc.add(new KnnFloatVectorField(F1, f1[i], EUCLIDEAN)); doc.add(new KnnFloatVectorField(F2, f2[i], EUCLIDEAN)); w.addDocument(doc); @@ -130,19 +126,8 @@ public void testTwoVectorFieldsPerDoc() throws Exception { try (DirectoryReader reader = DirectoryReader.open(w)) { LeafReader r = getOnlyLeafReader(reader); - FloatVectorValues values = r.getFloatVectorValues(F1); - assertNotNull(values); - assertEquals(R, values.size()); - for (int i = 0; i < R; i++) { - assertArrayEquals(f1[i], values.vectorValue(i), 0.0f); - } - - values = r.getFloatVectorValues(F2); - assertNotNull(values); - assertEquals(R, values.size()); - for (int i = 0; i < R; i++) { - assertArrayEquals(f2[i], values.vectorValue(i), 0.0f); - } + assertVectorsKeepTheirDocuments(r, F1, f1); + assertVectorsKeepTheirDocuments(r, F2, f2); } } } diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestUtils.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestUtils.java index e19e5ce64f..c96c17f493 100644 --- a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestUtils.java +++ b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestUtils.java @@ -4,16 +4,70 @@ */ package com.nvidia.cuvs.lucene; +import static org.junit.Assert.assertArrayEquals; +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertNotNull; +import static org.junit.Assert.assertTrue; + import java.io.IOException; import java.nio.ByteBuffer; import java.nio.ByteOrder; import java.nio.channels.FileChannel; import java.nio.file.Path; import java.nio.file.StandardOpenOption; +import java.util.HashSet; import java.util.Random; +import java.util.Set; +import org.apache.lucene.index.FloatVectorValues; +import org.apache.lucene.index.LeafReader; +import org.apache.lucene.index.StoredFields; public class TestUtils { + /** + * Asserts that every vector in {@code field} is still paired with the document that indexed it, + * matching on the document's stored {@code id} rather than on ordinal position. + * + *

Nothing fixes the order documents land in after a merge, and the randomized test + * framework's {@code MockRandomMergePolicy} disturbs it two independent ways: {@code + * findForcedMerges} shuffles the segments it is about to merge, and {@code + * MockRandomOneMerge.reorder} reverses doc IDs outright. The second applies even when there is + * only one segment, so a test that never commits between documents is no safer than one that + * does. Asserting {@code vectorValue(i) == expectedById[i]} therefore fails on some seeds with + * nothing wrong. What has to hold is that no document loses its own vector, which is what this + * checks. + * + * @param expectedById the vector indexed for each document id, indexed by that id + */ + public static void assertVectorsKeepTheirDocuments( + LeafReader reader, String field, float[][] expectedById) throws IOException { + FloatVectorValues values = reader.getFloatVectorValues(field); + assertNotNull("no vector values for field " + field, values); + assertEquals(expectedById.length, values.size()); + StoredFields storedFields = reader.storedFields(); + Set seen = new HashSet<>(); + int previousDoc = -1; + for (int ord = 0; ord < values.size(); ord++) { + int doc = values.ordToDoc(ord); + // Ordinals are assigned in docID order. Where every document has a vector -- every caller + // today -- ordToDoc is the hardcoded identity and this cannot fire. It earns its keep only + // if a caller passes a field that some documents lack, where the id lookup below would + // otherwise be free to agree with a garbled mapping. + assertTrue("ordToDoc went backwards at ordinal " + ord, doc > previousDoc); + previousDoc = doc; + String storedId = storedFields.document(doc).get("id"); + assertNotNull("document at ordinal " + ord + " has no stored id", storedId); + int id = Integer.parseInt(storedId); + assertTrue("document id " + id + " appeared twice", seen.add(id)); + assertArrayEquals( + "vector for document id " + id + " (doc " + doc + ", ordinal " + ord + ")", + expectedById[id], + values.vectorValue(ord), + 0.0f); + } + assertEquals("not every document was found", expectedById.length, seen.size()); + } + /** Writes {@code dataset} as an uncompressed {@code .fbin} file (little-endian float32 rows). */ public static void writeFbin(Path path, float[][] dataset) throws IOException { int numVectors = dataset.length; From e69606bc6e9be8048906156c42b1d8a12ffa2edf Mon Sep 17 00:00:00 2001 From: Igor Motov Date: Fri, 11 Sep 2026 12:53:52 -1000 Subject: [PATCH 4/7] Revert "Optimize cuvs-lucene CAGRA_HNSW index build; add example (#2481)" (#2597) This reverts commit 6a5bba01911fb66bcfe8490abe33d71b75010313. After some additional discussion, we concluded that while this change addresses a very specific use case well, there's an opportunity to make it more generic and more widely applicable. We'll be working on a solution that covers a broader set of use cases and open a follow-up PR soon. Authors: - Igor Motov (https://github.com/imotov) Approvers: - Corey J. Nolet (https://github.com/cjnolet) URL: https://github.com/NVIDIA/cuvs/pull/2597 --- examples/java/cuvs-lucene/README.md | 14 - .../OptimizedCagraHnswBuildExample.java | 268 ------- ...va-api-com-nvidia-cuvs-spi-cuvsprovider.md | 27 +- .../com/nvidia/cuvs/spi/CuVSProvider.java | 17 +- java/cuvs-lucene/pom.xml | 5 - .../cuvs/lucene/AcceleratedHNSWParams.java | 41 +- .../cuvs/lucene/AcceleratedHNSWUtils.java | 224 ++---- .../lucene/AcceleratedHnswGraphOutput.java | 214 ------ .../cuvs/lucene/CagraHnswBulkIndexWriter.java | 660 ------------------ .../cuvs/lucene/CagraIndexParamsFactory.java | 10 +- .../cuvs/lucene/CuVS2510GPUVectorsWriter.java | 6 +- .../nvidia/cuvs/lucene/FbinVectorSource.java | 213 ------ .../com/nvidia/cuvs/lucene/FieldWriter.java | 6 - .../nvidia/cuvs/lucene/GPUBuiltHnswGraph.java | 97 +-- .../lucene/Lucene101AcceleratedHNSWCodec.java | 28 +- .../Lucene99AcceleratedHNSWVectorsFormat.java | 29 +- .../Lucene99AcceleratedHNSWVectorsWriter.java | 226 ++++-- ...ratedHNSWBinaryQuantizedVectorsWriter.java | 15 +- ...ratedHNSWScalarQuantizedVectorsWriter.java | 15 +- .../nvidia/cuvs/lucene/NativeFieldWriter.java | 124 ---- .../NativeFlatBufferedHNSWVectorsWriter.java | 214 ------ .../cuvs/lucene/NativeFlatVectorsWriter.java | 199 ------ .../java/com/nvidia/cuvs/lucene/Utils.java | 62 +- .../com/nvidia/cuvs/lucene/VectorSource.java | 35 - .../TestAcceleratedHNSWParamsSurface.java | 47 -- .../lucene/TestCagraHnswBulkIndexWriter.java | 446 ------------ .../lucene/TestCagraIndexParamsFactory.java | 39 -- .../cuvs/lucene/TestFbinVectorSource.java | 120 ---- .../lucene/TestMergedGraphOrdinalBounds.java | 198 ------ .../TestNativeFlatBufferingGuardRails.java | 269 ------- ...TestNativeFlatBufferingIndexAndSearch.java | 220 ------ .../TestNativeFlatVectorsWriterRoundTrip.java | 279 -------- .../com/nvidia/cuvs/lucene/TestUtils.java | 32 - .../TestWriterThreadsGraphEquivalence.java | 156 ----- 34 files changed, 324 insertions(+), 4231 deletions(-) delete mode 100644 examples/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/examples/OptimizedCagraHnswBuildExample.java delete mode 100644 java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHnswGraphOutput.java delete mode 100644 java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CagraHnswBulkIndexWriter.java delete mode 100644 java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/FbinVectorSource.java delete mode 100644 java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/NativeFieldWriter.java delete mode 100644 java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/NativeFlatBufferedHNSWVectorsWriter.java delete mode 100644 java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/NativeFlatVectorsWriter.java delete mode 100644 java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/VectorSource.java delete mode 100644 java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedHNSWParamsSurface.java delete mode 100644 java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestCagraHnswBulkIndexWriter.java delete mode 100644 java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestFbinVectorSource.java delete mode 100644 java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestMergedGraphOrdinalBounds.java delete mode 100644 java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestNativeFlatBufferingGuardRails.java delete mode 100644 java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestNativeFlatBufferingIndexAndSearch.java delete mode 100644 java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestNativeFlatVectorsWriterRoundTrip.java delete mode 100644 java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestWriterThreadsGraphEquivalence.java diff --git a/examples/java/cuvs-lucene/README.md b/examples/java/cuvs-lucene/README.md index 4113b00adc..5a3c3675e0 100644 --- a/examples/java/cuvs-lucene/README.md +++ b/examples/java/cuvs-lucene/README.md @@ -32,17 +32,3 @@ To run the Index and Search on GPU example do: ```sh mvn clean install && java -Djava.util.logging.config.file=src/main/resources/logging.properties -cp target/examples-26.10.0-jar-with-merged-services.jar com.nvidia.cuvs.lucene.examples.IndexAndSearchonGPUExample ``` - -To run the optimized CAGRA-HNSW build example (reference pattern for efficiently building an -accelerated HNSW index from a large `.fbin` with every ingest-side knob on — open the file once and -stream sequential prefetched chunks that overlap the disk read with indexing, hold at most two chunks -in memory, reuse a single vector array, size a native flat buffer per segment, auto-select the CAGRA -graph-build algorithm, and optionally partition into K segments built sequentially or overlapped) do: - -```sh -mvn clean install && java -Djava.util.logging.config.file=src/main/resources/logging.properties -cp target/examples-26.10.0-jar-with-merged-services.jar com.nvidia.cuvs.lucene.examples.OptimizedCagraHnswBuildExample -``` - -With no arguments it generates and indexes a small demo `.fbin` as a single segment; pass a real file, -chunk size, segment count, and overlap flag as -`... OptimizedCagraHnswBuildExample `. diff --git a/examples/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/examples/OptimizedCagraHnswBuildExample.java b/examples/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/examples/OptimizedCagraHnswBuildExample.java deleted file mode 100644 index 6cfcfafa61..0000000000 --- a/examples/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/examples/OptimizedCagraHnswBuildExample.java +++ /dev/null @@ -1,268 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: Apache-2.0 - */ -package com.nvidia.cuvs.lucene.examples; - -import com.nvidia.cuvs.CagraIndexParams.CuvsDistanceType; -import com.nvidia.cuvs.lucene.AcceleratedHNSWParams; -import com.nvidia.cuvs.lucene.CagraHnswBulkIndexWriter; -import com.nvidia.cuvs.lucene.FbinVectorSource; -import com.nvidia.cuvs.spi.CuVSProvider; -import java.io.IOException; -import java.nio.ByteBuffer; -import java.nio.ByteOrder; -import java.nio.channels.FileChannel; -import java.nio.file.Files; -import java.nio.file.Path; -import java.nio.file.Paths; -import java.nio.file.StandardOpenOption; -import java.util.Random; -import java.util.UUID; -import java.util.logging.Logger; -import org.apache.commons.io.FileUtils; -import org.apache.lucene.document.Document; -import org.apache.lucene.document.Field; -import org.apache.lucene.document.KnnFloatVectorField; -import org.apache.lucene.document.StringField; -import org.apache.lucene.index.DirectoryReader; -import org.apache.lucene.index.IndexWriterConfig; -import org.apache.lucene.index.VectorSimilarityFunction; -import org.apache.lucene.search.IndexSearcher; -import org.apache.lucene.search.KnnFloatVectorQuery; -import org.apache.lucene.search.ScoreDoc; -import org.apache.lucene.search.TopDocs; -import org.apache.lucene.store.Directory; -import org.apache.lucene.store.FSDirectory; - -/** - * Reference usage of {@link CagraHnswBulkIndexWriter}: builds an accelerated HNSW index (whose - * graph is built on the GPU with CAGRA) from a {@code .fbin} vector file on local disk. - * - *

Demonstrates both ways to use {@link CagraHnswBulkIndexWriter}: - * - *

    - *
  • {@link #main} — the one-shot convenience path. All of the bulk-build mechanics — - * prefetched streaming reads, native flat buffering, the {@code IndexWriterConfig} tuning - * that guarantees a single unmerged segment per slice, K-segment partitioning, and combining - * the result by hardlink — are owned by {@link CagraHnswBulkIndexWriter} itself; see its - * Javadoc for how each of those works and the tradeoffs of {@code numSegments} and {@code - * overlap}. This example only wires up what's genuinely application-specific: where the - * vectors come from ({@link FbinVectorSource}, or your own {@link - * com.nvidia.cuvs.lucene.VectorSource} for a different data source), the graph-build quality - * knobs ({@link AcceleratedHNSWParams}), and — via a {@link - * CagraHnswBulkIndexWriter.FieldCallback} — any per-vector metadata to attach. - *
  • {@link #runManualExample} — the manual, direct-instance path: construct a {@link - * CagraHnswBulkIndexWriter} yourself and drive {@code addDocument}/{@code close} exactly - * like a plain Lucene {@code IndexWriter}, building each {@link Document} (metadata included) - * yourself instead of going through a callback. - *
- * - *

Usage: {@code OptimizedCagraHnswBuildExample [] [] [] - * []}. With no arguments a small demo {@code .fbin} is generated and indexed as a - * single segment. - */ -public class OptimizedCagraHnswBuildExample { - - private static final Logger log = - Logger.getLogger(OptimizedCagraHnswBuildExample.class.getName()); - private static final String ID_FIELD = "id"; - private static final String CATEGORY_FIELD = "category"; - private static final String VECTOR_FIELD = "vector_field"; - - public static void main(String[] args) throws Exception { - // It is recommended to enable RMM allocation mode at application start, before constructing - // any CagraHnswBulkIndexWriter, to avoid device-wide sync from the default allocator. - CuVSProvider.provider().enableRMMAsyncMemory(); - - int chunkSizeMB = args.length >= 2 ? Integer.parseInt(args[1]) : 32; - int numSegments = args.length >= 3 ? Math.max(1, Integer.parseInt(args[2])) : 1; - boolean overlap = args.length >= 4 && Boolean.parseBoolean(args[3]); - Path indexDirPath = Paths.get(UUID.randomUUID().toString()); - - Path fbinPath; - boolean generated = false; - if (args.length >= 1) { - fbinPath = Paths.get(args[0]); - } else { - fbinPath = Paths.get("demo-" + UUID.randomUUID() + ".fbin"); - writeDemoFbin(fbinPath, 5000, 32, new Random(222)); - generated = true; - log.info("No .fbin provided; generated a demo file at " + fbinPath); - } - - try { - int dim; - try (FbinVectorSource probe = new FbinVectorSource(fbinPath, 1)) { - dim = probe.dimensions(); - } - - CagraHnswBulkIndexWriter.Config config = - CagraHnswBulkIndexWriter.Config.builder() - .field(VECTOR_FIELD, dim, VectorSimilarityFunction.EUCLIDEAN) - .idField(ID_FIELD) - .graphBuild( - new AcceleratedHNSWParams.Builder() - // HEURISTIC lets cuVS pick the build algorithm and auto-tune its parameters - // based on maxConn and beamWidth below. - .withStrategy(AcceleratedHNSWParams.Strategy.HEURISTIC) - // Primary recall/graph-size knobs. Higher values improve recall at the cost - // of a larger graph and longer build. Match to your dataset and recall - // target. - .withMaxConn(32) - .withBeamWidth(32) - // Must match the distance metric used when querying the index. - .withCuvsDistanceType(CuvsDistanceType.L2Expanded) - // Starting point: one thread per logical CPU. Profile and tune for your - // hardware. - .withWriterThreads(Runtime.getRuntime().availableProcessors()) - .build()) - .segments(numSegments, overlap) - .targetDirectory(indexDirPath) - .build(); - - log.info( - "Indexing " - + fbinPath - + " (" - + dim - + "-dim) into " - + numSegments - + " segment(s), " - + (overlap && numSegments > 1 ? "overlapped" : "sequential") - + " build, " - + chunkSizeMB - + " MB prefetched chunks"); - - // FieldCallback lets the one-shot path attach metadata per vector: indexFbin/build build the - // id+vector fields internally (they own the loop), so this is how a caller reaches the - // Document to add anything else -- here, an illustrative "even"/"odd" category by id. - CagraHnswBulkIndexWriter.indexFbin( - fbinPath, - config, - (doc, id) -> - doc.add( - new StringField(CATEGORY_FIELD, id % 2 == 0 ? "even" : "odd", Field.Store.YES)), - chunkSizeMB); - log.info("Index build complete: " + indexDirPath); - - runSampleSearch(indexDirPath, fbinPath, 5); - } finally { - FileUtils.deleteDirectory(indexDirPath.toFile()); - if (generated) { - Files.deleteIfExists(fbinPath); - } - } - - runManualExample(); - } - - /** Runs one k-NN query using the first vector in the file to show the index is searchable. */ - private static void runSampleSearch(Path indexDirPath, Path fbinPath, int topK) throws Exception { - float[] queryVector; - try (FbinVectorSource reader = new FbinVectorSource(fbinPath, 1)) { - queryVector = reader.get(0); - } - try (Directory dir = FSDirectory.open(indexDirPath); - DirectoryReader reader = DirectoryReader.open(dir)) { - IndexSearcher searcher = new IndexSearcher(reader); - TopDocs results = - searcher.search(new KnnFloatVectorQuery(VECTOR_FIELD, queryVector, topK), topK); - log.info("Sample search returned " + results.scoreDocs.length + " hits:"); - for (int i = 0; i < results.scoreDocs.length; i++) { - ScoreDoc sd = results.scoreDocs[i]; - Document hit = searcher.storedFields().document(sd.doc); - log.info( - " rank " - + (i + 1) - + ": id=" - + hit.get(ID_FIELD) - + " category=" - + hit.get(CATEGORY_FIELD) - + " score=" - + sd.score); - } - } - } - - /** - * Short demonstration of the manual, direct-instance API: {@link CagraHnswBulkIndexWriter} is - * constructed directly and driven with {@code addDocument}/{@code close}, the same shape as a - * plain Lucene {@code IndexWriter} — the caller builds each {@link Document} itself, including - * whatever metadata it wants, with no callback needed since it already owns the loop. Unlike the - * one-shot path above, this only ever builds a single segment; K-segment partitioning and - * overlap are only available via {@link CagraHnswBulkIndexWriter#indexFbin}/{@link - * CagraHnswBulkIndexWriter#build}. - */ - private static void runManualExample() throws Exception { - int numDocs = 200; - int dim = 16; - Random random = new Random(7); - Path manualIndexDirPath = Paths.get("manual-" + UUID.randomUUID()); - - try { - CagraHnswBulkIndexWriter.Config config = - CagraHnswBulkIndexWriter.Config.builder() - .field(VECTOR_FIELD, dim, VectorSimilarityFunction.EUCLIDEAN) - .graphBuild(new AcceleratedHNSWParams.Builder().build()) - .build(); - - float[][] vectors = new float[numDocs][dim]; - try (Directory dir = FSDirectory.open(manualIndexDirPath); - CagraHnswBulkIndexWriter writer = - new CagraHnswBulkIndexWriter(dir, new IndexWriterConfig(), config, numDocs)) { - for (int i = 0; i < numDocs; i++) { - for (int j = 0; j < dim; j++) { - vectors[i][j] = random.nextFloat() * 100; - } - Document doc = new Document(); - doc.add(new StringField(ID_FIELD, Integer.toString(i), Field.Store.YES)); - doc.add(new StringField(CATEGORY_FIELD, i % 2 == 0 ? "even" : "odd", Field.Store.YES)); - doc.add( - new KnnFloatVectorField( - VECTOR_FIELD, vectors[i], VectorSimilarityFunction.EUCLIDEAN)); - writer.addDocument(doc); // same call shape as a plain IndexWriter - } - } // close() runs the single native-buffered flush (the GPU CAGRA build happens here) - - try (Directory dir = FSDirectory.open(manualIndexDirPath); - DirectoryReader reader = DirectoryReader.open(dir)) { - IndexSearcher searcher = new IndexSearcher(reader); - TopDocs results = searcher.search(new KnnFloatVectorQuery(VECTOR_FIELD, vectors[0], 1), 1); - Document hit = searcher.storedFields().document(results.scoreDocs[0].doc); - log.info( - "Manual example: nearest neighbor of vector 0 is id=" - + hit.get(ID_FIELD) - + " category=" - + hit.get(CATEGORY_FIELD)); - } - } finally { - FileUtils.deleteDirectory(manualIndexDirPath.toFile()); - } - } - - /** Writes a small random {@code .fbin} so the example is runnable without external data. */ - private static void writeDemoFbin(Path path, int numVectors, int dim, Random random) - throws IOException { - ByteBuffer buf = - ByteBuffer.allocate(8 + numVectors * dim * Float.BYTES).order(ByteOrder.LITTLE_ENDIAN); - buf.putInt(numVectors); // .fbin header: [num_vectors int32][dimension int32] - buf.putInt(dim); - for (int i = 0; i < numVectors; i++) { - for (int j = 0; j < dim; j++) { - buf.putFloat(random.nextFloat() * 100); - } - } - buf.flip(); - try (FileChannel ch = - FileChannel.open( - path, - StandardOpenOption.CREATE, - StandardOpenOption.WRITE, - StandardOpenOption.TRUNCATE_EXISTING)) { - while (buf.hasRemaining()) { - ch.write(buf); - } - } - } -} diff --git a/fern/pages/java_api/java-api-com-nvidia-cuvs-spi-cuvsprovider.md b/fern/pages/java_api/java-api-com-nvidia-cuvs-spi-cuvsprovider.md index f19b50fb6b..1e5018c3e7 100644 --- a/fern/pages/java_api/java-api-com-nvidia-cuvs-spi-cuvsprovider.md +++ b/fern/pages/java_api/java-api-com-nvidia-cuvs-spi-cuvsprovider.md @@ -412,13 +412,6 @@ Switch RMM allocations (used internally by various cuVS algorithms and by the de `CuVSDeviceMatrix`) to use pooled memory. This operation has a global effect, and will affect all resources on the current device. -Ownership: this and the other `enableRMM*`/`resetRMMPooledMemory` methods are -not called automatically by any cuVS/cuvs-lucene class. Because the effect is global to the -current device, it is application code's responsibility to call one of them, if desired, once -during startup/initialization — before any GPU resource is created — rather than have an -individual class call it lazily on first use, which cannot be made race-free against -concurrent construction of that class from multiple threads. - **Parameters** | Name | Description | @@ -426,7 +419,7 @@ concurrent construction of that class from multiple threads. | `initialPoolSizePercent` | The initial pool size, in percentage of the total GPU memory | | `maxPoolSizePercent` | The maximum pool size, in percentage of the total GPU memory | -_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/CuVSProvider.java:250`_ +_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/CuVSProvider.java:243`_ ### enableRMMManagedPooledMemory @@ -438,8 +431,6 @@ Switch RMM allocations (used internally by various cuVS algorithms and by the de `CuVSDeviceMatrix`) to use pooled memory. This operation has a global effect, and will affect all resources on the current device. -Ownership: see `#enableRMMPooledMemory`. - **Parameters** | Name | Description | @@ -447,7 +438,7 @@ Ownership: see `#enableRMMPooledMemory`. | `initialPoolSizePercent` | The initial pool size, in percentage of the total GPU memory | | `maxPoolSizePercent` | The maximum pool size, in percentage of the total GPU memory | -_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/CuVSProvider.java:262`_ +_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/CuVSProvider.java:253`_ ### enableRMMAsyncMemory @@ -462,9 +453,7 @@ on deallocation. This is especially beneficial when multiple CAGRA searches run on separate CUDA streams, because internal workspace allocations no longer serialize kernel launches. This operation has a global effect and will affect all resources on the current device. -Ownership: see `#enableRMMPooledMemory`. - -_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/CuVSProvider.java:274`_ +_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/CuVSProvider.java:263`_ ### resetRMMPooledMemory @@ -474,9 +463,7 @@ void resetRMMPooledMemory() Disables pooled memory on the current device, reverting back to the default setting. -Ownership: see `#enableRMMPooledMemory`. - -_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/CuVSProvider.java:281`_ +_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/CuVSProvider.java:266`_ ### provider @@ -486,7 +473,7 @@ static CuVSProvider provider() Retrieves the system-wide provider. -_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/CuVSProvider.java:284`_ +_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/CuVSProvider.java:269`_ ### cagraIndexParamsFromHnswParams @@ -516,7 +503,7 @@ may be shifted along the curve right or left. See the heuristics descriptions fo A new CAGRA index parameters object -_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/CuVSProvider.java:304`_ +_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/CuVSProvider.java:289`_ ### cagraIndexParamsFromDataset @@ -540,6 +527,6 @@ Create CAGRA index parameters heuristically tuned for a dataset. A new CAGRA index parameters object -_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/CuVSProvider.java:322`_ +_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/CuVSProvider.java:307`_ _Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/CuVSProvider.java:18`_ diff --git a/java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/CuVSProvider.java b/java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/CuVSProvider.java index cb44f9b672..44706d0b6a 100644 --- a/java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/CuVSProvider.java +++ b/java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/CuVSProvider.java @@ -238,13 +238,6 @@ MultiPartitionSearchResults searchCagraMultiPartition( * {@link CuVSDeviceMatrix}) to use pooled memory. * This operation has a global effect, and will affect all resources on the current device. * - *

Ownership: this and the other {@code enableRMM*}/{@code resetRMMPooledMemory} methods are - * not called automatically by any cuVS/cuvs-lucene class. Because the effect is global to the - * current device, it is application code's responsibility to call one of them, if desired, once - * during startup/initialization — before any GPU resource is created — rather than have an - * individual class call it lazily on first use, which cannot be made race-free against - * concurrent construction of that class from multiple threads. - * * @param initialPoolSizePercent The initial pool size, in percentage of the total GPU memory * @param maxPoolSizePercent The maximum pool size, in percentage of the total GPU memory */ @@ -255,8 +248,6 @@ MultiPartitionSearchResults searchCagraMultiPartition( * {@link CuVSDeviceMatrix}) to use pooled memory. * This operation has a global effect, and will affect all resources on the current device. * - *

Ownership: see {@link #enableRMMPooledMemory}. - * * @param initialPoolSizePercent The initial pool size, in percentage of the total GPU memory * @param maxPoolSizePercent The maximum pool size, in percentage of the total GPU memory */ @@ -269,16 +260,10 @@ MultiPartitionSearchResults searchCagraMultiPartition( * on deallocation. This is especially beneficial when multiple CAGRA searches run concurrently * on separate CUDA streams, because internal workspace allocations no longer serialize kernel * launches. This operation has a global effect and will affect all resources on the current device. - * - *

Ownership: see {@link #enableRMMPooledMemory}. */ void enableRMMAsyncMemory(); - /** - * Disables pooled memory on the current device, reverting back to the default setting. - * - *

Ownership: see {@link #enableRMMPooledMemory}. - */ + /** Disables pooled memory on the current device, reverting back to the default setting. */ void resetRMMPooledMemory(); /** Retrieves the system-wide provider. */ diff --git a/java/cuvs-lucene/pom.xml b/java/cuvs-lucene/pom.xml index 925bd377da..cf7628137e 100644 --- a/java/cuvs-lucene/pom.xml +++ b/java/cuvs-lucene/pom.xml @@ -123,11 +123,6 @@ SPDX-License-Identifier: Apache-2.0 lucene-backward-codecs 10.2.0 - - org.apache.lucene - lucene-misc - 10.2.0 - org.apache.lucene lucene-test-framework diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java index b57614539a..a5f164b70b 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java @@ -59,7 +59,7 @@ public static enum Strategy { public static final int DEFAULT_MAX_CONN = 32; public static final int DEFAULT_BEAM_WIDTH = 32; public static final CagraGraphBuildAlgo DEFAULT_CAGRA_GRAPH_BUILD_ALGO = - CagraGraphBuildAlgo.AUTO_SELECT; + CagraGraphBuildAlgo.NN_DESCENT; public static final int DEFAULT_NUM_MERGE_WORKERS = 1; public static final Strategy DEFAULT_STRATEGY = Strategy.HEURISTIC; public static final CuvsDistanceType DEFAULT_CUVS_DISTANCE_TYPE = CuvsDistanceType.L2Expanded; @@ -97,22 +97,17 @@ public static enum Strategy { * * @param writerThreads Number of cuVS writer threads to use. * @param intermediateGraphDegree The intermediate graph degree while building the CAGRA index. - * Only consulted under the {@link Strategy#CUSTOM} strategy. - * @param graphdegree The graph degree to use while building the CAGRA index. Only consulted - * under the {@link Strategy#CUSTOM} strategy. + * @param graphdegree The graph degree to use while building the CAGRA index. * @param hnswLayers The number of HNSW layers to build in the HNSW index. * @param maxConn The max connection parameter used when building HNSW index with the fallback mechanism. * @param beamWidth The beam width parameter used when building HNSW index with the fallback mechanism. - * @param cagraGraphBuildAlgo The CAGRA graph build algorithm to use [NN_DESCENT, IVF_PQ]. Only - * consulted under the {@link Strategy#CUSTOM} strategy. + * @param cagraGraphBuildAlgo The CAGRA graph build algorithm to use [NN_DESCENT, IVF_PQ]. * @param cuVSIvfPqParams An instance of CuVSIvfPqParams containing IVF_PQ specific parameters. - * Only consulted under the {@link Strategy#CUSTOM} strategy. * @param numMergeWorkers The number of merge workers to use with the fallback mechanism. * @param mergeExec The instance of {@link ExecutorService} to use with the fallback mechanism. * @param strategy either HEURISTIC [Default] that delegates the CAGRA build parameters to cuVS (derived from the HNSW-equivalent maxConn and beamWidth) or CUSTOM that uses the parameters passed through this class. * @param cuvsDistanceType the cuvsDistanceType. The default option is L2Expanded. * @param nnDescentNumIterations the number of Iterations to run if building with NN_DESCENT. - * Only consulted under the {@link Strategy#CUSTOM} strategy. * @param hnswHeuristicType the heuristic cuVS applies when deriving the CAGRA build parameters from maxConn and beamWidth under the HEURISTIC strategy. */ private AcceleratedHNSWParams( @@ -157,7 +152,7 @@ public int getWriterThreads() { } /** - * Get the intermediate graph degree. Only consulted under the {@link Strategy#CUSTOM} strategy. + * Get the intermediate graph degree * * @return the graph degree parameter */ @@ -166,7 +161,7 @@ public int getIntermediateGraphDegree() { } /** - * Get the graph degree. Only consulted under the {@link Strategy#CUSTOM} strategy. + * Get the graph degree * * @return the graph degree parameter */ @@ -202,8 +197,7 @@ public int getBeamWidth() { } /** - * Get the CAGRA graph build algorithm. Only consulted under the {@link Strategy#CUSTOM} - * strategy; under {@link Strategy#HEURISTIC} the algorithm is chosen by cuVS. + * Get the CAGRA graph build algorithm * * @return the CAGRA graph build algorithm */ @@ -212,8 +206,7 @@ public CagraGraphBuildAlgo getCagraGraphBuildAlgo() { } /** - * Get the instance of {@link CuVSIvfPqParams}. Only consulted under the {@link - * Strategy#CUSTOM} strategy. + * Get the instance of {@link CuVSIvfPqParams} * * @return the instance of {@link CuVSIvfPqParams} */ @@ -261,8 +254,7 @@ public CuvsDistanceType getCuvsDistanceType() { } /** - * get the number of Iterations to run if building with NN_DESCENT. Only consulted under the - * {@link Strategy#CUSTOM} strategy. + * get the number of Iterations to run if building with NN_DESCENT * * @return the number of iterations for NN_DESCENT */ @@ -347,8 +339,7 @@ public Builder withWriterThreads(int writerThreads) { } /** - * Set the intermediate graph degree to use while building CAGRA index. Only consulted under - * the {@link Strategy#CUSTOM} strategy. + * Set the intermediate graph degree to use while building CAGRA index * Valid range - Minimum: {@value MIN_INT_GRAPH_DEG}, Maximum: {@value MAX_INT_GRAPH_DEG} * Default value - {@value DEFAULT_INT_GRAPH_DEGREE} * @@ -361,8 +352,7 @@ public Builder withIntermediateGraphDegree(int intermediateGraphDegree) { } /** - * Set the graph degree to use while building CAGRA index. Only consulted under the {@link - * Strategy#CUSTOM} strategy. + * Set the graph degree to use while building CAGRA index * Valid range - Minimum: {@value MIN_GRAPH_DEG}, Maximum: {@value MAX_GRAPH_DEG} * Default value - {@value DEFAULT_GRAPH_DEGREE} * @@ -414,9 +404,8 @@ public Builder withBeamWidth(int beamWidth) { } /** - * Set the CAGRA graph build algorithm to use. Only consulted under the {@link - * Strategy#CUSTOM} strategy; under {@link Strategy#HEURISTIC} the algorithm is chosen by cuVS. - * Default value - AUTO_SELECT + * Set the CAGRA graph build algorithm to use + * Default value - NN_DESCENT * * @param cagraGraphBuildAlgo * @return instance of {@link Builder} @@ -427,8 +416,7 @@ public Builder withCagraGraphBuildAlgo(CagraGraphBuildAlgo cagraGraphBuildAlgo) } /** - * Set the instance of {@link CuVSIvfPqParams}. Only consulted under the {@link - * Strategy#CUSTOM} strategy. + * Set the instance of {@link CuVSIvfPqParams} * * @param cuVSIvfPqParams * @return instance of {@link Builder} @@ -491,8 +479,7 @@ public Builder withCuvsDistanceType(CuvsDistanceType cuvsDistanceType) { } /** - * Set the number of Iterations to run if building with NN_DESCENT. Only consulted under the - * {@link Strategy#CUSTOM} strategy. + * Set the number of Iterations to run if building with NN_DESCENT * * Valid range - Minimum: {@value MIN_NN_DESCENT_NUM_ITERATIONS}, Maximum: {@value MAX_NN_DESCENT_NUM_ITERATIONS} * Default value - {@value DEFAULT_NN_DESCENT_NUM_ITERATIONS} diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java index fdd7dd441f..9c49c07fe0 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java @@ -19,14 +19,8 @@ import java.util.Random; import java.util.SortedSet; import java.util.TreeSet; -import java.util.concurrent.Callable; -import java.util.concurrent.ExecutorService; -import java.util.concurrent.Executors; import org.apache.lucene.index.FieldInfo; import org.apache.lucene.index.VectorSimilarityFunction; -import org.apache.lucene.search.TaskExecutor; -import org.apache.lucene.store.ByteBuffersDataOutput; -import org.apache.lucene.store.DataOutput; import org.apache.lucene.store.IndexOutput; import org.apache.lucene.util.InfoStream; import org.apache.lucene.util.hnsw.HnswGraph; @@ -75,7 +69,7 @@ public static GPUBuiltHnswGraph createSingleVectorHnswGraph(int size, int dimens layerAdjacencies.add(adjacencyMatrix); // Create the single-layer graph - return new GPUBuiltHnswGraph(size, dimensions, layerNodes, layerAdjacencies, 1); + return new GPUBuiltHnswGraph(size, dimensions, layerNodes, layerAdjacencies); } /** @@ -84,25 +78,18 @@ public static GPUBuiltHnswGraph createSingleVectorHnswGraph(int size, int dimens * (its column count). Ceil is used to accommodate odd graph degrees. * Each layer contains 1/M nodes from the previous layer * Creates layers until the highest layer has ≤ M nodes - *

- * Vectors for higher-layer subsets are read directly from the native matrix - * via {@link CuVSMatrix#getRow(long)} and {@link RowView#toArray(float[])}, - * avoiding any additional heap allocation of the full dataset. Used by both - * the flush and merge paths; the caller provides the vectors as a - * {@link CuVSMatrix}. */ public static GPUBuiltHnswGraph createMultiLayerHnswGraph( FieldInfo fieldInfo, + int size, int dimensions, CuVSMatrix adjacencyListMatrix, - CuVSMatrix vectorDataset, + List vectors, int hnswLayers, CagraIndexParams params, - QuantizationType quantization, - int numThreads) + QuantizationType quantization) throws Throwable { - int size = (int) vectorDataset.size(); int M = Math.ceilDiv((int) adjacencyListMatrix.columns(), 2); // Store all layers data @@ -132,7 +119,8 @@ public static GPUBuiltHnswGraph createMultiLayerHnswGraph( // Select from previous layer nodes int[] prevLayerNodes = layerNodes.get(layerNodes.size() - 1); while (selectedNodesSet.size() < nextLayerSize) { - selectedNodesSet.add(prevLayerNodes[random.nextInt(prevLayerNodes.length)]); + int idx = random.nextInt(prevLayerNodes.length); + selectedNodesSet.add(prevLayerNodes[idx]); } } @@ -143,22 +131,24 @@ public static GPUBuiltHnswGraph createMultiLayerHnswGraph( layerNodes.add(selectedNodes); if (quantization == QuantizationType.NONE) { - // Read only the sampled rows from the native matrix — no full-dataset heap copy - float[][] selectedVectors = new float[nextLayerSize][dimensions]; + // Extract vectors for selected nodes + float[][] selectedVectors = new float[nextLayerSize][]; for (int i = 0; i < nextLayerSize; i++) { - vectorDataset.getRow(selectedNodes[i]).toArray(selectedVectors[i]); + selectedVectors[i] = (float[]) vectors.get(selectedNodes[i]); } // Build CAGRA graph for this layer layerAdjacencies.add( buildCagraGraphForSubset( selectedVectors, selectedNodes, 0, params, dimensions, quantization)); + } else { - // Byte width comes from the matrix itself: binary packs 8 dims/byte, scalar is 1 byte/dim. - int bytesPerVector = (int) vectorDataset.columns(); - byte[][] selectedVectors = new byte[nextLayerSize][bytesPerVector]; + + // Extract vectors for selected nodes + int bytesPerVector = (dimensions + 7) / 8; + byte[][] selectedVectors = new byte[nextLayerSize][]; for (int i = 0; i < nextLayerSize; i++) { - vectorDataset.getRow(selectedNodes[i]).toArray(selectedVectors[i]); + selectedVectors[i] = (byte[]) vectors.get(selectedNodes[i]); } // Build CAGRA graph for this layer @@ -176,7 +166,7 @@ public static GPUBuiltHnswGraph createMultiLayerHnswGraph( } // Create the multi-layer graph with all layers - return new GPUBuiltHnswGraph(size, dimensions, layerNodes, layerAdjacencies, numThreads); + return new GPUBuiltHnswGraph(size, dimensions, layerNodes, layerAdjacencies); } /** @@ -194,9 +184,11 @@ private static CuVSMatrix buildCagraGraphForSubset( CuVSMatrix subsetDataset; if (quantization == QuantizationType.BINARY) { - subsetDataset = createByteMatrixFromArray((byte[][]) vectors, bytesPerVector); + subsetDataset = + createByteMatrixFromArray((byte[][]) vectors, bytesPerVector, getCuVSResourcesInstance()); } else if (quantization == QuantizationType.SCALAR) { - subsetDataset = createByteMatrixFromArray((byte[][]) vectors, dimensions); + subsetDataset = + createByteMatrixFromArray((byte[][]) vectors, dimensions, getCuVSResourcesInstance()); } else { subsetDataset = CuVSMatrix.ofArray((float[][]) vectors); } @@ -243,154 +235,56 @@ private static CuVSMatrix buildCagraGraphForSubset( * @return a 2D array of offsets * @throws IOException I/O Exceptions */ - public static int[][] writeGraph(GPUBuiltHnswGraph graph, IndexOutput vectorIndex, int numThreads) + public static int[][] writeGraph(GPUBuiltHnswGraph graph, IndexOutput vectorIndex) throws IOException { // write vectors' neighbors on each level into the vectorIndex file int countOnLevel0 = graph.size(); - int numLevels = graph.numLevels(); - int[][] offsets = new int[numLevels][]; - - // Level 0 holds all nodes and dominates serialization cost. Each node's delta/VInt block is - // independent, so encode level 0 in parallel and concatenate the per-thread buffers serially in - // node order, in memory-bounded waves. Higher levels are tiny and stay serial. The on-disk - // bytes - // are identical to the fully-serial path (blocks in node order, offsets = per-node byte - // lengths). - // graph.maxConn() scans every layer-0 adjacency row (O(graph size)); compute it once here - // rather than per level/per task below. - int maxConn = graph.maxConn(); - - int[] level0Nodes = NodesIterator.getSortedNodes(graph.getNodesOnLevel(0)); - offsets[0] = new int[level0Nodes.length]; - if (numThreads > 1 && level0Nodes.length >= PARALLEL_MIN_NODES) { - writeLevel0Parallel( - graph, vectorIndex, level0Nodes, offsets[0], countOnLevel0, maxConn, numThreads); - } else { - writeLevelSerial(graph, vectorIndex, 0, level0Nodes, offsets[0], countOnLevel0, maxConn); - } - - for (int level = 1; level < numLevels; level++) { + int[][] offsets = new int[graph.numLevels()][]; + int[] scratch = new int[graph.maxConn() * 2]; + for (int level = 0; level < graph.numLevels(); level++) { int[] sortedNodes = NodesIterator.getSortedNodes(graph.getNodesOnLevel(level)); offsets[level] = new int[sortedNodes.length]; - writeLevelSerial( - graph, vectorIndex, level, sortedNodes, offsets[level], countOnLevel0, maxConn); - } - // Return offsets (information written while writing the meta info) - return offsets; - } - - /** Node count below which parallel level-0 serialization is not worth the overhead. */ - private static final int PARALLEL_MIN_NODES = 1 << 16; - - /** Nodes per wave — bounds the transient encode buffer regardless of dataset size. */ - private static final int WAVE_NODES = 1 << 20; - - /** Serially encodes a level's nodes into {@code out}, recording per-node byte lengths. */ - private static void writeLevelSerial( - GPUBuiltHnswGraph graph, - IndexOutput out, - int level, - int[] sortedNodes, - int[] offsets, - int countOnLevel0, - int maxConn) - throws IOException { - int[] scratch = new int[maxConn * 2]; - int idx = 0; - for (int node : sortedNodes) { - long start = out.getFilePointer(); - encodeNode(graph.getNeighbors(level, node), scratch, out, countOnLevel0); - offsets[idx++] = Math.toIntExact(out.getFilePointer() - start); - } - } - - /** - * Encodes level 0 in parallel: within memory-bounded waves, threads encode contiguous node - * sub-ranges into per-thread buffers, which are then concatenated to {@code out} in node order - * (identical layout to the serial path). - */ - private static void writeLevel0Parallel( - GPUBuiltHnswGraph graph, - IndexOutput out, - int[] nodes, - int[] offsets, - int countOnLevel0, - int maxConn, - int numThreads) - throws IOException { - // invokeAll joins every task before it returns, including tasks still running when a sibling - // fails, so `buffers` is never concatenated while a worker might still be writing into it. - // It also runs one share of each wave on the calling thread instead of parking it, so the - // pool only has to cover the other ranges. - ExecutorService pool = Executors.newFixedThreadPool(Math.max(1, numThreads - 1)); - try { - TaskExecutor executor = new TaskExecutor(pool); - int n = nodes.length; - for (int waveStart = 0; waveStart < n; waveStart += WAVE_NODES) { - int waveEnd = Math.min(waveStart + WAVE_NODES, n); - int perThread = (waveEnd - waveStart + numThreads - 1) / numThreads; - - ByteBuffersDataOutput[] buffers = new ByteBuffersDataOutput[numThreads]; - List> tasks = new ArrayList<>(numThreads); - for (int t = 0; t < numThreads; t++) { - final int subStart = waveStart + t * perThread; - final int subEnd = Math.min(subStart + perThread, waveEnd); - final int slot = t; - if (subStart >= subEnd) { - continue; - } - tasks.add( - () -> { - ByteBuffersDataOutput buffer = new ByteBuffersDataOutput(); - int[] scratch = new int[maxConn * 2]; - for (int i = subStart; i < subEnd; i++) { - long before = buffer.size(); - encodeNode(graph.getNeighbors(0, nodes[i]), scratch, buffer, countOnLevel0); - offsets[i] = Math.toIntExact(buffer.size() - before); - } - buffers[slot] = buffer; - return null; - }); + int nodeOffsetId = 0; + + for (int node : sortedNodes) { + // Get node neighbors + NeighborArray neighbors = graph.getNeighbors(level, node); + // Get the size of the neighbor array + int size = neighbors.size(); + // Write size in VInt as the neighbors list is typically small + long offsetStart = vectorIndex.getFilePointer(); + // Get neighbors + int[] nnodes = neighbors.nodes(); + // Sort them + Arrays.sort(nnodes, 0, size); + // Now that we have sorted, do delta encoding to minimize the required bits to store the + // information + int actualSize = 0; + if (size > 0) { + scratch[0] = nnodes[0]; + actualSize = 1; } - executor.invokeAll(tasks); - // Concatenate in thread order (== node order), preserving the serial byte layout. - for (ByteBuffersDataOutput buffer : buffers) { - if (buffer != null) { - buffer.copyTo(out); + // De-duplication + for (int i = 1; i < size; i++) { + assert nnodes[i] < countOnLevel0 : "node too large: " + nnodes[i] + ">=" + countOnLevel0; + // Sorting step helps here + if (nnodes[i - 1] == nnodes[i]) { + continue; } + scratch[actualSize++] = nnodes[i] - nnodes[i - 1]; } - } - } finally { - pool.shutdown(); - } - } - - /** - * Sorts, delta-encodes and de-duplicates a node's neighbors and writes the block (VInt size + VInt - * deltas) to {@code out}. Shared by the serial and parallel paths so encoding is identical. - */ - private static void encodeNode( - NeighborArray neighbors, int[] scratch, DataOutput out, int countOnLevel0) - throws IOException { - int size = neighbors == null ? 0 : neighbors.size(); - int actualSize = 0; - if (size > 0) { - int[] nnodes = neighbors.nodes(); - Arrays.sort(nnodes, 0, size); - scratch[0] = nnodes[0]; - actualSize = 1; - for (int i = 1; i < size; i++) { - assert nnodes[i] < countOnLevel0 : "node too large: " + nnodes[i] + ">=" + countOnLevel0; - if (nnodes[i - 1] == nnodes[i]) { - continue; + // Write the size after duplicates are removed + vectorIndex.writeVInt(actualSize); + // Write de-duplicated neighbors + for (int i = 0; i < actualSize; i++) { + vectorIndex.writeVInt(scratch[i]); } - scratch[actualSize++] = nnodes[i] - nnodes[i - 1]; + offsets[level][nodeOffsetId++] = + Math.toIntExact(vectorIndex.getFilePointer() - offsetStart); } } - out.writeVInt(actualSize); - for (int i = 0; i < actualSize; i++) { - out.writeVInt(scratch[i]); - } + // Return offsets (information written while writing the meta info) + return offsets; } /** diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHnswGraphOutput.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHnswGraphOutput.java deleted file mode 100644 index e5317f4a01..0000000000 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHnswGraphOutput.java +++ /dev/null @@ -1,214 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: Apache-2.0 - */ -package com.nvidia.cuvs.lucene; - -import static com.nvidia.cuvs.lucene.AcceleratedHNSWUtils.createMultiLayerHnswGraph; -import static com.nvidia.cuvs.lucene.AcceleratedHNSWUtils.createSingleVectorHnswGraph; -import static com.nvidia.cuvs.lucene.AcceleratedHNSWUtils.writeEmpty; -import static com.nvidia.cuvs.lucene.AcceleratedHNSWUtils.writeGraph; -import static com.nvidia.cuvs.lucene.AcceleratedHNSWUtils.writeMeta; -import static com.nvidia.cuvs.lucene.Lucene99AcceleratedHNSWVectorsFormat.HNSW_INDEX_CODEC_NAME; -import static com.nvidia.cuvs.lucene.Lucene99AcceleratedHNSWVectorsFormat.HNSW_INDEX_EXT; -import static com.nvidia.cuvs.lucene.Lucene99AcceleratedHNSWVectorsFormat.HNSW_META_CODEC_EXT; -import static com.nvidia.cuvs.lucene.Lucene99AcceleratedHNSWVectorsFormat.HNSW_META_CODEC_NAME; -import static com.nvidia.cuvs.lucene.ThreadLocalCuVSResourcesProvider.getCuVSResourcesInstance; - -import com.nvidia.cuvs.CagraIndex; -import com.nvidia.cuvs.CagraIndexParams; -import com.nvidia.cuvs.CuVSMatrix; -import com.nvidia.cuvs.lucene.AcceleratedHNSWUtils.QuantizationType; -import java.io.Closeable; -import java.io.IOException; -import java.util.List; -import org.apache.lucene.codecs.CodecUtil; -import org.apache.lucene.index.FieldInfo; -import org.apache.lucene.index.IndexFileNames; -import org.apache.lucene.index.SegmentWriteState; -import org.apache.lucene.store.IndexOutput; -import org.apache.lucene.util.IOUtils; - -/** - * Owns the {@code .vem}/{@code .vex} outputs and builds/writes the CAGRA-derived HNSW graph for a - * field. Shared by both {@link Lucene99AcceleratedHNSWVectorsWriter} (heap-buffered flat path) - * and {@link NativeFlatBufferedHNSWVectorsWriter} (native flat-buffered path) — the only - * difference between those two writers is how the flat {@code .vec}/{@code .vemf} files and the - * dataset handed to {@link #writeField(FieldInfo, CuVSMatrix)} are produced. - */ -final class AcceleratedHnswGraphOutput implements Closeable { - - private static final LuceneProvider LUCENE_PROVIDER; - private static final Integer VERSION_CURRENT; - - static { - try { - LUCENE_PROVIDER = LuceneProvider.getInstance("99"); - VERSION_CURRENT = LUCENE_PROVIDER.getStaticIntParam("VERSION_CURRENT"); - } catch (Exception e) { - throw new ExceptionInInitializerError(e.getMessage()); - } - } - - private final AcceleratedHNSWParams acceleratedHNSWParams; - private final IndexOutput hnswMeta; - private final IndexOutput hnswVectorIndex; - private boolean finished; - - AcceleratedHnswGraphOutput(SegmentWriteState state, AcceleratedHNSWParams acceleratedHNSWParams) - throws IOException { - this.acceleratedHNSWParams = acceleratedHNSWParams; - String vemFileName = - IndexFileNames.segmentFileName( - state.segmentInfo.name, state.segmentSuffix, HNSW_META_CODEC_EXT); - String vexFileName = - IndexFileNames.segmentFileName(state.segmentInfo.name, state.segmentSuffix, HNSW_INDEX_EXT); - IndexOutput meta = null; - IndexOutput vectorIndex = null; - boolean success = false; - try { - meta = state.directory.createOutput(vemFileName, state.context); - vectorIndex = state.directory.createOutput(vexFileName, state.context); - CodecUtil.writeIndexHeader( - meta, - HNSW_META_CODEC_NAME, - VERSION_CURRENT, - state.segmentInfo.getId(), - state.segmentSuffix); - CodecUtil.writeIndexHeader( - vectorIndex, - HNSW_INDEX_CODEC_NAME, - VERSION_CURRENT, - state.segmentInfo.getId(), - state.segmentSuffix); - success = true; - } finally { - this.hnswMeta = meta; - this.hnswVectorIndex = vectorIndex; - if (success == false) { - IOUtils.closeWhileHandlingException(this); - } - } - } - - /** - * Flush/sorting path: builds a host matrix from the heap vectors, then delegates to {@link - * #writeField(FieldInfo, CuVSMatrix)}. - */ - void writeField(FieldInfo fieldInfo, List vectors) throws IOException { - if (vectors.size() == 0) { - writeEmpty(fieldInfo, hnswMeta); - return; - } - if (vectors.size() < 2) { - writeSingleVectorGraph(fieldInfo, vectors); - return; - } - CuVSMatrix dataset = Utils.createFloatMatrix(vectors, fieldInfo.getVectorDimension()); - writeField(fieldInfo, dataset); - } - - /** - * Builds the intermediate CAGRA index and builds and writes the HNSW index. Single - * implementation used by both the flush and merge paths of both writers. The dataset is a - * {@link CuVSMatrix} (host-backed on the merge/native paths) so the full set of vectors is - * never double-materialised on the Java heap. - */ - void writeField(FieldInfo fieldInfo, CuVSMatrix dataset) throws IOException { - try (dataset) { - int size = (int) dataset.size(); - if (size == 0) { - writeEmpty(fieldInfo, hnswMeta); - return; - } - if (size < 2) { - float[] buf = new float[fieldInfo.getVectorDimension()]; - dataset.getRow(0).toArray(buf); - writeSingleVectorGraph(fieldInfo, List.of(buf)); - return; - } - try { - CagraIndexParams params = - CagraIndexParamsFactory.create( - acceleratedHNSWParams, dataset.size(), dataset.columns()); - try (CagraIndex cagraIndex = - CagraIndex.newBuilder(getCuVSResourcesInstance()) - .withDataset(dataset) - .withIndexParams(params) - .build()) { - CuVSMatrix adjacencyListMatrix = cagraIndex.getGraph(); - int dimensions = fieldInfo.getVectorDimension(); - GPUBuiltHnswGraph hnswGraph = - createMultiLayerHnswGraph( - fieldInfo, - dimensions, - adjacencyListMatrix, - dataset, - acceleratedHNSWParams.getHnswLayers(), - params, - QuantizationType.NONE, - acceleratedHNSWParams.getWriterThreads()); - long vectorIndexOffset = hnswVectorIndex.getFilePointer(); - int[][] graphLevelNodeOffsets = - writeGraph(hnswGraph, hnswVectorIndex, acceleratedHNSWParams.getWriterThreads()); - long vectorIndexLength = hnswVectorIndex.getFilePointer() - vectorIndexOffset; - writeMeta( - hnswVectorIndex, - hnswMeta, - fieldInfo, - vectorIndexOffset, - vectorIndexLength, - size, - hnswGraph, - graphLevelNodeOffsets); - } - } catch (Throwable t) { - Utils.handleThrowable(t); - } - } - } - - private void writeSingleVectorGraph(FieldInfo fieldInfo, List vectors) - throws IOException { - try { - int size = 1; - int dimensions = fieldInfo.getVectorDimension(); - GPUBuiltHnswGraph hnswGraph = createSingleVectorHnswGraph(size, dimensions); - long vectorIndexOffset = hnswVectorIndex.getFilePointer(); - int[][] graphLevelNodeOffsets = - writeGraph(hnswGraph, hnswVectorIndex, acceleratedHNSWParams.getWriterThreads()); - long vectorIndexLength = hnswVectorIndex.getFilePointer() - vectorIndexOffset; - writeMeta( - hnswVectorIndex, - hnswMeta, - fieldInfo, - vectorIndexOffset, - vectorIndexLength, - size, - hnswGraph, - graphLevelNodeOffsets); - } catch (Throwable t) { - Utils.handleThrowable(t); - } - } - - void finish() throws IOException { - if (finished) { - throw new IllegalStateException("already finished"); - } - finished = true; - if (hnswMeta != null) { - // write end of fields marker - hnswMeta.writeInt(-1); - CodecUtil.writeFooter(hnswMeta); - } - if (hnswVectorIndex != null) { - CodecUtil.writeFooter(hnswVectorIndex); - } - } - - @Override - public void close() throws IOException { - IOUtils.close(hnswMeta, hnswVectorIndex); - } -} diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CagraHnswBulkIndexWriter.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CagraHnswBulkIndexWriter.java deleted file mode 100644 index c1cd5e71b4..0000000000 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CagraHnswBulkIndexWriter.java +++ /dev/null @@ -1,660 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: Apache-2.0 - */ -package com.nvidia.cuvs.lucene; - -import java.io.Closeable; -import java.io.IOException; -import java.nio.file.Files; -import java.nio.file.Path; -import java.util.ArrayList; -import java.util.Comparator; -import java.util.List; -import java.util.Objects; -import java.util.concurrent.ExecutorService; -import java.util.concurrent.Executors; -import java.util.concurrent.Future; -import java.util.concurrent.Semaphore; -import java.util.stream.Stream; -import org.apache.lucene.codecs.Codec; -import org.apache.lucene.document.Document; -import org.apache.lucene.document.Field; -import org.apache.lucene.document.KnnFloatVectorField; -import org.apache.lucene.document.StringField; -import org.apache.lucene.index.IndexWriter; -import org.apache.lucene.index.IndexWriterConfig; -import org.apache.lucene.index.IndexableField; -import org.apache.lucene.index.NoMergePolicy; -import org.apache.lucene.index.VectorSimilarityFunction; -import org.apache.lucene.misc.store.HardlinkCopyDirectoryWrapper; -import org.apache.lucene.store.Directory; -import org.apache.lucene.store.FSDirectory; - -/** - * Builds a GPU-accelerated CAGRA/HNSW Lucene index from data the caller already has ready on - * local disk, as few segments as possible, without going through a generically-configurable - * {@link org.apache.lucene.codecs.Codec}. - * - *

This is not a general-purpose Lucene extension point: an instance owns a single - * {@link IndexWriter} and its {@link IndexWriterConfig} so that the invariants the underlying GPU - * writer requires (single unmerged segment, no index sort, exact vector count) are guaranteed by - * this class rather than left to a caller to get right. If you need a general Lucene codec that - * any indexing pipeline (including ones that don't control their own {@code IndexWriter} - * lifecycle, e.g. Solr or Elasticsearch) can register, use {@link Lucene101AcceleratedHNSWCodec} - * directly instead. - * - *

Two ways to use this class: - * - *

    - *
  • Manual, single segment: construct an instance directly, call {@link #addDocument} - * per document exactly like a plain {@link IndexWriter}, then {@link #close}. You own the - * loop and the {@link Document} you build (any fields, not just the vector). - *
  • One-shot, one or many segments: {@link #indexFbin} / {@link #build(VectorSource, - * Config)} own the loop for you — they read vectors from a {@code .fbin} file or {@link - * VectorSource}, optionally split into {@code numSegments} partitions (sequential or - * overlapped), and combine the result. Since they build each row's {@link Document} - * internally, an optional {@link FieldCallback} lets you add extra fields to it. - *
- * - *

Scope: CAGRA_HNSW (GPU build, CPU search) only. This class builds indexes for {@link - * Lucene101AcceleratedHNSWCodec}. It does not support the GPU-search codec ({@code - * CuVS2510GPUSearchCodec}) — that writer does not (yet) have the native flat-buffering - * optimization this class relies on. Reading an existing index for CAGRA_SEARCH-style GPU search - * is unaffected by this class either way; only bulk-building one is out of scope for now. - */ -public final class CagraHnswBulkIndexWriter implements Closeable { - - private static final int DEFAULT_CHUNK_SIZE_MB = 32; - - private final IndexWriter writer; - private final int exactVectorCount; - private int documentsAdded; - private boolean closed; - - /** - * Opens a single-segment, native-flat-buffered writer. {@code exactVectorCount} must equal the - * number of {@link #addDocument} calls that will follow — the native buffer is pre-sized to it, - * so {@link #close} rejects a mismatched count rather than let the underlying writer produce a - * corrupt or incomplete segment. - * - *

{@code conf}'s {@code Analyzer}, {@code Similarity}, {@code InfoStream}, and {@code - * OpenMode} are honored; an explicit {@code IndexSort} is rejected ({@link - * IllegalArgumentException}), since native flat buffering does not support index-sorted - * segments. Codec, merge policy, and flush thresholds are always owned by this class regardless - * of what {@code conf} contains — {@link IndexWriterConfig} does not expose whether the caller - * explicitly set those or left them at Lucene's defaults, so there is no reliable way to - * validate-and-reject a caller-supplied value for them the way {@code IndexSort} can be; this - * class simply never reads them from {@code conf}. - * - *

{@code config.targetDirectory()}, {@code config.numSegments()}, {@code - * config.overlapped()}, and {@code config.pipelineDepth()} are not consulted here — {@code - * directory} is passed explicitly, and this constructor always builds exactly one segment. Those - * fields only matter to {@link #indexFbin} / {@link #build(VectorSource, Config)}. - */ - public CagraHnswBulkIndexWriter( - Directory directory, IndexWriterConfig conf, Config config, int exactVectorCount) - throws Exception { - Objects.requireNonNull(directory, "directory"); - Objects.requireNonNull(conf, "conf"); - Objects.requireNonNull(config, "config"); - if (exactVectorCount <= 0) { - throw new IllegalArgumentException("exactVectorCount must be > 0, got " + exactVectorCount); - } - if (conf.getIndexSort() != null) { - throw new IllegalArgumentException( - "CagraHnswBulkIndexWriter does not support an index-sorted segment (native flat" - + " buffering requires an unsorted single-segment build); leave" - + " IndexWriterConfig.indexSort unset"); - } - this.exactVectorCount = exactVectorCount; - - Codec codec = new Lucene101AcceleratedHNSWCodec(config.graphBuildParams(), exactVectorCount); - IndexWriterConfig ownedConf = - new IndexWriterConfig(conf.getAnalyzer()) - .setSimilarity(conf.getSimilarity()) - .setInfoStream(conf.getInfoStream()) - .setCodec(codec) - .setUseCompoundFile(false) - .setMaxBufferedDocs(Math.max(2, exactVectorCount + 1)) - .setRAMBufferSizeMB(IndexWriterConfig.DISABLE_AUTO_FLUSH) - .setMergePolicy(NoMergePolicy.INSTANCE) - .setOpenMode(conf.getOpenMode()); - this.writer = new IndexWriter(directory, ownedConf); - } - - /** - * Adds one document, exactly like {@link IndexWriter#addDocument}. {@code doc} may contain any - * fields — the vector field (matching {@link Config#fieldName()}) is routed into the native - * flat buffer automatically by the underlying codec, the same way any {@link - * KnnFloatVectorField} is for any Lucene codec; every other field is indexed normally. - */ - public long addDocument(Iterable doc) throws IOException { - if (closed) { - throw new IllegalStateException("addDocument called after close()"); - } - if (documentsAdded >= exactVectorCount) { - throw new IllegalStateException( - "addDocument called more than exactVectorCount (" + exactVectorCount + ") times"); - } - long seqNo = writer.addDocument(doc); - documentsAdded++; - return seqNo; - } - - /** - * Runs the single native-buffered flush (this is where the GPU CAGRA build happens) and closes - * the underlying writer. There is no separate {@code commit()}: unlike a plain {@link - * IndexWriter}, only one flush is ever valid for this instance, so folding it into {@code - * close()} removes any way to trigger it early with fewer than {@code exactVectorCount} vectors - * added. Throws {@link IllegalStateException} if {@link #addDocument} was called fewer times - * than {@code exactVectorCount} promised, discarding the buffered documents rather than - * flushing an under-filled segment. - */ - @Override - public void close() throws IOException { - if (closed) { - return; - } - closed = true; - if (documentsAdded != exactVectorCount) { - IllegalStateException mismatch = - new IllegalStateException( - "expected " + exactVectorCount + " documents, got " + documentsAdded); - // The native host matrix is preallocated for exactly exactVectorCount rows, so an - // under-filled buffer can only be discarded, never flushed. rollback() discards it and - // closes the IndexWriter, and is also what frees that preallocated memory: it aborts the - // indexing chain, which closes the vectors writer. Leaving without it leaks the buffer, - // which lives in a shared Arena and is never reclaimed by the GC. - try { - writer.rollback(); - } catch (Throwable t) { - mismatch.addSuppressed(t); - } - throw mismatch; - } - try (writer) { - writer.commit(); - } - } - - /** - * Discards everything added so far without producing a segment: no GPU build runs, nothing is - * committed, and the preallocated native buffer is released (see {@link #close()} for why that - * release is what matters). Unlike {@link #close()} this does not care how many documents were - * added, so it is the cleanup to use on a failure path -- a partially filled buffer is not - * worth building, and reporting its count would only bury the failure that caused it. - * Idempotent, and a no-op once {@link #close()} has run. - */ - void abort() throws IOException { - if (closed) { - return; - } - closed = true; - writer.rollback(); - } - - /** - * Callback invoked once per row by {@link #indexFbin} / {@link #build(VectorSource, Config, - * FieldCallback)}, right after the id and vector fields have been added to {@code document} and - * right before it is added to the index — lets the caller attach additional fields (metadata) - * per vector. Not used by the manual, direct-instance API, where the caller already builds the - * whole {@link Document} themselves. - */ - @FunctionalInterface - public interface FieldCallback { - void addFields(Document document, int id) throws IOException; - } - - /** - * Builds an index from {@code source} into {@code config.targetDirectory()}, partitioned into - * {@code config.numSegments()} contiguous slices as {@code source} is consumed front-to-back. - * - *

{@code config.overlapped()} is not supported here: a single {@link VectorSource} is - * forward-only/single-consumer (see its contract) and so cannot be read concurrently by - * multiple slice builders. Use {@link #indexFbin} for the overlapped pipeline, which knows how - * to open independent, per-slice sources over the same underlying file. - */ - public static void build(VectorSource source, Config config) throws Exception { - build(source, config, null); - } - - /** As {@link #build(VectorSource, Config)}, with a {@link FieldCallback} for extra fields. */ - public static void build(VectorSource source, Config config, FieldCallback callback) - throws Exception { - Objects.requireNonNull(source, "source"); - Objects.requireNonNull(config, "config"); - if (config.overlapped()) { - throw new IllegalArgumentException( - "Config.overlapped() is not supported by build(VectorSource, Config): a VectorSource is" - + " forward-only and cannot be read by multiple concurrent slice builders. Use" - + " indexFbin(...) for the overlapped multi-segment build."); - } - if (source.dimensions() != config.dimensions()) { - throw new IllegalArgumentException( - "source.dimensions() (" - + source.dimensions() - + ") does not match Config.dimensions() (" - + config.dimensions() - + ")"); - } - List slices = sliceEvenly(source.size(), config.numSegments()); - buildSequential(source, config, callback, slices); - } - - /** - * Convenience entry point: builds an index directly from a {@code .fbin} file ({@code - * [num_vectors int32][dim int32]} header, then contiguous little-endian float32 rows), using - * {@link FbinVectorSource} with {@link #DEFAULT_CHUNK_SIZE_MB}-sized prefetched chunks. - */ - public static void indexFbin(Path fbinPath, Config config) throws Exception { - indexFbin(fbinPath, config, null, DEFAULT_CHUNK_SIZE_MB); - } - - /** As {@link #indexFbin(Path, Config)}, with a {@link FieldCallback} for extra fields. */ - public static void indexFbin(Path fbinPath, Config config, FieldCallback callback) - throws Exception { - indexFbin(fbinPath, config, callback, DEFAULT_CHUNK_SIZE_MB); - } - - /** As {@link #indexFbin(Path, Config)}, with an explicit prefetch chunk size. */ - public static void indexFbin(Path fbinPath, Config config, int chunkSizeMB) throws Exception { - indexFbin(fbinPath, config, null, chunkSizeMB); - } - - /** - * As {@link #indexFbin(Path, Config)}, with an explicit prefetch chunk size and {@link - * FieldCallback}. When {@code config.numSegments() > 1} and {@code config.overlapped()}, builds - * up to {@code config.pipelineDepth()} segments concurrently, each over its own slice of {@code - * fbinPath}, then combines them by hardlink; otherwise builds sequentially over one shared - * reader. - */ - public static void indexFbin( - Path fbinPath, Config config, FieldCallback callback, int chunkSizeMB) throws Exception { - Objects.requireNonNull(fbinPath, "fbinPath"); - Objects.requireNonNull(config, "config"); - int total; - int dim; - try (FbinVectorSource probe = new FbinVectorSource(fbinPath, 1)) { - total = probe.size(); - dim = probe.dimensions(); - } - if (dim != config.dimensions()) { - throw new IllegalArgumentException( - "fbinPath dimension (" - + dim - + ") does not match Config.dimensions() (" - + config.dimensions() - + ")"); - } - List slices = sliceEvenly(total, config.numSegments()); - if (config.overlapped() && slices.size() > 1) { - buildOverlapped(fbinPath, config, callback, chunkSizeMB, slices, dim); - } else { - try (FbinVectorSource source = new FbinVectorSource(fbinPath, chunkSizeMB)) { - buildSequential(source, config, callback, slices); - } - } - } - - /** - * Sequential partitioned build: {@code source} is streamed front-to-back across all slices, - * each slice built as a single native-flat segment appended to the same directory (first slice - * {@code CREATE}, later slices {@code APPEND}). Peak host memory is one slice's native buffer. - */ - private static void buildSequential( - VectorSource source, Config config, FieldCallback callback, List slices) - throws Exception { - float[] scratch = new float[config.dimensions()]; - try (Directory dir = FSDirectory.open(config.targetDirectory())) { - for (int p = 0; p < slices.size(); p++) { - int[] slice = slices.get(p); - buildSegment( - dir, source, scratch, config, callback, slice[0], slice[0], slice[1], p == 0, null); - } - } - } - - /** - * Overlapped partitioned build: a bounded pool builds up to {@code config.pipelineDepth()} - * segments at once, each with its OWN {@link FbinVectorSource} over just its slice, so a - * segment's ingest overlaps a prior segment's GPU commit. The GPU build itself is serialized on - * a single permit. The finished per-segment indexes are combined into {@code - * config.targetDirectory()} by hardlinking their files (no bulk copy of the vector data). - */ - private static void buildOverlapped( - Path fbinPath, - Config config, - FieldCallback callback, - int chunkSizeMB, - List slices, - int dim) - throws Exception { - Path targetDir = config.targetDirectory(); - int depth = Math.min(slices.size(), config.pipelineDepth()); - List segDirs = new ArrayList<>(); - for (int p = 0; p < slices.size(); p++) { - segDirs.add(targetDir.resolveSibling(targetDir.getFileName() + "_p" + p)); - } - for (Path segDir : segDirs) { - deleteRecursivelyQuietly(segDir); - } - try { - Semaphore gpuPermit = new Semaphore(1); // serialize the GPU CAGRA build across segments - ExecutorService pool = Executors.newFixedThreadPool(depth); - List> futures = new ArrayList<>(); - for (int p = 0; p < slices.size(); p++) { - int[] slice = slices.get(p); - Path segDir = segDirs.get(p); - futures.add( - pool.submit( - () -> { - float[] scratch = new float[dim]; - try (FbinVectorSource source = - new FbinVectorSource(fbinPath, slice[0], slice[1], chunkSizeMB); - Directory d = FSDirectory.open(segDir)) { - // createNew=true: each segment is a fresh single-segment index in its own - // dir; sourceStart=0 since this source is already windowed to the slice. - buildSegment( - d, source, scratch, config, callback, 0, slice[0], slice[1], true, - gpuPermit); - } - return null; - })); - } - pool.shutdown(); - try { - for (Future f : futures) { - f.get(); // propagate any build failure - } - } catch (Exception e) { - throw new IOException("Overlapped bulk index build failed", e); - } finally { - pool.shutdownNow(); - } - combineByHardlink(targetDir, segDirs); - } finally { - for (Path segDir : segDirs) { - deleteRecursivelyQuietly(segDir); - } - } - } - - /** - * Builds one segment from {@code size} vectors: {@code source.get(sourceStart + i, ...)} for - * {@code i} in {@code [0, size)}, labelled with ids {@code idStart + i} (the vectors' absolute - * position in the overall build, regardless of whether {@code source} itself is windowed). Opens - * a {@link CagraHnswBulkIndexWriter} sized to {@code size} and drives it exactly like the manual - * API. When {@code gpuPermit} is non-null the close (which runs the GPU CAGRA build) is - * serialized on it while other segments' host-side ingest may proceed. - */ - private static void buildSegment( - Directory dir, - VectorSource source, - float[] scratch, - Config config, - FieldCallback callback, - int sourceStart, - int idStart, - int size, - boolean createNew, - Semaphore gpuPermit) - throws Exception { - IndexWriterConfig conf = - new IndexWriterConfig() - .setOpenMode( - createNew ? IndexWriterConfig.OpenMode.CREATE : IndexWriterConfig.OpenMode.APPEND); - CagraHnswBulkIndexWriter writer = new CagraHnswBulkIndexWriter(dir, conf, config, size); - try { - for (int i = 0; i < size; i++) { - source.get(sourceStart + i, scratch); - int id = idStart + i; - Document doc = new Document(); - if (config.idFieldName() != null) { - doc.add(new StringField(config.idFieldName(), Integer.toString(id), Field.Store.YES)); - } - doc.add(new KnnFloatVectorField(config.fieldName(), scratch, config.similarity())); - if (callback != null) { - callback.addFields(doc, id); - } - writer.addDocument(doc); // copies the vector -> 'scratch' is safe to reuse next iteration - } - // The single flush inside close() is where the GPU CAGRA build runs; serialize it if asked. - if (gpuPermit != null) { - gpuPermit.acquire(); - try { - writer.close(); - } finally { - gpuPermit.release(); - } - } else { - writer.close(); - } - } catch (Throwable t) { - // Cleanup must not become the reported failure. abort() discards the partially filled - // buffer rather than trying to build it, and anything it throws on the way out is attached - // to the original exception instead of replacing it. If the close() above is what failed, - // this is a no-op -- that writer has already released itself. - try { - writer.abort(); - } catch (Throwable cleanupFailure) { - t.addSuppressed(cleanupFailure); - } - throw t; - } - } - - /** - * Combines the per-segment indexes into {@code targetDir} by hardlinking their files (same - * filesystem) rather than copying the vector data. {@link HardlinkCopyDirectoryWrapper} falls - * back to a byte copy automatically if the segment dirs and the final dir are on different - * filesystems. - */ - private static void combineByHardlink(Path targetDir, List segDirs) throws IOException { - Directory[] sources = new Directory[segDirs.size()]; - try { - for (int i = 0; i < segDirs.size(); i++) { - sources[i] = FSDirectory.open(segDirs.get(i)); - } - IndexWriterConfig iwc = - new IndexWriterConfig().setMergePolicy(NoMergePolicy.INSTANCE); // keep segments separate - try (Directory target = new HardlinkCopyDirectoryWrapper(FSDirectory.open(targetDir)); - IndexWriter combiner = new IndexWriter(target, iwc)) { - combiner.addIndexes(sources); - } - } finally { - for (Directory s : sources) { - if (s != null) { - s.close(); - } - } - } - } - - /** Splits {@code total} into {@code k} contiguous [start, size] slices, spreading the remainder. */ - private static List sliceEvenly(int total, int k) { - List slices = new ArrayList<>(); - int base = total / k; - int rem = total % k; - int start = 0; - for (int p = 0; p < k; p++) { - int size = base + (p < rem ? 1 : 0); // spread the remainder over the first slices - if (size <= 0) { - continue; - } - slices.add(new int[] {start, size}); - start += size; - } - return slices; - } - - private static void deleteRecursivelyQuietly(Path path) { - if (!Files.exists(path)) { - return; - } - try (Stream walk = Files.walk(path)) { - walk.sorted(Comparator.reverseOrder()) - .forEach( - p -> { - try { - Files.deleteIfExists(p); - } catch (IOException ignored) { - // best-effort cleanup of a temp per-segment dir - } - }); - } catch (IOException ignored) { - // best-effort cleanup of a temp per-segment dir - } - } - - /** Immutable configuration for {@link CagraHnswBulkIndexWriter}. */ - public static final class Config { - private final String fieldName; - private final int dimensions; - private final VectorSimilarityFunction similarity; - private final String idFieldName; - private final AcceleratedHNSWParams graphBuildParams; - private final Path targetDirectory; - private final int numSegments; - private final boolean overlapped; - private final int pipelineDepth; - - private Config(Builder b) { - this.fieldName = b.fieldName; - this.dimensions = b.dimensions; - this.similarity = b.similarity; - this.idFieldName = b.idFieldName; - this.graphBuildParams = b.graphBuildParams; - this.targetDirectory = b.targetDirectory; - this.numSegments = b.numSegments; - this.overlapped = b.overlapped; - this.pipelineDepth = b.pipelineDepth; - } - - public String fieldName() { - return fieldName; - } - - public int dimensions() { - return dimensions; - } - - public VectorSimilarityFunction similarity() { - return similarity; - } - - public String idFieldName() { - return idFieldName; - } - - public AcceleratedHNSWParams graphBuildParams() { - return graphBuildParams; - } - - /** Only consulted by {@link #indexFbin} / {@link #build(VectorSource, Config)}. */ - public Path targetDirectory() { - return targetDirectory; - } - - /** Only consulted by {@link #indexFbin} / {@link #build(VectorSource, Config)}. */ - public int numSegments() { - return numSegments; - } - - /** Only consulted by {@link #indexFbin}. */ - public boolean overlapped() { - return overlapped; - } - - /** Only consulted by {@link #indexFbin}. */ - public int pipelineDepth() { - return pipelineDepth; - } - - public static Builder builder() { - return new Builder(); - } - - /** Builder for {@link Config}. */ - public static final class Builder { - private String fieldName = "vector"; - private int dimensions = -1; - private VectorSimilarityFunction similarity = VectorSimilarityFunction.EUCLIDEAN; - private String idFieldName = "id"; - private AcceleratedHNSWParams graphBuildParams; - private Path targetDirectory; - private int numSegments = 1; - private boolean overlapped = false; - private int pipelineDepth = 2; - - /** Sets the vector field name and dimensionality; required. */ - public Builder field(String fieldName, int dimensions, VectorSimilarityFunction similarity) { - this.fieldName = Objects.requireNonNull(fieldName, "fieldName"); - this.dimensions = dimensions; - this.similarity = Objects.requireNonNull(similarity, "similarity"); - return this; - } - - /** - * Sets the stored id field name (holding each document's absolute position in the build, as - * a string), added automatically by {@link #indexFbin} / {@link #build(VectorSource, - * Config)} before the {@link FieldCallback} runs. Defaults to {@code "id"}; pass {@code - * null} to disable. Not used by the manual, direct-instance API. - */ - public Builder idField(String idFieldName) { - this.idFieldName = idFieldName; - return this; - } - - /** Sets the CAGRA/HNSW graph-build parameters; required. */ - public Builder graphBuild(AcceleratedHNSWParams graphBuildParams) { - this.graphBuildParams = Objects.requireNonNull(graphBuildParams, "graphBuildParams"); - return this; - } - - /** Sets the directory the final index is written to; required for {@link #indexFbin}/{@link #build}. */ - public Builder targetDirectory(Path targetDirectory) { - this.targetDirectory = Objects.requireNonNull(targetDirectory, "targetDirectory"); - return this; - } - - /** - * Splits the build into {@code numSegments} contiguous slices, each a single native-flat - * segment. Peak host memory scales as {@code 1/numSegments}; the GPU build itself is always - * serialized across slices regardless of this setting. Default 1 (single segment). - * - * @param overlapped when {@code numSegments > 1} and building via {@link #indexFbin}, - * builds up to {@link #pipelineDepth} slices concurrently (ingest of one overlapping the - * GPU commit of another) instead of strictly sequentially. Ignored by {@link - * #build(VectorSource, Config)}, which always builds sequentially — see that method's - * Javadoc. - */ - public Builder segments(int numSegments, boolean overlapped) { - if (numSegments < 1) { - throw new IllegalArgumentException("numSegments must be >= 1, got " + numSegments); - } - this.numSegments = numSegments; - this.overlapped = overlapped; - return this; - } - - /** Max segments built concurrently in overlap mode; peak host memory is this many slice buffers. */ - public Builder pipelineDepth(int pipelineDepth) { - if (pipelineDepth < 1) { - throw new IllegalArgumentException("pipelineDepth must be >= 1, got " + pipelineDepth); - } - this.pipelineDepth = pipelineDepth; - return this; - } - - public Config build() { - if (dimensions <= 0) { - throw new IllegalStateException( - "field(...) must be called with a positive dimension count"); - } - Objects.requireNonNull(graphBuildParams, "graphBuild(...) must be called"); - return new Config(this); - } - } - } -} diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CagraIndexParamsFactory.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CagraIndexParamsFactory.java index ef0b82b3cf..8bb46faf61 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CagraIndexParamsFactory.java +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CagraIndexParamsFactory.java @@ -75,8 +75,7 @@ public static CagraIndexParams create( AcceleratedHNSWParams acceleratedHNSWParams, long rows, long dimension) { if (acceleratedHNSWParams.getStrategy().equals(AcceleratedHNSWParams.Strategy.HEURISTIC)) { // Delegate the derivation of the graph degrees, build algorithm and its parameters to cuVS, - // expressed in terms of the HNSW-equivalent maxConn/beamWidth. cagraGraphBuildAlgo is not - // consulted here; it only applies under CUSTOM. + // expressed in terms of the HNSW-equivalent maxConn/beamWidth. CagraIndexParams derived = CagraIndexParams.fromHnswParams( rows, @@ -85,9 +84,9 @@ public static CagraIndexParams create( acceleratedHNSWParams.getBeamWidth(), acceleratedHNSWParams.getHnswHeuristicType(), acceleratedHNSWParams.getCuvsDistanceType()); - // TODO: fromHnswParams has no writerThreads argument, so its result carries the cuVS - // default (not a heuristic value). We can rebuild the CagraIndexParams with the - // caller-supplied writerThreads for now but should fix this in cuVS in the future. + // TODO: fromHnswParams has no writerThreads argument, so its result carries the cuVS default + // (not a heuristic value). We can rebuild the CagraIndexParams with the caller-supplied + // writerThreads for now but should fix this in cuVS in the future. return new CagraIndexParams.Builder() .withGraphDegree(derived.getGraphDegree()) .withIntermediateGraphDegree(derived.getIntermediateGraphDegree()) @@ -98,7 +97,6 @@ public static CagraIndexParams create( .withNumWriterThreads(acceleratedHNSWParams.getWriterThreads()) .build(); } - // CUSTOM: forward the caller's algorithm and the parameters it consumes. return new CagraIndexParams.Builder() .withGraphDegree(acceleratedHNSWParams.getGraphdegree()) .withIntermediateGraphDegree(acceleratedHNSWParams.getIntermediateGraphDegree()) diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsWriter.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsWriter.java index 3a7eb7852e..5feb211afe 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsWriter.java +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsWriter.java @@ -202,7 +202,8 @@ private void writeFieldInternal(FieldInfo fieldInfo, List vectors) thro var cagraIndexOutputStream = new IndexOutputOutputStream(cuvsIndex); try { CuVSMatrix cagraDataset = - Utils.createFloatMatrix(vectors, fieldInfo.getVectorDimension()); + Utils.createFloatMatrix( + vectors, fieldInfo.getVectorDimension(), getCuVSResourcesInstance()); writeCagraIndex(cagraIndexOutputStream, cagraDataset); } catch (Throwable t) { // Fallback to brute force in a few cases, for now. @@ -223,7 +224,8 @@ private void writeFieldInternal(FieldInfo fieldInfo, List vectors) thro if (indexType.isBruteForce()) { var bruteForceIndexOutputStream = new IndexOutputOutputStream(cuvsIndex); CuVSMatrix bruteforceDataset = - Utils.createFloatMatrix(vectors, fieldInfo.getVectorDimension()); + Utils.createFloatMatrix( + vectors, fieldInfo.getVectorDimension(), getCuVSResourcesInstance()); writeBruteForceIndex(bruteForceIndexOutputStream, bruteforceDataset); bruteForceIndexLength = cuvsIndex.getFilePointer() - bruteForceIndexOffset; diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/FbinVectorSource.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/FbinVectorSource.java deleted file mode 100644 index 548874d9d6..0000000000 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/FbinVectorSource.java +++ /dev/null @@ -1,213 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: Apache-2.0 - */ -package com.nvidia.cuvs.lucene; - -import java.io.IOException; -import java.nio.ByteBuffer; -import java.nio.ByteOrder; -import java.nio.channels.FileChannel; -import java.nio.file.Path; -import java.nio.file.StandardOpenOption; -import java.util.concurrent.ArrayBlockingQueue; -import java.util.concurrent.BlockingQueue; - -/** - * Prefetching, double-buffered {@link VectorSource} over an uncompressed {@code .fbin} file - * ({@code [num_vectors int32][dim int32]} header, then contiguous little-endian float32 rows). - * - *

Opens the file ONCE. A background thread reads the (optionally sliced) range front-to-back - * into two reusable direct buffers: while the caller consumes the current chunk, the reader fills - * the next one, so the disk read overlaps with the caller's per-vector work. {@link #get(int, - * float[])} unpacks directly into a caller-supplied array (no per-vector allocation). - * - *

Forward-only, single-consumer. {@code get} must be called with strictly increasing - * indices from a single thread, relative to the window this instance was constructed over — index 0 is - * the first vector in that window, not necessarily the first vector in the file. To read only a - * slice of a larger file (e.g. one segment of a partitioned build), construct with a {@code - * [firstVector, count)} range so each instance streams just its own portion of the file. - */ -public final class FbinVectorSource implements VectorSource { - - private static final long HEADER_BYTES = 8; - - private static final class Chunk { - final ByteBuffer buf; - final long start; - final int len; - - Chunk(ByteBuffer buf, long start, int len) { - this.buf = buf; - this.start = start; - this.len = len; - } - } - - /** Sentinel placed on the ready queue once the reader has produced the final chunk. */ - private static final Chunk POISON = new Chunk(null, -1, 0); - - private final FileChannel channel; - private final int dimension; - private final int firstVector; // absolute index (in the file) of the first vector served - private final int windowSize; // number of vectors this instance serves - private final int vectorBytes; - private final int chunkVectors; - - private final BlockingQueue free = new ArrayBlockingQueue<>(2); - private final BlockingQueue ready = new ArrayBlockingQueue<>(2); - private final Thread reader; - private volatile IOException readerError; - - private Chunk current; // consumer-owned; the chunk currently being served - private long nextExpectedRelative; // forward-only guard, relative index - - /** Reads the whole file from index 0. */ - public FbinVectorSource(Path path, int chunkSizeMB) throws IOException { - this(path, 0, -1, chunkSizeMB); - } - - /** - * Reads the contiguous range {@code [firstVector, firstVector + count)} of {@code path}, or to - * end of file if {@code count <= 0}. {@link #get} then serves indices relative to {@code - * firstVector} (0 = {@code firstVector}). - */ - public FbinVectorSource(Path path, int firstVector, int count, int chunkSizeMB) - throws IOException { - this.channel = FileChannel.open(path, StandardOpenOption.READ); - ByteBuffer header = ByteBuffer.allocate((int) HEADER_BYTES).order(ByteOrder.LITTLE_ENDIAN); - readFully(header, 0); - header.flip(); - int numVectors = header.getInt(); - this.dimension = header.getInt(); - this.vectorBytes = dimension * Float.BYTES; - this.firstVector = firstVector; - int endVector = count > 0 ? (int) Math.min((long) firstVector + count, numVectors) : numVectors; - this.windowSize = Math.max(0, endVector - firstVector); - - long chunkBytes = (long) Math.max(1, chunkSizeMB) * 1024 * 1024; - int cap = (Integer.MAX_VALUE - 16) / vectorBytes; // keep chunkVectors * vectorBytes in an int - this.chunkVectors = (int) Math.max(1, Math.min(chunkBytes / vectorBytes, cap)); - - // Two reusable direct buffers: the reader fills one while the consumer drains the other. - for (int i = 0; i < 2; i++) { - free.add( - ByteBuffer.allocateDirect(chunkVectors * vectorBytes).order(ByteOrder.LITTLE_ENDIAN)); - } - - this.reader = new Thread(this::readLoop, "fbin-prefetch-reader"); - this.reader.setDaemon(true); - this.reader.start(); - } - - @Override - public int dimensions() { - return dimension; - } - - @Override - public int size() { - return windowSize; - } - - /** Reader thread: fill chunks front-to-back, blocking on a free buffer between chunks. */ - private void readLoop() { - long next = firstVector; - long endVector = firstVector + windowSize; - try { - while (next < endVector) { - ByteBuffer buf = free.take(); - int toRead = (int) Math.min(chunkVectors, endVector - next); - buf.clear(); - buf.limit(toRead * vectorBytes); - readFully(buf, HEADER_BYTES + next * (long) vectorBytes); - ready.put(new Chunk(buf, next, toRead)); - next += toRead; - } - ready.put(POISON); - } catch (InterruptedException e) { - Thread.currentThread().interrupt(); // close() requested; stop quietly - } catch (IOException e) { - readerError = e; - try { - ready.put(POISON); // unblock the consumer so it can observe the error - } catch (InterruptedException ignored) { - Thread.currentThread().interrupt(); - } - } - } - - private void advance() throws IOException { - if (current != null) { - try { - free.put(current.buf); // hand the drained buffer back to the reader - } catch (InterruptedException e) { - Thread.currentThread().interrupt(); - throw new IOException("Interrupted returning chunk buffer", e); - } - current = null; - } - Chunk next; - try { - next = ready.take(); - } catch (InterruptedException e) { - Thread.currentThread().interrupt(); - throw new IOException("Interrupted awaiting next chunk", e); - } - if (next == POISON) { - if (readerError != null) { - throw new IOException("Prefetch reader failed", readerError); - } - throw new IOException("No more chunks available (unexpected EOF in prefetch)"); - } - current = next; - } - - @Override - public void get(int index, float[] dst) throws IOException { - if (index < 0 || index >= windowSize) { - throw new IndexOutOfBoundsException( - "Index " + index + " out of bounds [0, " + windowSize + ")"); - } - if (index < nextExpectedRelative) { - throw new UnsupportedOperationException( - "FbinVectorSource requires forward-only sequential access; got index " - + index - + " before previously served index " - + (nextExpectedRelative - 1)); - } - nextExpectedRelative = index + 1; - long absolute = firstVector + index; - while (current == null || absolute >= current.start + current.len) { - advance(); - } - int base = (int) (absolute - current.start) * vectorBytes; - for (int i = 0; i < dimension; i++) { - dst[i] = current.buf.getFloat(base + i * Float.BYTES); - } - } - - /** Convenience allocating variant (e.g. for a one-off query vector). */ - public float[] get(int index) throws IOException { - float[] dst = new float[dimension]; - get(index, dst); - return dst; - } - - private void readFully(ByteBuffer buf, long position) throws IOException { - long pos = position; - while (buf.hasRemaining()) { - int n = channel.read(buf, pos); - if (n < 0) { - throw new IOException("Unexpected EOF reading at position " + pos); - } - pos += n; - } - } - - @Override - public void close() throws IOException { - reader.interrupt(); - channel.close(); - } -} diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/FieldWriter.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/FieldWriter.java index 1898f5b985..8e2d9a70e2 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/FieldWriter.java +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/FieldWriter.java @@ -23,7 +23,6 @@ public class FieldWriter extends KnnFieldVectorsWriter { RamUsageEstimator.shallowSizeOfInstance(FieldWriter.class); private final FieldInfo fieldInfo; - private final int dimension; private final FlatFieldVectorsWriter flatFieldVectorsWriter; private int lastDocID = -1; private QuantizationType quantizationType; @@ -35,7 +34,6 @@ public FieldWriter( FlatFieldVectorsWriter flatFieldVectorsWriter) { this.quantizationType = quantizationType; this.fieldInfo = fieldInfo; - this.dimension = fieldInfo.getVectorDimension(); this.flatFieldVectorsWriter = (FlatFieldVectorsWriter) flatFieldVectorsWriter; } @@ -73,10 +71,6 @@ DocsWithFieldSet getDocsWithFieldSet() { return flatFieldVectorsWriter.getDocsWithFieldSet(); } - int dimension() { - return dimension; - } - @Override public Object copyValue(Object vectorValue) { throw new UnsupportedOperationException(); diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java index 2007bc5c30..7e9f888e32 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java @@ -6,17 +6,10 @@ import static org.apache.lucene.search.DocIdSetIterator.NO_MORE_DOCS; -import com.nvidia.cuvs.CuVSDeviceMatrix; -import com.nvidia.cuvs.CuVSHostMatrix; import com.nvidia.cuvs.CuVSMatrix; import com.nvidia.cuvs.RowView; -import java.io.IOException; import java.util.ArrayList; import java.util.List; -import java.util.concurrent.Callable; -import java.util.concurrent.ExecutorService; -import java.util.concurrent.Executors; -import org.apache.lucene.search.TaskExecutor; import org.apache.lucene.util.hnsw.HnswGraph; import org.apache.lucene.util.hnsw.NeighborArray; @@ -45,15 +38,9 @@ public class GPUBuiltHnswGraph extends HnswGraph { * @param dimensions the vector dimension * @param layerNodes the nodes on the layer * @param layerAdjacencies adjacency list - * @param numThreads threads to use for materializing the adjacency (1 = serial) */ public GPUBuiltHnswGraph( - int size, - int dimensions, - List layerNodes, - List layerAdjacencies, - int numThreads) - throws IOException { + int size, int dimensions, List layerNodes, List layerAdjacencies) { this.size = size; this.dimensions = dimensions; @@ -63,102 +50,38 @@ public GPUBuiltHnswGraph( // Process Layer 0 (base layer with all nodes) CuVSMatrix layer0Adjacency = layerAdjacencies.get(0); - this.layer0Neighbors = fillNeighborArray(layer0Adjacency, size, numThreads); + this.layer0Neighbors = fillNeighborArray(layer0Adjacency, size); // Process higher layers (1 to numLevels-1) for (int level = 1; level < numLevels; level++) { int[] nodes = layerNodes.get(level); CuVSMatrix adjacency = layerAdjacencies.get(level); this.layerNodes.add(nodes); - this.layerNeighbors.add(fillNeighborArray(adjacency, nodes.length, numThreads)); + this.layerNeighbors.add(fillNeighborArray(adjacency, nodes.length)); } } - /** Node count below which parallel materialization is not worth the thread overhead. */ - private static final int PARALLEL_MIN_NODES = 1 << 16; - /** - * Materializes the adjacency matrix into on-heap {@link NeighborArray}s, one per node. - * - *

The serial path reads the adjacency directly (a device matrix's {@code getRow} is safe - * single-threaded). The parallel path cannot: the CAGRA layer-0 adjacency is a device matrix whose - * {@code getRow} uses a shared, stateful buffered reader that is not safe for concurrent access, so - * it is pulled to host once (a single bulk device->host copy) before materializing disjoint node - * ranges concurrently. Host matrices (the upper layers, built via {@link CuVSMatrix#ofArray}) are - * read directly in both paths. + * Fills the neighbor array using the adjacency matrix. * * @param adjacency instance of adjacency CuVSMatrix * @param size the number of nodes - * @param numThreads threads to use (1, or fewer than {@value #PARALLEL_MIN_NODES} nodes = serial) * @return the NeighborArray */ - private static NeighborArray[] fillNeighborArray(CuVSMatrix adjacency, int size, int numThreads) - throws IOException { + private NeighborArray[] fillNeighborArray(CuVSMatrix adjacency, int size) { NeighborArray[] neighbors = new NeighborArray[size]; - if (numThreads <= 1 || size < PARALLEL_MIN_NODES) { - fillNeighborRange(adjacency, neighbors, 0, size); - return neighbors; - } - CuVSMatrix source = adjacency; - CuVSHostMatrix hostCopy = null; - if (adjacency instanceof CuVSDeviceMatrix deviceAdjacency) { - hostCopy = deviceAdjacency.toHost(); - source = hostCopy; - } - try { - fillNeighborArrayParallel(source, neighbors, size, numThreads); - return neighbors; - } finally { - if (hostCopy != null) { - hostCopy.close(); - } - } - } - - /** - * Materializes disjoint node ranges concurrently. Each thread writes its own slots of {@code - * neighbors} and its own {@link NeighborArray} instances, so no synchronization is needed; {@code - * source} must be a host matrix (stateless {@code getRow}). - */ - private static void fillNeighborArrayParallel( - CuVSMatrix source, NeighborArray[] neighbors, int size, int numThreads) throws IOException { - ExecutorService pool = Executors.newFixedThreadPool(Math.max(1, numThreads - 1)); - try { - int perThread = (size + numThreads - 1) / numThreads; - List> tasks = new ArrayList<>(numThreads); - for (int t = 0; t < numThreads; t++) { - final int start = t * perThread; - final int end = Math.min(start + perThread, size); - if (start >= end) { - break; - } - tasks.add( - () -> { - fillNeighborRange(source, neighbors, start, end); - return null; - }); - } - new TaskExecutor(pool).invokeAll(tasks); - } finally { - pool.shutdown(); - } - } - - /** Fills {@code neighbors[start, end)} from the adjacency rows. */ - private static void fillNeighborRange( - CuVSMatrix source, NeighborArray[] neighbors, int start, int end) { - for (int i = start; i < end; i++) { - RowView rv = source.getRow(i); + for (int i = 0; i < size; i++) { + RowView rv = adjacency.getRow(i); if (rv != null && rv.size() > 0) { - NeighborArray na = new NeighborArray((int) rv.size(), true); + neighbors[i] = new NeighborArray((int) rv.size(), true); for (int j = 0; j < rv.size(); j++) { - na.addInOrder(rv.getAsInt(j), 1.0f - (j * 0.001f)); + neighbors[i].addInOrder(rv.getAsInt(j), 1.0f - (j * 0.001f)); } - neighbors[i] = na; } else { neighbors[i] = new NeighborArray(0, true); } } + return neighbors; } /** diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene101AcceleratedHNSWCodec.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene101AcceleratedHNSWCodec.java index f95487e953..e9f4d6fead 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene101AcceleratedHNSWCodec.java +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene101AcceleratedHNSWCodec.java @@ -53,25 +53,7 @@ public Lucene101AcceleratedHNSWCodec(String name, Codec delegate) { public Lucene101AcceleratedHNSWCodec(AcceleratedHNSWParams acceleratedHNSWParams) throws Exception { this(NAME, LuceneProvider.getCodec("101")); - initializeFormat(acceleratedHNSWParams, 0); - } - - /** - * Constructor for {@link Lucene101AcceleratedHNSWCodec} used for the native flat-buffered build - * path (see {@link NativeFlatBufferedHNSWVectorsWriter}). Package-private: only {@link - * CagraHnswBulkIndexWriter} can guarantee the invariants (single unmerged segment, no index - * sort, exact vector count) that native flat buffering requires, so this overload is not part - * of the public API. - * - * @param acceleratedHNSWParams instance of {@link AcceleratedHNSWParams} - * @param numInputVectors the exact number of vectors to be indexed, used to pre-size the native - * flat buffer - * @throws Exception exception - */ - Lucene101AcceleratedHNSWCodec(AcceleratedHNSWParams acceleratedHNSWParams, int numInputVectors) - throws Exception { - this(NAME, LuceneProvider.getCodec("101")); - initializeFormat(acceleratedHNSWParams, numInputVectors); + initializeFormat(acceleratedHNSWParams); } /** @@ -79,19 +61,17 @@ public Lucene101AcceleratedHNSWCodec(AcceleratedHNSWParams acceleratedHNSWParams * with an instance of {@link AcceleratedHNSWParams} with default parameter values. */ private void initializeFormatDefaultValues() { - initializeFormat(new AcceleratedHNSWParams.Builder().build(), 0); + initializeFormat(new AcceleratedHNSWParams.Builder().build()); } /** * Initialize an instance of {@link Lucene99AcceleratedHNSWVectorsFormat}. * * @param acceleratedHNSWParams instance of {@link AcceleratedHNSWParams} to use - * @param numInputVectors the exact number of vectors to be indexed, used to pre-size the native - * flat buffer (0 = disabled, the default heap-buffered path) */ - private void initializeFormat(AcceleratedHNSWParams acceleratedHNSWParams, int numInputVectors) { + private void initializeFormat(AcceleratedHNSWParams acceleratedHNSWParams) { try { - format = new Lucene99AcceleratedHNSWVectorsFormat(acceleratedHNSWParams, numInputVectors); + format = new Lucene99AcceleratedHNSWVectorsFormat(acceleratedHNSWParams); setKnnFormat(format); } catch (LibraryException ex) { log.log( diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsFormat.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsFormat.java index 59e012427d..1b1bc3cbd2 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsFormat.java +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsFormat.java @@ -31,7 +31,6 @@ public class Lucene99AcceleratedHNSWVectorsFormat extends KnnVectorsFormat { private static final FlatVectorsFormat FLAT_VECTORS_FORMAT; private static final int MAX_DIMENSIONS = 4096; private final AcceleratedHNSWParams acceleratedHNSWParams; - private final int numInputVectors; static final String HNSW_META_CODEC_NAME = "Lucene99HnswVectorsFormatMeta"; static final String HNSW_META_CODEC_EXT = "vem"; @@ -66,25 +65,8 @@ public Lucene99AcceleratedHNSWVectorsFormat() { * @param acceleratedHNSWParams An instance of {@link AcceleratedHNSWParams} */ public Lucene99AcceleratedHNSWVectorsFormat(AcceleratedHNSWParams acceleratedHNSWParams) { - this(acceleratedHNSWParams, 0); - } - - /** - * Initializes {@link Lucene99AcceleratedHNSWVectorsFormat} for the native flat-buffered build - * path (see {@link NativeFlatBufferedHNSWVectorsWriter}). Package-private: only {@link - * CagraHnswBulkIndexWriter} (via {@link Lucene101AcceleratedHNSWCodec}) can guarantee the - * invariants that native flat buffering requires, so this overload is not part of the public - * API. - * - * @param acceleratedHNSWParams An instance of {@link AcceleratedHNSWParams} - * @param numInputVectors the exact number of vectors to be indexed, used to pre-size the native - * flat buffer (0 = disabled, the default heap-buffered path) - */ - Lucene99AcceleratedHNSWVectorsFormat( - AcceleratedHNSWParams acceleratedHNSWParams, int numInputVectors) { super("Lucene99AcceleratedHNSWVectorsFormat"); this.acceleratedHNSWParams = acceleratedHNSWParams; - this.numInputVectors = numInputVectors; } /** @@ -92,23 +74,14 @@ public Lucene99AcceleratedHNSWVectorsFormat(AcceleratedHNSWParams acceleratedHNS */ @Override public KnnVectorsWriter fieldsWriter(SegmentWriteState state) throws IOException { - boolean nativeMode = isSupported() && numInputVectors > 0; + var flatWriter = FLAT_VECTORS_FORMAT.fieldsWriter(state); if (isSupported()) { - if (nativeMode) { - log.log(Level.FINE, "cuVS is supported so using the NativeFlatBufferedHNSWVectorsWriter"); - // In hint mode the accelerated writer owns the flat .vec/.vemf files, so the Lucene flat - // writer must not be created (it would open the same outputs). - return new NativeFlatBufferedHNSWVectorsWriter( - state, acceleratedHNSWParams, numInputVectors); - } log.log(Level.FINE, "cuVS is supported so using the Lucene99AcceleratedHNSWVectorsWriter"); - var flatWriter = FLAT_VECTORS_FORMAT.fieldsWriter(state); return new Lucene99AcceleratedHNSWVectorsWriter(state, acceleratedHNSWParams, flatWriter); } else { log.log( Level.WARNING, "GPU based indexing not supported, falling back to using the Lucene99HnswVectorsWriter"); - var flatWriter = FLAT_VECTORS_FORMAT.fieldsWriter(state); try { return LUCENE_PROVIDER.getLuceneHnswVectorsWriterInstance( state, diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java index bc27af0b06..13edc64975 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java @@ -4,31 +4,42 @@ */ package com.nvidia.cuvs.lucene; +import static com.nvidia.cuvs.lucene.AcceleratedHNSWUtils.createMultiLayerHnswGraph; +import static com.nvidia.cuvs.lucene.AcceleratedHNSWUtils.createSingleVectorHnswGraph; import static com.nvidia.cuvs.lucene.AcceleratedHNSWUtils.printInfoStream; +import static com.nvidia.cuvs.lucene.AcceleratedHNSWUtils.writeEmpty; +import static com.nvidia.cuvs.lucene.AcceleratedHNSWUtils.writeGraph; +import static com.nvidia.cuvs.lucene.AcceleratedHNSWUtils.writeMeta; +import static com.nvidia.cuvs.lucene.Lucene99AcceleratedHNSWVectorsFormat.HNSW_INDEX_CODEC_NAME; +import static com.nvidia.cuvs.lucene.Lucene99AcceleratedHNSWVectorsFormat.HNSW_INDEX_EXT; +import static com.nvidia.cuvs.lucene.Lucene99AcceleratedHNSWVectorsFormat.HNSW_META_CODEC_EXT; +import static com.nvidia.cuvs.lucene.Lucene99AcceleratedHNSWVectorsFormat.HNSW_META_CODEC_NAME; import static com.nvidia.cuvs.lucene.ThreadLocalCuVSResourcesProvider.closeCuVSResourcesInstance; +import static com.nvidia.cuvs.lucene.ThreadLocalCuVSResourcesProvider.getCuVSResourcesInstance; +import static com.nvidia.cuvs.lucene.Utils.createListFromMergedVectors; import static org.apache.lucene.index.VectorEncoding.FLOAT32; import static org.apache.lucene.util.RamUsageEstimator.shallowSizeOfInstance; -import com.nvidia.cuvs.CuVSHostMatrix; +import com.nvidia.cuvs.CagraIndex; +import com.nvidia.cuvs.CagraIndexParams; import com.nvidia.cuvs.CuVSMatrix; import com.nvidia.cuvs.lucene.AcceleratedHNSWUtils.QuantizationType; import java.io.IOException; import java.util.ArrayList; import java.util.List; import java.util.Objects; +import org.apache.lucene.codecs.CodecUtil; import org.apache.lucene.codecs.KnnFieldVectorsWriter; import org.apache.lucene.codecs.KnnVectorsWriter; import org.apache.lucene.codecs.hnsw.FlatVectorsWriter; import org.apache.lucene.index.DocsWithFieldSet; import org.apache.lucene.index.FieldInfo; -import org.apache.lucene.index.FloatVectorValues; -import org.apache.lucene.index.KnnVectorValues; +import org.apache.lucene.index.IndexFileNames; import org.apache.lucene.index.MergeState; import org.apache.lucene.index.SegmentWriteState; import org.apache.lucene.index.Sorter; import org.apache.lucene.index.Sorter.DocMap; -import org.apache.lucene.search.DocIdSetIterator; -import org.apache.lucene.util.Bits; +import org.apache.lucene.store.IndexOutput; import org.apache.lucene.util.IOUtils; import org.apache.lucene.util.InfoStream; @@ -43,13 +54,28 @@ public class Lucene99AcceleratedHNSWVectorsWriter extends KnnVectorsWriter { private static final long SHALLOW_RAM_BYTES_USED = shallowSizeOfInstance(Lucene99AcceleratedHNSWVectorsWriter.class); private static final String COMPONENT = "Lucene99AcceleratedHNSWVectorsWriter"; + private static final LuceneProvider LUCENE_PROVIDER; + private static final Integer VERSION_CURRENT; + private final AcceleratedHNSWParams acceleratedHNSWParams; private final FlatVectorsWriter flatVectorsWriter; private final List fields = new ArrayList<>(); private final InfoStream infoStream; - private AcceleratedHnswGraphOutput graphOutput; + private IndexOutput hnswMeta = null; + private IndexOutput hnswVectorIndex = null; + private String vemFileName; + private String vexFileName; private boolean finished; + static { + try { + LUCENE_PROVIDER = LuceneProvider.getInstance("99"); + VERSION_CURRENT = LUCENE_PROVIDER.getStaticIntParam("VERSION_CURRENT"); + } catch (Exception e) { + throw new ExceptionInInitializerError(e.getMessage()); + } + } + /** * Initializes {@link Lucene99AcceleratedHNSWVectorsWriter} * @@ -64,11 +90,30 @@ public Lucene99AcceleratedHNSWVectorsWriter( FlatVectorsWriter flatVectorsWriter) throws IOException { super(); - this.flatVectorsWriter = Objects.requireNonNull(flatVectorsWriter); + this.flatVectorsWriter = flatVectorsWriter; this.infoStream = state.infoStream; + this.acceleratedHNSWParams = acceleratedHNSWParams; + vemFileName = + IndexFileNames.segmentFileName( + state.segmentInfo.name, state.segmentSuffix, HNSW_META_CODEC_EXT); + vexFileName = + IndexFileNames.segmentFileName(state.segmentInfo.name, state.segmentSuffix, HNSW_INDEX_EXT); boolean success = false; try { - graphOutput = new AcceleratedHnswGraphOutput(state, acceleratedHNSWParams); + hnswMeta = state.directory.createOutput(vemFileName, state.context); + hnswVectorIndex = state.directory.createOutput(vexFileName, state.context); + CodecUtil.writeIndexHeader( + hnswMeta, + HNSW_META_CODEC_NAME, + VERSION_CURRENT, + state.segmentInfo.getId(), + state.segmentSuffix); + CodecUtil.writeIndexHeader( + hnswVectorIndex, + HNSW_INDEX_CODEC_NAME, + VERSION_CURRENT, + state.segmentInfo.getId(), + state.segmentSuffix); success = true; printInfoStream(infoStream, COMPONENT, "Lucene99AcceleratedHNSWVectorsWriter is initialized"); } finally { @@ -93,6 +138,66 @@ public KnnFieldVectorsWriter addField(FieldInfo fieldInfo) throws IOException return writer; } + /** + * Builds the intermediate CAGRA index and builds and writes the HNSW index. + * + * @param fieldInfo instance of FieldInfo that has the field description + * @param vectors vectors to index + * @throws IOException + */ + private void writeFieldInternal(FieldInfo fieldInfo, List vectors) throws IOException { + if (vectors.size() == 0) { + writeEmpty(fieldInfo, hnswMeta); + return; + } + if (vectors.size() < 2) { + writeSingleVectorGraph(fieldInfo, vectors); + return; + } + try { + CuVSMatrix dataset = + Utils.createFloatMatrix( + vectors, fieldInfo.getVectorDimension(), getCuVSResourcesInstance()); + + CagraIndexParams params = + CagraIndexParamsFactory.create(acceleratedHNSWParams, dataset.size(), dataset.columns()); + + CagraIndex cagraIndex = + CagraIndex.newBuilder(getCuVSResourcesInstance()) + .withDataset(dataset) + .withIndexParams(params) + .build(); + CuVSMatrix adjacencyListMatrix = cagraIndex.getGraph(); + int size = (int) dataset.size(); + int dimensions = fieldInfo.getVectorDimension(); + GPUBuiltHnswGraph hnswGraph = + createMultiLayerHnswGraph( + fieldInfo, + size, + dimensions, + adjacencyListMatrix, + vectors, + acceleratedHNSWParams.getHnswLayers(), + params, + QuantizationType.NONE); + long vectorIndexOffset = hnswVectorIndex.getFilePointer(); + int[][] graphLevelNodeOffsets = writeGraph(hnswGraph, hnswVectorIndex); + long vectorIndexLength = hnswVectorIndex.getFilePointer() - vectorIndexOffset; + writeMeta( + hnswVectorIndex, + hnswMeta, + fieldInfo, + vectorIndexOffset, + vectorIndexLength, + size, + hnswGraph, + graphLevelNodeOffsets); + cagraIndex.close(); + } catch (Throwable t) { + Utils.handleThrowable(t); + } + } + /** * Build the indexes and writes it to the disk. */ @@ -115,7 +220,7 @@ public void flush(int maxDoc, DocMap sortMap) throws IOException { * @throws IOException */ private void writeField(FieldWriter fieldData) throws IOException { - graphOutput.writeField(fieldData.fieldInfo(), fieldData.getFloatVectors()); + writeFieldInternal(fieldData.fieldInfo(), fieldData.getFloatVectors()); } /** @@ -134,68 +239,48 @@ private void writeSortingField(FieldWriter fieldData, Sorter.DocMap sortMap) thr for (int i = 0; i < floatVectors.size(); i++) { sortedVectors.add(floatVectors.get(new2OldOrd[i])); } - graphOutput.writeField(fieldData.fieldInfo(), sortedVectors); + writeFieldInternal(fieldData.fieldInfo(), sortedVectors); } /** - * Streams merged vectors directly into a native host-memory matrix (CuVSHostMatrix) - * without materialising a List on the Java heap, then calls graphOutput.writeField. - * This avoids the double-copy OOM (heap list + native matrix simultaneously) that - * occurs when force-merging large segments. + * Builds and writes a single vector graph. + * + * @param fieldInfo instance of FieldInfo + * @param vectors the list of float vectors + * @throws IOException I/O Exceptions */ - private void vectorBasedMerge(FieldInfo fieldInfo, MergeState mergeState) throws IOException { + private void writeSingleVectorGraph(FieldInfo fieldInfo, List vectors) + throws IOException { try { - // FloatVectorValues#size() on the merged view is the raw sum of every source segment's - // on-disk vector count (MergedVectorValues.MergedFloat32VectorValues computes it once at - // construction from each sub-reader's unfiltered size) -- NOT the number of live - // (non-deleted) vectors the iterator below will actually yield, which is what - // CuVSMatrix.hostBuilder needs since it preallocates a fixed-size native buffer. Using - // size() here under-fills that buffer whenever the merge drops deleted docs, leaving the - // graph built over more rows than were actually populated. - // - // size() IS trustworthy when no segment being merged has any deletions: per-segment vector - // counts already exclude docs without a value for this field (sparse fields are handled at - // the single-segment level, independent of deletions), so the raw sum equals the live count - // in that case and the extra counting pass below can be skipped. - boolean anySegmentHasDeletions = false; - for (Bits liveDocs : mergeState.liveDocs) { - if (liveDocs != null) { - anySegmentHasDeletions = true; - break; - } - } - - int size; - if (anySegmentHasDeletions) { - // Count the live vectors via a throwaway iteration first (mergeFloatVectorValues - // constructs a fresh, independent view each call, so this doesn't disturb the real build - // pass below). - size = 0; - FloatVectorValues counting = - KnnVectorsWriter.MergedVectorValues.mergeFloatVectorValues(fieldInfo, mergeState); - KnnVectorValues.DocIndexIterator countingIt = counting.iterator(); - for (int doc = countingIt.nextDoc(); - doc != DocIdSetIterator.NO_MORE_DOCS; - doc = countingIt.nextDoc()) { - size++; - } - } else { - size = - KnnVectorsWriter.MergedVectorValues.mergeFloatVectorValues(fieldInfo, mergeState) - .size(); - } + int size = 1; + int dimensions = fieldInfo.getVectorDimension(); + GPUBuiltHnswGraph hnswGraph = createSingleVectorHnswGraph(size, dimensions); + long vectorIndexOffset = hnswVectorIndex.getFilePointer(); + int[][] graphLevelNodeOffsets = writeGraph(hnswGraph, hnswVectorIndex); + long vectorIndexLength = hnswVectorIndex.getFilePointer() - vectorIndexOffset; + writeMeta( + hnswVectorIndex, + hnswMeta, + fieldInfo, + vectorIndexOffset, + vectorIndexLength, + size, + hnswGraph, + graphLevelNodeOffsets); + } catch (Throwable t) { + Utils.handleThrowable(t); + } + } - FloatVectorValues mergedVectors = - KnnVectorsWriter.MergedVectorValues.mergeFloatVectorValues(fieldInfo, mergeState); - int dims = fieldInfo.getVectorDimension(); - CuVSMatrix.Builder builder = - CuVSMatrix.hostBuilder(size, dims, CuVSMatrix.DataType.FLOAT); - KnnVectorValues.DocIndexIterator it = mergedVectors.iterator(); - for (int doc = it.nextDoc(); doc != DocIdSetIterator.NO_MORE_DOCS; doc = it.nextDoc()) { - builder.addVector(mergedVectors.vectorValue(it.index())); - } - CuVSHostMatrix dataset = builder.build(); - graphOutput.writeField(fieldInfo, dataset); + /** + * Create combined data set for the merged segment and call writeFieldInternal. + */ + private void vectorBasedMerge(FieldInfo fieldInfo, MergeState mergeState) throws IOException { + try { + List dataset = + createListFromMergedVectors( + KnnVectorsWriter.MergedVectorValues.mergeFloatVectorValues(fieldInfo, mergeState)); + writeFieldInternal(fieldInfo, dataset); } catch (Throwable t) { Utils.handleThrowable(t); } @@ -220,7 +305,14 @@ public void finish() throws IOException { } finished = true; flatVectorsWriter.finish(); - graphOutput.finish(); + if (hnswMeta != null) { + // write end of fields marker + hnswMeta.writeInt(-1); + CodecUtil.writeFooter(hnswMeta); + } + if (hnswVectorIndex != null) { + CodecUtil.writeFooter(hnswVectorIndex); + } } /** @@ -229,7 +321,7 @@ public void finish() throws IOException { @Override public void close() throws IOException { printInfoStream(infoStream, COMPONENT, "Closing resources"); - IOUtils.close(graphOutput, flatVectorsWriter); + IOUtils.close(hnswMeta, hnswVectorIndex, flatVectorsWriter); closeCuVSResourcesInstance(); } diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java index cf6d19f7ca..87907d2cbb 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java @@ -155,7 +155,8 @@ private void writeFieldInternal(FieldInfo fieldInfo, List vectors) throw int dimensions = fieldInfo.getVectorDimension(); int bytesPerVector = (dimensions + 7) / 8; - CuVSMatrix dataset = Utils.createByteMatrix(vectors, bytesPerVector); + CuVSMatrix dataset = + Utils.createByteMatrix(vectors, bytesPerVector, getCuVSResourcesInstance()); if (dataset.size() < 2) { writeSingleVectorGraph(fieldInfo, vectors); @@ -178,18 +179,17 @@ private void writeFieldInternal(FieldInfo fieldInfo, List vectors) throw GPUBuiltHnswGraph hnswGraph = createMultiLayerHnswGraph( fieldInfo, + size, dimensions, adjacencyListMatrix, - dataset, + vectors, acceleratedHNSWParams.getHnswLayers(), params, - QuantizationType.BINARY, - acceleratedHNSWParams.getWriterThreads()); + QuantizationType.BINARY); long vectorIndexOffset = hnswVectorIndex.getFilePointer(); // Write the graph to the vector index - int[][] graphLevelNodeOffsets = - writeGraph(hnswGraph, hnswVectorIndex, acceleratedHNSWParams.getWriterThreads()); + int[][] graphLevelNodeOffsets = writeGraph(hnswGraph, hnswVectorIndex); long vectorIndexLength = hnswVectorIndex.getFilePointer() - vectorIndexOffset; // Write metadata @@ -277,8 +277,7 @@ private void writeSingleVectorGraph(FieldInfo fieldInfo, List vectors) long vectorIndexOffset = hnswVectorIndex.getFilePointer(); // Write the graph to the vector index - int[][] graphLevelNodeOffsets = - writeGraph(hnswGraph, hnswVectorIndex, acceleratedHNSWParams.getWriterThreads()); + int[][] graphLevelNodeOffsets = writeGraph(hnswGraph, hnswVectorIndex); long vectorIndexLength = hnswVectorIndex.getFilePointer() - vectorIndexOffset; diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java index aebfbea7c6..7141af56ee 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java @@ -181,7 +181,8 @@ private void writeFieldInternal(FieldInfo fieldInfo, List vectors) throws IOE } // Create CuVSMatrix with BYTE data type (unsigned bytes) - CuVSMatrix dataset = Utils.createByteMatrix(unsignedVectors, dimensions); + CuVSMatrix dataset = + Utils.createByteMatrix(unsignedVectors, dimensions, getCuVSResourcesInstance()); if (dataset.size() < 2) { writeSingleVectorGraph(fieldInfo, unsignedVectors); @@ -203,19 +204,18 @@ private void writeFieldInternal(FieldInfo fieldInfo, List vectors) throws IOE GPUBuiltHnswGraph hnswGraph = createMultiLayerHnswGraph( fieldInfo, + size, dimensions, adjacencyListMatrix, - dataset, + unsignedVectors, acceleratedHNSWParams.getHnswLayers(), params, - QuantizationType.SCALAR, - acceleratedHNSWParams.getWriterThreads()); + QuantizationType.SCALAR); long vectorIndexOffset = hnswVectorIndex.getFilePointer(); // Write the graph to the vector index - int[][] graphLevelNodeOffsets = - writeGraph(hnswGraph, hnswVectorIndex, acceleratedHNSWParams.getWriterThreads()); + int[][] graphLevelNodeOffsets = writeGraph(hnswGraph, hnswVectorIndex); long vectorIndexLength = hnswVectorIndex.getFilePointer() - vectorIndexOffset; @@ -302,8 +302,7 @@ private void writeSingleVectorGraph(FieldInfo fieldInfo, List vectors) long vectorIndexOffset = hnswVectorIndex.getFilePointer(); // Write the graph to the vector index - int[][] graphLevelNodeOffsets = - writeGraph(hnswGraph, hnswVectorIndex, acceleratedHNSWParams.getWriterThreads()); + int[][] graphLevelNodeOffsets = writeGraph(hnswGraph, hnswVectorIndex); long vectorIndexLength = hnswVectorIndex.getFilePointer() - vectorIndexOffset; // Write metadata diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/NativeFieldWriter.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/NativeFieldWriter.java deleted file mode 100644 index a246efb69d..0000000000 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/NativeFieldWriter.java +++ /dev/null @@ -1,124 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: Apache-2.0 - */ - -package com.nvidia.cuvs.lucene; - -import com.nvidia.cuvs.CuVSHostMatrix; -import com.nvidia.cuvs.CuVSMatrix; -import java.io.IOException; -import org.apache.lucene.codecs.KnnFieldVectorsWriter; -import org.apache.lucene.index.DocsWithFieldSet; -import org.apache.lucene.index.FieldInfo; -import org.apache.lucene.util.RamUsageEstimator; - -/** - * Native-flat-buffered field writer, used exclusively by {@link - * NativeFlatBufferedHNSWVectorsWriter}: streams incoming vectors directly into a native host - * matrix (see {@link CuVSMatrix#hostBuilder}) instead of accumulating a {@code List} on - * the Java heap, so the full dataset is never held twice (heap list + native copy). The matrix is - * preallocated for exactly {@code numInputVectors} rows, so that count must equal the number of - * vectors actually added (validated by {@link NativeFlatBufferedHNSWVectorsWriter}). - * - *

Package-private and float32-only — quantization is not supported on this path. Distinct from - * the heap-buffered {@link FieldWriter} used by the other accelerated writers; the two share no - * state or behavior beyond both extending {@link KnnFieldVectorsWriter}. - */ -final class NativeFieldWriter extends KnnFieldVectorsWriter { - - private static final long SHALLOW_SIZE = - RamUsageEstimator.shallowSizeOfInstance(NativeFieldWriter.class); - - private final FieldInfo fieldInfo; - private final int numInputVectors; - private final CuVSMatrix.Builder hostMatrixBuilder; - private final DocsWithFieldSet docsWithField = new DocsWithFieldSet(); - private int lastDocID = -1; - private int count; - private CuVSHostMatrix builtMatrix; - - NativeFieldWriter(FieldInfo fieldInfo, int numInputVectors) { - this.fieldInfo = fieldInfo; - this.numInputVectors = numInputVectors; - // Preallocates one contiguous native region of numInputVectors * dimension * 4 bytes. - this.hostMatrixBuilder = - CuVSMatrix.hostBuilder( - numInputVectors, fieldInfo.getVectorDimension(), CuVSMatrix.DataType.FLOAT); - } - - @Override - public void addValue(int docID, Object vectorValue) throws IOException { - if (docID == lastDocID) { - throw new IllegalArgumentException( - "VectorValuesField \"" - + fieldInfo.name - + "\" appears more than once in this document (only one value is allowed per" - + " field)"); - } - if (count >= numInputVectors) { - throw new IllegalStateException( - "Buffered vectors (" - + (count + 1) - + ") exceed numInputVectors (" - + numInputVectors - + ") for field \"" - + fieldInfo.name - + "\". This usually means more vectors arrived than the numInputVectors hint" - + " promised (e.g. a merge or a later flush cycle reused this config); see" - + " CagraHnswBulkIndexWriter."); - } - // hostMatrixBuilder.addVector validates the dimension and performs the native row copy. - hostMatrixBuilder.addVector((float[]) vectorValue); - docsWithField.add(docID); - count++; - lastDocID = docID; - } - - FieldInfo fieldInfo() { - return fieldInfo; - } - - DocsWithFieldSet getDocsWithFieldSet() { - return docsWithField; - } - - /** - * The native host matrix holding the buffered vectors. Built once and cached; the caller owns - * closing it via {@link #releaseNativeBuffer()} once the CAGRA build has consumed it. - */ - CuVSHostMatrix getHostMatrix() { - if (builtMatrix == null) { - builtMatrix = hostMatrixBuilder.build(); - } - return builtMatrix; - } - - /** Number of vectors buffered so far. */ - int getNativeVectorCount() { - return count; - } - - /** - * Closes the native host matrix. Safe to call multiple times, including on a field whose - * buffer was never fully populated (e.g. after a count-mismatch or other flush failure): {@code - * build()} just returns the pre-allocated matrix regardless of how many rows were written, and - * {@code close()} on that matrix is itself idempotent. - */ - void releaseNativeBuffer() { - getHostMatrix().close(); - builtMatrix = null; - } - - @Override - public Object copyValue(Object vectorValue) { - throw new UnsupportedOperationException(); - } - - @Override - public long ramBytesUsed() { - // The native host matrix is off-heap and intentionally excluded from Lucene's heap RAM - // accounting. - return SHALLOW_SIZE; - } -} diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/NativeFlatBufferedHNSWVectorsWriter.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/NativeFlatBufferedHNSWVectorsWriter.java deleted file mode 100644 index 4f52d882d7..0000000000 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/NativeFlatBufferedHNSWVectorsWriter.java +++ /dev/null @@ -1,214 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: Apache-2.0 - */ -package com.nvidia.cuvs.lucene; - -import static com.nvidia.cuvs.lucene.AcceleratedHNSWUtils.printInfoStream; -import static com.nvidia.cuvs.lucene.ThreadLocalCuVSResourcesProvider.closeCuVSResourcesInstance; -import static org.apache.lucene.index.VectorEncoding.FLOAT32; -import static org.apache.lucene.util.RamUsageEstimator.shallowSizeOfInstance; - -import com.nvidia.cuvs.CuVSHostMatrix; -import java.io.IOException; -import java.util.ArrayList; -import java.util.List; -import org.apache.lucene.codecs.KnnFieldVectorsWriter; -import org.apache.lucene.codecs.KnnVectorsWriter; -import org.apache.lucene.index.FieldInfo; -import org.apache.lucene.index.MergeState; -import org.apache.lucene.index.SegmentWriteState; -import org.apache.lucene.index.Sorter.DocMap; -import org.apache.lucene.util.IOUtils; -import org.apache.lucene.util.InfoStream; - -/** - * Native-flat-buffered {@link KnnVectorsWriter}, streaming vectors directly into a native host - * matrix (see {@link NativeFieldWriter}) rather than a heap {@code List}, and writing the - * flat {@code .vec}/{@code .vemf} files itself via {@link NativeFlatVectorsWriter} instead of - * delegating to Lucene's {@link org.apache.lucene.codecs.hnsw.FlatVectorsWriter}. - * - *

This writer supports only the unsorted, single-segment flush path: merges and index-sorted - * flushes are rejected. It is package-private and constructed only by {@link - * Lucene99AcceleratedHNSWVectorsFormat#fieldsWriter} when {@code numInputVectors > 0}, which is - * reachable only via the package-private constructors on {@link Lucene99AcceleratedHNSWVectorsFormat} - * and {@link Lucene101AcceleratedHNSWCodec} used exclusively by {@link CagraHnswBulkIndexWriter} — - * the only caller that owns its {@link org.apache.lucene.index.IndexWriter} and can guarantee - * those invariants. - */ -final class NativeFlatBufferedHNSWVectorsWriter extends KnnVectorsWriter { - - private static final long SHALLOW_RAM_BYTES_USED = - shallowSizeOfInstance(NativeFlatBufferedHNSWVectorsWriter.class); - private static final String COMPONENT = "NativeFlatBufferedHNSWVectorsWriter"; - - private final int numInputVectors; - private final List fields = new ArrayList<>(); - private final InfoStream infoStream; - private NativeFlatVectorsWriter nativeFlat; - private AcceleratedHnswGraphOutput graphOutput; - private boolean finished; - - NativeFlatBufferedHNSWVectorsWriter( - SegmentWriteState state, AcceleratedHNSWParams acceleratedHNSWParams, int numInputVectors) - throws IOException { - super(); - if (numInputVectors <= 0) { - throw new IllegalArgumentException("numInputVectors must be > 0, got " + numInputVectors); - } - if (state.segmentInfo.getIndexSort() != null) { - throw new IllegalArgumentException( - "AcceleratedHNSWParams.numInputVectors (native flat buffering) does not support" - + " index-sorted segments; unset it (0) to use the heap-buffered path"); - } - this.infoStream = state.infoStream; - this.numInputVectors = numInputVectors; - boolean success = false; - try { - // In hint mode we own the flat files; the Lucene flat writer must be absent to avoid opening - // the same .vec/.vemf outputs. - nativeFlat = new NativeFlatVectorsWriter(state); - graphOutput = new AcceleratedHnswGraphOutput(state, acceleratedHNSWParams); - success = true; - printInfoStream(infoStream, COMPONENT, "NativeFlatBufferedHNSWVectorsWriter is initialized"); - } finally { - if (success == false) { - IOUtils.closeWhileHandlingException(this); - } - } - } - - /** - * Add new field for indexing. - */ - @Override - public KnnFieldVectorsWriter addField(FieldInfo fieldInfo) throws IOException { - var encoding = fieldInfo.getVectorEncoding(); - if (encoding != FLOAT32) { - throw new IllegalArgumentException("Expected float32, got:" + encoding); - } - // Buffer directly into a native host matrix; return the field writer itself so Lucene routes - // addValue() here rather than to a (nonexistent) Lucene flat field writer. - var cuvsFieldWriter = new NativeFieldWriter(fieldInfo, numInputVectors); - fields.add(cuvsFieldWriter); - return cuvsFieldWriter; - } - - /** - * Writes the flat {@code .vec}/{@code .vemf} from each field's native host matrix, builds the - * CAGRA/HNSW graph from the same matrix, then releases the matrix. - */ - @Override - public void flush(int maxDoc, DocMap sortMap) throws IOException { - if (sortMap != null) { - throw new UnsupportedOperationException( - "AcceleratedHNSWParams.numInputVectors (native flat buffering) does not support" - + " index-sorted segments; unset it (0) to enable the sorted flush path"); - } - for (var field : fields) { - writeFieldNative(field, maxDoc); - } - } - - private void writeFieldNative(NativeFieldWriter fieldData, int maxDoc) throws IOException { - int count = fieldData.getNativeVectorCount(); - if (count != numInputVectors) { - throw new IllegalStateException( - "numInputVectors (" - + numInputVectors - + ") must equal the number of vectors added (" - + count - + ") for field \"" - + fieldData.fieldInfo().name - + "\"; the native host matrix is sized for the hint exactly. This usually means" - + " IndexWriterConfig's auto-flush wasn't disabled (setMaxBufferedDocs /" - + " setRAMBufferSizeMB(DISABLE_AUTO_FLUSH)), so a flush landed before exactly" - + " numInputVectors vectors were added; see CagraHnswBulkIndexWriter."); - } - FieldInfo fieldInfo = fieldData.fieldInfo(); - try { - CuVSHostMatrix dataset = fieldData.getHostMatrix(); - nativeFlat.writeField(fieldInfo, dataset, maxDoc, fieldData.getDocsWithFieldSet()); - graphOutput.writeField(fieldInfo, dataset); - } finally { - fieldData.releaseNativeBuffer(); - } - } - - /** - * Native flat buffering supports only the unsorted single-segment flush path; merges are - * rejected. - */ - @Override - public void mergeOneField(FieldInfo fieldInfo, MergeState mergeState) throws IOException { - throw new UnsupportedOperationException( - "AcceleratedHNSWParams.numInputVectors (native flat buffering) supports only the" - + " unsorted single-segment flush path; unset it (0) to enable merges"); - } - - /** - * Called once at the end before close. - */ - @Override - public void finish() throws IOException { - if (finished) { - throw new IllegalStateException("already finished"); - } - finished = true; - nativeFlat.finish(); - graphOutput.finish(); - } - - /** - * Closes the resources. - */ - @Override - public void close() throws IOException { - printInfoStream(infoStream, COMPONENT, "Closing resources"); - try { - releaseAllNativeBuffers(); - } finally { - IOUtils.close(graphOutput, nativeFlat); - closeCuVSResourcesInstance(); - } - } - - /** - * Releases every field's native host matrix. {@link #writeFieldNative} already releases a - * field's buffer once its write succeeds, but that is not a reliable place to guarantee cleanup: - * a count mismatch throws before that field's buffer is ever touched, and any field failing - * aborts {@link #flush}'s loop before later fields are even reached. This is the guaranteed - * backstop, run unconditionally on close regardless of how flush exited. - * {@link NativeFieldWriter#releaseNativeBuffer} is idempotent, so re-releasing an - * already-released field here is a safe no-op. - */ - private void releaseAllNativeBuffers() { - RuntimeException firstFailure = null; - for (var field : fields) { - try { - field.releaseNativeBuffer(); - } catch (RuntimeException e) { - if (firstFailure == null) { - firstFailure = e; - } else { - firstFailure.addSuppressed(e); - } - } - } - if (firstFailure != null) { - throw firstFailure; - } - } - - /** - * Returns the memory usage of this object in bytes. - */ - @Override - public long ramBytesUsed() { - long total = SHALLOW_RAM_BYTES_USED; - for (var field : fields) { - total += field.ramBytesUsed(); - } - return total; - } -} diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/NativeFlatVectorsWriter.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/NativeFlatVectorsWriter.java deleted file mode 100644 index 4b37686590..0000000000 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/NativeFlatVectorsWriter.java +++ /dev/null @@ -1,199 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: Apache-2.0 - */ - -package com.nvidia.cuvs.lucene; - -import com.nvidia.cuvs.CuVSHostMatrix; -import java.io.Closeable; -import java.io.IOException; -import java.lang.foreign.MemorySegment; -import java.lang.foreign.ValueLayout; -import java.nio.ByteOrder; -import org.apache.lucene.codecs.CodecUtil; -import org.apache.lucene.codecs.lucene95.OrdToDocDISIReaderConfiguration; -import org.apache.lucene.index.DocsWithFieldSet; -import org.apache.lucene.index.FieldInfo; -import org.apache.lucene.index.IndexFileNames; -import org.apache.lucene.index.SegmentWriteState; -import org.apache.lucene.store.IndexOutput; -import org.apache.lucene.util.IOUtils; - -/** - * Writes the flat vector files ({@code .vec} data + {@code .vemf} meta) directly from a native host - * matrix, byte-for-byte compatible with Lucene's {@code Lucene99FlatVectorsWriter} so the stock - * {@code Lucene99FlatVectorsReader} can read them. - * - *

This is the hint-path counterpart to delegating to Lucene's {@code FlatVectorsWriter}: the - * accelerated writer streams vectors into a {@link CuVSHostMatrix} during indexing (see - * {@link NativeFieldWriter}) and never materialises the full dataset as a {@code List} - * on the Java heap, so the flat file is written here from that native matrix instead. - * - *

Ported code — pinned to the {@code Lucene99} format. The dense float32 layout (format - * constants, header, meta field order, and footer) is transcribed from {@code - * Lucene99FlatVectorsFormat}/{@code Lucene99FlatVectorsWriter} (tag {@code - * releases/lucene/10.2.0}): - * - *

    - *
  • {@code org.apache.lucene.codecs.lucene99.Lucene99FlatVectorsFormat} — format constants: - * https://github.com/apache/lucene/blob/releases/lucene/10.2.0/lucene/core/src/java/org/apache/lucene/codecs/lucene99/Lucene99FlatVectorsFormat.java - *
  • {@code org.apache.lucene.codecs.lucene99.Lucene99FlatVectorsWriter} — write sequence: - * https://github.com/apache/lucene/blob/releases/lucene/10.2.0/lucene/core/src/java/org/apache/lucene/codecs/lucene99/Lucene99FlatVectorsWriter.java - *
- * - *

{@code Lucene99} is a frozen, versioned codec name: its constants and byte layout stay fixed - * for the life of the codec, and files written to it stay readable by the stock {@code - * Lucene99FlatVectorsReader} — via {@code lucene-backward-codecs} once a newer default codec ships - * — for the rest of the Lucene 10.x major version. The constants below (codec names, extensions, - * {@code VERSION_CURRENT}) are the fixed handshake the reader string-matches ({@code - * CodecUtil.checkIndexHeader}) and file-extension-matches against, so they stay as transcribed here - * for the life of the {@code Lucene99} codec. - * - *

On a {@code lucene-core} version bump: run {@code TestNativeFlatVectorsWriterRoundTrip} - * to confirm the stock reader on the new classpath still accepts what this class writes; it builds - * a small index and asserts every vector round-trips byte-exact through {@code - * Lucene99FlatVectorsReader}. - * - *

Moving to a newer flat-vector format (e.g. once {@code Lucene99} support is dropped - * ahead of an 11.x major bump): diff the new format's writer against the one linked above and - * re-derive the {@code writeField}/{@code writeMeta} byte sequence (header, meta field order, - * footer) here. This class writes directly from native memory to avoid the per-vector {@code - * FloatVectorValues} indirection Lucene's own writer requires — that's the reason to keep - * hand-porting the format rather than delegating to it. Then update the codec name/extension/version - * constants below and the Lucene version and links in this javadoc. - */ -final class NativeFlatVectorsWriter implements Closeable { - - // Mirrors org.apache.lucene.codecs.lucene99.Lucene99FlatVectorsFormat (10.2.0) so the standard - // Lucene99FlatVectorsReader accepts the header/codec of the files written here. - private static final String META_CODEC_NAME = "Lucene99FlatVectorsFormatMeta"; - private static final String VECTOR_DATA_CODEC_NAME = "Lucene99FlatVectorsFormatData"; - private static final String META_EXTENSION = "vemf"; - private static final String VECTOR_DATA_EXTENSION = "vec"; - private static final int VERSION_CURRENT = 0; - private static final int DIRECT_MONOTONIC_BLOCK_SHIFT = 16; - - // Little-endian float layout matching Lucene's on-disk .vec byte order. Must be UNALIGNED: the - // destination is a heap byte[]-backed MemorySegment whose max alignment is 1 byte, so a 4-byte - // aligned JAVA_FLOAT layout is rejected with "incompatible with alignment constraints". - private static final ValueLayout.OfFloat LE_FLOAT = - ValueLayout.JAVA_FLOAT_UNALIGNED.withOrder(ByteOrder.LITTLE_ENDIAN); - - // Byte granularity for a single writeBytes call; bounds the transient encode buffer. - private static final int CHUNK_BYTES = 1 << 18; // 256 KiB - - private final IndexOutput meta; - private final IndexOutput vectorData; - private boolean finished; - - NativeFlatVectorsWriter(SegmentWriteState state) throws IOException { - String metaFileName = - IndexFileNames.segmentFileName(state.segmentInfo.name, state.segmentSuffix, META_EXTENSION); - String vectorDataFileName = - IndexFileNames.segmentFileName( - state.segmentInfo.name, state.segmentSuffix, VECTOR_DATA_EXTENSION); - boolean success = false; - try { - meta = state.directory.createOutput(metaFileName, state.context); - vectorData = state.directory.createOutput(vectorDataFileName, state.context); - CodecUtil.writeIndexHeader( - meta, META_CODEC_NAME, VERSION_CURRENT, state.segmentInfo.getId(), state.segmentSuffix); - CodecUtil.writeIndexHeader( - vectorData, - VECTOR_DATA_CODEC_NAME, - VERSION_CURRENT, - state.segmentInfo.getId(), - state.segmentSuffix); - success = true; - } finally { - if (success == false) { - IOUtils.closeWhileHandlingException(this); - } - } - } - - /** - * Writes one dense float32 field: the raw vectors to {@code .vec} and the field metadata (plus the - * ordinal-to-doc mapping) to {@code .vemf}. Vectors are read from {@code matrix} in ordinal order, - * which matches the ascending-docID order in which {@code docsWithField} was populated. - * - * @param field the field being written - * @param matrix the native host matrix holding {@code docsWithField.cardinality()} rows of {@code - * field.getVectorDimension()} floats each - * @param maxDoc the segment's maxDoc, used to build the ordinal-to-doc mapping - * @param docsWithField the set of docs that have a value for this field - */ - void writeField( - FieldInfo field, CuVSHostMatrix matrix, int maxDoc, DocsWithFieldSet docsWithField) - throws IOException { - // Mirrors Lucene99FlatVectorsWriter#writeField (see class-level version pin). - int count = docsWithField.cardinality(); - int dim = field.getVectorDimension(); - long vectorDataOffset = vectorData.alignFilePointer(Float.BYTES); - writeFloat32Vectors(matrix, count, dim); - long vectorDataLength = vectorData.getFilePointer() - vectorDataOffset; - writeMeta(field, maxDoc, count, vectorDataOffset, vectorDataLength, docsWithField); - } - - private void writeFloat32Vectors(CuVSHostMatrix matrix, int count, int dim) throws IOException { - int rowBytes = dim * Float.BYTES; - int chunkRows = Math.max(1, CHUNK_BYTES / rowBytes); - byte[] chunk = new byte[chunkRows * rowBytes]; - MemorySegment chunkSeg = MemorySegment.ofArray(chunk); - float[] rowBuf = new float[dim]; - int r = 0; - for (int ord = 0; ord < count; ord++) { - matrix.getRow(ord).toArray(rowBuf); // native -> heap float[] (bulk) - MemorySegment.copy(rowBuf, 0, chunkSeg, LE_FLOAT, (long) r * rowBytes, dim); // -> LE bytes - if (++r == chunkRows) { - vectorData.writeBytes(chunk, r * rowBytes); - r = 0; - } - } - if (r > 0) { - vectorData.writeBytes(chunk, r * rowBytes); - } - } - - private void writeMeta( - FieldInfo field, - int maxDoc, - int count, - long vectorDataOffset, - long vectorDataLength, - DocsWithFieldSet docsWithField) - throws IOException { - // Mirrors Lucene99FlatVectorsWriter#writeMeta (see class-level version pin); field order is - // load-bearing and must match Lucene's reader. - meta.writeInt(field.number); - meta.writeInt(field.getVectorEncoding().ordinal()); - meta.writeInt(field.getVectorSimilarityFunction().ordinal()); - meta.writeVLong(vectorDataOffset); - meta.writeVLong(vectorDataLength); - meta.writeVInt(field.getVectorDimension()); - meta.writeInt(count); - OrdToDocDISIReaderConfiguration.writeStoredMeta( - DIRECT_MONOTONIC_BLOCK_SHIFT, meta, vectorData, count, maxDoc, docsWithField); - } - - /** Writes the end-of-fields marker and footers. Mirrors {@code Lucene99FlatVectorsWriter.finish}. */ - void finish() throws IOException { - if (finished) { - throw new IllegalStateException("already finished"); - } - finished = true; - if (meta != null) { - meta.writeInt(-1); - CodecUtil.writeFooter(meta); - } - if (vectorData != null) { - CodecUtil.writeFooter(vectorData); - } - } - - @Override - public void close() throws IOException { - IOUtils.close(meta, vectorData); - } -} diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Utils.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Utils.java index 9b8028543c..034f8aa5d8 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Utils.java +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Utils.java @@ -51,50 +51,82 @@ static RuntimeException handleThrowable(Throwable t) throws IOException { } /** - * Builds a host-memory CuVSMatrix from a list of float vectors. + * A method to build a CuVSMatrix from a list of float vectors. * - *

Copies vectors directly into a native host matrix via {@link CuVSMatrix#hostBuilder}, - * without creating an intermediate {@code float[][]} on the heap. + * Uses CuVSMatrix.Builder to copy vectors directly to device memory + * without creating intermediate heap arrays. * * @param data The float vectors - * @param dimensions The number of float elements in each vector - * @return a host-memory CuVSMatrix + * @param dimensions The number float elements in each vector + * @param resources The CuVS resources for device matrix creation + * @return an instance of CuVSMatrix */ - static CuVSMatrix createFloatMatrix(List data, int dimensions) { + static CuVSMatrix createFloatMatrix(List data, int dimensions, CuVSResources resources) { + // Use Builder pattern to avoid intermediate float[][] allocation + // and copy directly from List to device memory CuVSMatrix.Builder builder = - CuVSMatrix.hostBuilder(data.size(), dimensions, CuVSMatrix.DataType.FLOAT); + CuVSMatrix.deviceBuilder( + resources, + data.size(), // rows (number of vectors) + dimensions, // columns (vector dimension) + CuVSMatrix.DataType.FLOAT); + + // Add vectors one by one - builder copies directly to device memory for (float[] vector : data) { builder.addVector(vector); } + return builder.build(); } /** - * Builds a host-memory CuVSMatrix from a list of byte vectors (e.g. quantized vectors). + * A method to build a CuVSMatrix from a list of byte vectors (for binary quantized vectors). + * + * Uses CuVSMatrix.Builder to copy vectors directly to device memory + * without creating intermediate heap arrays. * * @param data The byte vectors (packed bits for binary quantization) * @param bytesPerVector The number of bytes in each vector - * @return a host-memory CuVSMatrix with BYTE data type + * @param resources The CuVS resources for device matrix creation + * @return an instance of CuVSMatrix with BYTE data type */ - static CuVSMatrix createByteMatrix(List data, int bytesPerVector) { + static CuVSMatrix createByteMatrix( + List data, int bytesPerVector, CuVSResources resources) { + // Use Builder pattern to avoid intermediate byte[][] allocation + // and copy directly from List to device memory CuVSMatrix.Builder builder = - CuVSMatrix.hostBuilder(data.size(), bytesPerVector, CuVSMatrix.DataType.BYTE); + CuVSMatrix.deviceBuilder( + resources, + data.size(), // rows (number of vectors) + bytesPerVector, // columns (bytes per vector) + CuVSMatrix.DataType.BYTE); + + // Add vectors one by one - builder copies directly to device memory for (byte[] vector : data) { builder.addVector(vector); } + return builder.build(); } /** - * Builds a host-memory CuVSMatrix from a 2D byte array (e.g. quantized vectors). + * A method to build a CuVSMatrix from a 2D byte array (for binary quantized vectors). * * @param data The 2D byte array (packed bits for binary quantization) * @param bytesPerVector The number of bytes in each vector - * @return a host-memory CuVSMatrix with BYTE data type + * @param resources The CuVS resources for device matrix creation + * @return an instance of CuVSMatrix with BYTE data type */ - static CuVSMatrix createByteMatrixFromArray(byte[][] data, int bytesPerVector) { + static CuVSMatrix createByteMatrixFromArray( + byte[][] data, int bytesPerVector, CuVSResources resources) { CuVSMatrix.Builder builder = - CuVSMatrix.hostBuilder(data.length, bytesPerVector, CuVSMatrix.DataType.BYTE); + CuVSMatrix.deviceBuilder( + resources, + data.length, // rows (number of vectors) + bytesPerVector, // columns (bytes per vector) + CuVSMatrix.DataType.BYTE); + + // Add vectors one by one - builder copies directly to device memory for (byte[] vector : data) { builder.addVector(vector); } diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/VectorSource.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/VectorSource.java deleted file mode 100644 index eeab85cf1b..0000000000 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/VectorSource.java +++ /dev/null @@ -1,35 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: Apache-2.0 - */ -package com.nvidia.cuvs.lucene; - -import java.io.Closeable; -import java.io.IOException; - -/** - * A forward-only, single-consumer source of float vectors for {@link CagraHnswBulkIndexWriter}. - * - *

Implementations own how the underlying data is fetched (a local file, a database cursor, an - * object-store stream, ...); {@link CagraHnswBulkIndexWriter} only depends on this contract, not on - * any particular storage. - * - *

{@link #get(int, float[])} must be called with strictly increasing indices from a single - * thread -- repeating the previous call's index is not permitted. This mirrors how a bulk indexer - * consumes its input (front-to-back, one pass, exactly once per vector) and lets implementations - * use a simple prefetch/streaming strategy instead of arbitrary random access. - */ -public interface VectorSource extends Closeable { - - /** Number of dimensions of every vector returned by this source. */ - int dimensions(); - - /** Number of vectors available, i.e. the exclusive upper bound on indices passed to {@link #get}. */ - int size(); - - /** - * Fills {@code dst} with the vector at {@code index} (no allocation). {@code index} must be - * strictly greater than the index passed to the previous call, if any. - */ - void get(int index, float[] dst) throws IOException; -} diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedHNSWParamsSurface.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedHNSWParamsSurface.java deleted file mode 100644 index ac027e2701..0000000000 --- a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedHNSWParamsSurface.java +++ /dev/null @@ -1,47 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: Apache-2.0 - */ -package com.nvidia.cuvs.lucene; - -import static org.junit.Assert.assertFalse; - -import java.lang.reflect.Constructor; -import java.lang.reflect.Modifier; -import org.junit.Test; - -/** - * Regression guard for the native-flat-buffering surface narrowing: the {@code numInputVectors} - * constructor overloads on {@link Lucene101AcceleratedHNSWCodec} and {@link - * Lucene99AcceleratedHNSWVectorsFormat} must stay package-private, reachable only from within - * {@code com.nvidia.cuvs.lucene} (i.e. only by {@link CagraHnswBulkIndexWriter}), not from a - * generic external Lucene codec user. A future change that accidentally re-widens either - * constructor to {@code public} would reopen exactly the footgun this package's design is meant - * to close, without necessarily being caught by any functional test -- this test exists to catch - * that specific mistake directly. - */ -public class TestAcceleratedHNSWParamsSurface { - - @Test - public void testCodecNumInputVectorsConstructorIsNotPublic() throws NoSuchMethodException { - Constructor ctor = - Lucene101AcceleratedHNSWCodec.class.getDeclaredConstructor( - AcceleratedHNSWParams.class, int.class); - assertFalse( - "Lucene101AcceleratedHNSWCodec(AcceleratedHNSWParams, int) must not be public; it should" - + " only be reachable from within com.nvidia.cuvs.lucene (CagraHnswBulkIndexWriter)", - Modifier.isPublic(ctor.getModifiers())); - } - - @Test - public void testFormatNumInputVectorsConstructorIsNotPublic() throws NoSuchMethodException { - Constructor ctor = - Lucene99AcceleratedHNSWVectorsFormat.class.getDeclaredConstructor( - AcceleratedHNSWParams.class, int.class); - assertFalse( - "Lucene99AcceleratedHNSWVectorsFormat(AcceleratedHNSWParams, int) must not be public; it" - + " should only be reachable from within com.nvidia.cuvs.lucene" - + " (CagraHnswBulkIndexWriter)", - Modifier.isPublic(ctor.getModifiers())); - } -} diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestCagraHnswBulkIndexWriter.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestCagraHnswBulkIndexWriter.java deleted file mode 100644 index 2a299a6440..0000000000 --- a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestCagraHnswBulkIndexWriter.java +++ /dev/null @@ -1,446 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: Apache-2.0 - */ -package com.nvidia.cuvs.lucene; - -import static com.nvidia.cuvs.lucene.TestUtils.generateDataset; -import static com.nvidia.cuvs.lucene.ThreadLocalCuVSResourcesProvider.isSupported; -import static org.apache.lucene.index.VectorSimilarityFunction.EUCLIDEAN; - -import com.nvidia.cuvs.spi.CuVSProvider; -import java.io.File; -import java.io.IOException; -import java.nio.file.Path; -import java.nio.file.Paths; -import java.util.Random; -import java.util.UUID; -import org.apache.commons.io.FileUtils; -import org.apache.lucene.document.Document; -import org.apache.lucene.document.Field; -import org.apache.lucene.document.KnnFloatVectorField; -import org.apache.lucene.document.StringField; -import org.apache.lucene.index.DirectoryReader; -import org.apache.lucene.index.IndexWriterConfig; -import org.apache.lucene.search.IndexSearcher; -import org.apache.lucene.search.KnnFloatVectorQuery; -import org.apache.lucene.search.Sort; -import org.apache.lucene.search.SortField; -import org.apache.lucene.search.TopDocs; -import org.apache.lucene.store.Directory; -import org.apache.lucene.store.FSDirectory; -import org.apache.lucene.tests.util.LuceneTestCase; -import org.apache.lucene.tests.util.LuceneTestCase.SuppressSysoutChecks; -import org.junit.After; -import org.junit.Before; -import org.junit.BeforeClass; -import org.junit.Test; - -/** - * Functional coverage for {@link CagraHnswBulkIndexWriter}: the manual/direct-instance {@link - * CagraHnswBulkIndexWriter#addDocument} API (including its safety checks) and the one-shot {@link - * CagraHnswBulkIndexWriter#indexFbin}/{@link CagraHnswBulkIndexWriter#build(VectorSource, Config)} - * convenience entry points (single-segment, K-segment sequential, K-segment overlapped, - * {@link CagraHnswBulkIndexWriter.FieldCallback} metadata, rejection of {@code overlapped} for - * {@code build}). Guard-rail behavior of the underlying native-buffered writer itself is already - * covered by {@link TestNativeFlatBufferingGuardRails}; this class covers this class's own - * orchestration and its own safety checks layered on top. - */ -@SuppressSysoutChecks(bugUrl = "") -public class TestCagraHnswBulkIndexWriter extends LuceneTestCase { - - private static final String ID_FIELD = "id"; - private static final String VECTOR_FIELD = "vector_field"; - private static final String CATEGORY_FIELD = "category"; - - private Random random; - private Path indexDirPath; - private Path fbinPath; - - @BeforeClass - public static void beforeClass() { - // It is recommended to enable RMM allocation mode at application start, before constructing - // any CagraHnswBulkIndexWriter, to avoid device-wide sync from the default allocator. - try { - CuVSProvider.provider().enableRMMAsyncMemory(); - } catch (UnsupportedOperationException unsupported) { - assumeTrue("cuVS not supported: " + unsupported.getMessage(), false); - } - } - - @Before - public void beforeTest() { - assumeTrue("cuVS not supported", isSupported()); - random = new Random(222); - indexDirPath = Paths.get(UUID.randomUUID().toString()); - fbinPath = Paths.get(UUID.randomUUID() + ".fbin"); - } - - @After - public void afterTest() throws IOException { - if (indexDirPath != null) { - File dir = indexDirPath.toFile(); - if (dir.exists() && dir.isDirectory()) { - FileUtils.deleteDirectory(dir); - } - } - if (fbinPath != null) { - new File(fbinPath.toString()).delete(); - } - } - - @Test - public void testSingleSegmentBuildIsSearchable() throws Exception { - int numDocs = 300; - int dimension = 32; - float[][] dataset = generateDataset(random, numDocs, dimension); - TestUtils.writeFbin(fbinPath, dataset); - - CagraHnswBulkIndexWriter.indexFbin(fbinPath, configFor(dimension, 1, false)); - - assertSearchable(numDocs, dataset, /* expectedSegments= */ 1); - } - - @Test - public void testPartitionedSequentialBuildProducesKSegments() throws Exception { - int numDocs = 400; - int dimension = 24; - int k = 4; - float[][] dataset = generateDataset(random, numDocs, dimension); - TestUtils.writeFbin(fbinPath, dataset); - - CagraHnswBulkIndexWriter.indexFbin(fbinPath, configFor(dimension, k, false)); - - assertSearchable(numDocs, dataset, /* expectedSegments= */ k); - } - - @Test - public void testOverlappedBuildProducesKSegments() throws Exception { - int numDocs = 400; - int dimension = 24; - int k = 4; - float[][] dataset = generateDataset(random, numDocs, dimension); - TestUtils.writeFbin(fbinPath, dataset); - - CagraHnswBulkIndexWriter.indexFbin(fbinPath, configFor(dimension, k, true)); - - assertSearchable(numDocs, dataset, /* expectedSegments= */ k); - } - - @Test - public void testGenericBuildViaVectorSource() throws Exception { - int numDocs = 200; - int dimension = 16; - float[][] dataset = generateDataset(random, numDocs, dimension); - - CagraHnswBulkIndexWriter.build( - new InMemoryVectorSource(dataset), configFor(dimension, 1, false)); - - assertSearchable(numDocs, dataset, /* expectedSegments= */ 1); - } - - @Test - public void testGenericBuildRejectsOverlapped() throws Exception { - int dimension = 8; - float[][] dataset = generateDataset(random, 50, dimension); - CagraHnswBulkIndexWriter.Config config = configFor(dimension, 2, true); - try { - CagraHnswBulkIndexWriter.build(new InMemoryVectorSource(dataset), config); - fail( - "expected IllegalArgumentException: overlapped is not supported by build(VectorSource," - + " Config)"); - } catch (IllegalArgumentException expected) { - // expected - } - } - - @Test - public void testIdFieldCanBeDisabled() throws Exception { - int numDocs = 50; - int dimension = 8; - float[][] dataset = generateDataset(random, numDocs, dimension); - TestUtils.writeFbin(fbinPath, dataset); - - CagraHnswBulkIndexWriter.Config config = - CagraHnswBulkIndexWriter.Config.builder() - .field(VECTOR_FIELD, dimension, EUCLIDEAN) - .idField(null) - .graphBuild(new AcceleratedHNSWParams.Builder().build()) - .segments(1, false) - .targetDirectory(indexDirPath) - .build(); - CagraHnswBulkIndexWriter.indexFbin(fbinPath, config); - - try (Directory dir = FSDirectory.open(indexDirPath); - DirectoryReader reader = DirectoryReader.open(dir)) { - IndexSearcher searcher = new IndexSearcher(reader); - TopDocs results = searcher.search(new KnnFloatVectorQuery(VECTOR_FIELD, dataset[0], 5), 5); - assertEquals(5, results.scoreDocs.length); - assertNull( - "id field should not be stored when idField(null) is used", - searcher.storedFields().document(results.scoreDocs[0].doc).get(ID_FIELD)); - } - } - - /** {@link CagraHnswBulkIndexWriter.FieldCallback} lets the one-shot API attach metadata per row. */ - @Test - public void testFieldCallbackAttachesMetadata() throws Exception { - int numDocs = 60; - int dimension = 8; - float[][] dataset = generateDataset(random, numDocs, dimension); - TestUtils.writeFbin(fbinPath, dataset); - - CagraHnswBulkIndexWriter.indexFbin( - fbinPath, - configFor(dimension, 1, false), - (doc, id) -> - doc.add( - new StringField(CATEGORY_FIELD, id % 2 == 0 ? "even" : "odd", Field.Store.YES))); - - try (Directory dir = FSDirectory.open(indexDirPath); - DirectoryReader reader = DirectoryReader.open(dir)) { - IndexSearcher searcher = new IndexSearcher(reader); - TopDocs results = searcher.search(new KnnFloatVectorQuery(VECTOR_FIELD, dataset[0], 1), 1); - String category = - searcher.storedFields().document(results.scoreDocs[0].doc).get(CATEGORY_FIELD); - assertEquals("even", category); // id=0 is even - } - } - - /** - * The manual/direct-instance API: caller builds the {@link Document} (including arbitrary extra - * fields) and drives {@link CagraHnswBulkIndexWriter#addDocument} directly, same shape as a - * plain {@link org.apache.lucene.index.IndexWriter}. - */ - @Test - public void testManualAddDocumentWithMetadata() throws Exception { - int numDocs = 40; - int dimension = 12; - float[][] dataset = generateDataset(random, numDocs, dimension); - - CagraHnswBulkIndexWriter.Config config = - CagraHnswBulkIndexWriter.Config.builder() - .field(VECTOR_FIELD, dimension, EUCLIDEAN) - .graphBuild(new AcceleratedHNSWParams.Builder().build()) - .build(); - - try (Directory dir = FSDirectory.open(indexDirPath); - CagraHnswBulkIndexWriter writer = - new CagraHnswBulkIndexWriter(dir, new IndexWriterConfig(), config, numDocs)) { - for (int i = 0; i < numDocs; i++) { - Document doc = new Document(); - doc.add(new StringField(ID_FIELD, Integer.toString(i), Field.Store.YES)); - doc.add(new StringField(CATEGORY_FIELD, i % 2 == 0 ? "even" : "odd", Field.Store.YES)); - doc.add(new KnnFloatVectorField(VECTOR_FIELD, dataset[i], EUCLIDEAN)); - writer.addDocument(doc); - } - } - - try (Directory dir = FSDirectory.open(indexDirPath); - DirectoryReader reader = DirectoryReader.open(dir)) { - assertEquals(1, reader.leaves().size()); - IndexSearcher searcher = new IndexSearcher(reader); - TopDocs results = searcher.search(new KnnFloatVectorQuery(VECTOR_FIELD, dataset[0], 1), 1); - assertEquals( - "even", searcher.storedFields().document(results.scoreDocs[0].doc).get(CATEGORY_FIELD)); - } - } - - @Test - public void testCloseWithTooFewDocumentsThrows() throws Exception { - int dimension = 8; - CagraHnswBulkIndexWriter.Config config = - CagraHnswBulkIndexWriter.Config.builder() - .field(VECTOR_FIELD, dimension, EUCLIDEAN) - .graphBuild(new AcceleratedHNSWParams.Builder().build()) - .build(); - - try (Directory dir = FSDirectory.open(indexDirPath)) { - CagraHnswBulkIndexWriter writer = - new CagraHnswBulkIndexWriter(dir, new IndexWriterConfig(), config, 10); - Document doc = new Document(); - doc.add( - new KnnFloatVectorField( - VECTOR_FIELD, generateDataset(random, 1, dimension)[0], EUCLIDEAN)); - writer.addDocument(doc); // only 1 of the promised 10 - - try { - writer.close(); - fail("expected IllegalStateException: fewer documents added than exactVectorCount"); - } catch (IllegalStateException expected) { - // Specifically this class's mismatch error, not the underlying native writer's own count - // check: close() must roll the buffered documents back, not flush them. - assertEquals("expected 10 documents, got 1", expected.getMessage()); - } - // The rollback left nothing committed, so there is no segment to open. - assertFalse(DirectoryReader.indexExists(dir)); - } - } - - @Test - public void testSourceFailureMidBuildIsNotMaskedByCleanup() throws Exception { - int dimension = 8; - float[][] dataset = generateDataset(random, 200, dimension); - // Fails halfway through, so the writer is cleaned up holding 100 of the 200 promised vectors. - VectorSource failing = new FailingVectorSource(dataset, 100); - - try { - CagraHnswBulkIndexWriter.build(failing, configFor(dimension, 1, false)); - fail("expected the source's IOException to propagate"); - } catch (IOException expected) { - // The cleanup path must not report the count mismatch it necessarily sees instead of the - // failure that caused it. - assertEquals("source failed at 100", expected.getMessage()); - } - } - - @Test - public void testAddDocumentBeyondExactCountThrows() throws Exception { - int dimension = 8; - CagraHnswBulkIndexWriter.Config config = - CagraHnswBulkIndexWriter.Config.builder() - .field(VECTOR_FIELD, dimension, EUCLIDEAN) - .graphBuild(new AcceleratedHNSWParams.Builder().build()) - .build(); - - try (Directory dir = FSDirectory.open(indexDirPath)) { - CagraHnswBulkIndexWriter writer = - new CagraHnswBulkIndexWriter(dir, new IndexWriterConfig(), config, 1); - float[][] vectors = generateDataset(random, 2, dimension); - Document doc1 = new Document(); - doc1.add(new KnnFloatVectorField(VECTOR_FIELD, vectors[0], EUCLIDEAN)); - writer.addDocument(doc1); - - Document doc2 = new Document(); - doc2.add(new KnnFloatVectorField(VECTOR_FIELD, vectors[1], EUCLIDEAN)); - try { - writer.addDocument(doc2); // exactVectorCount was 1 - fail("expected IllegalStateException: more documents added than exactVectorCount"); - } catch (IllegalStateException expected) { - // expected -- the rejected call never reached the underlying writer or incremented the - // count, so it's still exactly 1 and close() below completes normally. - } finally { - writer.close(); - } - } - } - - @Test - public void testConstructorRejectsIndexSort() throws Exception { - int dimension = 8; - CagraHnswBulkIndexWriter.Config config = - CagraHnswBulkIndexWriter.Config.builder() - .field(VECTOR_FIELD, dimension, EUCLIDEAN) - .graphBuild(new AcceleratedHNSWParams.Builder().build()) - .build(); - IndexWriterConfig conf = - new IndexWriterConfig() - .setIndexSort(new Sort(new SortField(ID_FIELD, SortField.Type.STRING))); - - try (Directory dir = FSDirectory.open(indexDirPath)) { - try { - new CagraHnswBulkIndexWriter(dir, conf, config, 10); - fail("expected IllegalArgumentException: index-sorted segments are not supported"); - } catch (IllegalArgumentException expected) { - // expected - } - } - } - - private CagraHnswBulkIndexWriter.Config configFor( - int dimension, int numSegments, boolean overlapped) { - return CagraHnswBulkIndexWriter.Config.builder() - .field(VECTOR_FIELD, dimension, EUCLIDEAN) - .graphBuild(new AcceleratedHNSWParams.Builder().build()) - .segments(numSegments, overlapped) - .targetDirectory(indexDirPath) - .build(); - } - - private void assertSearchable(int numDocs, float[][] dataset, int expectedSegments) - throws Exception { - try (Directory dir = FSDirectory.open(indexDirPath); - DirectoryReader reader = DirectoryReader.open(dir)) { - assertEquals("unexpected segment count", expectedSegments, reader.leaves().size()); - - IndexSearcher searcher = new IndexSearcher(reader); - int topK = 10; - TopDocs results = - searcher.search(new KnnFloatVectorQuery(VECTOR_FIELD, dataset[0], topK), topK); - assertEquals(topK, results.scoreDocs.length); - boolean sawQueryVectorItself = false; - for (var scoreDoc : results.scoreDocs) { - String id = searcher.storedFields().document(scoreDoc.doc).get(ID_FIELD); - int idValue = Integer.parseInt(id); - assertTrue("returned id out of range: " + id, idValue >= 0 && idValue < numDocs); - sawQueryVectorItself |= idValue == 0; - } - // Querying with dataset[0] itself (an exact, zero-distance match) should reliably surface - // id=0 within topK for a graph this small, across all segments the query fans out over. - assertTrue( - "expected id=0 (the query vector itself) within topK results", sawQueryVectorItself); - } - } - - /** An {@link InMemoryVectorSource} that throws once a given row is reached. */ - private static final class FailingVectorSource implements VectorSource { - private final InMemoryVectorSource delegate; - private final int failAt; - - FailingVectorSource(float[][] dataset, int failAt) { - this.delegate = new InMemoryVectorSource(dataset); - this.failAt = failAt; - } - - @Override - public int dimensions() { - return delegate.dimensions(); - } - - @Override - public int size() { - return delegate.size(); - } - - @Override - public void get(int index, float[] dst) throws IOException { - if (index >= failAt) { - throw new IOException("source failed at " + failAt); - } - delegate.get(index, dst); - } - - @Override - public void close() { - delegate.close(); - } - } - - private static final class InMemoryVectorSource implements VectorSource { - private final float[][] dataset; - - InMemoryVectorSource(float[][] dataset) { - this.dataset = dataset; - } - - @Override - public int dimensions() { - return dataset.length == 0 ? 0 : dataset[0].length; - } - - @Override - public int size() { - return dataset.length; - } - - @Override - public void get(int index, float[] dst) { - System.arraycopy(dataset[index], 0, dst, 0, dst.length); - } - - @Override - public void close() { - // nothing to release - } - } -} diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestCagraIndexParamsFactory.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestCagraIndexParamsFactory.java index 302629d6c1..f559ea3808 100644 --- a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestCagraIndexParamsFactory.java +++ b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestCagraIndexParamsFactory.java @@ -195,43 +195,6 @@ public void testHnswHeuristicDelegatesToCuVS() { assertEquals(7, cagraParams.getNumWriterThreads()); } - /** - * {@code cagraGraphBuildAlgo} is only consulted under {@code CUSTOM}; an explicit override left - * on the builder while the strategy is {@code HEURISTIC} must not change the delegated-to-cuVS - * result. Requires the native cuVS library. - */ - @Test - public void testHnswHeuristicIgnoresCagraGraphBuildAlgoOverride() { - assumeTrue("cuVS not supported", isSupported()); - - AcceleratedHNSWParams withoutOverride = - new AcceleratedHNSWParams.Builder() - .withStrategy(AcceleratedHNSWParams.Strategy.HEURISTIC) - .withMaxConn(16) - .withBeamWidth(100) - .build(); - AcceleratedHNSWParams withOverride = - new AcceleratedHNSWParams.Builder() - .withStrategy(AcceleratedHNSWParams.Strategy.HEURISTIC) - .withCagraGraphBuildAlgo(CagraGraphBuildAlgo.NN_DESCENT) - .withMaxConn(16) - .withBeamWidth(100) - .build(); - - CagraIndexParams derivedWithoutOverride = - CagraIndexParamsFactory.create(withoutOverride, 10_000, 128); - CagraIndexParams derivedWithOverride = - CagraIndexParamsFactory.create(withOverride, 10_000, 128); - - assertEquals( - derivedWithoutOverride.getCagraGraphBuildAlgo(), - derivedWithOverride.getCagraGraphBuildAlgo()); - assertEquals(derivedWithoutOverride.getGraphDegree(), derivedWithOverride.getGraphDegree()); - assertEquals( - derivedWithoutOverride.getIntermediateGraphDegree(), - derivedWithOverride.getIntermediateGraphDegree()); - } - /** * The heuristic type must reach cuVS rather than being pinned to the default: under * SIMILAR_SEARCH_PERFORMANCE cuVS derives {@code graph_degree = 2 + maxConn * 2 / 3} instead of @@ -291,12 +254,10 @@ public void testStrategySpecificParamsRemainAccepted() { .withStrategy(AcceleratedHNSWParams.Strategy.HEURISTIC) .withGraphDegree(96) .withIntermediateGraphDegree(192) - .withCagraGraphBuildAlgo(CagraGraphBuildAlgo.NN_DESCENT) .build(); // Retained verbatim on the instance; CagraIndexParamsFactory is what declines to apply them. assertEquals(96, hnswParams.getGraphdegree()); assertEquals(192, hnswParams.getIntermediateGraphDegree()); - assertEquals(CagraGraphBuildAlgo.NN_DESCENT, hnswParams.getCagraGraphBuildAlgo()); GPUSearchParams gpuParams = new GPUSearchParams.Builder() diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestFbinVectorSource.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestFbinVectorSource.java deleted file mode 100644 index 41e7d4a081..0000000000 --- a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestFbinVectorSource.java +++ /dev/null @@ -1,120 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: Apache-2.0 - */ -package com.nvidia.cuvs.lucene; - -import static com.nvidia.cuvs.lucene.TestUtils.generateDataset; -import static com.nvidia.cuvs.lucene.TestUtils.writeFbin; - -import java.io.File; -import java.nio.file.Path; -import java.nio.file.Paths; -import java.util.Random; -import java.util.UUID; -import org.apache.lucene.tests.util.LuceneTestCase; -import org.junit.After; -import org.junit.Before; -import org.junit.Test; - -/** - * Correctness coverage for {@link FbinVectorSource}: header parsing, whole-file vs. windowed - * (sliced) reads, and the forward-only/single-consumer contract it shares with {@link - * VectorSource}. Does not require GPU support -- this is pure file I/O. - */ -public class TestFbinVectorSource extends LuceneTestCase { - - private Path fbinPath; - - @Before - public void beforeTest() { - fbinPath = Paths.get(UUID.randomUUID() + ".fbin"); - } - - @After - public void afterTest() { - if (fbinPath != null) { - new File(fbinPath.toString()).delete(); - } - } - - @Test - public void testWholeFileReadMatchesDataset() throws Exception { - Random random = new Random(1); - int numVectors = 250; - int dimension = 17; - float[][] dataset = generateDataset(random, numVectors, dimension); - writeFbin(fbinPath, dataset); - - // A small chunk size forces multiple prefetch chunks so this also exercises the chunk - // boundary/advance() path, not just a single-chunk read. - try (FbinVectorSource source = new FbinVectorSource(fbinPath, /* chunkSizeMB= */ 1)) { - assertEquals(dimension, source.dimensions()); - assertEquals(numVectors, source.size()); - - float[] scratch = new float[dimension]; - for (int i = 0; i < numVectors; i++) { - source.get(i, scratch); - assertArrayEquals("vector " + i + " mismatch", dataset[i], scratch, 0f); - } - } - } - - @Test - public void testWindowedReadServesOnlyItsSlice() throws Exception { - Random random = new Random(2); - int numVectors = 100; - int dimension = 8; - int sliceStart = 30; - int sliceSize = 25; - float[][] dataset = generateDataset(random, numVectors, dimension); - writeFbin(fbinPath, dataset); - - try (FbinVectorSource source = new FbinVectorSource(fbinPath, sliceStart, sliceSize, 1)) { - assertEquals(dimension, source.dimensions()); - assertEquals(sliceSize, source.size()); - - for (int i = 0; i < sliceSize; i++) { - float[] got = source.get(i); - assertArrayEquals( - "relative index " + i + " should map to absolute " + (sliceStart + i), - dataset[sliceStart + i], - got, - 0f); - } - } - } - - @Test - public void testOutOfOrderAccessIsRejected() throws Exception { - Random random = new Random(3); - float[][] dataset = generateDataset(random, 20, 4); - writeFbin(fbinPath, dataset); - - try (FbinVectorSource source = new FbinVectorSource(fbinPath, 1)) { - source.get(5); - try { - source.get(2); // going backwards must be rejected: forward-only contract - fail("expected UnsupportedOperationException for out-of-order access"); - } catch (UnsupportedOperationException expected) { - // expected - } - } - } - - @Test - public void testOutOfBoundsIndexIsRejected() throws Exception { - Random random = new Random(4); - float[][] dataset = generateDataset(random, 10, 4); - writeFbin(fbinPath, dataset); - - try (FbinVectorSource source = new FbinVectorSource(fbinPath, 1)) { - try { - source.get(10); // window is [0, 10) - fail("expected IndexOutOfBoundsException"); - } catch (IndexOutOfBoundsException expected) { - // expected - } - } - } -} diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestMergedGraphOrdinalBounds.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestMergedGraphOrdinalBounds.java deleted file mode 100644 index 12b00c6e12..0000000000 --- a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestMergedGraphOrdinalBounds.java +++ /dev/null @@ -1,198 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: Apache-2.0 - */ -package com.nvidia.cuvs.lucene; - -import static com.nvidia.cuvs.lucene.TestUtils.generateDataset; -import static com.nvidia.cuvs.lucene.ThreadLocalCuVSResourcesProvider.isSupported; -import static org.apache.lucene.index.VectorSimilarityFunction.EUCLIDEAN; -import static org.apache.lucene.search.DocIdSetIterator.NO_MORE_DOCS; - -import java.util.Random; -import org.apache.lucene.codecs.Codec; -import org.apache.lucene.codecs.KnnVectorsReader; -import org.apache.lucene.codecs.hnsw.HnswGraphProvider; -import org.apache.lucene.codecs.perfield.PerFieldKnnVectorsFormat; -import org.apache.lucene.document.Document; -import org.apache.lucene.document.Field; -import org.apache.lucene.document.KnnFloatVectorField; -import org.apache.lucene.document.StringField; -import org.apache.lucene.index.CodecReader; -import org.apache.lucene.index.DirectoryReader; -import org.apache.lucene.index.FloatVectorValues; -import org.apache.lucene.index.IndexWriter; -import org.apache.lucene.index.IndexWriterConfig; -import org.apache.lucene.index.LeafReader; -import org.apache.lucene.index.NoMergePolicy; -import org.apache.lucene.index.Term; -import org.apache.lucene.index.TieredMergePolicy; -import org.apache.lucene.store.Directory; -import org.apache.lucene.tests.util.LuceneTestCase; -import org.apache.lucene.tests.util.LuceneTestCase.SuppressSysoutChecks; -import org.apache.lucene.tests.util.TestUtil; -import org.apache.lucene.util.hnsw.HnswGraph; -import org.junit.Test; - -/** - * Repro for the CI-observed {@code EOFException} in {@code OffHeapFloatVectorValues} during - * concurrent KNN search over an accelerated-HNSW index built via {@link - * TestAcceleratedHNSWDeletedDocuments}/{@link TestCuVSAcceleratedHNSWDeletedDocuments} (both: - * deletions + a real merge, heap-buffered path, no {@code numInputVectors}). - * - *

Rather than relying on a random concurrent search happening to traverse a bad graph node - * (which only reproduced intermittently, on one CI node), this walks the entire merged - * HNSW graph directly and asserts every neighbor ordinal is within the merged segment's actual - * flat-vector count. This targets the suspected root cause: {@code - * Lucene99AcceleratedHNSWVectorsWriter#mergeOneField} derives the merged vector set twice, - * independently -- once via the real {@code flatVectorsWriter.mergeOneField} (the authoritative - * flat {@code .vec} file) and again via {@code vectorBasedMerge}'s own call to {@code - * KnnVectorsWriter.MergedVectorValues.mergeFloatVectorValues} to build the CAGRA/HNSW graph. If - * those two independently-derived views of "the merged, post-deletion vector set" ever disagree in - * count or ordinal order, the graph ends up referencing ordinals the flat file doesn't actually - * have, which is exactly what an out-of-bounds read (EOFException) during traversal would look - * like. - * - *

This test is deterministic: it fails on any disagreement between the graph and the flat file, - * rather than depending on a search happening to reach the bad node. - */ -@SuppressSysoutChecks(bugUrl = "") -public class TestMergedGraphOrdinalBounds extends LuceneTestCase { - - private static final String ID_FIELD = "id"; - private static final String FIELD = "vector"; - - @Test - public void testMergedGraphOrdinalsStayWithinFlatVectorBounds() throws Exception { - assumeTrue("cuVS not supported", isSupported()); - - Random random = new Random(1234); - int segmentSize = 300; - int dimension = 32; - // Interspersed deletions on both segments, so the merge must drop a scattered subset of - // ordinals from each -- not just a contiguous prefix/suffix -- when it re-derives the merged - // vector set. - int deleteEveryNth = 4; - - Codec codec = TestUtil.alwaysKnnVectorsFormat(new Lucene99AcceleratedHNSWVectorsFormat()); - IndexWriterConfig config = - new IndexWriterConfig().setCodec(codec).setMergePolicy(NoMergePolicy.INSTANCE); - - int expectedLiveVectors; - try (Directory dir = newDirectory(); - IndexWriter writer = new IndexWriter(dir, config)) { - int deletedFromSegment1 = - addSegmentWithInterspersedDeletions( - writer, 0, segmentSize, dimension, deleteEveryNth, random); - writer.commit(); // segment 1, alone - int deletedFromSegment2 = - addSegmentWithInterspersedDeletions( - writer, segmentSize, segmentSize, dimension, deleteEveryNth, random); - writer.commit(); // segment 2, alone - - expectedLiveVectors = 2 * segmentSize - deletedFromSegment1 - deletedFromSegment2; - - // NoMergePolicy blocks forced merges too, so swap it out now that the two segments (each - // with their own interspersed deletions already committed) are set up. - writer.getConfig().setMergePolicy(new TieredMergePolicy()); - writer.forceMerge(1); - writer.commit(); - - try (DirectoryReader reader = DirectoryReader.open(dir)) { - assertEquals( - "expected the forced merge to produce a single segment", 1, reader.leaves().size()); - LeafReader leaf = reader.leaves().get(0).reader(); - - FloatVectorValues flatValues = leaf.getFloatVectorValues(FIELD); - assertEquals( - "merged flat vector count should equal (added - deleted)", - expectedLiveVectors, - flatValues.size()); - - HnswGraph graph = graphOf(leaf); - int level0NodeCount = graph.getNodesOnLevel(0).size(); - assertEquals( - "HNSW graph's level-0 node count disagrees with the merged flat vector file's actual" - + " count -- the graph and the flat file were derived independently by" - + " vectorBasedMerge and flatVectorsWriter.mergeOneField and disagree", - flatValues.size(), - level0NodeCount); - - assertAllNeighborOrdinalsInBounds(graph, flatValues.size()); - } - } - } - - /** - * Every neighbor referenced anywhere in the graph, at every level, must be a valid ordinal into - * the merged segment's actual flat vector data -- otherwise a reader resolving that neighbor's - * vector (e.g. mid-search, to score it) reads past the end of the flat file. - */ - private static void assertAllNeighborOrdinalsInBounds(HnswGraph graph, int liveVectorCount) - throws Exception { - for (int level = 0; level < graph.numLevels(); level++) { - HnswGraph.NodesIterator nodes = graph.getNodesOnLevel(level); - while (nodes.hasNext()) { - int node = nodes.nextInt(); - assertTrue( - "node " - + node - + " at level " - + level - + " is itself out of bounds (live vectors: " - + liveVectorCount - + ")", - node >= 0 && node < liveVectorCount); - graph.seek(level, node); - for (int neighbor = graph.nextNeighbor(); - neighbor != NO_MORE_DOCS; - neighbor = graph.nextNeighbor()) { - assertTrue( - "node " - + node - + " at level " - + level - + " has a neighbor ordinal " - + neighbor - + " out of bounds for the merged segment's " - + liveVectorCount - + " live vectors", - neighbor >= 0 && neighbor < liveVectorCount); - } - } - } - } - - private static HnswGraph graphOf(LeafReader leaf) throws Exception { - KnnVectorsReader knnReader = ((CodecReader) leaf).getVectorReader(); - if (knnReader instanceof PerFieldKnnVectorsFormat.FieldsReader fieldsReader) { - knnReader = fieldsReader.getFieldReader(FIELD); - } - return ((HnswGraphProvider) knnReader).getGraph(FIELD); - } - - /** - * Adds {@code count} documents (global ids {@code [startId, startId + count)}), then deletes - * every {@code deleteEveryNth}-th one by id, scattering the deletions across the segment rather - * than leaving a contiguous surviving range. - * - * @return the number of documents deleted from this segment - */ - private static int addSegmentWithInterspersedDeletions( - IndexWriter writer, int startId, int count, int dimension, int deleteEveryNth, Random random) - throws Exception { - float[][] dataset = generateDataset(random, count, dimension); - for (int i = 0; i < count; i++) { - Document document = new Document(); - document.add(new StringField(ID_FIELD, Integer.toString(startId + i), Field.Store.YES)); - document.add(new KnnFloatVectorField(FIELD, dataset[i], EUCLIDEAN)); - writer.addDocument(document); - } - int deleted = 0; - for (int i = 0; i < count; i += deleteEveryNth) { - writer.deleteDocuments(new Term(ID_FIELD, Integer.toString(startId + i))); - deleted++; - } - return deleted; - } -} diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestNativeFlatBufferingGuardRails.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestNativeFlatBufferingGuardRails.java deleted file mode 100644 index e116e9f02a..0000000000 --- a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestNativeFlatBufferingGuardRails.java +++ /dev/null @@ -1,269 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: Apache-2.0 - */ -package com.nvidia.cuvs.lucene; - -import static com.nvidia.cuvs.lucene.TestUtils.generateDataset; -import static com.nvidia.cuvs.lucene.ThreadLocalCuVSResourcesProvider.isSupported; -import static org.apache.lucene.index.VectorSimilarityFunction.EUCLIDEAN; - -import java.io.File; -import java.nio.file.Path; -import java.nio.file.Paths; -import java.util.Random; -import java.util.UUID; -import org.apache.commons.io.FileUtils; -import org.apache.lucene.codecs.Codec; -import org.apache.lucene.document.Document; -import org.apache.lucene.document.Field; -import org.apache.lucene.document.KnnFloatVectorField; -import org.apache.lucene.document.NumericDocValuesField; -import org.apache.lucene.document.StringField; -import org.apache.lucene.index.IndexWriter; -import org.apache.lucene.index.IndexWriterConfig; -import org.apache.lucene.index.NoMergePolicy; -import org.apache.lucene.index.SerialMergeScheduler; -import org.apache.lucene.index.TieredMergePolicy; -import org.apache.lucene.search.Sort; -import org.apache.lucene.search.SortField; -import org.apache.lucene.store.Directory; -import org.apache.lucene.store.FSDirectory; -import org.apache.lucene.tests.util.LuceneTestCase; -import org.apache.lucene.tests.util.LuceneTestCase.SuppressSysoutChecks; -import org.junit.After; -import org.junit.Before; -import org.junit.Test; - -/** - * Negative-path coverage for the guard rails in {@code NativeFlatBufferedHNSWVectorsWriter} around - * {@code numInputVectors} (native flat buffering): a count mismatch, an index-sorted segment, a - * merge attempt, and that a count mismatch on one field doesn't leave another field's native - * buffer unreleased on close. - * - *

Positive-path coverage (does a natively-buffered index actually search correctly, tolerate - * deletions, etc.) lives separately in {@link TestNativeFlatBufferingIndexAndSearch}. - */ -@SuppressSysoutChecks(bugUrl = "") -public class TestNativeFlatBufferingGuardRails extends LuceneTestCase { - - private static final String ID_FIELD = "id"; - private static final String VECTOR_FIELD = "vector_field"; - - private Random random; - private Path indexDirPath; - - @Before - public void beforeTest() throws Exception { - assumeTrue("cuVS not supported", isSupported()); - random = new Random(222); - indexDirPath = Paths.get(UUID.randomUUID().toString()); - } - - @After - public void afterTest() throws Exception { - if (indexDirPath == null) { - return; - } - File indexDirPathFile = indexDirPath.toFile(); - if (indexDirPathFile.exists() && indexDirPathFile.isDirectory()) { - FileUtils.deleteDirectory(indexDirPathFile); - } - } - - /** - * The count-mismatch guard exists for a caller-side bookkeeping error: declaring {@code - * numInputVectors} against the pre-filter document count instead of the number of vectors that - * actually reach {@code addValue} (e.g. an ingest-time filter skips some documents' vector - * field). It is unrelated to, and not triggered by, Lucene-level deletion -- see the javadoc on - * {@link TestNativeFlatBufferingIndexAndSearch#testDeletedDocsAfterNativeFlatBufferedFlush}. - */ - @Test - public void testCountMismatchFromIngestTimeFilterIsRejected() throws Exception { - int declaredNumInputVectors = 100; - int actuallyIndexed = declaredNumInputVectors - 1; // one doc "filtered out" before addValue - int dimension = 32; - float[][] dataset = generateDataset(random, actuallyIndexed, dimension); - - AcceleratedHNSWParams params = new AcceleratedHNSWParams.Builder().build(); - Codec codec = new Lucene101AcceleratedHNSWCodec(params, declaredNumInputVectors); - IndexWriterConfig config = - new IndexWriterConfig() - .setCodec(codec) - .setUseCompoundFile(false) - .setMaxBufferedDocs(declaredNumInputVectors + 1) - .setRAMBufferSizeMB(IndexWriterConfig.DISABLE_AUTO_FLUSH) - .setMergePolicy(NoMergePolicy.INSTANCE); - - try (Directory dir = FSDirectory.open(indexDirPath); - IndexWriter writer = new IndexWriter(dir, config)) { - for (int i = 0; i < actuallyIndexed; i++) { - Document document = new Document(); - document.add(new StringField(ID_FIELD, Integer.toString(i), Field.Store.YES)); - document.add(new KnnFloatVectorField(VECTOR_FIELD, dataset[i], EUCLIDEAN)); - writer.addDocument(document); - } - IllegalStateException thrown = expectThrows(IllegalStateException.class, writer::commit); - assertTrue( - "unexpected message: " + thrown.getMessage(), - thrown.getMessage().contains("numInputVectors")); - } - } - - /** - * When one field's native buffer fails the count-mismatch guard, {@code flush}'s per-field loop - * throws immediately -- a later field, even if it was populated correctly, is never reached by - * {@code writeFieldNative} and so never gets its native buffer released there. Confirms {@code - * NativeFlatBufferedHNSWVectorsWriter#close} still releases every field's buffer as a backstop - * regardless of what flush reached, by asserting close does not throw even though the second - * field's buffer was never touched during the failed flush. - */ - @Test - public void testCloseAfterCountMismatchReleasesEveryFieldsBuffer() throws Exception { - int declaredNumInputVectors = 50; - int dimension = 16; - // fieldA is added to documents first, so it is flush()'s first (and only-attempted) field; - // fieldB is added second and, unlike fieldA, is fully and correctly filled on every document. - String fieldA = "vector_field_a"; - String fieldB = "vector_field_b"; - float[][] datasetA = generateDataset(random, declaredNumInputVectors - 1, dimension); - float[][] datasetB = generateDataset(random, declaredNumInputVectors, dimension); - - AcceleratedHNSWParams params = new AcceleratedHNSWParams.Builder().build(); - Codec codec = new Lucene101AcceleratedHNSWCodec(params, declaredNumInputVectors); - IndexWriterConfig config = - new IndexWriterConfig() - .setCodec(codec) - .setUseCompoundFile(false) - .setMaxBufferedDocs(declaredNumInputVectors + 1) - .setRAMBufferSizeMB(IndexWriterConfig.DISABLE_AUTO_FLUSH) - .setMergePolicy(NoMergePolicy.INSTANCE); - - try (Directory dir = FSDirectory.open(indexDirPath); - IndexWriter writer = new IndexWriter(dir, config)) { - for (int i = 0; i < declaredNumInputVectors; i++) { - Document document = new Document(); - document.add(new StringField(ID_FIELD, Integer.toString(i), Field.Store.YES)); - // Skip fieldA on the last doc so it ends up one short of declaredNumInputVectors. - if (i < declaredNumInputVectors - 1) { - document.add(new KnnFloatVectorField(fieldA, datasetA[i], EUCLIDEAN)); - } - document.add(new KnnFloatVectorField(fieldB, datasetB[i], EUCLIDEAN)); - writer.addDocument(document); - } - IllegalStateException thrown = expectThrows(IllegalStateException.class, writer::commit); - assertTrue( - "unexpected message: " + thrown.getMessage(), - thrown.getMessage().contains("numInputVectors")); - // The try-with-resources close() below must not throw even though fieldB's fully-populated - // native buffer was never reached by the aborted flush() loop. - } - } - - /** Native flat buffering pre-sizes a single flush's buffer and cannot support sorted flushes. */ - @Test - public void testIndexSortedSegmentIsRejected() throws Exception { - int numDocs = 50; - int dimension = 32; - float[][] dataset = generateDataset(random, numDocs, dimension); - - AcceleratedHNSWParams params = new AcceleratedHNSWParams.Builder().build(); - Codec codec = new Lucene101AcceleratedHNSWCodec(params, numDocs); - Sort indexSort = new Sort(new SortField("sort_key", SortField.Type.LONG)); - IndexWriterConfig config = - new IndexWriterConfig() - .setCodec(codec) - .setUseCompoundFile(false) - .setIndexSort(indexSort) - .setMaxBufferedDocs(numDocs + 1) - .setRAMBufferSizeMB(IndexWriterConfig.DISABLE_AUTO_FLUSH) - .setMergePolicy(NoMergePolicy.INSTANCE); - - try (Directory dir = FSDirectory.open(indexDirPath); - IndexWriter writer = new IndexWriter(dir, config)) { - // KnnVectorsFormat#fieldsWriter (and so the writer's index-sort check in its constructor) - // is invoked on the first addDocument() for the segment, not at commit() -- so the guard - // must be expected around the whole indexing loop, not just the flush. - IllegalArgumentException thrown = - expectThrows( - IllegalArgumentException.class, - () -> { - for (int i = 0; i < numDocs; i++) { - Document document = new Document(); - document.add(new StringField(ID_FIELD, Integer.toString(i), Field.Store.YES)); - document.add(new NumericDocValuesField("sort_key", numDocs - i)); - document.add(new KnnFloatVectorField(VECTOR_FIELD, dataset[i], EUCLIDEAN)); - writer.addDocument(document); - } - }); - assertTrue( - "unexpected message: " + thrown.getMessage(), - thrown.getMessage().contains("index-sorted")); - } - } - - /** - * Native flat buffering supports only the unsorted single-segment flush path; merging two - * natively-buffered segments must be rejected rather than silently mis-sizing the native buffer. - */ - @Test - public void testMergeIsRejected() throws Exception { - int segmentSize = 40; - int dimension = 32; - - AcceleratedHNSWParams params = new AcceleratedHNSWParams.Builder().build(); - Codec codec = new Lucene101AcceleratedHNSWCodec(params, segmentSize); - IndexWriterConfig config = - new IndexWriterConfig() - .setCodec(codec) - .setUseCompoundFile(false) - .setMaxBufferedDocs(segmentSize + 1) - .setRAMBufferSizeMB(IndexWriterConfig.DISABLE_AUTO_FLUSH) - // Keep the two flushes below as separate segments; NoMergePolicy also blocks forced - // merges, so it is swapped out before forceMerge() is called. - .setMergePolicy(NoMergePolicy.INSTANCE) - // Force merges to run synchronously on the calling thread, so the guard's exception - // (or whatever IndexWriter/SegmentMerger wraps it as) surfaces directly from - // forceMerge() instead of on a background merge thread. - .setMergeScheduler(new SerialMergeScheduler()); - - try (Directory dir = FSDirectory.open(indexDirPath); - IndexWriter writer = new IndexWriter(dir, config)) { - addSegment(writer, 0, segmentSize, dimension); - writer.commit(); // segment 1: exactly segmentSize vectors, matching numInputVectors - addSegment(writer, segmentSize, segmentSize, dimension); - writer.commit(); // segment 2: exactly segmentSize vectors, matching numInputVectors - - writer.getConfig().setMergePolicy(new TieredMergePolicy()); - Throwable thrown = expectThrows(Throwable.class, () -> writer.forceMerge(1)); - assertTrue( - "expected UnsupportedOperationException somewhere in the cause chain of: " + thrown, - causedBy(thrown, UnsupportedOperationException.class)); - } - } - - private void addSegment(IndexWriter writer, int startId, int count, int dimension) - throws Exception { - float[][] dataset = generateDataset(random, count, dimension); - for (int i = 0; i < count; i++) { - Document document = new Document(); - document.add(new StringField(ID_FIELD, Integer.toString(startId + i), Field.Store.YES)); - document.add(new KnnFloatVectorField(VECTOR_FIELD, dataset[i], EUCLIDEAN)); - writer.addDocument(document); - } - } - - private static boolean causedBy(Throwable t, Class type) { - for (Throwable cur = t; cur != null; cur = cur.getCause()) { - if (type.isInstance(cur)) { - return true; - } - for (Throwable suppressed : cur.getSuppressed()) { - if (causedBy(suppressed, type)) { - return true; - } - } - } - return false; - } -} diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestNativeFlatBufferingIndexAndSearch.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestNativeFlatBufferingIndexAndSearch.java deleted file mode 100644 index 0ba476c652..0000000000 --- a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestNativeFlatBufferingIndexAndSearch.java +++ /dev/null @@ -1,220 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: Apache-2.0 - */ -package com.nvidia.cuvs.lucene; - -import static com.nvidia.cuvs.lucene.TestUtils.generateDataset; -import static com.nvidia.cuvs.lucene.ThreadLocalCuVSResourcesProvider.isSupported; -import static org.apache.lucene.index.VectorSimilarityFunction.EUCLIDEAN; - -import com.nvidia.cuvs.CagraIndexParams.CagraGraphBuildAlgo; -import java.io.File; -import java.nio.file.Path; -import java.nio.file.Paths; -import java.util.Random; -import java.util.UUID; -import org.apache.commons.io.FileUtils; -import org.apache.lucene.codecs.Codec; -import org.apache.lucene.document.Document; -import org.apache.lucene.document.Field; -import org.apache.lucene.document.KnnFloatVectorField; -import org.apache.lucene.document.StringField; -import org.apache.lucene.index.DirectoryReader; -import org.apache.lucene.index.IndexWriter; -import org.apache.lucene.index.IndexWriterConfig; -import org.apache.lucene.index.NoMergePolicy; -import org.apache.lucene.index.Term; -import org.apache.lucene.search.IndexSearcher; -import org.apache.lucene.search.KnnFloatVectorQuery; -import org.apache.lucene.search.TopDocs; -import org.apache.lucene.store.Directory; -import org.apache.lucene.store.FSDirectory; -import org.apache.lucene.tests.util.LuceneTestCase; -import org.apache.lucene.tests.util.LuceneTestCase.SuppressSysoutChecks; -import org.junit.After; -import org.junit.Before; -import org.junit.Test; - -/** - * Positive-path functional coverage for native flat buffering ({@code - * AcceleratedHNSWParams.numInputVectors}) beyond {@code TestNativeFlatVectorsWriterRoundTrip}, - * which only checks that the flat {@code .vec} file round-trips -- not that the resulting index is - * actually searchable, tolerates deletions, or composes correctly with the odd-graph-degree fix. - * - *

The negative/guard-rail paths (count mismatch, index-sorted segments, merges) are covered - * separately in {@link TestNativeFlatBufferingGuardRails}. - */ -@SuppressSysoutChecks(bugUrl = "") -public class TestNativeFlatBufferingIndexAndSearch extends LuceneTestCase { - - private static final String ID_FIELD = "id"; - private static final String VECTOR_FIELD = "vector_field"; - - private Random random; - private Path indexDirPath; - - @Before - public void beforeTest() throws Exception { - assumeTrue("cuVS not supported", isSupported()); - random = new Random(222); - indexDirPath = Paths.get(UUID.randomUUID().toString()); - } - - @After - public void afterTest() throws Exception { - if (indexDirPath == null) { - return; - } - File indexDirPathFile = indexDirPath.toFile(); - if (indexDirPathFile.exists() && indexDirPathFile.isDirectory()) { - FileUtils.deleteDirectory(indexDirPathFile); - } - } - - /** A natively-buffered index must still be searchable through the normal Lucene KNN query API. */ - @Test - public void testIndexAndSearch() throws Exception { - int numDocs = 500; - int dimension = 32; - int topK = 10; - float[][] dataset = generateDataset(random, numDocs, dimension); - - buildNativeFlatBufferedIndex(numDocs, dimension, dataset); - - try (Directory dir = FSDirectory.open(indexDirPath); - DirectoryReader reader = DirectoryReader.open(dir)) { - assertEquals("expected a single native-flat-buffered segment", 1, reader.leaves().size()); - - IndexSearcher searcher = new IndexSearcher(reader); - float[] queryVector = generateDataset(random, 1, dimension)[0]; - TopDocs results = - searcher.search(new KnnFloatVectorQuery(VECTOR_FIELD, queryVector, topK), topK); - - assertEquals("expected topK results", topK, results.scoreDocs.length); - for (var scoreDoc : results.scoreDocs) { - String id = searcher.storedFields().document(scoreDoc.doc).get(ID_FIELD); - int idValue = Integer.parseInt(id); - assertTrue("returned id out of range: " + id, idValue >= 0 && idValue < numDocs); - } - } - } - - /** - * Deletion is orthogonal to native flat buffering: {@code IndexWriter.deleteDocuments} only - * updates Lucene's liveDocs bitset at search time -- it never touches {@code NativeFieldWriter} - * or the native host matrix, so it cannot trip (and isn't meant to be caught by) the count-mismatch - * guard rail tested in {@link TestNativeFlatBufferingGuardRails}. This test instead confirms that - * deletions applied after a natively-buffered flush are still honored correctly at search time. - */ - @Test - public void testDeletedDocsAfterNativeFlatBufferedFlush() throws Exception { - int numDocs = 300; - int dimension = 32; - float[][] dataset = generateDataset(random, numDocs, dimension); - - AcceleratedHNSWParams params = new AcceleratedHNSWParams.Builder().build(); - Codec codec = new Lucene101AcceleratedHNSWCodec(params, numDocs); - IndexWriterConfig config = - new IndexWriterConfig() - .setCodec(codec) - .setUseCompoundFile(false) - .setMaxBufferedDocs(numDocs + 1) - .setRAMBufferSizeMB(IndexWriterConfig.DISABLE_AUTO_FLUSH) - .setMergePolicy(NoMergePolicy.INSTANCE); - - try (Directory dir = FSDirectory.open(indexDirPath); - IndexWriter writer = new IndexWriter(dir, config)) { - for (int i = 0; i < numDocs; i++) { - Document document = new Document(); - document.add(new StringField(ID_FIELD, Integer.toString(i), Field.Store.YES)); - document.add(new KnnFloatVectorField(VECTOR_FIELD, dataset[i], EUCLIDEAN)); - writer.addDocument(document); - } - writer.commit(); // single natively-buffered flush: NativeFieldWriter's count matches numDocs - - // Delete every 3rd doc. No new vectors are added, so this does not trigger another flush of - // the vector field and cannot interact with the numInputVectors hint. - for (int i = 0; i < numDocs; i += 3) { - writer.deleteDocuments(new Term(ID_FIELD, Integer.toString(i))); - } - writer.commit(); - } - - try (Directory dir = FSDirectory.open(indexDirPath); - DirectoryReader reader = DirectoryReader.open(dir)) { - assertEquals( - "expected the deletions to land in the same single segment", 1, reader.leaves().size()); - assertTrue("expected some deleted docs", reader.numDeletedDocs() > 0); - - IndexSearcher searcher = new IndexSearcher(reader); - TopDocs results = - searcher.search(new KnnFloatVectorQuery(VECTOR_FIELD, dataset[0], numDocs), numDocs); - for (var scoreDoc : results.scoreDocs) { - String id = searcher.storedFields().document(scoreDoc.doc).get(ID_FIELD); - assertNotEquals( - "deleted doc id=" + id + " was still returned by search", 0, Integer.parseInt(id) % 3); - } - } - } - - /** - * The M = ceil(cagraGraphDegree / 2) fix ({@link TestAcceleratedHNSWOddGraphDegree}) must also - * hold on the native-flat-buffered write path ({@code writeFieldNative}), which is a distinct - * call path from the heap-buffered one that test exercises. - */ - @Test - public void testOddGraphDegreeWithNativeFlatBuffering() throws Exception { - int numDocs = 200; - int dimension = 32; - int oddGraphDegree = 63; - float[][] dataset = generateDataset(random, numDocs, dimension); - - AcceleratedHNSWParams params = - new AcceleratedHNSWParams.Builder() - .withStrategy(AcceleratedHNSWParams.Strategy.CUSTOM) - .withCagraGraphBuildAlgo(CagraGraphBuildAlgo.NN_DESCENT) - .withIntermediateGraphDegree(128) - .withGraphDegree(oddGraphDegree) - .build(); - - buildNativeFlatBufferedIndex(numDocs, dimension, dataset, params); - - try (Directory dir = FSDirectory.open(indexDirPath); - DirectoryReader reader = DirectoryReader.open(dir)) { - IndexSearcher searcher = new IndexSearcher(reader); - TopDocs results = searcher.search(new KnnFloatVectorQuery(VECTOR_FIELD, dataset[0], 5), 5); - assertEquals(5, results.scoreDocs.length); - } - } - - private void buildNativeFlatBufferedIndex(int numDocs, int dimension, float[][] dataset) - throws Exception { - buildNativeFlatBufferedIndex( - numDocs, dimension, dataset, new AcceleratedHNSWParams.Builder().build()); - } - - private void buildNativeFlatBufferedIndex( - int numDocs, int dimension, float[][] dataset, AcceleratedHNSWParams params) - throws Exception { - Codec codec = new Lucene101AcceleratedHNSWCodec(params, numDocs); - IndexWriterConfig config = - new IndexWriterConfig() - .setCodec(codec) - .setUseCompoundFile(false) - .setMaxBufferedDocs(numDocs + 1) - .setRAMBufferSizeMB(IndexWriterConfig.DISABLE_AUTO_FLUSH) - .setMergePolicy(NoMergePolicy.INSTANCE); - - try (Directory dir = FSDirectory.open(indexDirPath); - IndexWriter writer = new IndexWriter(dir, config)) { - for (int i = 0; i < numDocs; i++) { - Document document = new Document(); - document.add(new StringField(ID_FIELD, Integer.toString(i), Field.Store.YES)); - document.add(new KnnFloatVectorField(VECTOR_FIELD, dataset[i], EUCLIDEAN)); - writer.addDocument(document); - } - writer.commit(); - } - } -} diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestNativeFlatVectorsWriterRoundTrip.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestNativeFlatVectorsWriterRoundTrip.java deleted file mode 100644 index 027a086f88..0000000000 --- a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestNativeFlatVectorsWriterRoundTrip.java +++ /dev/null @@ -1,279 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: Apache-2.0 - */ -package com.nvidia.cuvs.lucene; - -import static com.nvidia.cuvs.lucene.TestUtils.generateDataset; -import static com.nvidia.cuvs.lucene.ThreadLocalCuVSResourcesProvider.isSupported; -import static org.apache.lucene.index.VectorSimilarityFunction.EUCLIDEAN; - -import java.io.File; -import java.io.IOException; -import java.nio.file.Path; -import java.nio.file.Paths; -import java.util.ArrayList; -import java.util.Collections; -import java.util.HashMap; -import java.util.HashSet; -import java.util.List; -import java.util.Map; -import java.util.Random; -import java.util.Set; -import java.util.UUID; -import java.util.function.IntFunction; -import org.apache.commons.io.FileUtils; -import org.apache.lucene.codecs.Codec; -import org.apache.lucene.document.Document; -import org.apache.lucene.document.Field; -import org.apache.lucene.document.KnnFloatVectorField; -import org.apache.lucene.document.StringField; -import org.apache.lucene.index.DirectoryReader; -import org.apache.lucene.index.FloatVectorValues; -import org.apache.lucene.index.IndexWriter; -import org.apache.lucene.index.IndexWriterConfig; -import org.apache.lucene.index.KnnVectorValues; -import org.apache.lucene.index.LeafReader; -import org.apache.lucene.index.NoMergePolicy; -import org.apache.lucene.search.DocIdSetIterator; -import org.apache.lucene.store.Directory; -import org.apache.lucene.store.FSDirectory; -import org.apache.lucene.tests.util.LuceneTestCase; -import org.apache.lucene.tests.util.LuceneTestCase.SuppressSysoutChecks; -import org.junit.After; -import org.junit.Before; -import org.junit.Test; - -/** - * Positive round-trip check for {@link NativeFlatVectorsWriter}: builds a single-segment index - * with {@code AcceleratedHNSWParams.numInputVectors} set (native flat buffering) and reads the - * {@code .vec}/{@code .vemf} files back through the stock {@code Lucene99FlatVectorsReader}, - * asserting every vector round-trips byte-exact. - * - *

This is the check called for by {@link NativeFlatVectorsWriter}'s "on a Lucene upgrade" class - * javadoc: it confirms the hand-transcribed format is still readable by Lucene's real reader, and - * should pass on every {@code lucene-core} version bump. - */ -@SuppressSysoutChecks(bugUrl = "") -public class TestNativeFlatVectorsWriterRoundTrip extends LuceneTestCase { - - private static final String ID_FIELD = "id"; - private static final String VECTOR_FIELD = "vector_field"; - - private Random random; - private Path indexDirPath; - - @Before - public void beforeTest() throws Exception { - assumeTrue("cuVS not supported", isSupported()); - random = new Random(222); - indexDirPath = Paths.get(UUID.randomUUID().toString()); - } - - @After - public void afterTest() throws Exception { - if (indexDirPath == null) { - return; - } - File indexDirPathFile = indexDirPath.toFile(); - if (indexDirPathFile.exists() && indexDirPathFile.isDirectory()) { - FileUtils.deleteDirectory(indexDirPathFile); - } - } - - @Test - public void vectorsRoundTripThroughStockLucene99FlatVectorsReader() throws Exception { - int numDocs = 500; - int dimension = 32; - float[][] dataset = generateDataset(random, numDocs, dimension); - - AcceleratedHNSWParams params = new AcceleratedHNSWParams.Builder().build(); - Codec codec = new Lucene101AcceleratedHNSWCodec(params, numDocs); - - // Force everything into a single unsorted, unmerged flush: native flat buffering requires - // numInputVectors to equal the exact number of vectors landing in that one flush. - IndexWriterConfig config = - new IndexWriterConfig() - .setCodec(codec) - .setUseCompoundFile(false) - .setMaxBufferedDocs(numDocs + 1) - .setRAMBufferSizeMB(IndexWriterConfig.DISABLE_AUTO_FLUSH) - .setMergePolicy(NoMergePolicy.INSTANCE); - - try (Directory dir = FSDirectory.open(indexDirPath); - IndexWriter writer = new IndexWriter(dir, config)) { - for (int i = 0; i < numDocs; i++) { - Document document = new Document(); - document.add(new StringField(ID_FIELD, Integer.toString(i), Field.Store.YES)); - document.add(new KnnFloatVectorField(VECTOR_FIELD, dataset[i], EUCLIDEAN)); - writer.addDocument(document); - } - writer.commit(); - } - - try (Directory dir = FSDirectory.open(indexDirPath); - DirectoryReader reader = DirectoryReader.open(dir)) { - assertEquals("expected a single native-flat-buffered segment", 1, reader.leaves().size()); - LeafReader leafReader = reader.leaves().get(0).reader(); - assertFieldRoundTrips(leafReader, VECTOR_FIELD, dimension, numDocs, docId -> dataset[docId]); - } - } - - /** - * Cardinality (8000) exceeds {@code IndexedDISI.MAX_ARRAY_LENGTH} (4095), so {@code - * OrdToDocDISIReaderConfiguration} picks the DENSE (bitset) {@code docsWithField} encoding for - * this field. - */ - @Test - public void vectorsRoundTripWithDenseDocsWithFieldEncoding() throws Exception { - roundTripPartialField(10000, 8000, 32); - } - - /** - * Cardinality (800) stays under {@code IndexedDISI.MAX_ARRAY_LENGTH} (4095), so {@code - * OrdToDocDISIReaderConfiguration} picks the SPARSE (array) {@code docsWithField} encoding - * instead. - */ - @Test - public void vectorsRoundTripWithSparseDocsWithFieldEncoding() throws Exception { - roundTripPartialField(1000, 800, 32); - } - - /** - * Two vector fields written to the same segment, exercising the boundary between consecutive - * per-field records in {@code .vemf}: the second field's header must be found where the first - * field's record actually ends. - */ - @Test - public void vectorsRoundTripAcrossMultipleFieldsInSameSegment() throws Exception { - int numDocs = 500; - int dimensionA = 32; - int dimensionB = 16; - String fieldB = "vector_field_two"; - float[][] datasetA = generateDataset(random, numDocs, dimensionA); - float[][] datasetB = generateDataset(random, numDocs, dimensionB); - - AcceleratedHNSWParams params = new AcceleratedHNSWParams.Builder().build(); - Codec codec = new Lucene101AcceleratedHNSWCodec(params, numDocs); - - IndexWriterConfig config = - new IndexWriterConfig() - .setCodec(codec) - .setUseCompoundFile(false) - .setMaxBufferedDocs(numDocs + 1) - .setRAMBufferSizeMB(IndexWriterConfig.DISABLE_AUTO_FLUSH) - .setMergePolicy(NoMergePolicy.INSTANCE); - - try (Directory dir = FSDirectory.open(indexDirPath); - IndexWriter writer = new IndexWriter(dir, config)) { - for (int i = 0; i < numDocs; i++) { - Document document = new Document(); - document.add(new StringField(ID_FIELD, Integer.toString(i), Field.Store.YES)); - document.add(new KnnFloatVectorField(VECTOR_FIELD, datasetA[i], EUCLIDEAN)); - document.add(new KnnFloatVectorField(fieldB, datasetB[i], EUCLIDEAN)); - writer.addDocument(document); - } - writer.commit(); - } - - try (Directory dir = FSDirectory.open(indexDirPath); - DirectoryReader reader = DirectoryReader.open(dir)) { - assertEquals("expected a single native-flat-buffered segment", 1, reader.leaves().size()); - LeafReader leafReader = reader.leaves().get(0).reader(); - assertFieldRoundTrips( - leafReader, VECTOR_FIELD, dimensionA, numDocs, docId -> datasetA[docId]); - assertFieldRoundTrips(leafReader, fieldB, dimensionB, numDocs, docId -> datasetB[docId]); - } - } - - /** - * Builds a segment of {@code numDocs} documents where only a random {@code numDocsWithVector} - * of them carry {@code VECTOR_FIELD}, then verifies the round trip. - */ - private void roundTripPartialField(int numDocs, int numDocsWithVector, int dimension) - throws Exception { - Set docsWithVector = randomDocSubset(numDocs, numDocsWithVector); - float[][] dataset = generateDataset(random, numDocsWithVector, dimension); - - AcceleratedHNSWParams params = new AcceleratedHNSWParams.Builder().build(); - Codec codec = new Lucene101AcceleratedHNSWCodec(params, numDocsWithVector); - - IndexWriterConfig config = - new IndexWriterConfig() - .setCodec(codec) - .setUseCompoundFile(false) - .setMaxBufferedDocs(numDocs + 1) - .setRAMBufferSizeMB(IndexWriterConfig.DISABLE_AUTO_FLUSH) - .setMergePolicy(NoMergePolicy.INSTANCE); - - Map docIdToVectorIndex = new HashMap<>(); - try (Directory dir = FSDirectory.open(indexDirPath); - IndexWriter writer = new IndexWriter(dir, config)) { - int vectorIndex = 0; - for (int i = 0; i < numDocs; i++) { - Document document = new Document(); - document.add(new StringField(ID_FIELD, Integer.toString(i), Field.Store.YES)); - if (docsWithVector.contains(i)) { - document.add(new KnnFloatVectorField(VECTOR_FIELD, dataset[vectorIndex], EUCLIDEAN)); - docIdToVectorIndex.put(i, vectorIndex); - vectorIndex++; - } - writer.addDocument(document); - } - writer.commit(); - } - - try (Directory dir = FSDirectory.open(indexDirPath); - DirectoryReader reader = DirectoryReader.open(dir)) { - assertEquals("expected a single native-flat-buffered segment", 1, reader.leaves().size()); - LeafReader leafReader = reader.leaves().get(0).reader(); - assertFieldRoundTrips( - leafReader, - VECTOR_FIELD, - dimension, - numDocsWithVector, - docId -> dataset[docIdToVectorIndex.get(docId)]); - } - } - - /** Picks {@code count} distinct doc ids out of {@code [0, numDocs)}. */ - private Set randomDocSubset(int numDocs, int count) { - List allDocs = new ArrayList<>(numDocs); - for (int i = 0; i < numDocs; i++) { - allDocs.add(i); - } - Collections.shuffle(allDocs, random); - return new HashSet<>(allDocs.subList(0, count)); - } - - private void assertFieldRoundTrips( - LeafReader leafReader, - String fieldName, - int dimension, - int expectedCount, - IntFunction expectedVectorForDocId) - throws IOException { - FloatVectorValues values = leafReader.getFloatVectorValues(fieldName); - assertNotNull(values); - assertEquals(expectedCount, values.size()); - assertEquals(dimension, values.dimension()); - - int seen = 0; - KnnVectorValues.DocIndexIterator it = values.iterator(); - for (int doc = it.nextDoc(); doc != DocIdSetIterator.NO_MORE_DOCS; doc = it.nextDoc()) { - String id = leafReader.storedFields().document(doc).get(ID_FIELD); - float[] roundTripped = values.vectorValue(it.index()); - assertArrayEquals( - "vector for field=" - + fieldName - + " id=" - + id - + " did not round-trip byte-exact through the stock Lucene99FlatVectorsReader", - expectedVectorForDocId.apply(Integer.parseInt(id)), - roundTripped, - 0f); - seen++; - } - assertEquals("did not visit every vector for field=" + fieldName, expectedCount, seen); - } -} diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestUtils.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestUtils.java index c96c17f493..22cd509e7b 100644 --- a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestUtils.java +++ b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestUtils.java @@ -10,11 +10,6 @@ import static org.junit.Assert.assertTrue; import java.io.IOException; -import java.nio.ByteBuffer; -import java.nio.ByteOrder; -import java.nio.channels.FileChannel; -import java.nio.file.Path; -import java.nio.file.StandardOpenOption; import java.util.HashSet; import java.util.Random; import java.util.Set; @@ -68,33 +63,6 @@ public static void assertVectorsKeepTheirDocuments( assertEquals("not every document was found", expectedById.length, seen.size()); } - /** Writes {@code dataset} as an uncompressed {@code .fbin} file (little-endian float32 rows). */ - public static void writeFbin(Path path, float[][] dataset) throws IOException { - int numVectors = dataset.length; - int dimension = numVectors == 0 ? 0 : dataset[0].length; - ByteBuffer buf = - ByteBuffer.allocate(8 + numVectors * dimension * Float.BYTES) - .order(ByteOrder.LITTLE_ENDIAN); - buf.putInt(numVectors); - buf.putInt(dimension); - for (float[] vector : dataset) { - for (float v : vector) { - buf.putFloat(v); - } - } - buf.flip(); - try (FileChannel ch = - FileChannel.open( - path, - StandardOpenOption.CREATE, - StandardOpenOption.WRITE, - StandardOpenOption.TRUNCATE_EXISTING)) { - while (buf.hasRemaining()) { - ch.write(buf); - } - } - } - public static float[][] generateDataset(Random random, int size, int dimensions) { float[][] dataset = new float[size][dimensions]; for (int i = 0; i < size; i++) { diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestWriterThreadsGraphEquivalence.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestWriterThreadsGraphEquivalence.java deleted file mode 100644 index d95c60ae6b..0000000000 --- a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestWriterThreadsGraphEquivalence.java +++ /dev/null @@ -1,156 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: Apache-2.0 - */ -package com.nvidia.cuvs.lucene; - -import static com.nvidia.cuvs.lucene.ThreadLocalCuVSResourcesProvider.isSupported; -import static org.apache.lucene.search.DocIdSetIterator.NO_MORE_DOCS; - -import com.nvidia.cuvs.CuVSMatrix; -import java.io.IOException; -import java.util.ArrayList; -import java.util.Arrays; -import java.util.List; -import java.util.Random; -import org.apache.lucene.store.ByteBuffersDirectory; -import org.apache.lucene.store.Directory; -import org.apache.lucene.store.IOContext; -import org.apache.lucene.store.IndexInput; -import org.apache.lucene.store.IndexOutput; -import org.apache.lucene.tests.util.LuceneTestCase; -import org.apache.lucene.tests.util.LuceneTestCase.SuppressSysoutChecks; -import org.apache.lucene.util.hnsw.HnswGraph; -import org.apache.lucene.util.hnsw.HnswGraph.NodesIterator; -import org.junit.Test; - -/** - * Verifies that {@code writerThreads > 1} produces the same result as the serial path for the - * two parallelizations gated on {@code AcceleratedHNSWUtils}/{@code GPUBuiltHnswGraph}'s {@code - * PARALLEL_MIN_NODES} threshold: materializing the CAGRA adjacency into {@code NeighborArray}s - * (the {@code GPUBuiltHnswGraph} constructor), and encoding level 0 to disk ({@code - * AcceleratedHNSWUtils#writeGraph}). - * - *

Both tests use a synthetic adjacency ({@link CuVSMatrix#ofArray(int[][])}, the same host-matrix - * construction the higher-layer subset builder already uses) rather than a real CAGRA build, so - * that the comparison isolates these two parallelizations from CAGRA's own build-to-build - * variance -- a real GPU build is not guaranteed to produce the identical graph twice even with the - * same input and thread count, which would make an end-to-end build comparison unreliable for this - * purpose. This needs cuVS/GPU only to allocate the host matrix, not to run a build. - */ -@SuppressSysoutChecks(bugUrl = "") -public class TestWriterThreadsGraphEquivalence extends LuceneTestCase { - - // Must be >= PARALLEL_MIN_NODES (1 << 16) in both AcceleratedHNSWUtils and - // GPUBuiltHnswGraph, or the "parallel" runs below silently fall through to the serial branch and - // the test would pass without exercising anything. - private static final int NUM_NODES = (1 << 16) + 1000; - private static final int DEGREE = 12; - private static final int NUM_THREADS = 4; - - @Test - public void fillNeighborArrayParallelMatchesSerial() throws Exception { - assumeTrue("cuVS not supported", isSupported()); - int[][] adjacency = randomAdjacency(NUM_NODES, DEGREE, new Random(1)); - - try (CuVSMatrix matrix = CuVSMatrix.ofArray(adjacency)) { - GPUBuiltHnswGraph serial = newSingleLayerGraph(matrix, 1); - GPUBuiltHnswGraph parallel = newSingleLayerGraph(matrix, NUM_THREADS); - assertGraphsEqual(serial, parallel); - } - } - - @Test - public void writeGraphParallelMatchesSerial() throws Exception { - assumeTrue("cuVS not supported", isSupported()); - int[][] adjacency = randomAdjacency(NUM_NODES, DEGREE, new Random(2)); - - try (CuVSMatrix matrix = CuVSMatrix.ofArray(adjacency); - Directory dir = new ByteBuffersDirectory()) { - // Materialize once, serially, so any difference found below is attributable only to - // writeGraph's own parallelization, not to fillNeighborArray's. - GPUBuiltHnswGraph graph = newSingleLayerGraph(matrix, 1); - - int[][] serialOffsets; - try (IndexOutput out = dir.createOutput("serial", IOContext.DEFAULT)) { - serialOffsets = AcceleratedHNSWUtils.writeGraph(graph, out, 1); - } - int[][] parallelOffsets; - try (IndexOutput out = dir.createOutput("parallel", IOContext.DEFAULT)) { - parallelOffsets = AcceleratedHNSWUtils.writeGraph(graph, out, NUM_THREADS); - } - - assertEquals(serialOffsets.length, parallelOffsets.length); - for (int level = 0; level < serialOffsets.length; level++) { - assertArrayEquals( - "per-node byte-length offsets differ for level " + level, - serialOffsets[level], - parallelOffsets[level]); - } - - assertArrayEquals( - "writeGraph's parallel level-0 encoding produced different bytes than the serial path", - readAllBytes(dir, "serial"), - readAllBytes(dir, "parallel")); - } - } - - private static GPUBuiltHnswGraph newSingleLayerGraph(CuVSMatrix layer0Adjacency, int numThreads) - throws IOException { - // A single layer (layer 0 only): the constructor never consults layerNodes in that case, so - // the placeholder null entry mirrors the convention used elsewhere for "layer 0 needs no node - // list" without actually being read. - return new GPUBuiltHnswGraph( - NUM_NODES, - /* dimensions= */ 4, - Arrays.asList((int[]) null), - List.of(layer0Adjacency), - numThreads); - } - - /** Every node/level's in-order arc list must match exactly between the two graphs. */ - private static void assertGraphsEqual(HnswGraph a, HnswGraph b) throws Exception { - assertEquals(a.numLevels(), b.numLevels()); - for (int level = 0; level < a.numLevels(); level++) { - int[] nodes = NodesIterator.getSortedNodes(a.getNodesOnLevel(level)); - for (int node : nodes) { - assertArrayEquals( - "node " + node + " at level " + level + " has different neighbors", - arcsOf(a, level, node), - arcsOf(b, level, node)); - } - } - } - - private static int[] arcsOf(HnswGraph graph, int level, int node) throws Exception { - graph.seek(level, node); - List arcs = new ArrayList<>(); - for (int n = graph.nextNeighbor(); n != NO_MORE_DOCS; n = graph.nextNeighbor()) { - arcs.add(n); - } - return arcs.stream().mapToInt(Integer::intValue).toArray(); - } - - private static byte[] readAllBytes(Directory dir, String name) throws Exception { - try (IndexInput in = dir.openInput(name, IOContext.DEFAULT)) { - byte[] bytes = new byte[(int) in.length()]; - in.readBytes(bytes, 0, bytes.length); - return bytes; - } - } - - /** - * A deterministic, seeded pseudo-adjacency. It doesn't need to be a real CAGRA graph -- only a - * realistic shape (fixed degree, valid node ids) -- since {@code GPUBuiltHnswGraph} and {@code - * AcceleratedHNSWUtils#writeGraph} don't interpret the neighbor ids semantically. - */ - private static int[][] randomAdjacency(int numNodes, int degree, Random random) { - int[][] adjacency = new int[numNodes][degree]; - for (int[] row : adjacency) { - for (int j = 0; j < degree; j++) { - row[j] = random.nextInt(numNodes); - } - } - return adjacency; - } -} From 96be825839e6627c19e365b1855948fee973292d Mon Sep 17 00:00:00 2001 From: James Lamb Date: Fri, 11 Sep 2026 21:36:33 -0500 Subject: [PATCH 5/7] Containerize standalone tarball build (#2595) Replaces https://github.com/NVIDIA/cuvs/pull/2443 Containerizes the process of building the standalone C library tarballs The starting point for this PR is @cjnolet 's work in #2443, where he wrote this description: > _This PR is based on @msarahan's original POC, with the logic of the ci sript moved into build.sh and invoked through the CI script. The readme is also moved into the build and install guide in the docs._ Changes relative to that PR: * threads `PARALLEL_LEVEL` throught so builds are parallelized * enables the build cluster (`sccache-dist`) in CI * simplifies configuration flow (e.g. reduced duplication of default values, removal of unnecessary flexibility) * makes flow of AWS creds for `sccache` more secure * removes unnecessary configuration * removes unnecessary package installs, consolidates package installs * removes an unnecessary `git clone` of https://github.com/rapidsai/spdx-license-builder * enables `sccache` in CI and with a pattern that could work locally (will share details privately) * updates GitHub Actions third-party actions to their latest versions ## Notes for Reviewers ### How I tested this Locally tried each of the new commands added to `build.md`, with `sccache` enabled. Tried just `build.sh` without `sccache` enabled (that takes a lot longer to run).

code for flipping between those (click me) ```shell # enabling sccache export CI=true # (private steps setting up 'sccache' profile) AWS_ACCESS_KEY_ID=$( aws configure get aws_access_key_id \ --profile sccache ) AWS_SECRET_ACCESS_KEY=$( aws configure get aws_secret_access_key \ --profile sccache ) AWS_SESSION_TOKEN=$( aws configure get aws_session_token \ --profile sccache ) export AWS_ACCESS_KEY_ID AWS_SECRET_ACCESS_KEY AWS_SESSION_TOKEN # testing without 'sccache' unset AWS_ACCESS_KEY_ID AWS_SECRET_ACCESS_KEY AWS_SESSION_TOKEN CI ```
```shell clean() { sudo rm -rf ./{build,c/build/,cpp/build,libcuvs_c.tar.gz} } # all defaults clean ./build.sh tarball # customizing base image clean CUVS_TARBALL_CUDA_VERSION=12.9.2 \ CUVS_TARBALL_PYTHON_VERSION=3.11 \ ./build.sh tarball # customizing output directory, building tests clean CUVS_TARBALL_BUILD_OUTPUT_DIR="${PWD}/dist" ./build.sh tarball --tarball-build-tests # manual run (no build.sh) clean docker build \ -f Dockerfile.standalone \ --build-arg CUDA_VERSION="13.3.0" \ --build-arg PYTHON_VERSION="3.14" \ --build-arg RAPIDS_VERSION="$(head -1 ./VERSION | cut -d. -f1,2 )" \ -t cuvs-standalone-c:local \ . mkdir -p "${PWD}/dist" docker run --rm \ -v "${PWD}:/workspace" \ -v "${PWD}/dist:/build" \ cuvs-standalone-c:local --tarball-build-tests ``` Saw high cache hit rates from `sccache` and everything working as expected. Also looked at CI logs and saw that fully-cached jobs take around 30 minutes, pretty similar to the timings in CI today. Authors: - James Lamb (https://github.com/jameslamb) Approvers: - Mike Sarahan (https://github.com/msarahan) - Corey J. Nolet (https://github.com/cjnolet) URL: https://github.com/NVIDIA/cuvs/pull/2595 --- .github/workflows/build.yaml | 58 ++++++---- .github/workflows/pr.yaml | 45 +++++--- .gitignore | 2 + Dockerfile.standalone | 23 ++++ README.md | 4 + build.sh | 57 +++++++++- ci/{ => standalone_c}/build_standalone_c.sh | 61 ++++------- ci/standalone_c/dlsym_smoke.c | 36 +++++++ ci/test_standalone_c.sh | 22 +++- fern/pages/build.md | 112 ++++++++++++++++++++ 10 files changed, 342 insertions(+), 78 deletions(-) create mode 100644 Dockerfile.standalone rename ci/{ => standalone_c}/build_standalone_c.sh (64%) create mode 100644 ci/standalone_c/dlsym_smoke.c diff --git a/.github/workflows/build.yaml b/.github/workflows/build.yaml index 6c18e9f97a..13f70c1fb2 100644 --- a/.github/workflows/build.yaml +++ b/.github/workflows/build.yaml @@ -61,34 +61,54 @@ jobs: contents: read uses: rapidsai/shared-workflows/.github/workflows/compute-matrix.yaml@release/26.10 with: - build_type: pull-request + build_type: ${{ inputs.build_type || 'branch' }} matrix_name: conda-cpp-build rocky8-clib-standalone-build: - needs: [build-details, rocky8-clib-standalone-build-matrix] + needs: [rocky8-clib-standalone-build-matrix] + runs-on: linux-${{ matrix.ARCH }}-cpu16 + env: + RAPIDS_BUILD_TYPE: ${{ inputs.build_type || 'branch' }} permissions: actions: read contents: read id-token: write packages: read pull-requests: read - secrets: inherit # zizmor: ignore[secrets-inherit] - uses: rapidsai/shared-workflows/.github/workflows/custom-job.yaml@release/26.10 strategy: fail-fast: false matrix: ${{ fromJSON(needs.rocky8-clib-standalone-build-matrix.outputs.matrix) }} - with: - build_type: ${{ inputs.build_type || 'branch' }} - build-datetime: ${{ needs.build-details.outputs.build-datetime }} - branch: ${{ inputs.branch }} - arch: "${{matrix.ARCH}}" - date: ${{ inputs.date }} - container_image: "rapidsai/ci-wheel:26.10-cuda${{ matrix.CUDA_VER }}-${{ matrix.LINUX_VER }}-py${{ matrix.PY_VER }}" - node_type: "cpu16" - requires_license_builder: true - script: "ci/build_standalone_c.sh" - artifact-name: "libcuvs_c_${{ matrix.CUDA_VER }}_${{ matrix.ARCH }}.tar.gz" - file_to_upload: "libcuvs_c.tar.gz" - sha: ${{ inputs.sha }} + steps: + - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 + with: + ref: ${{ inputs.sha || github.sha }} + fetch-depth: 0 + persist-credentials: false + - name: Standardize repository information + uses: rapidsai/shared-actions/rapids-github-info@main + with: + branch: ${{ inputs.branch }} + date: ${{ inputs.date }} + sha: ${{ inputs.sha }} + - uses: aws-actions/configure-aws-credentials@cbe3b392738ccf3f987d68400dafcf4b0624a56c # v6.2.4 + with: + role-to-assume: ${{ vars.AWS_ROLE_ARN }} + aws-region: ${{ vars.AWS_REGION }} + role-duration-seconds: 10800 + - name: Setup sccache-dist + uses: rapidsai/shared-actions/setup-sccache-dist@main + with: + request-timeout: 7140 + - name: Build standalone C tarball + env: + CUVS_TARBALL_CUDA_VERSION: ${{ matrix.CUDA_VER }} + CUVS_TARBALL_PYTHON_VERSION: ${{ matrix.PY_VER }} + run: | + ./build.sh tarball --tarball-build-tests + - name: Upload artifacts + uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1 + with: + name: libcuvs_c_${{ matrix.CUDA_VER }}_${{ matrix.ARCH }}.tar.gz + path: libcuvs_c.tar.gz rust-build-matrix: needs: cpp-build permissions: @@ -128,7 +148,7 @@ jobs: contents: read uses: rapidsai/shared-workflows/.github/workflows/compute-matrix.yaml@release/26.10 with: - build_type: pull-request + build_type: ${{ inputs.build_type || 'branch' }} matrix_name: conda-cpp-build matrix_filter: map(select(.ARCH == "amd64")) go-build: @@ -161,7 +181,7 @@ jobs: contents: read uses: rapidsai/shared-workflows/.github/workflows/compute-matrix.yaml@release/26.10 with: - build_type: pull-request + build_type: ${{ inputs.build_type || 'branch' }} matrix_name: conda-cpp-build matrix_filter: map(select(.ARCH == "amd64")) java-build: diff --git a/.github/workflows/pr.yaml b/.github/workflows/pr.yaml index c4b576d6a4..10b12303e6 100644 --- a/.github/workflows/pr.yaml +++ b/.github/workflows/pr.yaml @@ -563,30 +563,45 @@ jobs: build_type: pull-request matrix_name: conda-cpp-build rocky8-clib-standalone-build: - needs: [build-details, rocky8-clib-standalone-build-matrix] + needs: [rocky8-clib-standalone-build-matrix] + runs-on: linux-${{ matrix.ARCH }}-cpu16 + env: + RAPIDS_BUILD_TYPE: pull-request permissions: actions: read contents: read id-token: write packages: read pull-requests: read - secrets: inherit # zizmor: ignore[secrets-inherit] - uses: rapidsai/shared-workflows/.github/workflows/custom-job.yaml@release/26.10 strategy: fail-fast: false matrix: ${{ fromJSON(needs.rocky8-clib-standalone-build-matrix.outputs.matrix) }} - with: - build_type: pull-request - build-datetime: ${{ needs.build-details.outputs.build-datetime }} - arch: "${{matrix.arch}}" - date: ${{ inputs.date }}_c - container_image: "rapidsai/ci-wheel:26.10-cuda${{ matrix.CUDA_VER }}-${{ matrix.LINUX_VER }}-py${{ matrix.PY_VER }}" - node_type: "cpu16" - requires_license_builder: true - script: "ci/build_standalone_c.sh --build-tests" - artifact-name: "libcuvs_c_${{ matrix.CUDA_VER }}_${{ matrix.ARCH }}.tar.gz" - file_to_upload: "libcuvs_c.tar.gz" - sha: ${{ inputs.sha }} + steps: + - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 + with: + ref: ${{ inputs.sha || github.sha }} + fetch-depth: 0 + persist-credentials: false + - uses: aws-actions/configure-aws-credentials@cbe3b392738ccf3f987d68400dafcf4b0624a56c # v6.2.4 + with: + role-to-assume: ${{ vars.AWS_ROLE_ARN }} + aws-region: ${{ vars.AWS_REGION }} + role-duration-seconds: 10800 + - name: Setup sccache-dist + uses: rapidsai/shared-actions/setup-sccache-dist@main + with: + request-timeout: 7140 + - name: Build standalone C tarball + env: + CUVS_TARBALL_CUDA_VERSION: ${{ matrix.CUDA_VER }} + CUVS_TARBALL_PYTHON_VERSION: ${{ matrix.PY_VER }} + run: | + ./build.sh tarball --tarball-build-tests + - name: Upload artifacts + uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1 + with: + name: libcuvs_c_${{ matrix.CUDA_VER }}_${{ matrix.ARCH }}.tar.gz + path: libcuvs_c.tar.gz rocky8-clib-tests-matrix: needs: [rocky8-clib-standalone-build, changed-files] permissions: diff --git a/.gitignore b/.gitignore index 3627558ff5..7357b2a81e 100644 --- a/.gitignore +++ b/.gitignore @@ -19,6 +19,7 @@ dist/ python/**/**/*.cpp python/cuvs/record.txt log +infrequent_licenses/ .ipynb_checkpoints .DS_Store dask-worker-space/ @@ -29,6 +30,7 @@ temporary_*.json rust/target/ rust/Cargo.lock rmm_log.txt +*.tar.gz ## example notebooks notebooks/simplewiki-2020-11-01-nq-distilbert-base-v1.pt diff --git a/Dockerfile.standalone b/Dockerfile.standalone new file mode 100644 index 0000000000..11a4580d81 --- /dev/null +++ b/Dockerfile.standalone @@ -0,0 +1,23 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Environment and runner for the standalone C build. +# +# End-user instructions: https://docs.nvidia.com/cuvs/installation#build-the-standalone-c-library-with-docker + +ARG CUDA_VERSION=notset +ARG PYTHON_VERSION=notset +ARG RAPIDS_VERSION=notset +FROM rapidsai/ci-wheel:${RAPIDS_VERSION}-cuda${CUDA_VERSION}-rockylinux8-py${PYTHON_VERSION} + +# Output directory for the standalone archive. Bind-mount a host folder here +# (e.g. -v $(pwd)/build:/build) so libcuvs_c.tar.gz is written to the host. +ENV CUVS_TARBALL_BUILD_OUTPUT_DIR=/build +ENV CUVS_TARBALL_IN_CONTAINER=1 + +# Run from repo root; the repo is expected to be bind-mounted at /workspace. +WORKDIR /workspace + +# Run the build script, which writes the standalone tarball to the mounted output dir. +# Example: docker run -v $(pwd):/workspace -v $(pwd)/build:/build cuvs-standalone-c [--tarball-build-tests] +ENTRYPOINT ["/bin/bash", "-c", "ci/standalone_c/build_standalone_c.sh \"$@\"", "--"] diff --git a/README.md b/README.md index 77d2466c6a..fa25907ca6 100755 --- a/README.md +++ b/README.md @@ -84,6 +84,10 @@ cuVS comes with pre-built packages that can be installed through [conda](https:/ Please see the [Build and Install Guide](https://docs.rapids.ai/api/cuvs/nightly/build/) for more information on installing the available cuVS packages and building from source. +### Standalone C library (Docker) + +To build the standalone C library tarball (`libcuvs_c.tar.gz`) for use in your own C/C++ projects, see **[Building the standalone C library with Docker](https://docs.nvidia.com/cuvs/installation#build-the-standalone-c-library-with-docker)**. + ## Getting Started The following code snippets train an approximate nearest neighbors index for the CAGRA algorithm in the various different languages supported by cuVS. diff --git a/build.sh b/build.sh index da8adb793f..dadbba9ce7 100755 --- a/build.sh +++ b/build.sh @@ -19,7 +19,7 @@ ARGS=$* # scripts, and that this script resides in the repo dir! REPODIR=$(cd "$(dirname "$0")"; pwd) -VALIDARGS="clean libcuvs python rust go java lucene docs tests bench-ann examples --uninstall -v -g -n --allgpuarch --no-mg --mnmg-tests --no-cpu --cpu-only --no-shared-libs --no-nvtx --show_depr_warn --incl-cache-stats --time -h --run-java-tests --build-java-examples" +VALIDARGS="clean libcuvs python rust go java lucene docs tests bench-ann examples tarball --tarball-build-tests --uninstall -v -g -n --allgpuarch --no-mg --mnmg-tests --no-cpu --cpu-only --no-shared-libs --no-nvtx --show_depr_warn --incl-cache-stats --time -h --run-java-tests --build-java-examples" HELP="$0 [ ...] [ ...] [--cmake-args=\"\"] [--cache-tool=] [--limit-tests=] [--limit-bench-ann=] [--build-metrics=] where is: clean - remove all existing build artifacts and configuration (start over) @@ -34,6 +34,7 @@ HELP="$0 [ ...] [ ...] [--cmake-args=\"\"] [--cache-tool= is: -v - verbose build mode @@ -59,6 +60,7 @@ HELP="$0 [ ...] [ ...] [--cmake-args=\"\"] [--cache-tool=\\\" - pass arbitrary list of CMake configuration options (escape all quotes in argument) @@ -607,3 +609,56 @@ if hasArg examples; then ./build.sh popd fi + +################################################################################ +# Build the standalone C library tarball (if requested) + +if hasArg tarball; then + if [[ "${CUVS_TARBALL_IN_CONTAINER:-0}" == "1" ]]; then + CUVS_TARBALL_BUILD_OUTPUT_DIR="${CUVS_TARBALL_BUILD_OUTPUT_DIR:-${REPODIR}}" + tar czf "${CUVS_TARBALL_BUILD_OUTPUT_DIR}/libcuvs_c.tar.gz" -C "${REPODIR}/c/build/install" . + else + CUVS_TARBALL_CUDA_VERSION="${CUVS_TARBALL_CUDA_VERSION:-13.3.0}" + CUVS_TARBALL_PYTHON_VERSION="${CUVS_TARBALL_PYTHON_VERSION:-3.14}" + CUVS_TARBALL_BUILD_OUTPUT_DIR="${CUVS_TARBALL_BUILD_OUTPUT_DIR:-${REPODIR}/build}" + CUVS_TARBALL_IMAGE_NAME="nvidia/cuvs-standalone-c:local-cuda${CUVS_TARBALL_CUDA_VERSION}-py${CUVS_TARBALL_PYTHON_VERSION}" + + mkdir -p "${CUVS_TARBALL_BUILD_OUTPUT_DIR}" + BUILD_OUTPUT_DIR_ABS=$(realpath "${CUVS_TARBALL_BUILD_OUTPUT_DIR}") + + echo "Building Docker image ${CUVS_TARBALL_IMAGE_NAME} (CUDA ${CUVS_TARBALL_CUDA_VERSION}, Python ${CUVS_TARBALL_PYTHON_VERSION})..." + docker build \ + -f "${REPODIR}/Dockerfile.standalone" \ + --build-arg CUDA_VERSION="${CUVS_TARBALL_CUDA_VERSION}" \ + --build-arg PYTHON_VERSION="${CUVS_TARBALL_PYTHON_VERSION}" \ + --build-arg RAPIDS_VERSION="${RAPIDS_VERSION_MAJOR_MINOR}" \ + -t "${CUVS_TARBALL_IMAGE_NAME}" \ + "${REPODIR}" + + # optionally pass additional arguments through to the container's entrypoint + DOCKER_ENTRYPOINT_ARGS=() + if hasArg --tarball-build-tests; then + DOCKER_ENTRYPOINT_ARGS+=(--tarball-build-tests) + fi + + echo "Running standalone C build in container..." + # NOTE: the '--env-file' trick for AWS credentials keeps them out of 'ps aux' / 'docker ps' output + docker run \ + --rm \ + -v "${REPODIR}:/workspace:rw" \ + -v "${BUILD_OUTPUT_DIR_ABS}:/build:rw" \ + --env CI="${CI:-false}" \ + --env PARALLEL_LEVEL="${PARALLEL_LEVEL}" \ + --env RAPIDS_BUILD_TYPE="${RAPIDS_BUILD_TYPE:-}" \ + --env-file <(env | grep -E '^AWS_(ACCESS_KEY_ID|SECRET_ACCESS_KEY|SESSION_TOKEN)=') \ + --env-file <(env | grep -E '^SCCACHE_.*=') \ + "${CUVS_TARBALL_IMAGE_NAME}" \ + "${DOCKER_ENTRYPOINT_ARGS[@]}" + + cp -v "${BUILD_OUTPUT_DIR_ABS}/libcuvs_c.tar.gz" "${REPODIR}/libcuvs_c.tar.gz" + echo "Copied libcuvs_c.tar.gz to ${REPODIR}/libcuvs_c.tar.gz" + fi + + # print contents of the tarball + ls -lh "${CUVS_TARBALL_BUILD_OUTPUT_DIR}/libcuvs_c.tar.gz" +fi diff --git a/ci/build_standalone_c.sh b/ci/standalone_c/build_standalone_c.sh similarity index 64% rename from ci/build_standalone_c.sh rename to ci/standalone_c/build_standalone_c.sh index e1267f71b7..9b7a4d4a47 100755 --- a/ci/build_standalone_c.sh +++ b/ci/standalone_c/build_standalone_c.sh @@ -1,67 +1,47 @@ #!/bin/bash # SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 +# +# Build script for the standalone C library. +# +# Use 'Dockerfile.standalone' to build an image with all the prerequisites +# and run this in a container. set -euo pipefail +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +REPO_ROOT="$(cd "${SCRIPT_DIR}/../.." && pwd)" + TOOLSET_VERSION=14 -NINJA_VERSION=v1.13.1 BUILD_C_LIB_TESTS="OFF" -if [[ "${1:-}" == "--build-tests" ]]; then +if [[ "${1:-}" == "--tarball-build-tests" ]]; then BUILD_C_LIB_TESTS="ON" fi -dnf install -y \ - patch \ - tar \ - unzip \ - wget - -if ! command -V ninja >/dev/null 2>&1; then - case "$(uname -m)" in - x86_64) - wget --no-hsts -q -O /tmp/ninja-linux.zip "https://github.com/ninja-build/ninja/releases/download/${NINJA_VERSION}/ninja-linux.zip"; - ;; - aarch64) - wget --no-hsts -q -O /tmp/ninja-linux.zip "https://github.com/ninja-build/ninja/releases/download/${NINJA_VERSION}/ninja-linux-aarch64.zip"; - ;; - *) - echo "Unrecognized platform '$(uname -m)'" >&2 - exit 1 - ;; - esac - unzip -d /usr/bin /tmp/ninja-linux.zip - chmod +x /usr/bin/ninja - rm /tmp/ninja-linux.zip -fi - source rapids-install-sccache source rapids-configure-sccache -source rapids-datetime-string -rapids-pip-retry install cmake +PIP_PACKAGES=( + 'cmake>=4.0' + 'git+https://github.com/rapidsai/spdx-license-builder.git' + 'ninja>=1.13' +) RAPIDS_CUDA_MAJOR="${RAPIDS_CUDA_VERSION%%.*}" if [[ "${RAPIDS_CUDA_MAJOR}" == "13" ]]; then - rapids-pip-retry install cuda-tile "cuda-toolkit[tileiras]==${RAPIDS_CUDA_VERSION%.*}.*" + PIP_PACKAGES+=( + cuda-tile + ) fi +rapids-pip-retry install "${PIP_PACKAGES[@]}" pyenv rehash -rapids-print-env - rapids-logger "Begin cpp build" sccache --stop-server 2>/dev/null || true -RAPIDS_PACKAGE_VERSION=$(rapids-generate-version) -export RAPIDS_PACKAGE_VERSION - -RAPIDS_ARTIFACTS_DIR=${RAPIDS_ARTIFACTS_DIR:-"${PWD}/artifacts"} -mkdir -p "${RAPIDS_ARTIFACTS_DIR}" -export RAPIDS_ARTIFACTS_DIR - scl enable gcc-toolset-${TOOLSET_VERSION} -- \ cmake -S cpp -B cpp/build/ -GNinja \ -DCMAKE_CUDA_HOST_COMPILER=/opt/rh/gcc-toolset-${TOOLSET_VERSION}/root/usr/bin/gcc \ @@ -107,11 +87,8 @@ if [ "${BUILD_C_LIB_TESTS}" != "OFF" ]; then cmake --install c/build --prefix c/build/install --component testing fi - rapids-logger "Begin gathering licenses" -rapids-pip-retry install git+https://github.com/rapidsai/spdx-license-builder.git license-builder . --output-json c/build/install/licenses.json --output-txt c/build/install/LICENSE rapids-logger "Begin c tarball creation" -tar czf libcuvs_c.tar.gz -C c/build/install/ . -ls -lh libcuvs_c.tar.gz +"${REPO_ROOT}/build.sh" tarball diff --git a/ci/standalone_c/dlsym_smoke.c b/ci/standalone_c/dlsym_smoke.c new file mode 100644 index 0000000000..e9148a5012 --- /dev/null +++ b/ci/standalone_c/dlsym_smoke.c @@ -0,0 +1,36 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include +#include + +int main(int argc, char** argv) +{ + if (argc != 2) { + fprintf(stderr, "Usage: %s \n", argv[0]); + return 2; + } + + void* handle = dlopen(argv[1], RTLD_NOW | RTLD_LOCAL); + if (handle == NULL) { + fprintf(stderr, "Failed to load %s: %s\n", argv[1], dlerror()); + return 1; + } + + dlerror(); + void* symbol = dlsym(handle, "cuvsResourcesCreate"); + const char* error = dlerror(); + if (error != NULL || symbol == NULL) { + fprintf(stderr, + "Failed to resolve cuvsResourcesCreate: %s\n", + error != NULL ? error : "unknown error"); + dlclose(handle); + return 1; + } + + printf("Successfully loaded %s and resolved cuvsResourcesCreate\n", argv[1]); + dlclose(handle); + return 0; +} diff --git a/ci/test_standalone_c.sh b/ci/test_standalone_c.sh index 894d687bff..519fa2470d 100755 --- a/ci/test_standalone_c.sh +++ b/ci/test_standalone_c.sh @@ -1,5 +1,5 @@ #!/bin/bash -# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 set -euo pipefail @@ -7,6 +7,7 @@ set -euo pipefail rapids-pip-retry install cmake pyenv rehash +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" INSTALL_PREFIX="${PWD}/libcuvs_c_install" mkdir -p "${INSTALL_PREFIX}" @@ -24,6 +25,25 @@ DOWNLOAD_LOCATION=$(rapids-download-from-github "${payload_name}") # Extract the artifact to a staging directory tar -xf "${DOWNLOAD_LOCATION}/${pkg_name}" -C "${INSTALL_PREFIX}" +rapids-logger "Validate C API shared library" +C_API_LIBRARY="" +for C_API_LIBRARY_DIR in "${INSTALL_PREFIX}/lib" "${INSTALL_PREFIX}/lib64"; do + if [[ -f "${C_API_LIBRARY_DIR}/libcuvs_c.so" ]]; then + C_API_LIBRARY="${C_API_LIBRARY_DIR}/libcuvs_c.so" + break + fi +done + +if [[ -z "${C_API_LIBRARY}" ]]; then + echo "Error: C API shared library not found under ${INSTALL_PREFIX}/lib or ${INSTALL_PREFIX}/lib64" >&2 + exit 1 +fi + +C_API_SMOKE_TEST="${INSTALL_PREFIX}/bin/cuvs_c_dlsym_smoke" +"${CC:-cc}" -std=c11 -Wall -Wextra -Werror \ + "${SCRIPT_DIR}/standalone_c/dlsym_smoke.c" -ldl -o "${C_API_SMOKE_TEST}" +LD_LIBRARY_PATH="$(dirname "${C_API_LIBRARY}"):${LD_LIBRARY_PATH:-}" \ + "${C_API_SMOKE_TEST}" "${C_API_LIBRARY}" rapids-logger "Run C API tests" ls -l "${INSTALL_PREFIX}" diff --git a/fern/pages/build.md b/fern/pages/build.md index e8fabd8d27..4cd49bba2b 100644 --- a/fern/pages/build.md +++ b/fern/pages/build.md @@ -41,6 +41,118 @@ conda activate cuvs You may prefer `mamba` over `conda` for faster environment solves. The `conda/environments` directory also contains language-specific environment YAML files for narrower development environments. Conda is not required, but if you do not use it, install all required build dependencies explicitly before running `build.sh`. +## Build the Standalone C Library with Docker + + +The standalone tarball is built with Docker so that it uses the supported toolchain versions and the build remains portable and reproducible across all supported installation platforms. + + +Use the standalone Docker build when you want a `libcuvs_c.tar.gz` archive that you can unpack and use to build your own C or C++ binaries for deployment or integration. + +### Prerequisites + +- Docker with support for the target platform: x86_64 or aarch64. +- At least 16 GB of memory and 20 GB of free disk space available to Docker. +- NVIDIA Container Toolkit and a GPU if you want to run GPU-dependent steps. The image is based on CUDA and may require GPU support at runtime. + +### Use the Helper Script + +From the repository root, run: + +```bash +# (optional) clean old build directories +rm -rf ./{build,c/build/,cpp/build} + +# build +./build.sh tarball +``` + +The script builds the Docker image, runs the build in a container, writes the tarball to `./build/libcuvs_c.tar.gz`, and copies it to `./libcuvs_c.tar.gz` for CI artifact upload and convenience. + +The helper accepts the following environment variables: + +| Variable | Default | Purpose | +| --- | --- | --- | +| `CUVS_TARBALL_CUDA_VERSION` | `13.3.0` | CUDA version for the `rapidsai/ci-wheel` base image. | +| `CUVS_TARBALL_PYTHON_VERSION` | `3.14` | Python version for the `rapidsai/ci-wheel` base image. | +| `CUVS_TARBALL_BUILD_OUTPUT_DIR` | `./build` | Host directory where the tarball is written. | + +CUDA and Python versions should match an existing `rapidsai/ci-wheel` tag. +See https://hub.docker.com/r/rapidsai/ci-wheel/tags + +For example: + +```bash +CUVS_TARBALL_CUDA_VERSION=13.3.0 \ +CUVS_TARBALL_PYTHON_VERSION=3.14 \ + ./build.sh tarball +``` + +To write the tarball to another directory, set `CUVS_TARBALL_BUILD_OUTPUT_DIR`: + +```console +$ CUVS_TARBALL_BUILD_OUTPUT_DIR="${PWD}/dist" ./build.sh tarball +$ find . -name 'libcuvs_c.tar.gz' +./dist/libcuvs_c.tar.gz +./libcuvs_c.tar.gz +``` + +To build and install the C library tests in the archive, pass `--tarball-build-tests`: + +```bash +./build.sh tarball --tarball-build-tests +``` + +### Tarball Contents + +The archive contains the headers, libraries, CMake configuration, and license information needed to compile and link C or C++ applications against the standalone NVIDIA cuVS libraries. + +### Build and Run the Docker Image Manually + +If you do not want to use the helper script, build the image directly from the repository root: + +```bash +docker build \ + -f Dockerfile.standalone \ + --build-arg CUDA_VERSION="13.3.0" \ + --build-arg PYTHON_VERSION="3.14" \ + --build-arg RAPIDS_VERSION="$(head -1 ./VERSION | cut -d. -f1,2 )" \ + -t cuvs-standalone-c:local \ + . +``` + +This command builds a local image and tags it as `cuvs-standalone-c:local`. + +Run the build in a container using that image and mount the repository plus an output directory: + +```bash +mkdir -p build +docker run --rm \ + -v "${PWD}:/workspace" \ + -v "${PWD}/build:/build" \ + cuvs-standalone-c:local +``` + +Mount another host directory at `/build` to change the output location: + +```bash +mkdir -p "${PWD}/dist" +docker run --rm \ + -v "${PWD}:/workspace" \ + -v "${PWD}/dist:/build" \ + cuvs-standalone-c:local +``` + +Pass `--tarball-build-tests` to include the C library tests: + +```bash +mkdir -p build +docker run --rm \ + -v "${PWD}:/workspace" \ + -v "${PWD}/build:/build" \ + cuvs-standalone-c:local --tarball-build-tests +``` + ## Documentation Preview The NVIDIA cuVS documentation is a Fern project in the repository's `fern` directory. Fern requires Node.js 22 or newer. If the docs fail with an error such as `SyntaxError: Unexpected token '.'`, check `node --version` and activate a newer Node.js runtime. From 23d4e42fb171f2a00af6b7bec6a217e8b9f9f897 Mon Sep 17 00:00:00 2001 From: Igor Motov Date: Fri, 11 Sep 2026 17:26:54 -1000 Subject: [PATCH 6/7] Enable ARM (aarch64) support for cuvs-java and cuvs-lucene (#2541) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The C++ core already builds and is tested on both amd64 and arm64; cuvs-java had a hardcoded amd64-only gate that cuvs-lucene inherited. Java bytecode is portable, but cuvs-java isn't pure Java — it uses the Panama FFM API to call into native libcuvs_c.so, and those bindings are generated by jextract, which bakes struct/function ABI layouts (computed by parsing the C headers with clang for the host's target triple) into the generated .class files at build time. Reusing an amd64-generated binding on aarch64 should work, since no arch-conditional logic in the jextract-visible header surface changes any type layout, struct offset, or function signature between x86_64 and aarch64 Linux. We therefore expect the plain (no-native-bundled) jar to be arch-portable in practice. To verify that, the aarch64 cuvs-lucene CI job installs the amd64-built jar and runs cuvs-lucene's full test suite against it on aarch64 hardware with an aarch64-native libcuvs_c.so. Closes #1236. Authors: - Igor Motov (https://github.com/imotov) Approvers: - James Lamb (https://github.com/jameslamb) - MithunR (https://github.com/mythrocks) URL: https://github.com/NVIDIA/cuvs/pull/2541 --- .github/workflows/build.yaml | 3 +- .github/workflows/pr.yaml | 59 +++++ .github/workflows/test.yaml | 61 +++++ ci/test_java_prebuilt.sh | 83 +++++++ ci/test_lucene_prebuilt.sh | 112 +++++++++ java/cuvs-java/pom.xml | 212 ++++++++++++++++++ .../nvidia/cuvs/spi/CuVSServiceProvider.java | 5 +- .../nvidia/cuvs/BruteForceAndSearchIT.java | 2 +- .../nvidia/cuvs/BruteForceRandomizedIT.java | 4 +- .../nvidia/cuvs/CagraAceBuildAndSearchIT.java | 2 +- .../nvidia/cuvs/CagraBuildAndSearchIT.java | 2 +- .../cuvs/CagraMultiThreadStabilityIT.java | 2 +- .../com/nvidia/cuvs/CagraRandomizedIT.java | 2 +- .../java/com/nvidia/cuvs/CuVSMatrixIT.java | 4 +- .../java/com/nvidia/cuvs/CuVSResourcesIT.java | 2 +- .../java/com/nvidia/cuvs/CuVSTestCase.java | 5 +- .../com/nvidia/cuvs/FilterBitsetHandleIT.java | 2 +- .../test/java/com/nvidia/cuvs/GPUInfoIT.java | 4 +- .../nvidia/cuvs/HnswAceBuildAndSearchIT.java | 2 +- .../com/nvidia/cuvs/HnswBuildAndSearchIT.java | 2 +- .../com/nvidia/cuvs/HnswRandomizedIT.java | 2 +- .../cuvs/MemoryTrackingResourcesIT.java | 2 +- .../cuvs/MultiPartitionCagraSearchIT.java | 2 +- .../java/com/nvidia/cuvs/TieredIndexIT.java | 4 +- .../nvidia/cuvs/internal/common/UtilIT.java | 4 +- .../com/nvidia/cuvs/spi/CuVSProviderIT.java | 4 +- java/cuvs-lucene/pom.xml | 33 +++ java/panama-bindings/generate-bindings.sh | 21 +- 28 files changed, 611 insertions(+), 31 deletions(-) create mode 100755 ci/test_java_prebuilt.sh create mode 100755 ci/test_lucene_prebuilt.sh diff --git a/.github/workflows/build.yaml b/.github/workflows/build.yaml index 13f70c1fb2..af13b353e7 100644 --- a/.github/workflows/build.yaml +++ b/.github/workflows/build.yaml @@ -210,7 +210,8 @@ jobs: file_to_upload: "java/cuvs-java/target/" sha: ${{ inputs.sha }} lucene-build: - # Depends on the Java job for the cuvs-java artifact that cuvs-lucene builds against. + # Depends on the Java job for the cuvs-java artifact that cuvs-lucene builds against, and + # reuses its matrix so the two always agree on the set of CUDA versions. needs: [java-build-matrix, java-build] permissions: actions: read diff --git a/.github/workflows/pr.yaml b/.github/workflows/pr.yaml index 10b12303e6..e42c1d7d61 100644 --- a/.github/workflows/pr.yaml +++ b/.github/workflows/pr.yaml @@ -27,7 +27,10 @@ jobs: - rocky8-clib-tests - conda-java-build-and-tests-matrix - conda-java-build-and-tests + - conda-java-tests-other-arch-matrix + - conda-java-tests-other-arch - conda-lucene-build-and-tests + - conda-lucene-tests-other-arch - rust-build-matrix - rust-build - go-build-matrix @@ -665,6 +668,39 @@ jobs: script: "ci/test_java.sh" artifact-name: "cuvs-java-cuda${{ matrix.CUDA_VER }}" file_to_upload: "java/cuvs-java/target/" + conda-java-tests-other-arch-matrix: + # All architectures the C++ core builds for, except amd64. + needs: [conda-cpp-build, changed-files] + permissions: + contents: read + uses: rapidsai/shared-workflows/.github/workflows/compute-matrix.yaml@release/26.10 + if: fromJSON(needs.changed-files.outputs.changed_file_groups).test_java || fromJSON(needs.changed-files.outputs.changed_file_groups).test_cpp + with: + build_type: pull-request + matrix_name: conda-cpp-build + matrix_filter: map(select(.ARCH != "amd64")) + conda-java-tests-other-arch: + # Verifies that the amd64-built cuvs-java jar (jextract-generated Panama bindings and + # all) also works correctly on every non-amd64 architecture, by running its own IT test + # suite against the amd64-built classes without recompiling -- see ci/test_java_prebuilt.sh. + needs: [conda-java-tests-other-arch-matrix, conda-java-build-and-tests] + permissions: + actions: read + contents: read + id-token: write + packages: read + pull-requests: read + secrets: inherit # zizmor: ignore[secrets-inherit] + uses: rapidsai/shared-workflows/.github/workflows/custom-job.yaml@release/26.10 + strategy: + fail-fast: false + matrix: ${{ fromJSON(needs.conda-java-tests-other-arch-matrix.outputs.matrix) }} + with: + build_type: pull-request + node_type: "gpu-l4-latest-1" + arch: "${{ matrix.ARCH }}" + container_image: "rapidsai/ci-conda:26.10-cuda${{ matrix.CUDA_VER }}-${{ matrix.LINUX_VER }}-py${{ matrix.PY_VER }}" + script: "ci/test_java_prebuilt.sh ${{ matrix.ARCH }} cuvs-java-cuda${{ matrix.CUDA_VER }}" conda-lucene-build-and-tests: # Depends on the Java job for the cuvs-java artifact that cuvs-lucene builds against, and # reuses its matrix so the two always agree on the set of CUDA versions. @@ -690,6 +726,29 @@ jobs: script: "ci/test_lucene.sh cuvs-java-cuda${{ matrix.CUDA_VER }}" artifact-name: "cuvs-lucene-cuda${{ matrix.CUDA_VER }}" file_to_upload: "java/cuvs-lucene/target/" + conda-lucene-tests-other-arch: + # Verifies that cuvs-lucene, built against the amd64-built cuvs-java jar, also works + # correctly on every non-amd64 architecture, by running its test suite against the + # amd64-built classes without recompiling -- see ci/test_lucene_prebuilt.sh. Reuses + # conda-java-tests-other-arch-matrix so the two always agree on which architectures to test. + needs: [conda-java-tests-other-arch-matrix, conda-java-build-and-tests, conda-lucene-build-and-tests] + permissions: + actions: read + contents: read + id-token: write + packages: read + pull-requests: read + secrets: inherit # zizmor: ignore[secrets-inherit] + uses: rapidsai/shared-workflows/.github/workflows/custom-job.yaml@release/26.10 + strategy: + fail-fast: false + matrix: ${{ fromJSON(needs.conda-java-tests-other-arch-matrix.outputs.matrix) }} + with: + build_type: pull-request + node_type: "gpu-l4-latest-1" + arch: "${{ matrix.ARCH }}" + container_image: "rapidsai/ci-conda:26.10-cuda${{ matrix.CUDA_VER }}-${{ matrix.LINUX_VER }}-py${{ matrix.PY_VER }}" + script: "ci/test_lucene_prebuilt.sh ${{ matrix.ARCH }} cuvs-java-cuda${{ matrix.CUDA_VER }} cuvs-lucene-cuda${{ matrix.CUDA_VER }}" rust-build-matrix: needs: [conda-cpp-build, changed-files] permissions: diff --git a/.github/workflows/test.yaml b/.github/workflows/test.yaml index 16d967973b..98e51d8365 100644 --- a/.github/workflows/test.yaml +++ b/.github/workflows/test.yaml @@ -104,6 +104,41 @@ jobs: script: "ci/test_java.sh" artifact-name: "cuvs-java-cuda${{ matrix.CUDA_VER }}" file_to_upload: "java/cuvs-java/target/" + conda-java-tests-other-arch-matrix: + # All architectures the C++ core builds for, except amd64 -- so a newly added + # architecture is automatically picked up here without a workflow change. + permissions: + contents: read + uses: rapidsai/shared-workflows/.github/workflows/compute-matrix.yaml@release/26.10 + with: + build_type: ${{ inputs.build_type }} + matrix_name: conda-cpp-build + matrix_filter: map(select(.ARCH != "amd64")) + conda-java-tests-other-arch: + # Verifies that the amd64-built cuvs-java jar (jextract-generated Panama bindings and + # all) also works correctly on every non-amd64 architecture, by running its own IT test + # suite against the amd64-built classes without recompiling -- see ci/test_java_prebuilt.sh. + needs: [conda-java-tests-other-arch-matrix, conda-java-build-and-tests] + permissions: + actions: read + contents: read + id-token: write + packages: read + pull-requests: read + secrets: inherit # zizmor: ignore[secrets-inherit] + uses: rapidsai/shared-workflows/.github/workflows/custom-job.yaml@release/26.10 + strategy: + fail-fast: false + matrix: ${{ fromJSON(needs.conda-java-tests-other-arch-matrix.outputs.matrix) }} + with: + build_type: ${{ inputs.build_type }} + branch: ${{ inputs.branch }} + date: ${{ inputs.date }} + sha: ${{ inputs.sha }} + node_type: "gpu-l4-latest-1" + arch: "${{ matrix.ARCH }}" + container_image: "rapidsai/ci-conda:26.10-cuda${{ matrix.CUDA_VER }}-${{ matrix.LINUX_VER }}-py${{ matrix.PY_VER }}" + script: "ci/test_java_prebuilt.sh ${{ matrix.ARCH }} cuvs-java-cuda${{ matrix.CUDA_VER }}" conda-lucene-build-and-tests: # Depends on the Java job for the cuvs-java artifact that cuvs-lucene builds against, and # reuses its matrix so the two always agree on the set of CUDA versions. @@ -132,6 +167,32 @@ jobs: script: "ci/test_lucene.sh cuvs-java-cuda${{ matrix.CUDA_VER }}" artifact-name: "cuvs-lucene-cuda${{ matrix.CUDA_VER }}" file_to_upload: "java/cuvs-lucene/target/" + conda-lucene-tests-other-arch: + # Verifies that cuvs-lucene, built against the amd64-built cuvs-java jar, also works + # correctly on every non-amd64 architecture, by running its test suite against the + # amd64-built classes without recompiling -- see ci/test_lucene_prebuilt.sh. Reuses + # conda-java-tests-other-arch-matrix so the two always agree on which architectures to test. + needs: [conda-java-tests-other-arch-matrix, conda-java-build-and-tests, conda-lucene-build-and-tests] + permissions: + actions: read + contents: read + id-token: write + packages: read + pull-requests: read + secrets: inherit # zizmor: ignore[secrets-inherit] + uses: rapidsai/shared-workflows/.github/workflows/custom-job.yaml@release/26.10 + strategy: + fail-fast: false + matrix: ${{ fromJSON(needs.conda-java-tests-other-arch-matrix.outputs.matrix) }} + with: + build_type: ${{ inputs.build_type }} + branch: ${{ inputs.branch }} + date: ${{ inputs.date }} + sha: ${{ inputs.sha }} + node_type: "gpu-l4-latest-1" + arch: "${{ matrix.ARCH }}" + container_image: "rapidsai/ci-conda:26.10-cuda${{ matrix.CUDA_VER }}-${{ matrix.LINUX_VER }}-py${{ matrix.PY_VER }}" + script: "ci/test_lucene_prebuilt.sh ${{ matrix.ARCH }} cuvs-java-cuda${{ matrix.CUDA_VER }} cuvs-lucene-cuda${{ matrix.CUDA_VER }}" wheel-tests-cuvs: permissions: actions: read diff --git a/ci/test_java_prebuilt.sh b/ci/test_java_prebuilt.sh new file mode 100755 index 0000000000..5f811aa9b5 --- /dev/null +++ b/ci/test_java_prebuilt.sh @@ -0,0 +1,83 @@ +#!/bin/bash +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +set -euo pipefail + +# Runs cuvs-java's own IT test suite against the classes from a prior amd64 build job, +# without recompiling (or running jextract) on this host. cuvs-java is always built on +# amd64; this verifies that the resulting jextract-generated Panama bindings also work +# correctly against a native libcuvs_c.so on the given (non-amd64) architecture. +# +# Takes the target architecture (as used by the GitHub Actions job matrix, e.g. "arm64") +# and the name of the cuvs-java artifact uploaded by the amd64 Java job. + +ARCH="${1:?Usage: $0 }" +CUVS_JAVA_ARTIFACT="${2:?Usage: $0 }" + +rapids-logger "Testing the amd64-built cuvs-java artifact on ${ARCH}" + +case "${ARCH}" in + amd64) CONDA_ARCH="x86_64" ;; + arm64) CONDA_ARCH="aarch64" ;; + *) CONDA_ARCH="${ARCH}" ;; +esac + +if [ -e "/opt/conda/etc/profile.d/conda.sh" ]; then + . /opt/conda/etc/profile.d/conda.sh +fi + +rapids-logger "Check GPU usage" +nvidia-smi + +rapids-logger "Configuring conda strict channel priority" +conda config --set channel_priority strict + +rapids-logger "Downloading artifacts from previous jobs" +CPP_CHANNEL=$(rapids-download-from-github "$(rapids-artifact-name conda_cpp libcuvs cuvs --cuda "$RAPIDS_CUDA_VERSION")") +CUVS_JAVA_DIR=$(rapids-download-from-github "${CUVS_JAVA_ARTIFACT}") + +rapids-logger "Generate Java testing dependencies" + +ENV_YAML_DIR="$(mktemp -d)" + +rapids-dependency-file-generator \ + --output conda \ + --file-key java \ + --prepend-channel "${CPP_CHANNEL}" \ + --matrix "cuda=${RAPIDS_CUDA_VERSION%.*};arch=${CONDA_ARCH}" | tee "${ENV_YAML_DIR}/env.yaml" + +rapids-mamba-retry env create --yes -f "${ENV_YAML_DIR}/env.yaml" -n java + +# Temporarily allow unbound variables for conda activation. +set +u +conda activate java +set -u + +rapids-print-env + +# libcuvs comes from the conda environment here, matching the architecture this script is +# running on. +export LD_LIBRARY_PATH="${CONDA_PREFIX}/lib${LD_LIBRARY_PATH:+:${LD_LIBRARY_PATH}}" + +rapids-logger "Restore the amd64-built target/ directory so Maven can run the tests without recompiling" + +rm -rf java/cuvs-java/target +mkdir -p java/cuvs-java/target +cp -a "${CUVS_JAVA_DIR}/." java/cuvs-java/target/ + +EXITCODE=0 +trap "EXITCODE=1" ERR +set +e + +rapids-logger "Run cuvs-java IT tests against the amd64-built classes" + +# -Dskip.compile activates the pom's "skip-compile" profile, which disables all +# compilation for this run, forcing test to use amd64-compiled jar instead of +# local code. +pushd java/cuvs-java +mvn --batch-mode verify -Dskip.compile=true +popd + +rapids-logger "Test script exiting with value: $EXITCODE" +exit ${EXITCODE} diff --git a/ci/test_lucene_prebuilt.sh b/ci/test_lucene_prebuilt.sh new file mode 100755 index 0000000000..477ff75ad3 --- /dev/null +++ b/ci/test_lucene_prebuilt.sh @@ -0,0 +1,112 @@ +#!/bin/bash +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +set -euo pipefail + +# Runs cuvs-lucene's test suite against the classes from a prior amd64 build job, without +# recompiling on this host. cuvs-lucene is always built on amd64, against the amd64-built +# cuvs-java jar; this verifies that pairing also works correctly on the given (non-amd64) +# architecture. Since cuvs-lucene depends on the plain (no bundled natives) cuvs-java jar, +# this doubles as a cross-arch check that the jextract-generated Panama bindings baked +# into the amd64 jar work unmodified against a native libcuvs_c.so on that architecture. +# +# Takes the target architecture (as used by the GitHub Actions job matrix, e.g. "arm64") +# and the names of the cuvs-java and cuvs-lucene artifacts uploaded by the amd64 jobs. + +ARCH="${1:?Usage: $0 }" +CUVS_JAVA_ARTIFACT="${2:?Usage: $0 }" +CUVS_LUCENE_ARTIFACT="${3:?Usage: $0 }" + +rapids-logger "Testing the amd64-built cuvs-java/cuvs-lucene artifacts on ${ARCH}" + +case "${ARCH}" in + amd64) CONDA_ARCH="x86_64" ;; + arm64) CONDA_ARCH="aarch64" ;; + *) CONDA_ARCH="${ARCH}" ;; +esac + +if [ -e "/opt/conda/etc/profile.d/conda.sh" ]; then + . /opt/conda/etc/profile.d/conda.sh +fi + +rapids-logger "Check GPU usage" +nvidia-smi + +rapids-logger "Configuring conda strict channel priority" +conda config --set channel_priority strict + +rapids-logger "Downloading artifacts from previous jobs" +CPP_CHANNEL=$(rapids-download-from-github "$(rapids-artifact-name conda_cpp libcuvs cuvs --cuda "$RAPIDS_CUDA_VERSION")") +CUVS_JAVA_DIR=$(rapids-download-from-github "${CUVS_JAVA_ARTIFACT}") +CUVS_LUCENE_DIR=$(rapids-download-from-github "${CUVS_LUCENE_ARTIFACT}") + +rapids-logger "Generate Java testing dependencies" + +ENV_YAML_DIR="$(mktemp -d)" + +rapids-dependency-file-generator \ + --output conda \ + --file-key java \ + --prepend-channel "${CPP_CHANNEL}" \ + --matrix "cuda=${RAPIDS_CUDA_VERSION%.*};arch=${CONDA_ARCH}" | tee "${ENV_YAML_DIR}/env.yaml" + +rapids-mamba-retry env create --yes -f "${ENV_YAML_DIR}/env.yaml" -n java + +# Temporarily allow unbound variables for conda activation. +set +u +conda activate java +set -u + +rapids-print-env + +# libcuvs comes from the conda environment here, matching the architecture this script is +# running on. cuvs-lucene depends on the plain cuvs-java jar, which bundles no native +# libraries, so the JVM resolves libcuvs_c.so through the dynamic loader. +export LD_LIBRARY_PATH="${CONDA_PREFIX}/lib${LD_LIBRARY_PATH:+:${LD_LIBRARY_PATH}}" + +rapids-logger "Install the amd64-built cuvs-java artifact into the local Maven repository" + +# cuvs-lucene resolves cuvs-java from the local Maven repository. Its pom.xml travels with +# the artifact and supplies the coordinates, so no version needs to be hardcoded. +CUVS_JAVA_POM="${CUVS_JAVA_DIR}/pom.xml" +if [ ! -f "${CUVS_JAVA_POM}" ]; then + echo "Could not find pom.xml in the cuvs-java artifact at ${CUVS_JAVA_DIR}" >&2 + exit 1 +fi + +# The artifact also carries the per-architecture native jar and the sources/javadoc/test jars; +# cuvs-lucene depends on the plain one. +mapfile -t CUVS_JAVA_JARS < <(find "${CUVS_JAVA_DIR}" -maxdepth 1 -name 'cuvs-java-*.jar' \ + ! -name '*-sources.jar' ! -name '*-javadoc.jar' ! -name '*-tests.jar' ! -name '*-cuda*.jar') +if [ "${#CUVS_JAVA_JARS[@]}" -ne 1 ]; then + echo "Expected exactly one cuvs-java jar in ${CUVS_JAVA_DIR}, found: ${CUVS_JAVA_JARS[*]:-none}" >&2 + exit 1 +fi + +# Install cuvs-java jar into .m2, cd is needed to pick up pom.xml in order to avoid rate limit of main maven repo +pushd java/cuvs-lucene +mvn --batch-mode install:install-file -Dfile="${CUVS_JAVA_JARS[0]}" -DpomFile="${CUVS_JAVA_POM}" +popd + +rapids-logger "Restore the amd64-built cuvs-lucene target/ directory so Maven can run the tests without recompiling" + +rm -rf java/cuvs-lucene/target +mkdir -p java/cuvs-lucene/target +cp -a "${CUVS_LUCENE_DIR}/." java/cuvs-lucene/target/ + +EXITCODE=0 +trap "EXITCODE=1" ERR +set +e + +rapids-logger "Run cuvs-lucene tests against the amd64-built classes" + +# -Dskip.compile activates the pom's "skip-compile" profile, which disables all +# compilation for this run, forcing test to use amd64-compiled jar instead of +# local code. +pushd java/cuvs-lucene +mvn --batch-mode test -Dskip.compile=true +popd + +rapids-logger "Test script exiting with value: $EXITCODE" +exit ${EXITCODE} diff --git a/java/cuvs-java/pom.xml b/java/cuvs-java/pom.xml index b8e969d07b..d57cb3ccac 100644 --- a/java/cuvs-java/pom.xml +++ b/java/cuvs-java/pom.xml @@ -285,9 +285,46 @@ SPDX-License-Identifier: Apache-2.0 + + + skip-compile + + + skip.compile + + + + + + org.apache.maven.plugins + maven-compiler-plugin + + + default-compile + none + + + compile-java-22 + none + + + default-testCompile + none + + + + + + x86_64-cuda12 + + amd64 + cuda.version 12 @@ -371,6 +408,9 @@ SPDX-License-Identifier: Apache-2.0 x86_64-cuda13 + + amd64 + cuda.version 13 @@ -450,5 +490,177 @@ SPDX-License-Identifier: Apache-2.0 + + + aarch64-cuda12 + + + aarch64 + + + cuda.version + 12 + + + + aarch64-cuda12 + ${native.build.path}/cuda12 + + + + + org.apache.maven.plugins + maven-resources-plugin + 3.3.1 + + + copy-native-libs + prepare-package + + copy-resources + + + true + ${project.build.directory}/native-libs/${os.arch}/${os.name} + + + ${native.lib.path} + + libcuvs.so + libcuvs_c.so + + + + ${native.lib.path}/_deps/rmm-build + + librmm.so + + + + ${native.lib.path}/_deps/rapids_logger-build + + librapids_logger.so + + + + + + + + + org.apache.maven.plugins + maven-assembly-plugin + 3.4.2 + + + src/assembly/native-with-deps.xml + + + + true + true + 12 + + + + + + assemble-native + package + + single + + + + + + + + + + aarch64-cuda13 + + + aarch64 + + + cuda.version + 13 + + + + aarch64-cuda13 + ${native.build.path}/cuda13 + + + + + org.apache.maven.plugins + maven-resources-plugin + 3.3.1 + + + copy-native-libs + prepare-package + + copy-resources + + + true + ${project.build.directory}/native-libs/${os.arch}/${os.name} + + + ${native.lib.path} + + libcuvs.so + libcuvs_c.so + + + + ${native.lib.path}/_deps/rmm-build + + librmm.so + + + + ${native.lib.path}/_deps/rapids_logger-build + + librapids_logger.so + + + + + + + + + org.apache.maven.plugins + maven-assembly-plugin + 3.4.2 + + + src/assembly/native-with-deps.xml + + + + true + true + 13 + + + + + + assemble-native + package + + single + + + + + + + diff --git a/java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/CuVSServiceProvider.java b/java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/CuVSServiceProvider.java index 874f35ee20..b4f16a084c 100644 --- a/java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/CuVSServiceProvider.java +++ b/java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/CuVSServiceProvider.java @@ -39,7 +39,7 @@ static CuVSProvider builtinProvider() { var osArch = System.getProperty("os.arch"); var supportedJavaRuntime = javaRuntimeVersion > 21; var supportedOs = osName.startsWith("Linux"); - var supportedArchitecture = osArch.equals("amd64"); + var supportedArchitecture = osArch.equals("amd64") || osArch.equals("aarch64"); if (supportedJavaRuntime && supportedOs && supportedArchitecture) { try { var cls = Class.forName("com.nvidia.cuvs.spi.JDKProvider"); @@ -63,7 +63,8 @@ static CuVSProvider builtinProvider() { unsupportedReasons.add("cuvs-java supports only Linux, but os.name is " + osName); } if (!supportedArchitecture) { - unsupportedReasons.add("cuvs-java supports only x86-64 (amd64), but os.arch is " + osArch); + unsupportedReasons.add( + "cuvs-java supports only x86-64 (amd64) and arm64 (aarch64), but os.arch is " + osArch); } return new UnsupportedProvider(String.join("; ", unsupportedReasons)); diff --git a/java/cuvs-java/src/test/java/com/nvidia/cuvs/BruteForceAndSearchIT.java b/java/cuvs-java/src/test/java/com/nvidia/cuvs/BruteForceAndSearchIT.java index ab59395dd2..1f28c45950 100644 --- a/java/cuvs-java/src/test/java/com/nvidia/cuvs/BruteForceAndSearchIT.java +++ b/java/cuvs-java/src/test/java/com/nvidia/cuvs/BruteForceAndSearchIT.java @@ -23,7 +23,7 @@ public class BruteForceAndSearchIT extends CuVSTestCase { @Before public void setup() { - assumeTrue("not supported on " + System.getProperty("os.name"), isLinuxAmd64()); + assumeTrue("not supported on " + System.getProperty("os.name"), isLinuxSupportedArch()); CuVSProvider.provider().enableRMMPooledMemory(10, 60); } diff --git a/java/cuvs-java/src/test/java/com/nvidia/cuvs/BruteForceRandomizedIT.java b/java/cuvs-java/src/test/java/com/nvidia/cuvs/BruteForceRandomizedIT.java index 22af6fb4c3..96c10a0936 100644 --- a/java/cuvs-java/src/test/java/com/nvidia/cuvs/BruteForceRandomizedIT.java +++ b/java/cuvs-java/src/test/java/com/nvidia/cuvs/BruteForceRandomizedIT.java @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ package com.nvidia.cuvs; @@ -25,7 +25,7 @@ public class BruteForceRandomizedIT extends CuVSTestCase { @Before public void setup() { - assumeTrue(isLinuxAmd64()); + assumeTrue(isLinuxSupportedArch()); initializeRandom(); log.trace("Random context initialized for test."); CuVSProvider.provider().enableRMMPooledMemory(10, 60); diff --git a/java/cuvs-java/src/test/java/com/nvidia/cuvs/CagraAceBuildAndSearchIT.java b/java/cuvs-java/src/test/java/com/nvidia/cuvs/CagraAceBuildAndSearchIT.java index f762334619..423c58db35 100644 --- a/java/cuvs-java/src/test/java/com/nvidia/cuvs/CagraAceBuildAndSearchIT.java +++ b/java/cuvs-java/src/test/java/com/nvidia/cuvs/CagraAceBuildAndSearchIT.java @@ -36,7 +36,7 @@ public class CagraAceBuildAndSearchIT extends CuVSTestCase { @Before public void setup() { - assumeTrue("not supported on " + System.getProperty("os.name"), isLinuxAmd64()); + assumeTrue("not supported on " + System.getProperty("os.name"), isLinuxSupportedArch()); initializeRandom(); log.trace("Random context initialized for test."); } diff --git a/java/cuvs-java/src/test/java/com/nvidia/cuvs/CagraBuildAndSearchIT.java b/java/cuvs-java/src/test/java/com/nvidia/cuvs/CagraBuildAndSearchIT.java index d90b975846..a4f08c47d4 100644 --- a/java/cuvs-java/src/test/java/com/nvidia/cuvs/CagraBuildAndSearchIT.java +++ b/java/cuvs-java/src/test/java/com/nvidia/cuvs/CagraBuildAndSearchIT.java @@ -47,7 +47,7 @@ public class CagraBuildAndSearchIT extends CuVSTestCase { @Before public void setup() { - assumeTrue("not supported on " + System.getProperty("os.name"), isLinuxAmd64()); + assumeTrue("not supported on " + System.getProperty("os.name"), isLinuxSupportedArch()); initializeRandom(); log.trace("Random context initialized for test."); } diff --git a/java/cuvs-java/src/test/java/com/nvidia/cuvs/CagraMultiThreadStabilityIT.java b/java/cuvs-java/src/test/java/com/nvidia/cuvs/CagraMultiThreadStabilityIT.java index 371e842973..17790eb2fc 100644 --- a/java/cuvs-java/src/test/java/com/nvidia/cuvs/CagraMultiThreadStabilityIT.java +++ b/java/cuvs-java/src/test/java/com/nvidia/cuvs/CagraMultiThreadStabilityIT.java @@ -44,7 +44,7 @@ private interface ResourcesSupplier { @Before public void setup() { - assumeTrue("not supported on " + System.getProperty("os.name"), isLinuxAmd64()); + assumeTrue("not supported on " + System.getProperty("os.name"), isLinuxSupportedArch()); initializeRandom(); log.trace("Multi-threaded stability test initialized"); } diff --git a/java/cuvs-java/src/test/java/com/nvidia/cuvs/CagraRandomizedIT.java b/java/cuvs-java/src/test/java/com/nvidia/cuvs/CagraRandomizedIT.java index 2dcb8b1cd0..9110f6ddc2 100644 --- a/java/cuvs-java/src/test/java/com/nvidia/cuvs/CagraRandomizedIT.java +++ b/java/cuvs-java/src/test/java/com/nvidia/cuvs/CagraRandomizedIT.java @@ -24,7 +24,7 @@ public class CagraRandomizedIT extends CuVSTestCase { @Before public void setup() { - assumeTrue("not supported on " + System.getProperty("os.name"), isLinuxAmd64()); + assumeTrue("not supported on " + System.getProperty("os.name"), isLinuxSupportedArch()); initializeRandom(); log.trace("Random context initialized for test."); } diff --git a/java/cuvs-java/src/test/java/com/nvidia/cuvs/CuVSMatrixIT.java b/java/cuvs-java/src/test/java/com/nvidia/cuvs/CuVSMatrixIT.java index cd3cf43e55..954a7d1e4f 100644 --- a/java/cuvs-java/src/test/java/com/nvidia/cuvs/CuVSMatrixIT.java +++ b/java/cuvs-java/src/test/java/com/nvidia/cuvs/CuVSMatrixIT.java @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ package com.nvidia.cuvs; @@ -25,7 +25,7 @@ public class CuVSMatrixIT extends CuVSTestCase { @Before public void setup() { - assumeTrue("not supported on " + System.getProperty("os.name"), isLinuxAmd64()); + assumeTrue("not supported on " + System.getProperty("os.name"), isLinuxSupportedArch()); initializeRandom(); } diff --git a/java/cuvs-java/src/test/java/com/nvidia/cuvs/CuVSResourcesIT.java b/java/cuvs-java/src/test/java/com/nvidia/cuvs/CuVSResourcesIT.java index 00753b42c4..632eefc6a1 100644 --- a/java/cuvs-java/src/test/java/com/nvidia/cuvs/CuVSResourcesIT.java +++ b/java/cuvs-java/src/test/java/com/nvidia/cuvs/CuVSResourcesIT.java @@ -20,7 +20,7 @@ public class CuVSResourcesIT extends CuVSTestCase { @Before public void setup() { - assumeTrue("not supported on " + System.getProperty("os.name"), isLinuxAmd64()); + assumeTrue("not supported on " + System.getProperty("os.name"), isLinuxSupportedArch()); } @Test diff --git a/java/cuvs-java/src/test/java/com/nvidia/cuvs/CuVSTestCase.java b/java/cuvs-java/src/test/java/com/nvidia/cuvs/CuVSTestCase.java index 62e636eb82..53cbe47392 100644 --- a/java/cuvs-java/src/test/java/com/nvidia/cuvs/CuVSTestCase.java +++ b/java/cuvs-java/src/test/java/com/nvidia/cuvs/CuVSTestCase.java @@ -155,9 +155,10 @@ public boolean equals(Object o) { assertEquals(sortedExpected, sortedActual); } - protected static boolean isLinuxAmd64() { + protected static boolean isLinuxSupportedArch() { String name = System.getProperty("os.name"); - return (name.startsWith("Linux")) && System.getProperty("os.arch").equals("amd64"); + String arch = System.getProperty("os.arch"); + return (name.startsWith("Linux")) && (arch.equals("amd64") || arch.equals("aarch64")); } protected static int[][] createIntMatrix() { diff --git a/java/cuvs-java/src/test/java/com/nvidia/cuvs/FilterBitsetHandleIT.java b/java/cuvs-java/src/test/java/com/nvidia/cuvs/FilterBitsetHandleIT.java index 2abb153207..d7464ca97e 100644 --- a/java/cuvs-java/src/test/java/com/nvidia/cuvs/FilterBitsetHandleIT.java +++ b/java/cuvs-java/src/test/java/com/nvidia/cuvs/FilterBitsetHandleIT.java @@ -52,7 +52,7 @@ public class FilterBitsetHandleIT extends CuVSTestCase { @Before public void setup() { - assumeTrue("not supported on " + System.getProperty("os.name"), isLinuxAmd64()); + assumeTrue("not supported on " + System.getProperty("os.name"), isLinuxSupportedArch()); initializeRandom(); } diff --git a/java/cuvs-java/src/test/java/com/nvidia/cuvs/GPUInfoIT.java b/java/cuvs-java/src/test/java/com/nvidia/cuvs/GPUInfoIT.java index fbfab24518..4b9fc045a3 100644 --- a/java/cuvs-java/src/test/java/com/nvidia/cuvs/GPUInfoIT.java +++ b/java/cuvs-java/src/test/java/com/nvidia/cuvs/GPUInfoIT.java @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ package com.nvidia.cuvs; @@ -19,7 +19,7 @@ public class GPUInfoIT extends CuVSTestCase { @Before public void setup() { - assumeTrue("not supported on " + System.getProperty("os.name"), isLinuxAmd64()); + assumeTrue("not supported on " + System.getProperty("os.name"), isLinuxSupportedArch()); } @Test diff --git a/java/cuvs-java/src/test/java/com/nvidia/cuvs/HnswAceBuildAndSearchIT.java b/java/cuvs-java/src/test/java/com/nvidia/cuvs/HnswAceBuildAndSearchIT.java index 835fcd6e23..351873ce9a 100644 --- a/java/cuvs-java/src/test/java/com/nvidia/cuvs/HnswAceBuildAndSearchIT.java +++ b/java/cuvs-java/src/test/java/com/nvidia/cuvs/HnswAceBuildAndSearchIT.java @@ -40,7 +40,7 @@ public class HnswAceBuildAndSearchIT extends CuVSTestCase { @Before public void setup() { - assumeTrue("not supported on " + System.getProperty("os.name"), isLinuxAmd64()); + assumeTrue("not supported on " + System.getProperty("os.name"), isLinuxSupportedArch()); initializeRandom(); log.trace("Random context initialized for test."); } diff --git a/java/cuvs-java/src/test/java/com/nvidia/cuvs/HnswBuildAndSearchIT.java b/java/cuvs-java/src/test/java/com/nvidia/cuvs/HnswBuildAndSearchIT.java index 7cb4f0c3fa..78f86bd964 100644 --- a/java/cuvs-java/src/test/java/com/nvidia/cuvs/HnswBuildAndSearchIT.java +++ b/java/cuvs-java/src/test/java/com/nvidia/cuvs/HnswBuildAndSearchIT.java @@ -28,7 +28,7 @@ public class HnswBuildAndSearchIT extends CuVSTestCase { @Before public void setup() { - assumeTrue("not supported on " + System.getProperty("os.name"), isLinuxAmd64()); + assumeTrue("not supported on " + System.getProperty("os.name"), isLinuxSupportedArch()); } // Sample data and query diff --git a/java/cuvs-java/src/test/java/com/nvidia/cuvs/HnswRandomizedIT.java b/java/cuvs-java/src/test/java/com/nvidia/cuvs/HnswRandomizedIT.java index fe5ca88073..4ace8fc92e 100644 --- a/java/cuvs-java/src/test/java/com/nvidia/cuvs/HnswRandomizedIT.java +++ b/java/cuvs-java/src/test/java/com/nvidia/cuvs/HnswRandomizedIT.java @@ -28,7 +28,7 @@ public class HnswRandomizedIT extends CuVSTestCase { @Before public void setup() { - assumeTrue("not supported on " + System.getProperty("os.name"), isLinuxAmd64()); + assumeTrue("not supported on " + System.getProperty("os.name"), isLinuxSupportedArch()); initializeRandom(); log.trace("Random context initialized for test."); } diff --git a/java/cuvs-java/src/test/java/com/nvidia/cuvs/MemoryTrackingResourcesIT.java b/java/cuvs-java/src/test/java/com/nvidia/cuvs/MemoryTrackingResourcesIT.java index b80a789025..1d9d40ab27 100644 --- a/java/cuvs-java/src/test/java/com/nvidia/cuvs/MemoryTrackingResourcesIT.java +++ b/java/cuvs-java/src/test/java/com/nvidia/cuvs/MemoryTrackingResourcesIT.java @@ -18,7 +18,7 @@ public class MemoryTrackingResourcesIT extends CuVSTestCase { @Before public void setup() { - assumeTrue("not supported on " + System.getProperty("os.name"), isLinuxAmd64()); + assumeTrue("not supported on " + System.getProperty("os.name"), isLinuxSupportedArch()); } @Test diff --git a/java/cuvs-java/src/test/java/com/nvidia/cuvs/MultiPartitionCagraSearchIT.java b/java/cuvs-java/src/test/java/com/nvidia/cuvs/MultiPartitionCagraSearchIT.java index 62e13b733e..2d1e33484a 100644 --- a/java/cuvs-java/src/test/java/com/nvidia/cuvs/MultiPartitionCagraSearchIT.java +++ b/java/cuvs-java/src/test/java/com/nvidia/cuvs/MultiPartitionCagraSearchIT.java @@ -40,7 +40,7 @@ public class MultiPartitionCagraSearchIT extends CuVSTestCase { @Before public void setup() { - assumeTrue("not supported on " + System.getProperty("os.name"), isLinuxAmd64()); + assumeTrue("not supported on " + System.getProperty("os.name"), isLinuxSupportedArch()); initializeRandom(); log.trace("Random context initialized for test."); } diff --git a/java/cuvs-java/src/test/java/com/nvidia/cuvs/TieredIndexIT.java b/java/cuvs-java/src/test/java/com/nvidia/cuvs/TieredIndexIT.java index 7965ca34b6..6a47e506ca 100644 --- a/java/cuvs-java/src/test/java/com/nvidia/cuvs/TieredIndexIT.java +++ b/java/cuvs-java/src/test/java/com/nvidia/cuvs/TieredIndexIT.java @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ package com.nvidia.cuvs; @@ -31,7 +31,7 @@ public class TieredIndexIT extends CuVSTestCase { @Before public void setup() { - assumeTrue("not supported on " + System.getProperty("os.name"), isLinuxAmd64()); + assumeTrue("not supported on " + System.getProperty("os.name"), isLinuxSupportedArch()); initializeRandom(); CuVSProvider.provider().enableRMMPooledMemory(10, 60); log.debug("Random context initialized for test"); diff --git a/java/cuvs-java/src/test/java/com/nvidia/cuvs/internal/common/UtilIT.java b/java/cuvs-java/src/test/java/com/nvidia/cuvs/internal/common/UtilIT.java index b6a5664831..3bd8d81464 100644 --- a/java/cuvs-java/src/test/java/com/nvidia/cuvs/internal/common/UtilIT.java +++ b/java/cuvs-java/src/test/java/com/nvidia/cuvs/internal/common/UtilIT.java @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ package com.nvidia.cuvs.internal.common; @@ -22,7 +22,7 @@ public class UtilIT extends CuVSTestCase { @Before public void setup() { - assumeTrue("not supported on " + System.getProperty("os.name"), isLinuxAmd64()); + assumeTrue("not supported on " + System.getProperty("os.name"), isLinuxSupportedArch()); } @Test diff --git a/java/cuvs-java/src/test/java/com/nvidia/cuvs/spi/CuVSProviderIT.java b/java/cuvs-java/src/test/java/com/nvidia/cuvs/spi/CuVSProviderIT.java index c40c42204e..076db40553 100644 --- a/java/cuvs-java/src/test/java/com/nvidia/cuvs/spi/CuVSProviderIT.java +++ b/java/cuvs-java/src/test/java/com/nvidia/cuvs/spi/CuVSProviderIT.java @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ package com.nvidia.cuvs.spi; @@ -18,7 +18,7 @@ public class CuVSProviderIT extends CuVSTestCase { @Before public void setup() { - assumeTrue("not supported on " + System.getProperty("os.name"), isLinuxAmd64()); + assumeTrue("not supported on " + System.getProperty("os.name"), isLinuxSupportedArch()); // Clear sysprop from previous runs/command line System.clearProperty("cuvs.max_version"); } diff --git a/java/cuvs-lucene/pom.xml b/java/cuvs-lucene/pom.xml index cf7628137e..ef31cd37af 100644 --- a/java/cuvs-lucene/pom.xml +++ b/java/cuvs-lucene/pom.xml @@ -251,4 +251,37 @@ SPDX-License-Identifier: Apache-2.0 + + + + + skip-compile + + + skip.compile + + + + + + org.apache.maven.plugins + maven-compiler-plugin + + + default-compile + none + + + default-testCompile + none + + + + + + + diff --git a/java/panama-bindings/generate-bindings.sh b/java/panama-bindings/generate-bindings.sh index cdefd1ae1c..315b58f89f 100755 --- a/java/panama-bindings/generate-bindings.sh +++ b/java/panama-bindings/generate-bindings.sh @@ -10,7 +10,16 @@ REPODIR=$(cd "$(dirname "$0")"; cd ../../ ; pwd) CURDIR=$(cd "$(dirname "$0")"; pwd) TARGET_PACKAGE="com.nvidia.cuvs.internal.panama" -TARGET_DIR="targets/x86_64-linux/include" +ARCH=$(uname -m) +case "${ARCH}" in + x86_64) CUDA_TARGET_ARCH="x86_64-linux" ;; + aarch64) CUDA_TARGET_ARCH="sbsa-linux" ;; + *) + echo "Unsupported architecture for CUDA include directory lookup: ${ARCH}" + exit 1 + ;; +esac +TARGET_DIR="targets/${CUDA_TARGET_ARCH}/include" if [ -n "${CONDA_PREFIX:-}" ] && [ -d "${CONDA_PREFIX}/${TARGET_DIR}" ]; then CUDA_INCLUDE_DIR="${CONDA_PREFIX}/${TARGET_DIR}" elif [ -d "/usr/local/cuda/${TARGET_DIR}" ]; then @@ -25,7 +34,15 @@ export PATH if [[ $(command -v jextract) == "" ]]; then - JEXTRACT_FILENAME="openjdk-22-jextract+6-47_linux-x64_bin.tar.gz" + case "${ARCH}" in + x86_64) JEXTRACT_ARCH="linux-x64" ;; + aarch64) JEXTRACT_ARCH="linux-aarch64" ;; + *) + echo "Unsupported architecture for jextract download: ${ARCH}" + exit 1 + ;; + esac + JEXTRACT_FILENAME="openjdk-22-jextract+6-47_${JEXTRACT_ARCH}_bin.tar.gz" JEXTRACT_DOWNLOAD_URL="https://download.java.net/java/early_access/jextract/22/6/${JEXTRACT_FILENAME}" echo "jextract doesn't exist. Downloading it from $JEXTRACT_DOWNLOAD_URL."; wget -c $JEXTRACT_DOWNLOAD_URL From cdfb7b895640821296e5d0ee9494806871d48f22 Mon Sep 17 00:00:00 2001 From: Bradley Dice Date: Sun, 13 Sep 2026 23:24:31 -0500 Subject: [PATCH 7/7] Migrate stream APIs from rmm::cuda_stream_view to cuda::stream_ref (#2521) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Summary Track the coordinated migration of stream APIs and call sites from `rmm::cuda_stream_view` to CCCL's `cuda::stream_ref`. This propagates `cuda::stream_ref` through RMM containers and memory resources, RAFT resource and handle APIs, downstream C++ interfaces, Python/Cython bindings, benchmarks, tests, and documentation. This migrates affected cuVS API signatures and internal call sites while extracting raw stream handles only where CUDA, generated/JIT, or legacy APIs require them. Depends on rapidsai/rmm#2372 and NVIDIA/raft#3129. Tracked in rapidsai/build-planning#318. ## Migrations - Pass `cuda::stream_ref` through stream pools, resource accessors, conditionals, and downstream APIs without converting to `rmm::cuda_stream_view` - Use `cuda::stream_ref` constructions for default/legacy/per-thread streams - `rmm::cuda_stream_default` ➡️ `cuda::stream_ref{cudaStream_t{cudaStreamDefault}}` - `rmm::cuda_stream_legacy` ➡️ `cuda::stream_ref{cudaStreamLegacy}` - `rmm::cuda_stream_per_thread` ➡️ `cuda::stream_ref{cudaStreamPerThread}` - Use `.get()` when calling an API that requires a raw `cudaStream_t`, including CUDA runtime, library, CUB, and legacy API boundaries (previously `rmm::cuda_stream_view` used `value()`) - Use `.sync()` when synchronizing a `cuda::stream_ref` (previously `rmm::cuda_stream_view` used `synchronize()`) - Update Cython declarations and call sites to pass stream references directly where supported Authors: - Bradley Dice (https://github.com/bdice) Approvers: - Corey J. Nolet (https://github.com/cjnolet) - Divye Gala (https://github.com/divyegala) URL: https://github.com/NVIDIA/cuvs/pull/2521 --- c/src/core/c_api.cpp | 4 +- cpp/bench/ann/src/common/util.hpp | 2 +- cpp/bench/ann/src/cuvs/cuvs_ann_bench_utils.h | 6 +- .../ann/src/cuvs/cuvs_brute_force_knn.cu | 4 +- .../ann/src/cuvs/cuvs_cagra_hnswlib_wrapper.h | 2 +- cpp/bench/ann/src/cuvs/cuvs_cagra_wrapper.h | 2 +- .../ann/src/cuvs/cuvs_ivf_flat_wrapper.h | 2 +- cpp/bench/ann/src/cuvs/cuvs_ivf_pq_wrapper.h | 2 +- .../ann/src/cuvs/cuvs_ivf_rabitq_wrapper.h | 2 +- cpp/bench/ann/src/cuvs/cuvs_ivf_sq_wrapper.h | 2 +- .../ann/src/cuvs/cuvs_mg_cagra_wrapper.h | 2 +- .../ann/src/cuvs/cuvs_mg_ivf_flat_wrapper.h | 4 +- .../ann/src/cuvs/cuvs_mg_ivf_pq_wrapper.h | 4 +- cpp/bench/ann/src/cuvs/cuvs_vamana_wrapper.h | 4 +- cpp/bench/ann/src/cuvs/cuvs_wrapper.h | 4 +- .../patches/faiss-1.14-cuda-stream-ref.diff | 68 +++++++++++++++++++ cpp/cmake/patches/faiss_override.json | 5 ++ cpp/include/cuvs/neighbors/scann.hpp | 3 +- cpp/include/cuvs/neighbors/vamana.hpp | 1 - .../cuvs_internal/neighbors/refine_helper.cuh | 6 +- cpp/src/cluster/detail/kmeans_balanced.cuh | 1 - .../all_neighbors/all_neighbors_batched.cuh | 2 - cpp/src/neighbors/detail/ann_utils.cuh | 31 ++++----- .../neighbors/detail/cagra/cagra_search.cuh | 1 - .../detail/cagra/compute_distance.hpp | 14 ++-- .../cagra/compute_distance_standard-impl.cuh | 35 +++++----- .../cagra/compute_distance_vpq-impl.cuh | 4 +- .../neighbors/detail/cagra/search_plan.cuh | 9 +-- cpp/src/neighbors/detail/tiered_index.cuh | 1 - cpp/src/neighbors/ivf_common.cu | 5 +- cpp/src/neighbors/ivf_common.cuh | 7 +- cpp/src/neighbors/ivf_flat/ivf_flat_build.cuh | 2 - ...vf_flat_interleaved_scan_explicit_inst.cuh | 5 +- .../ivf_flat_interleaved_scan_ext.cuh | 7 +- .../ivf_flat_interleaved_scan_jit.cuh | 6 +- .../ivf_pq_compute_similarity_run_inst.cu.in | 5 +- .../detail/ivf_pq_contiguous_list_data.cu | 7 +- .../ivf_pq/detail/ivf_pq_list_data.cu | 7 +- cpp/src/neighbors/ivf_pq/ivf_pq_build.cuh | 7 +- .../ivf_pq/ivf_pq_compute_similarity.cuh | 8 +-- .../ivf_pq/ivf_pq_compute_similarity_impl.cuh | 4 +- .../ivf_pq/ivf_pq_contiguous_list_data.cuh | 7 +- .../ivf_pq_contiguous_list_data_impl.cuh | 5 +- cpp/src/neighbors/ivf_pq/ivf_pq_list_data.hpp | 7 +- .../ivf_pq/ivf_pq_list_data_impl.cuh | 5 +- cpp/src/neighbors/ivf_pq/ivf_pq_transform.cuh | 1 - .../ivf_rabitq/gpu_index/initializer_gpu.cuh | 3 +- .../neighbors/ivf_rabitq/gpu_index/ivf_gpu.cu | 3 +- .../ivf_rabitq/gpu_index/ivf_gpu.cuh | 6 +- .../ivf_rabitq/gpu_index/quantizer_gpu.cuh | 5 +- .../ivf_rabitq/gpu_index/searcher_gpu.cu | 3 +- .../ivf_rabitq/gpu_index/searcher_gpu.cuh | 8 +-- cpp/src/neighbors/ivf_sq/ivf_sq_search.cuh | 5 +- cpp/src/neighbors/scann/detail/scann_avq.cuh | 1 - .../preprocessing/quantize/detail/scalar.cuh | 4 -- cpp/tests/cluster/kmeans_mnmg.cu | 4 +- cpp/tests/neighbors/all_neighbors.cuh | 3 +- cpp/tests/neighbors/ann_brute_force.cuh | 3 +- cpp/tests/neighbors/ann_cagra.cuh | 13 ++-- cpp/tests/neighbors/ann_hnsw_ace.cuh | 3 +- cpp/tests/neighbors/ann_ivf_flat.cuh | 3 +- cpp/tests/neighbors/ann_ivf_pq.cuh | 5 +- cpp/tests/neighbors/ann_ivf_rabitq.cuh | 5 +- cpp/tests/neighbors/ann_ivf_sq.cuh | 3 +- cpp/tests/neighbors/ann_nn_descent.cuh | 7 +- cpp/tests/neighbors/ann_scann.cuh | 3 +- cpp/tests/neighbors/ann_vamana.cuh | 3 +- cpp/tests/neighbors/distance_nn.cu | 3 +- cpp/tests/neighbors/hnsw.cu | 3 +- cpp/tests/neighbors/refine.cu | 6 +- cpp/tests/neighbors/refine_helper.cuh | 6 +- cpp/tests/neighbors/tiered_index.cu | 3 +- .../sparse/neighbors/cross_component_nn.cu | 2 - examples/cpp/src/dynamic_batching_example.cu | 5 +- ...pes-copy-serialization-and-utility-apis.md | 16 ++--- ...pp-api-common-types-execution-resources.md | 36 +++++----- fern/pages/other/resources.md | 4 +- fern/scripts/generate_api_reference.py | 52 +++++++------- 78 files changed, 318 insertions(+), 233 deletions(-) create mode 100644 cpp/cmake/patches/faiss-1.14-cuda-stream-ref.diff diff --git a/c/src/core/c_api.cpp b/c/src/core/c_api.cpp index 5d431afcf8..4b5c2c351c 100644 --- a/c/src/core/c_api.cpp +++ b/c/src/core/c_api.cpp @@ -15,7 +15,7 @@ #include #include #include -#include +#include #include #include #include @@ -136,7 +136,7 @@ extern "C" cuvsError_t cuvsStreamSet(cuvsResources_t res, cudaStream_t stream) { return cuvs::core::translate_exceptions([=] { auto res_ptr = reinterpret_cast(res); - raft::resource::set_cuda_stream(*res_ptr, static_cast(stream)); + raft::resource::set_cuda_stream(*res_ptr, static_cast(stream)); }); } diff --git a/cpp/bench/ann/src/common/util.hpp b/cpp/bench/ann/src/common/util.hpp index 069d257717..3f23e74a7a 100644 --- a/cpp/bench/ann/src/common/util.hpp +++ b/cpp/bench/ann/src/common/util.hpp @@ -171,7 +171,7 @@ inline auto get_stream_from_global_pool() -> cudaStream_t detail::global_stream_pool.emplace_back(rmm::cuda_stream::flags::non_blocking); } } - return detail::global_stream_pool[benchmark_thread_id].view(); + return detail::global_stream_pool[benchmark_thread_id].value(); #else return nullptr; #endif diff --git a/cpp/bench/ann/src/cuvs/cuvs_ann_bench_utils.h b/cpp/bench/ann/src/cuvs/cuvs_ann_bench_utils.h index f699bdad1a..472957d049 100644 --- a/cpp/bench/ann/src/cuvs/cuvs_ann_bench_utils.h +++ b/cpp/bench/ann/src/cuvs/cuvs_ann_bench_utils.h @@ -19,8 +19,8 @@ #include #include +#include #include -#include #include #include #include @@ -124,8 +124,8 @@ class configured_raft_resources { */ explicit configured_raft_resources(const std::shared_ptr& shared_res) : shared_res_{shared_res}, - res_{std::make_unique( - rmm::cuda_stream_view(get_stream_from_global_pool()))} + res_{ + std::make_unique(cuda::stream_ref(get_stream_from_global_pool()))} { raft::resource::set_large_workspace_resource( *res_, raft::mr::device_resource{shared_res_->get_large_memory_resource()}); diff --git a/cpp/bench/ann/src/cuvs/cuvs_brute_force_knn.cu b/cpp/bench/ann/src/cuvs/cuvs_brute_force_knn.cu index 9aa491191a..3c8bf94f1f 100644 --- a/cpp/bench/ann/src/cuvs/cuvs_brute_force_knn.cu +++ b/cpp/bench/ann/src/cuvs/cuvs_brute_force_knn.cu @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2024, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ #include @@ -72,7 +72,7 @@ template class BruteForceKNNBenchmark { public: BruteForceKNNBenchmark(const RandomKNNInputs& params, const std::string& type_str) - : stream_(raft::resource::get_cuda_stream(handle_)), + : stream_(raft::resource::get_cuda_stream(handle_).get()), params_(params), type_str_(type_str), database(params_.num_db_vecs * params_.dim, stream_), diff --git a/cpp/bench/ann/src/cuvs/cuvs_cagra_hnswlib_wrapper.h b/cpp/bench/ann/src/cuvs/cuvs_cagra_hnswlib_wrapper.h index e24802a298..81c0e60ecf 100644 --- a/cpp/bench/ann/src/cuvs/cuvs_cagra_hnswlib_wrapper.h +++ b/cpp/bench/ann/src/cuvs/cuvs_cagra_hnswlib_wrapper.h @@ -45,7 +45,7 @@ class cuvs_cagra_hnswlib : public algo, public algo_gpu { [[nodiscard]] auto get_sync_stream() const noexcept -> cudaStream_t override { - return handle_.get_sync_stream(); + return handle_.get_sync_stream().get(); } [[nodiscard]] auto uses_stream() const noexcept -> bool override diff --git a/cpp/bench/ann/src/cuvs/cuvs_cagra_wrapper.h b/cpp/bench/ann/src/cuvs/cuvs_cagra_wrapper.h index ed067fac39..2507fc3de1 100644 --- a/cpp/bench/ann/src/cuvs/cuvs_cagra_wrapper.h +++ b/cpp/bench/ann/src/cuvs/cuvs_cagra_wrapper.h @@ -205,7 +205,7 @@ class cuvs_cagra : public algo, public algo_gpu { [[nodiscard]] auto get_sync_stream() const noexcept -> cudaStream_t override { - return handle_.get_sync_stream(); + return handle_.get_sync_stream().get(); } [[nodiscard]] auto uses_stream() const noexcept -> bool override diff --git a/cpp/bench/ann/src/cuvs/cuvs_ivf_flat_wrapper.h b/cpp/bench/ann/src/cuvs/cuvs_ivf_flat_wrapper.h index 8ed5adee26..b9eba9fcb3 100644 --- a/cpp/bench/ann/src/cuvs/cuvs_ivf_flat_wrapper.h +++ b/cpp/bench/ann/src/cuvs/cuvs_ivf_flat_wrapper.h @@ -57,7 +57,7 @@ class cuvs_ivf_flat : public algo, public algo_gpu { [[nodiscard]] auto get_sync_stream() const noexcept -> cudaStream_t override { - return handle_.get_sync_stream(); + return handle_.get_sync_stream().get(); } // to enable dataset access from GPU memory diff --git a/cpp/bench/ann/src/cuvs/cuvs_ivf_pq_wrapper.h b/cpp/bench/ann/src/cuvs/cuvs_ivf_pq_wrapper.h index 363d065e91..1783801c95 100644 --- a/cpp/bench/ann/src/cuvs/cuvs_ivf_pq_wrapper.h +++ b/cpp/bench/ann/src/cuvs/cuvs_ivf_pq_wrapper.h @@ -70,7 +70,7 @@ class cuvs_ivf_pq : public algo, public algo_gpu { [[nodiscard]] auto get_sync_stream() const noexcept -> cudaStream_t override { - return handle_.get_sync_stream(); + return handle_.get_sync_stream().get(); } // to enable dataset access from GPU memory diff --git a/cpp/bench/ann/src/cuvs/cuvs_ivf_rabitq_wrapper.h b/cpp/bench/ann/src/cuvs/cuvs_ivf_rabitq_wrapper.h index 0c41ee1c96..0050be1704 100644 --- a/cpp/bench/ann/src/cuvs/cuvs_ivf_rabitq_wrapper.h +++ b/cpp/bench/ann/src/cuvs/cuvs_ivf_rabitq_wrapper.h @@ -55,7 +55,7 @@ class cuvs_ivf_rabitq : public algo, public algo_gpu { [[nodiscard]] auto get_sync_stream() const noexcept -> cudaStream_t override { - return handle_.get_sync_stream(); + return handle_.get_sync_stream().get(); } // to enable dataset access from GPU memory diff --git a/cpp/bench/ann/src/cuvs/cuvs_ivf_sq_wrapper.h b/cpp/bench/ann/src/cuvs/cuvs_ivf_sq_wrapper.h index 5bf0098eaa..6acd91586b 100644 --- a/cpp/bench/ann/src/cuvs/cuvs_ivf_sq_wrapper.h +++ b/cpp/bench/ann/src/cuvs/cuvs_ivf_sq_wrapper.h @@ -52,7 +52,7 @@ class cuvs_ivf_sq : public algo, public algo_gpu { [[nodiscard]] auto get_sync_stream() const noexcept -> cudaStream_t override { - return handle_.get_sync_stream(); + return handle_.get_sync_stream().get(); } [[nodiscard]] auto get_preference() const -> algo_property override diff --git a/cpp/bench/ann/src/cuvs/cuvs_mg_cagra_wrapper.h b/cpp/bench/ann/src/cuvs/cuvs_mg_cagra_wrapper.h index 6d94e9495f..7a89b68394 100644 --- a/cpp/bench/ann/src/cuvs/cuvs_mg_cagra_wrapper.h +++ b/cpp/bench/ann/src/cuvs/cuvs_mg_cagra_wrapper.h @@ -56,7 +56,7 @@ class cuvs_mg_cagra : public algo, public algo_gpu { [[nodiscard]] auto get_sync_stream() const noexcept -> cudaStream_t override { auto stream = raft::resource::get_cuda_stream(clique_); - return stream; + return stream.get(); } // to enable dataset access from GPU memory diff --git a/cpp/bench/ann/src/cuvs/cuvs_mg_ivf_flat_wrapper.h b/cpp/bench/ann/src/cuvs/cuvs_mg_ivf_flat_wrapper.h index 496b89dbb9..a1dd955efd 100644 --- a/cpp/bench/ann/src/cuvs/cuvs_mg_ivf_flat_wrapper.h +++ b/cpp/bench/ann/src/cuvs/cuvs_mg_ivf_flat_wrapper.h @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2024, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ @@ -52,7 +52,7 @@ class cuvs_mg_ivf_flat : public algo, public algo_gpu { [[nodiscard]] auto get_sync_stream() const noexcept -> cudaStream_t override { auto stream = raft::resource::get_cuda_stream(clique_); - return stream; + return stream.get(); } [[nodiscard]] auto uses_stream() const noexcept -> bool override { return false; } diff --git a/cpp/bench/ann/src/cuvs/cuvs_mg_ivf_pq_wrapper.h b/cpp/bench/ann/src/cuvs/cuvs_mg_ivf_pq_wrapper.h index 735f2a2cbc..e0a39c7ce4 100644 --- a/cpp/bench/ann/src/cuvs/cuvs_mg_ivf_pq_wrapper.h +++ b/cpp/bench/ann/src/cuvs/cuvs_mg_ivf_pq_wrapper.h @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2024, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ @@ -52,7 +52,7 @@ class cuvs_mg_ivf_pq : public algo, public algo_gpu { [[nodiscard]] auto get_sync_stream() const noexcept -> cudaStream_t override { auto stream = raft::resource::get_cuda_stream(clique_); - return stream; + return stream.get(); } [[nodiscard]] auto uses_stream() const noexcept -> bool override { return false; } diff --git a/cpp/bench/ann/src/cuvs/cuvs_vamana_wrapper.h b/cpp/bench/ann/src/cuvs/cuvs_vamana_wrapper.h index 7d7292bd50..9bc492dc92 100644 --- a/cpp/bench/ann/src/cuvs/cuvs_vamana_wrapper.h +++ b/cpp/bench/ann/src/cuvs/cuvs_vamana_wrapper.h @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ #pragma once @@ -37,7 +37,7 @@ class cuvs_vamana : public algo, public algo_gpu { [[nodiscard]] auto get_sync_stream() const noexcept -> cudaStream_t override { - return handle_.get_sync_stream(); + return handle_.get_sync_stream().get(); } // to enable dataset access from GPU memory diff --git a/cpp/bench/ann/src/cuvs/cuvs_wrapper.h b/cpp/bench/ann/src/cuvs/cuvs_wrapper.h index 584d64d939..b5470dedae 100644 --- a/cpp/bench/ann/src/cuvs/cuvs_wrapper.h +++ b/cpp/bench/ann/src/cuvs/cuvs_wrapper.h @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2023-2024, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ #pragma once @@ -65,7 +65,7 @@ class cuvs_gpu : public algo, public algo_gpu { } [[nodiscard]] auto get_sync_stream() const noexcept -> cudaStream_t override { - return handle_.get_sync_stream(); + return handle_.get_sync_stream().get(); } void set_search_dataset(const T* dataset, size_t nrow) override; void save(const std::string& file) const override; diff --git a/cpp/cmake/patches/faiss-1.14-cuda-stream-ref.diff b/cpp/cmake/patches/faiss-1.14-cuda-stream-ref.diff new file mode 100644 index 0000000000..c4577e7f0d --- /dev/null +++ b/cpp/cmake/patches/faiss-1.14-cuda-stream-ref.diff @@ -0,0 +1,68 @@ +diff --git a/faiss/gpu/GpuIndexIVFFlat.cu b/faiss/gpu/GpuIndexIVFFlat.cu +index aef0b26..d8a6f9c 100644 +--- a/faiss/gpu/GpuIndexIVFFlat.cu ++++ b/faiss/gpu/GpuIndexIVFFlat.cu +@@ -267,7 +267,7 @@ void GpuIndexIVFFlat::train(idx_t n, const float* x) { + // transfer centroids to host + auto host_centroids = toHost( + cuvs_ivfflat_index.value().centers().data_handle(), +- raft_handle.get_stream(), ++ raft_handle.get_stream().get(), + {idx_t(nlist), this->d}); + quantizer->train(nlist, host_centroids.data()); + quantizer->add(nlist, host_centroids.data()); +diff --git a/faiss/gpu/GpuIndexIVFPQ.cu b/faiss/gpu/GpuIndexIVFPQ.cu +index a4d849c..fae7df1 100644 +--- a/faiss/gpu/GpuIndexIVFPQ.cu ++++ b/faiss/gpu/GpuIndexIVFPQ.cu +@@ -431,7 +431,7 @@ void GpuIndexIVFPQ::train(idx_t n, const float* x) { + // transfer centroids to host + auto host_centroids = toHost( + cluster_centers.data_handle(), +- raft_handle.get_stream(), ++ raft_handle.get_stream().get(), + {idx_t(nlist), this->d}); + quantizer->train(nlist, host_centroids.data()); + quantizer->add(nlist, host_centroids.data()); +diff --git a/faiss/gpu/impl/CuvsIVFFlat.cu b/faiss/gpu/impl/CuvsIVFFlat.cu +index 59ccb21..2497891 100644 +--- a/faiss/gpu/impl/CuvsIVFFlat.cu ++++ b/faiss/gpu/impl/CuvsIVFFlat.cu +@@ -459,7 +459,7 @@ void CuvsIVFFlat::copyInvertedListsFrom(const InvertedLists* ivf) { + cuvs_index->centers().data_handle(), + cuvs_index->dim(), + (uint32_t)nlist, +- raft_handle.get_stream()); ++ raft_handle.get_stream().get()); + } + } + +diff --git a/faiss/gpu/utils/CuvsFilterConvert.cu b/faiss/gpu/utils/CuvsFilterConvert.cu +index 8a10aee..bbc5f8b 100644 +--- a/faiss/gpu/utils/CuvsFilterConvert.cu ++++ b/faiss/gpu/utils/CuvsFilterConvert.cu +@@ -91,13 +91,13 @@ void convert_to_bitset_range( + (n_elements_to_set + threads_per_block - 1) / threads_per_block; + + if (nbits == original_nbits) { +- set_range_kernel<<>>( ++ set_range_kernel<<>>( + (uint32_t*)bitset.data(), imin, imax, n_elements_to_set); + } else if (original_nbits == 8) { +- set_range_kernel<<>>( ++ set_range_kernel<<>>( + (uint8_t*)bitset.data(), imin, imax, n_elements_to_set); + } else if (original_nbits == 64) { +- set_range_kernel<<>>( ++ set_range_kernel<<>>( + (uint64_t*)bitset.data(), imin, imax, n_elements_to_set); + } else { + throw std::invalid_argument("Unsupported original_nbits"); +@@ -175,6 +175,6 @@ void convert_to_bitset_bitmap( + const int threads_per_block = 256; + const int blocks = (n_elements + threads_per_block - 1) / threads_per_block; + +- set_bitmap_kernel<<>>( ++ set_bitmap_kernel<<>>( + bitset.data(), d_bitmap_ptr, n_elements, bitset_original_nbits); + } diff --git a/cpp/cmake/patches/faiss_override.json b/cpp/cmake/patches/faiss_override.json index fb76b0d53b..f209f97a72 100644 --- a/cpp/cmake/patches/faiss_override.json +++ b/cpp/cmake/patches/faiss_override.json @@ -9,6 +9,11 @@ "file" : "${current_json_dir}/faiss-1.14-cuvs-26.08.diff", "issue" : "Multiple fixes for cuVS compatibility. Update Faiss cuVS to be compatible with new Dataset API: update_device_dataset_same_layout now takes dataset_view and make_padded_dataset_view must be called beforehand. Loading an index built from a user-provided KNN graph passes dataset_view into cagra::index, not raw mdspan.", "fixed_in" : "" + }, + { + "file" : "${current_json_dir}/faiss-1.14-cuda-stream-ref.diff", + "issue" : "Unwrap cuda::stream_ref at CUDA kernel launch, Faiss toHost, and legacy RAFT rowNorm boundaries.", + "fixed_in" : "" } ] } diff --git a/cpp/include/cuvs/neighbors/scann.hpp b/cpp/include/cuvs/neighbors/scann.hpp index de0487265e..f186ef21e2 100644 --- a/cpp/include/cuvs/neighbors/scann.hpp +++ b/cpp/include/cuvs/neighbors/scann.hpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ @@ -17,7 +17,6 @@ #include #include #include -#include #include #include diff --git a/cpp/include/cuvs/neighbors/vamana.hpp b/cpp/include/cuvs/neighbors/vamana.hpp index 398c8cd6e5..517e1f76e1 100644 --- a/cpp/include/cuvs/neighbors/vamana.hpp +++ b/cpp/include/cuvs/neighbors/vamana.hpp @@ -16,7 +16,6 @@ #include #include #include -#include #include #include diff --git a/cpp/internal/cuvs_internal/neighbors/refine_helper.cuh b/cpp/internal/cuvs_internal/neighbors/refine_helper.cuh index d5c43aaa34..f5cfe6b79d 100644 --- a/cpp/internal/cuvs_internal/neighbors/refine_helper.cuh +++ b/cpp/internal/cuvs_internal/neighbors/refine_helper.cuh @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2022-2023, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ #pragma once @@ -15,7 +15,7 @@ #include #include -#include +#include #include namespace cuvs::neighbors { @@ -127,7 +127,7 @@ class RefineHelper { public: RefineInputs p; const raft::resources& handle_; - rmm::cuda_stream_view stream_; + cuda::stream_ref stream_; raft::device_matrix dataset; raft::device_matrix queries; diff --git a/cpp/src/cluster/detail/kmeans_balanced.cuh b/cpp/src/cluster/detail/kmeans_balanced.cuh index 3eb4f9275a..7ae70a7a04 100644 --- a/cpp/src/cluster/detail/kmeans_balanced.cuh +++ b/cpp/src/cluster/detail/kmeans_balanced.cuh @@ -898,7 +898,6 @@ void build_clusters(const raft::resources& handle, rmm::device_async_resource_ref device_memory, const MathT* dataset_norm = nullptr) { - auto stream = raft::resource::get_cuda_stream(handle); // "randomly" initialize labels auto labels_view = raft::make_device_vector_view(cluster_labels, n_rows); raft::linalg::map_offset( diff --git a/cpp/src/neighbors/all_neighbors/all_neighbors_batched.cuh b/cpp/src/neighbors/all_neighbors/all_neighbors_batched.cuh index 6ba5b9395d..6d5581310d 100644 --- a/cpp/src/neighbors/all_neighbors/all_neighbors_batched.cuh +++ b/cpp/src/neighbors/all_neighbors/all_neighbors_batched.cuh @@ -131,8 +131,6 @@ void single_gpu_assign_clusters( std::optional> norms_view; cuvs::neighbors::brute_force::index brute_force_index(res, centroids, norms_view, metric); - auto stream = resource::get_cuda_stream(res); - for (size_t i = 0; i < num_batches; i++) { size_t row_offset = n_rows_per_batch * i + base_row_offset; size_t n_rows_of_current_batch = std::min(n_rows_per_batch, num_rows - row_offset); diff --git a/cpp/src/neighbors/detail/ann_utils.cuh b/cpp/src/neighbors/detail/ann_utils.cuh index b3c5b11f27..deddc47dcb 100644 --- a/cpp/src/neighbors/detail/ann_utils.cuh +++ b/cpp/src/neighbors/detail/ann_utils.cuh @@ -18,7 +18,7 @@ #include #include -#include +#include #include #include #include @@ -224,7 +224,7 @@ HDI constexpr auto mapping::operator()(const float& x) const -> int8_t * @param[in] n_bytes */ template -inline void memzero(T* ptr, IdxT n_elems, rmm::cuda_stream_view stream) +inline void memzero(T* ptr, IdxT n_elems, cuda::stream_ref stream) { switch (check_pointer_residency(ptr)) { case pointer_residency::host_and_device: @@ -298,7 +298,7 @@ void block_copy(const IdxT* in_offsets, const T* in_data, T* out_data, IdxT n_mult, - rmm::cuda_stream_view stream) + cuda::stream_ref stream) { IdxT in_size; update_host(&in_size, in_offsets + n_blocks, 1, stream); @@ -325,7 +325,7 @@ void block_copy(const IdxT* in_offsets, * @param stream */ template -void outer_add(const T* a, IdxT len_a, const T* b, IdxT len_b, T* c, rmm::cuda_stream_view stream) +void outer_add(const T* a, IdxT len_a, const T* b, IdxT len_b, T* c, cuda::stream_ref stream) { dim3 threads(128, 1, 1); dim3 blocks(raft::ceildiv(len_a * len_b, threads.x), 1, 1); @@ -370,7 +370,7 @@ void copy_selected(IdxT n_rows, IdxT ld_src, T* dst, IdxT ld_dst, - rmm::cuda_stream_view stream) + cuda::stream_ref stream) { switch (check_pointer_residency(src, dst, row_ids)) { case pointer_residency::host_and_device: @@ -403,8 +403,7 @@ void copy_selected(IdxT n_rows, * the main stream itself is returned with `false`, and the caller should treat prefetch as a * no-op (no overlap is possible on a single stream). */ -inline auto get_prefetch_stream(raft::resources const& res) - -> std::pair +inline auto get_prefetch_stream(raft::resources const& res) -> std::pair { if (res.has_resource_factory(raft::resource::resource_type::CUDA_STREAM_POOL) && raft::resource::get_stream_pool_size(res) >= 1) { @@ -600,7 +599,7 @@ struct batch_load_iterator { batch(raft::resources const& res, MdspanT input_view, size_type batch_size, - rmm::cuda_stream_view copy_stream, + cuda::stream_ref copy_stream, rmm::device_async_resource_ref mr, bool prefetch, bool initialize, @@ -806,7 +805,7 @@ struct batch_load_iterator { copy_stream_.get())); } - rmm::cuda_stream_view copy_stream_; + cuda::stream_ref copy_stream_; raft::resources const* res_; MdspanT input_view_; element_type* source_; @@ -860,7 +859,7 @@ struct batch_load_iterator { batch_load_iterator(raft::resources const& res, MdspanT input_view, size_type batch_size, - rmm::cuda_stream_view copy_stream, + cuda::stream_ref copy_stream, rmm::device_async_resource_ref mr, bool prefetch = false, bool initialize = true, @@ -876,7 +875,7 @@ struct batch_load_iterator { batch_load_iterator(raft::resources const& res, MdspanT input_view, size_type batch_size, - rmm::cuda_stream_view copy_stream, + cuda::stream_ref copy_stream, bool prefetch = false, bool initialize = true, bool host_writeback = false) @@ -1025,7 +1024,7 @@ class batch_load_iterator_dyn { IdxT n_rows, IdxT row_width, size_type batch_size, - rmm::cuda_stream_view copy_stream, + cuda::stream_ref copy_stream, rmm::device_async_resource_ref mr, bool prefetch = false, bool initialize = true, @@ -1050,7 +1049,7 @@ class batch_load_iterator_dyn { IdxT n_rows, IdxT row_width, size_type batch_size, - rmm::cuda_stream_view copy_stream, + cuda::stream_ref copy_stream, bool prefetch = false, bool initialize = true, bool host_writeback = false) @@ -1156,7 +1155,7 @@ class batch_load_iterator_dyn { IdxT n_rows, IdxT row_width, size_type batch_size, - rmm::cuda_stream_view copy_stream, + cuda::stream_ref copy_stream, rmm::device_async_resource_ref mr, bool prefetch, bool initialize, @@ -1219,7 +1218,7 @@ auto make_batch_load_iterator(raft::resources const& res, detail::type_identity_t n_rows, detail::type_identity_t row_width, size_t batch_size, - rmm::cuda_stream_view copy_stream, + cuda::stream_ref copy_stream, rmm::device_async_resource_ref mr, bool prefetch = false, bool initialize = true, @@ -1244,7 +1243,7 @@ auto make_batch_load_iterator(raft::resources const& res, detail::type_identity_t n_rows, detail::type_identity_t row_width, size_t batch_size, - rmm::cuda_stream_view copy_stream, + cuda::stream_ref copy_stream, bool prefetch = false, bool initialize = true, bool host_writeback = false) -> batch_load_iterator_dyn diff --git a/cpp/src/neighbors/detail/cagra/cagra_search.cuh b/cpp/src/neighbors/detail/cagra/cagra_search.cuh index 40645ffb09..362596e782 100644 --- a/cpp/src/neighbors/detail/cagra/cagra_search.cuh +++ b/cpp/src/neighbors/detail/cagra/cagra_search.cuh @@ -281,7 +281,6 @@ void search_main(raft::resources const& res, cuvs::spatial::knn::detail::utils::config::kDivisor; if (index.metric() == cuvs::distance::DistanceType::CosineExpanded) { - auto stream = raft::resource::get_cuda_stream(res); auto query_norms = raft::make_device_vector(res, queries.extent(0)); // first scale the queries and then compute norms diff --git a/cpp/src/neighbors/detail/cagra/compute_distance.hpp b/cpp/src/neighbors/detail/cagra/compute_distance.hpp index 55fe39d585..c1d0033a01 100644 --- a/cpp/src/neighbors/detail/cagra/compute_distance.hpp +++ b/cpp/src/neighbors/detail/cagra/compute_distance.hpp @@ -9,6 +9,7 @@ #include "hashmap.hpp" #include "utils.hpp" +#include #include #include #include @@ -211,9 +212,8 @@ struct dataset_descriptor_host { // Codebook type is determined by DataT for VPQ (always half for now) struct state { - using ready_t = std::tuple; - using init_f = - std::tuple, size_t>; + using ready_t = std::tuple; + using init_f = std::tuple, size_t>; std::mutex mutex; std::atomic ready; // Not sure if std::holds_alternative is thread-safe @@ -235,7 +235,7 @@ struct dataset_descriptor_host { RAFT_CUDA_TRY_NO_THROW(cudaEventDestroy(ready_event)); } - void eval(rmm::cuda_stream_view stream) + void eval(cuda::stream_ref stream) { std::lock_guard lock(mutex); if (std::holds_alternative(value)) { @@ -249,7 +249,7 @@ struct dataset_descriptor_host { } } - auto get(rmm::cuda_stream_view stream) -> dev_descriptor_t* + auto get(cuda::stream_ref stream) -> dev_descriptor_t* { if (!ready.load(std::memory_order_acquire)) { eval(stream); } // value is immutable at this point. @@ -288,12 +288,12 @@ struct dataset_descriptor_host { /** * Return the device pointer, possibly evaluating it in the given thread. */ - [[nodiscard]] auto dev_ptr(rmm::cuda_stream_view stream) const -> const dev_descriptor_t* + [[nodiscard]] auto dev_ptr(cuda::stream_ref stream) const -> const dev_descriptor_t* { return value_->get(stream); } - [[nodiscard]] auto dev_ptr(rmm::cuda_stream_view stream) -> dev_descriptor_t* + [[nodiscard]] auto dev_ptr(cuda::stream_ref stream) -> dev_descriptor_t* { return value_->get(stream); } diff --git a/cpp/src/neighbors/detail/cagra/compute_distance_standard-impl.cuh b/cpp/src/neighbors/detail/cagra/compute_distance_standard-impl.cuh index 06cdb375da..b9061f8796 100644 --- a/cpp/src/neighbors/detail/cagra/compute_distance_standard-impl.cuh +++ b/cpp/src/neighbors/detail/cagra/compute_distance_standard-impl.cuh @@ -6,6 +6,7 @@ #include "compute_distance_standard.hpp" +#include #include #include @@ -131,23 +132,23 @@ standard_descriptor_spec* dev_ptr, - rmm::cuda_stream_view stream) { - standard_dataset_descriptor_init_kernel - <<<1, 1, 0, stream.get()>>>(dev_ptr, ptr, size, dim, ld, dataset_norms); - RAFT_CUDA_TRY(cudaPeekAtLastError()); - }, - Metric, - DatasetBlockDim, - false, // is_vpq - 0, // pq_bits - 0}; // pq_len + return host_type{ + desc_type{ptr, size, dim, ld, dataset_norms}, + [=](dataset_descriptor_base_t* dev_ptr, cuda::stream_ref stream) { + standard_dataset_descriptor_init_kernel + <<<1, 1, 0, stream.get()>>>(dev_ptr, ptr, size, dim, ld, dataset_norms); + RAFT_CUDA_TRY(cudaPeekAtLastError()); + }, + Metric, + DatasetBlockDim, + false, // is_vpq + 0, // pq_bits + 0}; // pq_len } } // namespace cuvs::neighbors::cagra::detail diff --git a/cpp/src/neighbors/detail/cagra/compute_distance_vpq-impl.cuh b/cpp/src/neighbors/detail/cagra/compute_distance_vpq-impl.cuh index d51727be05..3dbd17859b 100644 --- a/cpp/src/neighbors/detail/cagra/compute_distance_vpq-impl.cuh +++ b/cpp/src/neighbors/detail/cagra/compute_distance_vpq-impl.cuh @@ -8,6 +8,7 @@ #include "compute_distance_vpq.hpp" #include "packed_type.hpp" +#include #include #include @@ -221,8 +222,7 @@ vpq_descriptor_spec* dev_ptr, - rmm::cuda_stream_view stream) { + [=](dataset_descriptor_base_t* dev_ptr, cuda::stream_ref stream) { vpq_dataset_descriptor_init_kernel #include #include #include @@ -38,8 +39,8 @@ namespace cuvs::neighbors::cagra::detail { template struct lightweight_uvector { private: - using raft_res_type = const raft::resources*; - using rmm_res_type = std::tuple; + using raft_res_type = const raft::resources*; + using rmm_res_type = std::tuple; static constexpr size_t kAlign = 256; std::variant res_; @@ -75,14 +76,14 @@ struct lightweight_uvector { size_ = new_size; } - void resize(size_t new_size, rmm::cuda_stream_view stream) + void resize(size_t new_size, cuda::stream_ref stream) { if (new_size == size_) { return; } if (std::holds_alternative(res_)) { auto& h = std::get(res_); res_ = rmm_res_type{raft::resource::get_workspace_resource_ref(*h), stream}; } else { - std::get(std::get(res_)) = stream; + std::get(std::get(res_)) = stream; } resize(new_size); } diff --git a/cpp/src/neighbors/detail/tiered_index.cuh b/cpp/src/neighbors/detail/tiered_index.cuh index 8db27d7d6c..c2f0f4517b 100644 --- a/cpp/src/neighbors/detail/tiered_index.cuh +++ b/cpp/src/neighbors/detail/tiered_index.cuh @@ -257,7 +257,6 @@ struct index_state { } // merge results from ann_index/bfknn together, translating the bfknn ids - auto stream = raft::resource::get_cuda_stream(res); int64_t host_translations[2] = {0, static_cast(ann_rows())}; auto device_translations = raft::make_device_vector(res, 2); raft::copy( diff --git a/cpp/src/neighbors/ivf_common.cu b/cpp/src/neighbors/ivf_common.cu index 7551ef1214..a7789a2ffa 100644 --- a/cpp/src/neighbors/ivf_common.cu +++ b/cpp/src/neighbors/ivf_common.cu @@ -5,6 +5,7 @@ #include "ivf_common.cuh" +#include #include #include @@ -53,7 +54,7 @@ void calc_chunk_indices::configured::operator()(const uint32_t* cluster_sizes, const uint32_t* clusters_to_probe, uint32_t* chunk_indices, uint32_t* n_samples, - rmm::cuda_stream_view stream) + cuda::stream_ref stream) { void* kernel = nullptr; switch (block_dim.x) { @@ -78,7 +79,7 @@ void calc_chunk_indices::configured::operator()(const uint32_t* cluster_sizes, void sort_cluster_sizes_descending(uint32_t* input, uint32_t* output, uint32_t n_lists, - rmm::cuda_stream_view stream, + cuda::stream_ref stream, rmm::device_async_resource_ref tmp_res) { int begin_bit = 0; diff --git a/cpp/src/neighbors/ivf_common.cuh b/cpp/src/neighbors/ivf_common.cuh index fb1513065d..c445e3b6b7 100644 --- a/cpp/src/neighbors/ivf_common.cuh +++ b/cpp/src/neighbors/ivf_common.cuh @@ -5,6 +5,7 @@ #pragma once +#include #include #include #include @@ -19,7 +20,7 @@ namespace cuvs::neighbors::ivf::detail { void sort_cluster_sizes_descending(uint32_t* input, uint32_t* output, uint32_t n_lists, - rmm::cuda_stream_view stream, + cuda::stream_ref stream, rmm::device_async_resource_ref tmp_res); /** @@ -57,7 +58,7 @@ struct calc_chunk_indices { const uint32_t* clusters_to_probe, uint32_t* chunk_indices, uint32_t* n_samples, - rmm::cuda_stream_view stream); + cuda::stream_ref stream); }; static inline auto configure(uint32_t n_probes, uint32_t n_queries) -> configured @@ -153,7 +154,7 @@ void postprocess_neighbors(IdxT* neighbors_out, // [n_queries, to uint32_t n_queries, uint32_t n_probes, uint32_t topk, - rmm::cuda_stream_view stream) + cuda::stream_ref stream) { constexpr int kPNThreads = 256; const int pn_blocks = raft::div_rounding_up_unsafe(n_queries * topk, kPNThreads); diff --git a/cpp/src/neighbors/ivf_flat/ivf_flat_build.cuh b/cpp/src/neighbors/ivf_flat/ivf_flat_build.cuh index a702212dab..294ae91f5b 100644 --- a/cpp/src/neighbors/ivf_flat/ivf_flat_build.cuh +++ b/cpp/src/neighbors/ivf_flat/ivf_flat_build.cuh @@ -45,8 +45,6 @@ namespace detail { template auto clone(const raft::resources& res, const index& source) -> index { - auto stream = raft::resource::get_cuda_stream(res); - // Allocate the new index index target(res, source.metric(), diff --git a/cpp/src/neighbors/ivf_flat/ivf_flat_interleaved_scan_explicit_inst.cuh b/cpp/src/neighbors/ivf_flat/ivf_flat_interleaved_scan_explicit_inst.cuh index 052b7bfe9a..3dd41c6fda 100644 --- a/cpp/src/neighbors/ivf_flat/ivf_flat_interleaved_scan_explicit_inst.cuh +++ b/cpp/src/neighbors/ivf_flat/ivf_flat_interleaved_scan_explicit_inst.cuh @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ @@ -8,6 +8,7 @@ #include "../detail/ann_utils.cuh" #include "ivf_flat_interleaved_scan_jit.cuh" #include +#include #include #include #include @@ -36,7 +37,7 @@ uint32_t* neighbors, \ float* distances, \ uint32_t& grid_dim_x, \ - rmm::cuda_stream_view stream, \ + cuda::stream_ref stream, \ const std::optional& metric_udf); #define COMMA , diff --git a/cpp/src/neighbors/ivf_flat/ivf_flat_interleaved_scan_ext.cuh b/cpp/src/neighbors/ivf_flat/ivf_flat_interleaved_scan_ext.cuh index 3a782822b4..199303d151 100644 --- a/cpp/src/neighbors/ivf_flat/ivf_flat_interleaved_scan_ext.cuh +++ b/cpp/src/neighbors/ivf_flat/ivf_flat_interleaved_scan_ext.cuh @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ @@ -11,6 +11,7 @@ #include #include "../detail/ann_utils.cuh" +#include #include #include #include @@ -35,7 +36,7 @@ void ivfflat_interleaved_scan(const index& index, uint32_t* neighbors, float* distances, uint32_t& grid_dim_x, - rmm::cuda_stream_view stream, + cuda::stream_ref stream, const std::optional& metric_udf) RAFT_EXPLICIT; #define CUVS_INST_IVF_FLAT_INTERLEAVED_SCAN(T, IdxT, SampleFilterT) \ @@ -58,7 +59,7 @@ void ivfflat_interleaved_scan(const index& index, uint32_t* neighbors, \ float* distances, \ uint32_t& grid_dim_x, \ - rmm::cuda_stream_view stream, \ + cuda::stream_ref stream, \ const std::optional& metric_udf); CUVS_INST_IVF_FLAT_INTERLEAVED_SCAN(float, int64_t, cuvs::neighbors::filtering::none_sample_filter); diff --git a/cpp/src/neighbors/ivf_flat/ivf_flat_interleaved_scan_jit.cuh b/cpp/src/neighbors/ivf_flat/ivf_flat_interleaved_scan_jit.cuh index cf4fface36..48aa62501c 100644 --- a/cpp/src/neighbors/ivf_flat/ivf_flat_interleaved_scan_jit.cuh +++ b/cpp/src/neighbors/ivf_flat/ivf_flat_interleaved_scan_jit.cuh @@ -25,7 +25,7 @@ #include // RAFT_CUDA_TRY #include -#include +#include namespace cuvs::neighbors::ivf_flat::detail { static constexpr int kThreadsPerBlock = 128; @@ -151,7 +151,7 @@ void launch_kernel(const index& index, uint32_t* neighbors, float* distances, uint32_t& grid_dim_x, - rmm::cuda_stream_view stream, + cuda::stream_ref stream, const std::optional& metric_udf) { RAFT_EXPECTS(Veclen == index.veclen(), @@ -435,7 +435,7 @@ void ivfflat_interleaved_scan(const index& index, uint32_t* neighbors, float* distances, uint32_t& grid_dim_x, - rmm::cuda_stream_view stream, + cuda::stream_ref stream, const std::optional& metric_udf) { const uint32_t n_probes_clamped = std::min(n_probes, index.n_lists()); diff --git a/cpp/src/neighbors/ivf_pq/detail/ivf_pq_compute_similarity_run_inst.cu.in b/cpp/src/neighbors/ivf_pq/detail/ivf_pq_compute_similarity_run_inst.cu.in index 05892b9333..30de9d57de 100644 --- a/cpp/src/neighbors/ivf_pq/detail/ivf_pq_compute_similarity_run_inst.cu.in +++ b/cpp/src/neighbors/ivf_pq/detail/ivf_pq_compute_similarity_run_inst.cu.in @@ -1,8 +1,9 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ +#include #include #include @@ -17,7 +18,7 @@ namespace cuvs::neighbors::ivf_pq::detail { template void cuvs::neighbors::ivf_pq::detail::compute_similarity_run( cuvs::neighbors::ivf_pq::detail::selected s, - rmm::cuda_stream_view stream, + cuda::stream_ref stream, uint32_t dim, uint32_t n_probes, uint32_t pq_dim, diff --git a/cpp/src/neighbors/ivf_pq/detail/ivf_pq_contiguous_list_data.cu b/cpp/src/neighbors/ivf_pq/detail/ivf_pq_contiguous_list_data.cu index 5fac491532..6900dd5e53 100644 --- a/cpp/src/neighbors/ivf_pq/detail/ivf_pq_contiguous_list_data.cu +++ b/cpp/src/neighbors/ivf_pq/detail/ivf_pq_contiguous_list_data.cu @@ -1,9 +1,10 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ #include "../ivf_pq_contiguous_list_data_impl.cuh" +#include #include namespace cuvs::neighbors::ivf_pq::detail { @@ -16,7 +17,7 @@ void unpack_contiguous_list_data( uint32_t pq_dim, std::variant offset_or_indices, uint32_t pq_bits, - rmm::cuda_stream_view stream) + cuda::stream_ref stream) { unpack_contiguous_list_data_impl( codes, list_data, n_rows, pq_dim, offset_or_indices, pq_bits, stream); @@ -31,7 +32,7 @@ void pack_contiguous_list_data( uint32_t pq_dim, std::variant offset_or_indices, uint32_t pq_bits, - rmm::cuda_stream_view stream) + cuda::stream_ref stream) { pack_contiguous_list_data_impl( list_data, codes, n_rows, pq_dim, offset_or_indices, pq_bits, stream); diff --git a/cpp/src/neighbors/ivf_pq/detail/ivf_pq_list_data.cu b/cpp/src/neighbors/ivf_pq/detail/ivf_pq_list_data.cu index 34a7244a67..1fae9b28d7 100644 --- a/cpp/src/neighbors/ivf_pq/detail/ivf_pq_list_data.cu +++ b/cpp/src/neighbors/ivf_pq/detail/ivf_pq_list_data.cu @@ -1,9 +1,10 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ #include "../ivf_pq_list_data_impl.cuh" +#include #include namespace cuvs::neighbors::ivf_pq::detail { @@ -13,7 +14,7 @@ void unpack_list_data(raft::device_matrix_view list_data, std::variant offset_or_indices, uint32_t pq_bits, - rmm::cuda_stream_view stream) + cuda::stream_ref stream) { unpack_list_data_impl(codes, list_data, offset_or_indices, pq_bits, stream); }; @@ -24,7 +25,7 @@ void pack_list_data(raft::device_mdspan codes, std::variant offset_or_indices, uint32_t pq_bits, - rmm::cuda_stream_view stream) + cuda::stream_ref stream) { pack_list_data_impl(list_data, codes, offset_or_indices, pq_bits, stream); }; diff --git a/cpp/src/neighbors/ivf_pq/ivf_pq_build.cuh b/cpp/src/neighbors/ivf_pq/ivf_pq_build.cuh index 78177b7acf..85024d2de2 100644 --- a/cpp/src/neighbors/ivf_pq/ivf_pq_build.cuh +++ b/cpp/src/neighbors/ivf_pq/ivf_pq_build.cuh @@ -50,7 +50,7 @@ #include #include -#include +#include #include #include @@ -235,7 +235,7 @@ auto calculate_offsets_and_indices(IdxT n_rows, const uint32_t* cluster_sizes, IdxT* cluster_offsets, IdxT* data_indices, - rmm::cuda_stream_view stream) -> uint32_t + cuda::stream_ref stream) -> uint32_t { auto exec_policy = rmm::exec_policy(stream); // Calculate the offsets @@ -304,7 +304,6 @@ void transpose_pq_centers(const raft::resources& handle, owning_impl* impl, const float* pq_centers_source) { - auto stream = raft::resource::get_cuda_stream(handle); auto extents = impl->pq_centers().extents(); static_assert(extents.rank() == 3); auto extents_source = @@ -962,8 +961,6 @@ void erase_list(raft::resources const& res, index* index, uint32_t label) template auto clone(const raft::resources& res, const index& source) -> index { - auto stream = raft::resource::get_cuda_stream(res); - // Create owning_impl directly to get mutable access for copying auto impl = std::make_unique>(res, source.metric(), diff --git a/cpp/src/neighbors/ivf_pq/ivf_pq_compute_similarity.cuh b/cpp/src/neighbors/ivf_pq/ivf_pq_compute_similarity.cuh index c55e3cdbc0..d313749840 100644 --- a/cpp/src/neighbors/ivf_pq/ivf_pq_compute_similarity.cuh +++ b/cpp/src/neighbors/ivf_pq/ivf_pq_compute_similarity.cuh @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2022-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ @@ -9,12 +9,12 @@ #include "ivf_pq_fp_8bit.cuh" // cuvs::neighbors::ivf_pq::detail::fp_8bit #include "ivf_pq_compute_similarity.hpp" // cuvs::neighbors::ivf_pq::detail::selected +#include // cuda::stream_ref #include #include // cuvs::distance::DistanceType #include #include // cuvs::neighbors::ivf_pq::codebook_gen #include // RAFT_WEAK_FUNCTION -#include // rmm::cuda_stream_view #include // __half @@ -28,7 +28,7 @@ auto RAFT_WEAK_FUNCTION is_local_topk_feasible(uint32_t k, uint32_t n_probes, ui template void compute_similarity_run(selected s, - rmm::cuda_stream_view stream, + cuda::stream_ref stream, uint32_t dim, uint32_t n_probes, uint32_t pq_dim, @@ -138,7 +138,7 @@ auto compute_similarity_select(const cudaDeviceProp& dev_props, \ extern template void cuvs::neighbors::ivf_pq::detail::compute_similarity_run( \ cuvs::neighbors::ivf_pq::detail::selected s, \ - rmm::cuda_stream_view stream, \ + cuda::stream_ref stream, \ uint32_t dim, \ uint32_t n_probes, \ uint32_t pq_dim, \ diff --git a/cpp/src/neighbors/ivf_pq/ivf_pq_compute_similarity_impl.cuh b/cpp/src/neighbors/ivf_pq/ivf_pq_compute_similarity_impl.cuh index 19f6bb6dc4..37bfa550d4 100644 --- a/cpp/src/neighbors/ivf_pq/ivf_pq_compute_similarity_impl.cuh +++ b/cpp/src/neighbors/ivf_pq/ivf_pq_compute_similarity_impl.cuh @@ -21,7 +21,7 @@ #include // raft::Pow2 #include // raft::TxN_t -#include // rmm::cuda_stream_view +#include // cuda::stream_ref namespace cuvs::neighbors::ivf_pq::detail { @@ -265,7 +265,7 @@ struct occupancy_t { template void compute_similarity_run(selected s, - rmm::cuda_stream_view stream, + cuda::stream_ref stream, uint32_t dim, uint32_t n_probes, uint32_t pq_dim, diff --git a/cpp/src/neighbors/ivf_pq/ivf_pq_contiguous_list_data.cuh b/cpp/src/neighbors/ivf_pq/ivf_pq_contiguous_list_data.cuh index 87c4c4ea79..59e5ff777c 100644 --- a/cpp/src/neighbors/ivf_pq/ivf_pq_contiguous_list_data.cuh +++ b/cpp/src/neighbors/ivf_pq/ivf_pq_contiguous_list_data.cuh @@ -1,10 +1,11 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ #pragma once #include +#include #include #include #include @@ -20,7 +21,7 @@ void unpack_contiguous_list_data( uint32_t pq_dim, std::variant offset_or_indices, uint32_t pq_bits, - rmm::cuda_stream_view stream); + cuda::stream_ref stream); template void unpack_contiguous_list_data(raft::resources const& res, @@ -53,7 +54,7 @@ void pack_contiguous_list_data( uint32_t pq_dim, std::variant offset_or_indices, uint32_t pq_bits, - rmm::cuda_stream_view stream); + cuda::stream_ref stream); template void pack_contiguous_list_data(raft::resources const& res, diff --git a/cpp/src/neighbors/ivf_pq/ivf_pq_contiguous_list_data_impl.cuh b/cpp/src/neighbors/ivf_pq/ivf_pq_contiguous_list_data_impl.cuh index a473324b60..0e347526b4 100644 --- a/cpp/src/neighbors/ivf_pq/ivf_pq_contiguous_list_data_impl.cuh +++ b/cpp/src/neighbors/ivf_pq/ivf_pq_contiguous_list_data_impl.cuh @@ -6,6 +6,7 @@ #pragma once #include "ivf_pq_codepacking.cuh" #include +#include #include #include #include @@ -72,7 +73,7 @@ inline void unpack_contiguous_list_data_impl( uint32_t pq_dim, std::variant offset_or_indices, uint32_t pq_bits, - rmm::cuda_stream_view stream) + cuda::stream_ref stream) { if (n_rows == 0) { return; } @@ -154,7 +155,7 @@ inline void pack_contiguous_list_data_impl( uint32_t pq_dim, std::variant offset_or_indices, uint32_t pq_bits, - rmm::cuda_stream_view stream) + cuda::stream_ref stream) { if (n_rows == 0) { return; } diff --git a/cpp/src/neighbors/ivf_pq/ivf_pq_list_data.hpp b/cpp/src/neighbors/ivf_pq/ivf_pq_list_data.hpp index ac6bbec662..7a06444732 100644 --- a/cpp/src/neighbors/ivf_pq/ivf_pq_list_data.hpp +++ b/cpp/src/neighbors/ivf_pq/ivf_pq_list_data.hpp @@ -1,10 +1,11 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ #pragma once #include +#include #include #include #include @@ -18,7 +19,7 @@ void unpack_list_data(raft::device_matrix_view list_data, std::variant offset_or_indices, uint32_t pq_bits, - rmm::cuda_stream_view stream); + cuda::stream_ref stream); void pack_list_data(raft::device_mdspan::list_extents, @@ -26,7 +27,7 @@ void pack_list_data(raft::device_mdspan codes, std::variant offset_or_indices, uint32_t pq_bits, - rmm::cuda_stream_view stream); + cuda::stream_ref stream); /** Unpack the list data; see the public interface for the api and usage. */ template diff --git a/cpp/src/neighbors/ivf_pq/ivf_pq_list_data_impl.cuh b/cpp/src/neighbors/ivf_pq/ivf_pq_list_data_impl.cuh index aad5f428b1..88240756af 100644 --- a/cpp/src/neighbors/ivf_pq/ivf_pq_list_data_impl.cuh +++ b/cpp/src/neighbors/ivf_pq/ivf_pq_list_data_impl.cuh @@ -5,6 +5,7 @@ #pragma once #include "ivf_pq_codepacking.cuh" +#include #include #include #include @@ -68,7 +69,7 @@ inline void unpack_list_data_impl( raft::row_major> list_data, std::variant offset_or_indices, uint32_t pq_bits, - rmm::cuda_stream_view stream) + cuda::stream_ref stream) { auto n_rows = codes.extent(0); if (n_rows == 0) { return; } @@ -143,7 +144,7 @@ inline void pack_list_data_impl( raft::device_matrix_view codes, std::variant offset_or_indices, uint32_t pq_bits, - rmm::cuda_stream_view stream) + cuda::stream_ref stream) { auto n_rows = codes.extent(0); if (n_rows == 0) { return; } diff --git a/cpp/src/neighbors/ivf_pq/ivf_pq_transform.cuh b/cpp/src/neighbors/ivf_pq/ivf_pq_transform.cuh index 0894287ade..38f16cabeb 100644 --- a/cpp/src/neighbors/ivf_pq/ivf_pq_transform.cuh +++ b/cpp/src/neighbors/ivf_pq/ivf_pq_transform.cuh @@ -120,7 +120,6 @@ void transform(raft::resources const& res, // The cluster centers in the index are stored padded, which is not acceptable by // the kmeans_balanced::predict. Thus, we need the restructuring raft::copy. - auto stream = raft::resource::get_cuda_stream(res); const auto n_clusters = index.n_lists(); auto cluster_centers = diff --git a/cpp/src/neighbors/ivf_rabitq/gpu_index/initializer_gpu.cuh b/cpp/src/neighbors/ivf_rabitq/gpu_index/initializer_gpu.cuh index d2bbbcf345..3d76778290 100644 --- a/cpp/src/neighbors/ivf_rabitq/gpu_index/initializer_gpu.cuh +++ b/cpp/src/neighbors/ivf_rabitq/gpu_index/initializer_gpu.cuh @@ -13,6 +13,7 @@ #include +#include #include #include #include @@ -81,7 +82,7 @@ class InitializerGPU { size_t D; // Dimension size_t K; // Num of Centroids raft::resources const& handle_; // reusable resource handle - rmm::cuda_stream_view stream_ = + cuda::stream_ref stream_ = raft::resource::get_cuda_stream(handle_); // CUDA stream obtained from handle_ }; diff --git a/cpp/src/neighbors/ivf_rabitq/gpu_index/ivf_gpu.cu b/cpp/src/neighbors/ivf_rabitq/gpu_index/ivf_gpu.cu index ca047b3246..1f24f3715c 100644 --- a/cpp/src/neighbors/ivf_rabitq/gpu_index/ivf_gpu.cu +++ b/cpp/src/neighbors/ivf_rabitq/gpu_index/ivf_gpu.cu @@ -12,6 +12,7 @@ #include "../utils/reductions.cuh" #include "ivf_gpu.cuh" #include "searcher_gpu.cuh" +#include #include #include @@ -1037,7 +1038,7 @@ void IVFGPU::PrepareClusterSearchInputs( raft::device_vector& d_G_kbxSumq) { raft::resources const& searcher_handle = searcher.get_handle(); - rmm::cuda_stream_view searcher_stream = searcher.get_stream(); + cuda::stream_ref searcher_stream = searcher.get_stream(); const size_t batch_size = queries.extent(0); // Compute ||q - c||^2 = -2 * q . c + ||q||^2 + ||c||^2 into centroid_distances: diff --git a/cpp/src/neighbors/ivf_rabitq/gpu_index/ivf_gpu.cuh b/cpp/src/neighbors/ivf_rabitq/gpu_index/ivf_gpu.cuh index 75ddaee865..0c0c1695d7 100644 --- a/cpp/src/neighbors/ivf_rabitq/gpu_index/ivf_gpu.cuh +++ b/cpp/src/neighbors/ivf_rabitq/gpu_index/ivf_gpu.cuh @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ @@ -20,7 +20,7 @@ #include #include -#include +#include #include @@ -342,7 +342,7 @@ class IVFGPU { void AllocateHostMemory(); raft::resources const& handle_; // reusable resource handle - rmm::cuda_stream_view stream_ = + cuda::stream_ref stream_ = raft::resource::get_cuda_stream(handle_); // CUDA stream obtained from handle_ // Device pointers for each data array. diff --git a/cpp/src/neighbors/ivf_rabitq/gpu_index/quantizer_gpu.cuh b/cpp/src/neighbors/ivf_rabitq/gpu_index/quantizer_gpu.cuh index 90ef71ed12..39e389dd9c 100644 --- a/cpp/src/neighbors/ivf_rabitq/gpu_index/quantizer_gpu.cuh +++ b/cpp/src/neighbors/ivf_rabitq/gpu_index/quantizer_gpu.cuh @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ @@ -11,6 +11,7 @@ #include "../defines.hpp" #include "rotator_gpu.cuh" +#include #include #include @@ -155,7 +156,7 @@ class DataQuantizerGPU { // RAFT resources raft::resources const& handle_; // reusable resource handle - rmm::cuda_stream_view stream_ = + cuda::stream_ref stream_ = raft::resource::get_cuda_stream(handle_); // CUDA stream obtained from handle_ // Device temporary buffers for quantization diff --git a/cpp/src/neighbors/ivf_rabitq/gpu_index/searcher_gpu.cu b/cpp/src/neighbors/ivf_rabitq/gpu_index/searcher_gpu.cu index af5dfddb61..1e8608da3f 100644 --- a/cpp/src/neighbors/ivf_rabitq/gpu_index/searcher_gpu.cu +++ b/cpp/src/neighbors/ivf_rabitq/gpu_index/searcher_gpu.cu @@ -16,6 +16,7 @@ #include "searcher_gpu.cuh" #include "searcher_gpu_common.cuh" +#include #include #include @@ -115,7 +116,7 @@ void launchPrecomputeLUTs(const float* d_query, float* d_lut_for_queries, size_t num_queries, size_t D, - rmm::cuda_stream_view stream) + cuda::stream_ref stream) { // Launch precompute kernel dim3 gridDim(num_queries, 1, 1); diff --git a/cpp/src/neighbors/ivf_rabitq/gpu_index/searcher_gpu.cuh b/cpp/src/neighbors/ivf_rabitq/gpu_index/searcher_gpu.cuh index 73557f57ea..3d5cd8c350 100644 --- a/cpp/src/neighbors/ivf_rabitq/gpu_index/searcher_gpu.cuh +++ b/cpp/src/neighbors/ivf_rabitq/gpu_index/searcher_gpu.cuh @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ @@ -17,7 +17,7 @@ #include #include -#include +#include namespace cuvs::neighbors::ivf_rabitq::detail { @@ -52,7 +52,7 @@ class SearcherGPU { // Getter methods std::string const& get_mode() { return mode_; } raft::resources const& get_handle() const { return handle_; } - rmm::cuda_stream_view get_stream() const { return stream_; } + cuda::stream_ref get_stream() const { return stream_; } float* get_centroid_distances() { return centroid_distances_.data_handle(); } float* get_q_norms() { return q_norms_.data_handle(); } @@ -110,7 +110,7 @@ class SearcherGPU { private: raft::resources const& handle_; // reusable resource handle - rmm::cuda_stream_view stream_ = + cuda::stream_ref stream_ = raft::resource::get_cuda_stream(handle_); // CUDA stream obtained from handle_ size_t D; // number of dimension const float* query_ = nullptr; // rotated query (non-owning) diff --git a/cpp/src/neighbors/ivf_sq/ivf_sq_search.cuh b/cpp/src/neighbors/ivf_sq/ivf_sq_search.cuh index 1fc32205df..716e7b733c 100644 --- a/cpp/src/neighbors/ivf_sq/ivf_sq_search.cuh +++ b/cpp/src/neighbors/ivf_sq/ivf_sq_search.cuh @@ -16,6 +16,7 @@ #include #include +#include #include #include #include @@ -149,7 +150,7 @@ void launch_kernel(const index& idx, float* out_distances, uint32_t* out_indices, uint32_t& grid_dim_x, - rmm::cuda_stream_view stream, + cuda::stream_ref stream, IvfSampleFilterT sample_filter) { static_assert(std::is_same_v, "IVF-SQ JIT-LTO scan only supports CodeT=uint8_t"); @@ -291,7 +292,7 @@ void ivf_sq_scan(raft::resources const& handle, uint32_t* out_indices, IvfSampleFilterT sample_filter, uint32_t& grid_dim_x, - rmm::cuda_stream_view stream) + cuda::stream_ref stream) { // Determine the fused top-k capacity (0 = disabled / fallback to materialization) int capacity = is_local_topk_feasible(k) ? raft::bound_by_power_of_two(int(k)) : 0; diff --git a/cpp/src/neighbors/scann/detail/scann_avq.cuh b/cpp/src/neighbors/scann/detail/scann_avq.cuh index b6d9a88989..2823f8f722 100644 --- a/cpp/src/neighbors/scann/detail/scann_avq.cuh +++ b/cpp/src/neighbors/scann/detail/scann_avq.cuh @@ -582,7 +582,6 @@ void apply_avq(raft::resources const& res, { // Compute clusters - cudaStream_t stream = raft::resource::get_cuda_stream(res).get(); auto cluster_offsets = raft::make_device_vector(res, centroids_view.extent(0)); auto clusters = raft::make_device_vector(res, dataset.extent(0)); int64_t max_cluster_size = 0; diff --git a/cpp/src/preprocessing/quantize/detail/scalar.cuh b/cpp/src/preprocessing/quantize/detail/scalar.cuh index 9b16b656d9..13f1a2adb2 100644 --- a/cpp/src/preprocessing/quantize/detail/scalar.cuh +++ b/cpp/src/preprocessing/quantize/detail/scalar.cuh @@ -168,8 +168,6 @@ void transform(raft::resources const& res, raft::device_matrix_view dataset, raft::device_matrix_view out) { - cudaStream_t stream = raft::resource::get_cuda_stream(res).get(); - raft::linalg::map(res, out, quantize_op(quantizer.min_, quantizer.max_), @@ -200,8 +198,6 @@ void inverse_transform(raft::resources const& res, raft::device_matrix_view dataset, raft::device_matrix_view out) { - cudaStream_t stream = raft::resource::get_cuda_stream(res).get(); - raft::linalg::map(res, out, quantize_op(quantizer.min_, quantizer.max_), diff --git a/cpp/tests/cluster/kmeans_mnmg.cu b/cpp/tests/cluster/kmeans_mnmg.cu index 3db41ee03f..4fddb99be6 100644 --- a/cpp/tests/cluster/kmeans_mnmg.cu +++ b/cpp/tests/cluster/kmeans_mnmg.cu @@ -468,9 +468,9 @@ class KmeansMGNcclTest : public ::testing::TestWithParam> raft::make_host_scalar_view(&pred_inertia_sg)); ari_vs_ref_ = raft::stats::adjusted_rand_index( - d_labels_ref.data(), d_labels_mg.data(), n_samples, sg_stream); + d_labels_ref.data(), d_labels_mg.data(), n_samples, sg_stream.get()); ari_vs_sg_ = raft::stats::adjusted_rand_index( - d_labels_sg.data(), d_labels_mg.data(), n_samples, sg_stream); + d_labels_sg.data(), d_labels_mg.data(), n_samples, sg_stream.get()); mg_inertia_ = mg_inertia; mg_n_iter_ = mg_n_iter; diff --git a/cpp/tests/neighbors/all_neighbors.cuh b/cpp/tests/neighbors/all_neighbors.cuh index 0b43023eff..8a5b435982 100644 --- a/cpp/tests/neighbors/all_neighbors.cuh +++ b/cpp/tests/neighbors/all_neighbors.cuh @@ -10,6 +10,7 @@ #include "ann_utils.cuh" #include "naive_knn.cuh" #include +#include #include #include #include @@ -259,7 +260,7 @@ class AllNeighborsTest : public ::testing::TestWithParam { private: raft::device_resources_snmg handle_; - rmm::cuda_stream_view stream_; + cuda::stream_ref stream_; AllNeighborsInputs ps; rmm::device_uvector database; }; diff --git a/cpp/tests/neighbors/ann_brute_force.cuh b/cpp/tests/neighbors/ann_brute_force.cuh index 3da3759804..c3feaaec69 100644 --- a/cpp/tests/neighbors/ann_brute_force.cuh +++ b/cpp/tests/neighbors/ann_brute_force.cuh @@ -11,6 +11,7 @@ #include +#include #include #include @@ -158,7 +159,7 @@ class AnnBruteForceTest : public ::testing::TestWithParam ps; rmm::device_uvector database; rmm::device_uvector search_queries; diff --git a/cpp/tests/neighbors/ann_cagra.cuh b/cpp/tests/neighbors/ann_cagra.cuh index 96ad03c4e9..381bddb989 100644 --- a/cpp/tests/neighbors/ann_cagra.cuh +++ b/cpp/tests/neighbors/ann_cagra.cuh @@ -7,6 +7,7 @@ #include "../test_utils.cuh" #include "ann_utils.cuh" #include "vpq_utils.cuh" +#include #include #include "cagra_padded_build_helpers.cuh" @@ -567,7 +568,7 @@ class AnnCagraTest : public ::testing::TestWithParam { private: raft::resources handle_; - rmm::cuda_stream_view stream_; + cuda::stream_ref stream_; AnnCagraInputs ps; rmm::device_uvector database; rmm::device_uvector search_queries; @@ -772,7 +773,7 @@ class AnnCagraAddNodesTest : public ::testing::TestWithParam { private: raft::resources handle_; - rmm::cuda_stream_view stream_; + cuda::stream_ref stream_; AnnCagraInputs ps; rmm::device_uvector database; rmm::device_uvector search_queries; @@ -1153,7 +1154,7 @@ class AnnCagraFilterTest : public ::testing::TestWithParam { private: raft::resources handle_; - rmm::cuda_stream_view stream_; + cuda::stream_ref stream_; AnnCagraInputs ps; rmm::device_uvector database; rmm::device_uvector search_queries; @@ -1402,7 +1403,7 @@ class AnnCagraIndexFilteredMergeTest : public ::testing::TestWithParam database; rmm::device_uvector search_queries; @@ -1650,7 +1651,7 @@ class AnnCagraIndexMergeTest : public ::testing::TestWithParam { private: raft::resources handle_; - rmm::cuda_stream_view stream_; + cuda::stream_ref stream_; AnnCagraInputs ps; rmm::device_uvector database; rmm::device_uvector search_queries; @@ -2420,7 +2421,7 @@ class AnnCagraMultiPartitionTest : public ::testing::TestWithParam database; rmm::device_uvector search_queries; diff --git a/cpp/tests/neighbors/ann_hnsw_ace.cuh b/cpp/tests/neighbors/ann_hnsw_ace.cuh index 03e9746b9b..f953da0c55 100644 --- a/cpp/tests/neighbors/ann_hnsw_ace.cuh +++ b/cpp/tests/neighbors/ann_hnsw_ace.cuh @@ -9,6 +9,7 @@ #include #include +#include #include #include @@ -660,7 +661,7 @@ class AnnHnswAceTest : public ::testing::TestWithParam { private: raft::resources handle_; - rmm::cuda_stream_view stream_; + cuda::stream_ref stream_; AnnHnswAceInputs ps; rmm::device_uvector database_dev; rmm::device_uvector search_queries; diff --git a/cpp/tests/neighbors/ann_ivf_flat.cuh b/cpp/tests/neighbors/ann_ivf_flat.cuh index a6f72e6d84..47a721bf2f 100644 --- a/cpp/tests/neighbors/ann_ivf_flat.cuh +++ b/cpp/tests/neighbors/ann_ivf_flat.cuh @@ -8,6 +8,7 @@ #include "ann_utils.cuh" #include "naive_knn.cuh" +#include #include #include #include @@ -515,7 +516,7 @@ class AnnIVFFlatTest : public ::testing::TestWithParam> { private: raft::resources handle_; - rmm::cuda_stream_view stream_; + cuda::stream_ref stream_; AnnIvfFlatInputs ps; rmm::device_uvector database; rmm::device_uvector search_queries; diff --git a/cpp/tests/neighbors/ann_ivf_pq.cuh b/cpp/tests/neighbors/ann_ivf_pq.cuh index 20e5231b71..85ec62678f 100644 --- a/cpp/tests/neighbors/ann_ivf_pq.cuh +++ b/cpp/tests/neighbors/ann_ivf_pq.cuh @@ -11,6 +11,7 @@ #include #include +#include #include #include #include @@ -705,7 +706,7 @@ class ivf_pq_test : public ::testing::TestWithParam { private: raft::resources handle_; - rmm::cuda_stream_view stream_; + cuda::stream_ref stream_; ivf_pq_inputs ps; // NOLINT rmm::device_uvector database; // NOLINT rmm::device_uvector search_queries; // NOLINT @@ -867,7 +868,7 @@ class ivf_pq_filter_test : public ::testing::TestWithParam { private: raft::resources handle_; - rmm::cuda_stream_view stream_; + cuda::stream_ref stream_; ivf_pq_inputs ps; // NOLINT rmm::device_uvector database; // NOLINT rmm::device_uvector search_queries; // NOLINT diff --git a/cpp/tests/neighbors/ann_ivf_rabitq.cuh b/cpp/tests/neighbors/ann_ivf_rabitq.cuh index 938f41f846..cdf0165b15 100644 --- a/cpp/tests/neighbors/ann_ivf_rabitq.cuh +++ b/cpp/tests/neighbors/ann_ivf_rabitq.cuh @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ #pragma once @@ -7,6 +7,7 @@ #include "../test_utils.cuh" #include "ann_utils.cuh" #include "naive_knn.cuh" +#include #include #include @@ -247,7 +248,7 @@ class ivf_rabitq_test : public ::testing::TestWithParam { private: raft::resources handle_; - rmm::cuda_stream_view stream_; + cuda::stream_ref stream_; ivf_rabitq_inputs ps; // NOLINT rmm::device_uvector database; // NOLINT rmm::device_uvector search_queries; // NOLINT diff --git a/cpp/tests/neighbors/ann_ivf_sq.cuh b/cpp/tests/neighbors/ann_ivf_sq.cuh index 2bd800493a..5299efd920 100644 --- a/cpp/tests/neighbors/ann_ivf_sq.cuh +++ b/cpp/tests/neighbors/ann_ivf_sq.cuh @@ -8,6 +8,7 @@ #include "ann_utils.cuh" #include "naive_knn.cuh" +#include #include #include #include @@ -373,7 +374,7 @@ class AnnIVFSQTest : public ::testing::TestWithParam> { } raft::resources handle_; - rmm::cuda_stream_view stream_; + cuda::stream_ref stream_; AnnIvfSqInputs ps; rmm::device_uvector database; rmm::device_uvector search_queries; diff --git a/cpp/tests/neighbors/ann_nn_descent.cuh b/cpp/tests/neighbors/ann_nn_descent.cuh index a4474a9382..7fd03b2f43 100644 --- a/cpp/tests/neighbors/ann_nn_descent.cuh +++ b/cpp/tests/neighbors/ann_nn_descent.cuh @@ -14,6 +14,7 @@ #include #include +#include #include #include #include @@ -181,7 +182,7 @@ class AnnNNDescentTest : public ::testing::TestWithParam { private: raft::resources handle_; - rmm::cuda_stream_view stream_; + cuda::stream_ref stream_; AnnNNDescentInputs ps; rmm::device_uvector database; }; @@ -339,7 +340,7 @@ class AnnNNDescentDistEpiTest : public ::testing::TestWithParam database; }; @@ -457,7 +458,7 @@ class AnnNNDescentBatchTest : public ::testing::TestWithParam database; }; diff --git a/cpp/tests/neighbors/ann_scann.cuh b/cpp/tests/neighbors/ann_scann.cuh index 81ef21c8e2..8f6deb30de 100644 --- a/cpp/tests/neighbors/ann_scann.cuh +++ b/cpp/tests/neighbors/ann_scann.cuh @@ -12,6 +12,7 @@ #include #include +#include #include #include #include @@ -282,7 +283,7 @@ class scann_test : public ::testing::TestWithParam { private: raft::resources handle_; - rmm::cuda_stream_view stream_; + cuda::stream_ref stream_; scann_inputs ps; // NOLINT rmm::device_uvector database; // NOLINT }; diff --git a/cpp/tests/neighbors/ann_vamana.cuh b/cpp/tests/neighbors/ann_vamana.cuh index d1b9b34864..ada130634a 100644 --- a/cpp/tests/neighbors/ann_vamana.cuh +++ b/cpp/tests/neighbors/ann_vamana.cuh @@ -7,6 +7,7 @@ #include "../test_utils.cuh" #include "ann_utils.cuh" +#include #include #include "cagra_padded_build_helpers.cuh" @@ -299,7 +300,7 @@ class AnnVamanaTest : public ::testing::TestWithParam { private: raft::resources handle_; - rmm::cuda_stream_view stream_; + cuda::stream_ref stream_; AnnVamanaInputs ps; rmm::device_uvector database; rmm::device_uvector search_queries; diff --git a/cpp/tests/neighbors/distance_nn.cu b/cpp/tests/neighbors/distance_nn.cu index 7c90c515d1..c7018b77c9 100644 --- a/cpp/tests/neighbors/distance_nn.cu +++ b/cpp/tests/neighbors/distance_nn.cu @@ -9,6 +9,7 @@ #include "../../src/distance/fused_distance_nn.cuh" #include "../../src/distance/unfused_distance_nn.cuh" +#include #include #include #include @@ -164,7 +165,7 @@ class NNTest : public ::testing::TestWithParam> { private: raft::resources handle; - rmm::cuda_stream_view stream; + cuda::stream_ref stream; NNInputs params_; ComparisonSummary summary; IdxT m; diff --git a/cpp/tests/neighbors/hnsw.cu b/cpp/tests/neighbors/hnsw.cu index ac86c19123..01a30950d0 100644 --- a/cpp/tests/neighbors/hnsw.cu +++ b/cpp/tests/neighbors/hnsw.cu @@ -8,6 +8,7 @@ #include "cagra_padded_build_helpers.cuh" #include +#include #include #include #include @@ -159,7 +160,7 @@ class AnnHNSWTest : public ::testing::TestWithParam { private: raft::resources handle_; - rmm::cuda_stream_view stream_; + cuda::stream_ref stream_; AnnHNSWInputs ps; rmm::device_uvector database; rmm::device_uvector queries; diff --git a/cpp/tests/neighbors/refine.cu b/cpp/tests/neighbors/refine.cu index bf8a2bd98e..1033fd6221 100644 --- a/cpp/tests/neighbors/refine.cu +++ b/cpp/tests/neighbors/refine.cu @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2022-2025, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ @@ -14,7 +14,7 @@ #include #include -#include +#include #include @@ -86,7 +86,7 @@ class RefineTest : public ::testing::TestWithParam> { public: raft::resources handle_; - rmm::cuda_stream_view stream_; + cuda::stream_ref stream_; RefineHelper data; }; diff --git a/cpp/tests/neighbors/refine_helper.cuh b/cpp/tests/neighbors/refine_helper.cuh index 2d82021d34..1680d9b5cf 100644 --- a/cpp/tests/neighbors/refine_helper.cuh +++ b/cpp/tests/neighbors/refine_helper.cuh @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2022-2024, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ #pragma once @@ -15,7 +15,7 @@ #include "naive_knn.cuh" -#include +#include #include namespace cuvs::neighbors { @@ -127,7 +127,7 @@ class RefineHelper { public: RefineInputs p; const raft::resources& handle_; - rmm::cuda_stream_view stream_; + cuda::stream_ref stream_; raft::device_matrix dataset; raft::device_matrix queries; diff --git a/cpp/tests/neighbors/tiered_index.cu b/cpp/tests/neighbors/tiered_index.cu index 4d1fd9c18d..5c71c31319 100644 --- a/cpp/tests/neighbors/tiered_index.cu +++ b/cpp/tests/neighbors/tiered_index.cu @@ -7,6 +7,7 @@ #include "ann_utils.cuh" #include +#include #include #include #include @@ -220,7 +221,7 @@ class ANNTieredIndexTest : public ::testing::TestWithParam private: raft::resources handle_; - rmm::cuda_stream_view stream_; + cuda::stream_ref stream_; AnnTieredIndexInputs ps; rmm::device_uvector database; rmm::device_uvector queries; diff --git a/cpp/tests/sparse/neighbors/cross_component_nn.cu b/cpp/tests/sparse/neighbors/cross_component_nn.cu index bf658a8335..76fdafff1d 100644 --- a/cpp/tests/sparse/neighbors/cross_component_nn.cu +++ b/cpp/tests/sparse/neighbors/cross_component_nn.cu @@ -388,8 +388,6 @@ class ConnectComponentsEdgesTest { raft::resources handle; - auto stream = raft::resource::get_cuda_stream(handle); - params = ::testing::TestWithParam< ConnectComponentsMutualReachabilityInputs>::GetParam(); diff --git a/examples/cpp/src/dynamic_batching_example.cu b/examples/cpp/src/dynamic_batching_example.cu index 0a2115eebe..2745b5e4a7 100644 --- a/examples/cpp/src/dynamic_batching_example.cu +++ b/examples/cpp/src/dynamic_batching_example.cu @@ -57,8 +57,9 @@ struct cuda_work_completion_promise { cuda_work_completion_promise(const raft::resources& res) { auto* promise = new std::promise; - RAFT_CUDA_TRY(cudaLaunchHostFunc( - raft::resource::get_cuda_stream(res), completion_callback, reinterpret_cast(promise))); + RAFT_CUDA_TRY(cudaLaunchHostFunc(raft::resource::get_cuda_stream(res).get(), + completion_callback, + reinterpret_cast(promise))); value_ = promise->get_future(); } diff --git a/fern/pages/cpp_api/cpp-api-common-types-copy-serialization-and-utility-apis.md b/fern/pages/cpp_api/cpp-api-common-types-copy-serialization-and-utility-apis.md index eeffd82ae6..7349cf13d7 100644 --- a/fern/pages/cpp_api/cpp-api-common-types-copy-serialization-and-utility-apis.md +++ b/fern/pages/cpp_api/cpp-api-common-types-copy-serialization-and-utility-apis.md @@ -14,7 +14,7 @@ Asynchronously copies elements between compatible memory locations. ```cpp template void copy(OutputIterator dst, InputIterator src, SizeType n, - rmm::cuda_stream_view stream); + cuda::stream_ref stream); ``` **Parameters** @@ -24,7 +24,7 @@ void copy(OutputIterator dst, InputIterator src, SizeType n, | `dst` | `OutputIterator` | Destination pointer or iterator. | | `src` | `InputIterator` | Source pointer or iterator. | | `n` | `SizeType` | Number of elements to copy. | -| `stream` | `rmm::cuda_stream_view` | CUDA stream used for the copy. | +| `stream` | `cuda::stream_ref` | CUDA stream used for the copy. | **Returns** @@ -39,7 +39,7 @@ Copies a dense matrix between compatible matrix views. ```cpp template -void copy_matrix(OutputView dst, InputView src, rmm::cuda_stream_view stream); +void copy_matrix(OutputView dst, InputView src, cuda::stream_ref stream); ``` **Parameters** @@ -48,7 +48,7 @@ void copy_matrix(OutputView dst, InputView src, rmm::cuda_stream_view stream); | --- | --- | --- | | `dst` | `OutputView` | Destination matrix view. | | `src` | `InputView` | Source matrix view. | -| `stream` | `rmm::cuda_stream_view` | CUDA stream used for the copy. | +| `stream` | `cuda::stream_ref` | CUDA stream used for the copy. | **Returns** @@ -64,7 +64,7 @@ Convenience helper for copying host data to device memory. ```cpp template void update_device(DevicePointer dst, HostPointer src, SizeType n, - rmm::cuda_stream_view stream); + cuda::stream_ref stream); ``` **Parameters** @@ -74,7 +74,7 @@ void update_device(DevicePointer dst, HostPointer src, SizeType n, | `dst` | `DevicePointer` | Destination device pointer. | | `src` | `HostPointer` | Source host pointer. | | `n` | `SizeType` | Number of elements to copy. | -| `stream` | `rmm::cuda_stream_view` | CUDA stream used for the copy. | +| `stream` | `cuda::stream_ref` | CUDA stream used for the copy. | **Returns** @@ -90,7 +90,7 @@ Convenience helper for copying device data to host memory. ```cpp template void update_host(HostPointer dst, DevicePointer src, SizeType n, - rmm::cuda_stream_view stream); + cuda::stream_ref stream); ``` **Parameters** @@ -100,7 +100,7 @@ void update_host(HostPointer dst, DevicePointer src, SizeType n, | `dst` | `HostPointer` | Destination host pointer. | | `src` | `DevicePointer` | Source device pointer. | | `n` | `SizeType` | Number of elements to copy. | -| `stream` | `rmm::cuda_stream_view` | CUDA stream used for the copy. | +| `stream` | `cuda::stream_ref` | CUDA stream used for the copy. | **Returns** diff --git a/fern/pages/cpp_api/cpp-api-common-types-execution-resources.md b/fern/pages/cpp_api/cpp-api-common-types-execution-resources.md index 90fa3eba43..dbc603a6a8 100644 --- a/fern/pages/cpp_api/cpp-api-common-types-execution-resources.md +++ b/fern/pages/cpp_api/cpp-api-common-types-execution-resources.md @@ -25,7 +25,7 @@ _Source header: `raft/core/resource/cuda_stream.hpp`_ Returns the CUDA stream associated with a resources object. ```cpp -rmm::cuda_stream_view get_cuda_stream(raft::resources const& res); +cuda::stream_ref get_cuda_stream(raft::resources const& res); ``` **Parameters** @@ -36,7 +36,7 @@ rmm::cuda_stream_view get_cuda_stream(raft::resources const& res); **Returns** -`rmm::cuda_stream_view` +`cuda::stream_ref` #### raft::resource::sync_stream @@ -47,7 +47,7 @@ Synchronizes the CUDA stream associated with a resources object. ```cpp void sync_stream(raft::resources const& res); -void sync_stream(raft::resources const& res, rmm::cuda_stream_view stream); +void sync_stream(raft::resources const& res, cuda::stream_ref stream); ``` **Parameters** @@ -55,7 +55,7 @@ void sync_stream(raft::resources const& res, rmm::cuda_stream_view stream); | Name | Type | Description | | --- | --- | --- | | `res` | `raft::resources const&` | Resources object to synchronize. | -| `stream` | `rmm::cuda_stream_view` | Optional stream to synchronize instead of the main stream. | +| `stream` | `cuda::stream_ref` | Optional stream to synchronize instead of the main stream. | **Returns** @@ -92,7 +92,7 @@ _Source header: `raft/core/resource/cuda_stream_pool.hpp`_ Returns a stream from the configured stream pool. ```cpp -rmm::cuda_stream_view get_stream_from_stream_pool(raft::resources const& res); +cuda::stream_ref get_stream_from_stream_pool(raft::resources const& res); ``` **Parameters** @@ -103,7 +103,7 @@ rmm::cuda_stream_view get_stream_from_stream_pool(raft::resources const& res); **Returns** -`rmm::cuda_stream_view` +`cuda::stream_ref` #### raft::resource::sync_stream_pool @@ -210,7 +210,7 @@ Constructs a single-GPU resources object. ```cpp device_resources( - rmm::cuda_stream_view stream_view = cuda::stream_ref{cudaStreamPerThread}, + cuda::stream_ref stream_view = cuda::stream_ref{cudaStreamPerThread}, std::shared_ptr stream_pool = nullptr, std::shared_ptr workspace_resource = nullptr, std::optional allocation_limit = std::nullopt); @@ -220,7 +220,7 @@ device_resources( | Name | Type | Description | | --- | --- | --- | -| `stream_view` | `rmm::cuda_stream_view` | Default CUDA stream used by algorithms. | +| `stream_view` | `cuda::stream_ref` | Default CUDA stream used by algorithms. | | `stream_pool` | `std::shared_ptr` | Optional CUDA stream pool. | | `workspace_resource` | `std::shared_ptr` | Optional workspace memory resource. | | `allocation_limit` | `std::optional` | Optional temporary workspace allocation limit in bytes. | @@ -232,14 +232,14 @@ Synchronizes either the main stream or a specific CUDA stream. ```cpp void sync_stream() const; -void sync_stream(rmm::cuda_stream_view stream) const; +void sync_stream(cuda::stream_ref stream) const; ``` **Parameters** | Name | Type | Description | | --- | --- | --- | -| `stream` | `rmm::cuda_stream_view` | Stream to synchronize. Omit to synchronize the main stream. | +| `stream` | `cuda::stream_ref` | Stream to synchronize. Omit to synchronize the main stream. | **Returns** @@ -251,12 +251,12 @@ void sync_stream(rmm::cuda_stream_view stream) const; Returns the main CUDA stream associated with the resources object. ```cpp -rmm::cuda_stream_view get_stream() const; +cuda::stream_ref get_stream() const; ``` **Returns** -`rmm::cuda_stream_view` +`cuda::stream_ref` #### raft::device_resources::is_stream_pool_initialized @@ -290,8 +290,8 @@ rmm::cuda_stream_pool const& get_stream_pool() const; Returns a stream from the configured CUDA stream pool. ```cpp -rmm::cuda_stream_view get_stream_from_stream_pool() const; -rmm::cuda_stream_view get_stream_from_stream_pool(std::size_t stream_idx) const; +cuda::stream_ref get_stream_from_stream_pool() const; +cuda::stream_ref get_stream_from_stream_pool(std::size_t stream_idx) const; ``` **Parameters** @@ -302,7 +302,7 @@ rmm::cuda_stream_view get_stream_from_stream_pool(std::size_t stream_idx) const; **Returns** -`rmm::cuda_stream_view` +`cuda::stream_ref` #### raft::device_resources::get_next_usable_stream @@ -310,8 +310,8 @@ rmm::cuda_stream_view get_stream_from_stream_pool(std::size_t stream_idx) const; Returns a stream from the pool when one exists; otherwise returns the main stream. ```cpp -rmm::cuda_stream_view get_next_usable_stream() const; -rmm::cuda_stream_view get_next_usable_stream(std::size_t stream_idx) const; +cuda::stream_ref get_next_usable_stream() const; +cuda::stream_ref get_next_usable_stream(std::size_t stream_idx) const; ``` **Parameters** @@ -322,7 +322,7 @@ rmm::cuda_stream_view get_next_usable_stream(std::size_t stream_idx) const; **Returns** -`rmm::cuda_stream_view` +`cuda::stream_ref` #### raft::device_resources::sync_stream_pool diff --git a/fern/pages/other/resources.md b/fern/pages/other/resources.md index 9cc45196b5..5b3c808e95 100644 --- a/fern/pages/other/resources.md +++ b/fern/pages/other/resources.md @@ -362,14 +362,14 @@ cudaStreamDestroy(stream); ```cpp #include -#include +#include #include cudaStream_t stream; cudaStreamCreate(&stream); -raft::device_resources resources{rmm::cuda_stream_view{stream}}; +raft::device_resources resources{cuda::stream_ref{stream}}; // cuVS C++ calls using resources are enqueued on stream. diff --git a/fern/scripts/generate_api_reference.py b/fern/scripts/generate_api_reference.py index 9a02eb1d0d..93a0b2b9dd 100755 --- a/fern/scripts/generate_api_reference.py +++ b/fern/scripts/generate_api_reference.py @@ -820,9 +820,9 @@ def add_symbol( "raft-resource-get-cuda-stream", "raft::resource::get_cuda_stream", "Returns the CUDA stream associated with a resources object.", - "rmm::cuda_stream_view get_cuda_stream(raft::resources const& res);", + "cuda::stream_ref get_cuda_stream(raft::resources const& res);", [("res", "raft::resources const&", "Resources object to query.")], - "rmm::cuda_stream_view", + "cuda::stream_ref", nested=True, ) add_symbol( @@ -832,7 +832,7 @@ def add_symbol( "Synchronizes the CUDA stream associated with a resources object.", ( "void sync_stream(raft::resources const& res);\n" - "void sync_stream(raft::resources const& res, rmm::cuda_stream_view stream);" + "void sync_stream(raft::resources const& res, cuda::stream_ref stream);" ), [ ( @@ -842,7 +842,7 @@ def add_symbol( ), ( "stream", - "rmm::cuda_stream_view", + "cuda::stream_ref", "Optional stream to synchronize instead of the main stream.", ), ], @@ -881,9 +881,9 @@ def add_symbol( "raft-resource-get-stream-from-stream-pool", "raft::resource::get_stream_from_stream_pool", "Returns a stream from the configured stream pool.", - "rmm::cuda_stream_view get_stream_from_stream_pool(raft::resources const& res);", + "cuda::stream_ref get_stream_from_stream_pool(raft::resources const& res);", [("res", "raft::resources const&", "Resources object to query.")], - "rmm::cuda_stream_view", + "cuda::stream_ref", nested=True, ) add_symbol( @@ -967,7 +967,7 @@ def add_symbol( "Constructs a single-GPU resources object.", ( "device_resources(\n" - " rmm::cuda_stream_view stream_view = cuda::stream_ref{cudaStreamPerThread},\n" + " cuda::stream_ref stream_view = cuda::stream_ref{cudaStreamPerThread},\n" " std::shared_ptr stream_pool = nullptr,\n" " std::shared_ptr workspace_resource = nullptr,\n" " std::optional allocation_limit = std::nullopt);" @@ -975,7 +975,7 @@ def add_symbol( [ ( "stream_view", - "rmm::cuda_stream_view", + "cuda::stream_ref", "Default CUDA stream used by algorithms.", ), ( @@ -1002,12 +1002,12 @@ def add_symbol( "Synchronizes either the main stream or a specific CUDA stream.", ( "void sync_stream() const;\n" - "void sync_stream(rmm::cuda_stream_view stream) const;" + "void sync_stream(cuda::stream_ref stream) const;" ), [ ( "stream", - "rmm::cuda_stream_view", + "cuda::stream_ref", "Stream to synchronize. Omit to synchronize the main stream.", ) ], @@ -1018,8 +1018,8 @@ def add_symbol( "raft-device-resources-get-stream", "raft::device_resources::get_stream", "Returns the main CUDA stream associated with the resources object.", - "rmm::cuda_stream_view get_stream() const;", - returns="rmm::cuda_stream_view", + "cuda::stream_ref get_stream() const;", + returns="cuda::stream_ref", ) add_symbol( lines, @@ -1043,8 +1043,8 @@ def add_symbol( "raft::device_resources::get_stream_from_stream_pool", "Returns a stream from the configured CUDA stream pool.", ( - "rmm::cuda_stream_view get_stream_from_stream_pool() const;\n" - "rmm::cuda_stream_view get_stream_from_stream_pool(std::size_t stream_idx) const;" + "cuda::stream_ref get_stream_from_stream_pool() const;\n" + "cuda::stream_ref get_stream_from_stream_pool(std::size_t stream_idx) const;" ), [ ( @@ -1053,7 +1053,7 @@ def add_symbol( "Optional index of the stream in the stream pool.", ) ], - "rmm::cuda_stream_view", + "cuda::stream_ref", ) add_symbol( lines, @@ -1064,8 +1064,8 @@ def add_symbol( "the main stream." ), ( - "rmm::cuda_stream_view get_next_usable_stream() const;\n" - "rmm::cuda_stream_view get_next_usable_stream(std::size_t stream_idx) const;" + "cuda::stream_ref get_next_usable_stream() const;\n" + "cuda::stream_ref get_next_usable_stream(std::size_t stream_idx) const;" ), [ ( @@ -1074,7 +1074,7 @@ def add_symbol( "Optional stream pool index to use when a stream pool is configured.", ) ], - "rmm::cuda_stream_view", + "cuda::stream_ref", ) add_symbol( lines, @@ -2105,7 +2105,7 @@ def add_symbol( ( "template \n" "void copy(OutputIterator dst, InputIterator src, SizeType n,\n" - " rmm::cuda_stream_view stream);" + " cuda::stream_ref stream);" ), [ ("dst", "OutputIterator", "Destination pointer or iterator."), @@ -2113,7 +2113,7 @@ def add_symbol( ("n", "SizeType", "Number of elements to copy."), ( "stream", - "rmm::cuda_stream_view", + "cuda::stream_ref", "CUDA stream used for the copy.", ), ], @@ -2124,13 +2124,13 @@ def add_symbol( "raft-copy-matrix", "raft::copy_matrix", "Copies a dense matrix between compatible matrix views.", - "template \nvoid copy_matrix(OutputView dst, InputView src, rmm::cuda_stream_view stream);", + "template \nvoid copy_matrix(OutputView dst, InputView src, cuda::stream_ref stream);", [ ("dst", "OutputView", "Destination matrix view."), ("src", "InputView", "Source matrix view."), ( "stream", - "rmm::cuda_stream_view", + "cuda::stream_ref", "CUDA stream used for the copy.", ), ], @@ -2144,7 +2144,7 @@ def add_symbol( ( "template \n" "void update_device(DevicePointer dst, HostPointer src, SizeType n,\n" - " rmm::cuda_stream_view stream);" + " cuda::stream_ref stream);" ), [ ("dst", "DevicePointer", "Destination device pointer."), @@ -2152,7 +2152,7 @@ def add_symbol( ("n", "SizeType", "Number of elements to copy."), ( "stream", - "rmm::cuda_stream_view", + "cuda::stream_ref", "CUDA stream used for the copy.", ), ], @@ -2166,7 +2166,7 @@ def add_symbol( ( "template \n" "void update_host(HostPointer dst, DevicePointer src, SizeType n,\n" - " rmm::cuda_stream_view stream);" + " cuda::stream_ref stream);" ), [ ("dst", "HostPointer", "Destination host pointer."), @@ -2174,7 +2174,7 @@ def add_symbol( ("n", "SizeType", "Number of elements to copy."), ( "stream", - "rmm::cuda_stream_view", + "cuda::stream_ref", "CUDA stream used for the copy.", ), ],