From 74186884f1e5665c0ffcacc3805b34f73dba36f1 Mon Sep 17 00:00:00 2001 From: Jake Awe Date: Thu, 17 Jul 2025 12:20:44 -0700 Subject: [PATCH 01/93] DOC v25.10 Updates [skip ci] --- .../cuda12.9-conda/devcontainer.json | 6 ++-- .devcontainer/cuda12.9-pip/devcontainer.json | 8 ++--- .github/workflows/build.yaml | 22 ++++++------- .github/workflows/pr.yaml | 32 +++++++++---------- .github/workflows/test.yaml | 10 +++--- .../trigger-breaking-change-alert.yaml | 2 +- README.md | 6 ++-- VERSION | 2 +- .../all_cuda-129_arch-aarch64.yaml | 4 +-- .../all_cuda-129_arch-x86_64.yaml | 4 +-- .../bench_ann_cuda-129_arch-aarch64.yaml | 8 ++--- .../bench_ann_cuda-129_arch-x86_64.yaml | 8 ++--- .../go_cuda-129_arch-aarch64.yaml | 4 +-- .../environments/go_cuda-129_arch-x86_64.yaml | 4 +-- .../rust_cuda-129_arch-aarch64.yaml | 4 +-- .../rust_cuda-129_arch-x86_64.yaml | 4 +-- dependencies.yaml | 24 +++++++------- docs/source/developer_guide.md | 4 +-- examples/go/README.md | 2 +- java/benchmarks/pom.xml | 4 +-- java/build.sh | 2 +- java/cuvs-java/pom.xml | 2 +- java/examples/README.md | 6 ++-- java/examples/pom.xml | 4 +-- python/cuvs/pyproject.toml | 10 +++--- python/cuvs_bench/cuvs_bench/plot/__main__.py | 2 +- python/cuvs_bench/pyproject.toml | 2 +- python/libcuvs/pyproject.toml | 8 ++--- rust/Cargo.toml | 2 +- rust/cuvs/Cargo.toml | 2 +- 30 files changed, 101 insertions(+), 101 deletions(-) diff --git a/.devcontainer/cuda12.9-conda/devcontainer.json b/.devcontainer/cuda12.9-conda/devcontainer.json index ce601856f1..ef5383cd80 100644 --- a/.devcontainer/cuda12.9-conda/devcontainer.json +++ b/.devcontainer/cuda12.9-conda/devcontainer.json @@ -5,17 +5,17 @@ "args": { "CUDA": "12.9", "PYTHON_PACKAGE_MANAGER": "conda", - "BASE": "rapidsai/devcontainers:25.08-cpp-mambaforge-ubuntu22.04" + "BASE": "rapidsai/devcontainers:25.10-cpp-mambaforge-ubuntu22.04" } }, "runArgs": [ "--rm", "--name", - "${localEnv:USER:anon}-rapids-${localWorkspaceFolderBasename}-25.08-cuda12.9-conda" + "${localEnv:USER:anon}-rapids-${localWorkspaceFolderBasename}-25.10-cuda12.9-conda" ], "hostRequirements": {"gpu": "optional"}, "features": { - "ghcr.io/rapidsai/devcontainers/features/rapids-build-utils:25.8": {} + "ghcr.io/rapidsai/devcontainers/features/rapids-build-utils:25.10": {} }, "overrideFeatureInstallOrder": [ "ghcr.io/rapidsai/devcontainers/features/rapids-build-utils" diff --git a/.devcontainer/cuda12.9-pip/devcontainer.json b/.devcontainer/cuda12.9-pip/devcontainer.json index be81ad6c64..5902c948f3 100644 --- a/.devcontainer/cuda12.9-pip/devcontainer.json +++ b/.devcontainer/cuda12.9-pip/devcontainer.json @@ -5,24 +5,24 @@ "args": { "CUDA": "12.9", "PYTHON_PACKAGE_MANAGER": "pip", - "BASE": "rapidsai/devcontainers:25.08-cpp-cuda12.9-ucx1.18.0-openmpi-ubuntu22.04" + "BASE": "rapidsai/devcontainers:25.10-cpp-cuda12.9-ucx1.18.0-openmpi-ubuntu22.04" } }, "runArgs": [ "--rm", "--name", - "${localEnv:USER:anon}-rapids-${localWorkspaceFolderBasename}-25.08-cuda12.9-pip" + "${localEnv:USER:anon}-rapids-${localWorkspaceFolderBasename}-25.10-cuda12.9-pip" ], "hostRequirements": {"gpu": "optional"}, "features": { - "ghcr.io/rapidsai/devcontainers/features/cuda:25.8": { + "ghcr.io/rapidsai/devcontainers/features/cuda:25.10": { "version": "12.9", "installcuBLAS": true, "installcuSOLVER": true, "installcuRAND": true, "installcuSPARSE": true }, - "ghcr.io/rapidsai/devcontainers/features/rapids-build-utils:25.8": {} + "ghcr.io/rapidsai/devcontainers/features/rapids-build-utils:25.10": {} }, "overrideFeatureInstallOrder": [ "ghcr.io/rapidsai/devcontainers/features/ucx", diff --git a/.github/workflows/build.yaml b/.github/workflows/build.yaml index 9c487f2bdf..d0836565a9 100644 --- a/.github/workflows/build.yaml +++ b/.github/workflows/build.yaml @@ -34,7 +34,7 @@ concurrency: jobs: cpp-build: secrets: inherit - uses: rapidsai/shared-workflows/.github/workflows/conda-cpp-build.yaml@branch-25.08 + uses: rapidsai/shared-workflows/.github/workflows/conda-cpp-build.yaml@branch-25.10 with: build_type: ${{ inputs.build_type || 'branch' }} branch: ${{ inputs.branch }} @@ -44,7 +44,7 @@ jobs: rust-build: needs: cpp-build secrets: inherit - uses: rapidsai/shared-workflows/.github/workflows/custom-job.yaml@branch-25.08 + uses: rapidsai/shared-workflows/.github/workflows/custom-job.yaml@branch-25.10 with: build_type: ${{ inputs.build_type || 'branch' }} branch: ${{ inputs.branch }} @@ -57,7 +57,7 @@ jobs: go-build: needs: cpp-build secrets: inherit - uses: rapidsai/shared-workflows/.github/workflows/custom-job.yaml@branch-25.08 + uses: rapidsai/shared-workflows/.github/workflows/custom-job.yaml@branch-25.10 with: build_type: ${{ inputs.build_type || 'branch' }} branch: ${{ inputs.branch }} @@ -70,7 +70,7 @@ jobs: java-build: needs: cpp-build secrets: inherit - uses: rapidsai/shared-workflows/.github/workflows/custom-job.yaml@branch-25.08 + uses: rapidsai/shared-workflows/.github/workflows/custom-job.yaml@branch-25.10 with: build_type: ${{ inputs.build_type || 'branch' }} branch: ${{ inputs.branch }} @@ -83,7 +83,7 @@ jobs: python-build: needs: [cpp-build] secrets: inherit - uses: rapidsai/shared-workflows/.github/workflows/conda-python-build.yaml@branch-25.08 + uses: rapidsai/shared-workflows/.github/workflows/conda-python-build.yaml@branch-25.10 with: build_type: ${{ inputs.build_type || 'branch' }} branch: ${{ inputs.branch }} @@ -93,7 +93,7 @@ jobs: upload-conda: needs: [cpp-build, python-build] secrets: inherit - uses: rapidsai/shared-workflows/.github/workflows/conda-upload-packages.yaml@branch-25.08 + uses: rapidsai/shared-workflows/.github/workflows/conda-upload-packages.yaml@branch-25.10 with: build_type: ${{ inputs.build_type || 'branch' }} branch: ${{ inputs.branch }} @@ -108,7 +108,7 @@ jobs: if: github.ref_type == 'branch' needs: python-build secrets: inherit - uses: rapidsai/shared-workflows/.github/workflows/custom-job.yaml@branch-25.08 + uses: rapidsai/shared-workflows/.github/workflows/custom-job.yaml@branch-25.10 with: arch: "amd64" branch: ${{ inputs.branch }} @@ -120,7 +120,7 @@ jobs: sha: ${{ inputs.sha }} wheel-build-libcuvs: secrets: inherit - uses: rapidsai/shared-workflows/.github/workflows/wheels-build.yaml@branch-25.08 + uses: rapidsai/shared-workflows/.github/workflows/wheels-build.yaml@branch-25.10 with: build_type: ${{ inputs.build_type || 'branch' }} branch: ${{ inputs.branch }} @@ -134,7 +134,7 @@ jobs: wheel-publish-libcuvs: needs: wheel-build-libcuvs secrets: inherit - uses: rapidsai/shared-workflows/.github/workflows/wheels-publish.yaml@branch-25.08 + uses: rapidsai/shared-workflows/.github/workflows/wheels-publish.yaml@branch-25.10 with: build_type: ${{ inputs.build_type || 'branch' }} branch: ${{ inputs.branch }} @@ -145,7 +145,7 @@ jobs: wheel-build-cuvs: needs: wheel-build-libcuvs secrets: inherit - uses: rapidsai/shared-workflows/.github/workflows/wheels-build.yaml@branch-25.08 + uses: rapidsai/shared-workflows/.github/workflows/wheels-build.yaml@branch-25.10 with: build_type: ${{ inputs.build_type || 'branch' }} branch: ${{ inputs.branch }} @@ -157,7 +157,7 @@ jobs: wheel-publish-cuvs: needs: wheel-build-cuvs secrets: inherit - uses: rapidsai/shared-workflows/.github/workflows/wheels-publish.yaml@branch-25.08 + uses: rapidsai/shared-workflows/.github/workflows/wheels-publish.yaml@branch-25.10 with: build_type: ${{ inputs.build_type || 'branch' }} branch: ${{ inputs.branch }} diff --git a/.github/workflows/pr.yaml b/.github/workflows/pr.yaml index 55d936b204..54c0144b0e 100644 --- a/.github/workflows/pr.yaml +++ b/.github/workflows/pr.yaml @@ -30,7 +30,7 @@ jobs: - devcontainer - telemetry-setup secrets: inherit - uses: rapidsai/shared-workflows/.github/workflows/pr-builder.yaml@branch-25.08 + uses: rapidsai/shared-workflows/.github/workflows/pr-builder.yaml@branch-25.10 if: always() with: needs: ${{ toJSON(needs) }} @@ -57,7 +57,7 @@ jobs: changed-files: needs: telemetry-setup secrets: inherit - uses: rapidsai/shared-workflows/.github/workflows/changed-files.yaml@branch-25.08 + uses: rapidsai/shared-workflows/.github/workflows/changed-files.yaml@branch-25.10 with: files_yaml: | test_cpp: @@ -109,14 +109,14 @@ jobs: checks: needs: telemetry-setup secrets: inherit - uses: rapidsai/shared-workflows/.github/workflows/checks.yaml@branch-25.08 + uses: rapidsai/shared-workflows/.github/workflows/checks.yaml@branch-25.10 with: enable_check_generated_files: false ignored_pr_jobs: "telemetry-summarize" conda-cpp-build: needs: checks secrets: inherit - uses: rapidsai/shared-workflows/.github/workflows/conda-cpp-build.yaml@branch-25.08 + uses: rapidsai/shared-workflows/.github/workflows/conda-cpp-build.yaml@branch-25.10 with: build_type: pull-request node_type: cpu16 @@ -124,7 +124,7 @@ jobs: conda-cpp-tests: needs: [conda-cpp-build, changed-files] secrets: inherit - uses: rapidsai/shared-workflows/.github/workflows/conda-cpp-tests.yaml@branch-25.08 + uses: rapidsai/shared-workflows/.github/workflows/conda-cpp-tests.yaml@branch-25.10 if: fromJSON(needs.changed-files.outputs.changed_file_groups).test_cpp with: build_type: pull-request @@ -132,21 +132,21 @@ jobs: conda-cpp-checks: needs: conda-cpp-build secrets: inherit - uses: rapidsai/shared-workflows/.github/workflows/conda-cpp-post-build-checks.yaml@branch-25.08 + uses: rapidsai/shared-workflows/.github/workflows/conda-cpp-post-build-checks.yaml@branch-25.10 with: build_type: pull-request symbol_exclusions: (void (thrust::|cub::)) conda-python-build: needs: conda-cpp-build secrets: inherit - uses: rapidsai/shared-workflows/.github/workflows/conda-python-build.yaml@branch-25.08 + uses: rapidsai/shared-workflows/.github/workflows/conda-python-build.yaml@branch-25.10 with: build_type: pull-request script: ci/build_python.sh conda-python-tests: needs: [conda-python-build, changed-files] secrets: inherit - uses: rapidsai/shared-workflows/.github/workflows/conda-python-tests.yaml@branch-25.08 + uses: rapidsai/shared-workflows/.github/workflows/conda-python-tests.yaml@branch-25.10 if: fromJSON(needs.changed-files.outputs.changed_file_groups).test_python with: build_type: pull-request @@ -154,7 +154,7 @@ jobs: conda-java-tests: needs: [conda-cpp-build, changed-files] secrets: inherit - uses: rapidsai/shared-workflows/.github/workflows/custom-job.yaml@branch-25.08 + uses: rapidsai/shared-workflows/.github/workflows/custom-job.yaml@branch-25.10 if: fromJSON(needs.changed-files.outputs.changed_file_groups).test_java with: build_type: pull-request @@ -166,7 +166,7 @@ jobs: docs-build: needs: conda-python-build secrets: inherit - uses: rapidsai/shared-workflows/.github/workflows/custom-job.yaml@branch-25.08 + uses: rapidsai/shared-workflows/.github/workflows/custom-job.yaml@branch-25.10 with: build_type: pull-request node_type: "gpu-l4-latest-1" @@ -176,7 +176,7 @@ jobs: rust-build: needs: conda-cpp-build secrets: inherit - uses: rapidsai/shared-workflows/.github/workflows/custom-job.yaml@branch-25.08 + uses: rapidsai/shared-workflows/.github/workflows/custom-job.yaml@branch-25.10 with: build_type: pull-request node_type: "gpu-l4-latest-1" @@ -186,7 +186,7 @@ jobs: go-build: needs: conda-cpp-build secrets: inherit - uses: rapidsai/shared-workflows/.github/workflows/custom-job.yaml@branch-25.08 + uses: rapidsai/shared-workflows/.github/workflows/custom-job.yaml@branch-25.10 with: build_type: pull-request node_type: "gpu-l4-latest-1" @@ -196,7 +196,7 @@ jobs: wheel-build-libcuvs: needs: checks secrets: inherit - uses: rapidsai/shared-workflows/.github/workflows/wheels-build.yaml@branch-25.08 + uses: rapidsai/shared-workflows/.github/workflows/wheels-build.yaml@branch-25.10 with: build_type: pull-request script: ci/build_wheel_libcuvs.sh @@ -207,7 +207,7 @@ jobs: wheel-build-cuvs: needs: wheel-build-libcuvs secrets: inherit - uses: rapidsai/shared-workflows/.github/workflows/wheels-build.yaml@branch-25.08 + uses: rapidsai/shared-workflows/.github/workflows/wheels-build.yaml@branch-25.10 with: build_type: pull-request script: ci/build_wheel_cuvs.sh @@ -216,7 +216,7 @@ jobs: wheel-tests-cuvs: needs: [wheel-build-cuvs, changed-files] secrets: inherit - uses: rapidsai/shared-workflows/.github/workflows/wheels-test.yaml@branch-25.08 + uses: rapidsai/shared-workflows/.github/workflows/wheels-test.yaml@branch-25.10 if: fromJSON(needs.changed-files.outputs.changed_file_groups).test_python with: build_type: pull-request @@ -224,7 +224,7 @@ jobs: devcontainer: secrets: inherit needs: telemetry-setup - uses: rapidsai/shared-workflows/.github/workflows/build-in-devcontainer.yaml@branch-25.08 + uses: rapidsai/shared-workflows/.github/workflows/build-in-devcontainer.yaml@branch-25.10 with: arch: '["amd64"]' cuda: '["12.9"]' diff --git a/.github/workflows/test.yaml b/.github/workflows/test.yaml index 7eaf54f7e5..64153d1809 100644 --- a/.github/workflows/test.yaml +++ b/.github/workflows/test.yaml @@ -25,7 +25,7 @@ on: jobs: conda-cpp-checks: secrets: inherit - uses: rapidsai/shared-workflows/.github/workflows/conda-cpp-post-build-checks.yaml@branch-25.08 + uses: rapidsai/shared-workflows/.github/workflows/conda-cpp-post-build-checks.yaml@branch-25.10 with: build_type: ${{ inputs.build_type }} branch: ${{ inputs.branch }} @@ -34,7 +34,7 @@ jobs: symbol_exclusions: (void (thrust::|cub::)) conda-cpp-tests: secrets: inherit - uses: rapidsai/shared-workflows/.github/workflows/conda-cpp-tests.yaml@branch-25.08 + uses: rapidsai/shared-workflows/.github/workflows/conda-cpp-tests.yaml@branch-25.10 with: build_type: ${{ inputs.build_type }} branch: ${{ inputs.branch }} @@ -43,7 +43,7 @@ jobs: sha: ${{ inputs.sha }} conda-python-tests: secrets: inherit - uses: rapidsai/shared-workflows/.github/workflows/conda-python-tests.yaml@branch-25.08 + uses: rapidsai/shared-workflows/.github/workflows/conda-python-tests.yaml@branch-25.10 with: build_type: ${{ inputs.build_type }} branch: ${{ inputs.branch }} @@ -52,7 +52,7 @@ jobs: sha: ${{ inputs.sha }} conda-java-tests: secrets: inherit - uses: rapidsai/shared-workflows/.github/workflows/custom-job.yaml@branch-25.08 + uses: rapidsai/shared-workflows/.github/workflows/custom-job.yaml@branch-25.10 with: build_type: ${{ inputs.build_type }} branch: ${{ inputs.branch }} @@ -64,7 +64,7 @@ jobs: script: "ci/test_java.sh" wheel-tests-cuvs: secrets: inherit - uses: rapidsai/shared-workflows/.github/workflows/wheels-test.yaml@branch-25.08 + uses: rapidsai/shared-workflows/.github/workflows/wheels-test.yaml@branch-25.10 with: build_type: ${{ inputs.build_type }} branch: ${{ inputs.branch }} diff --git a/.github/workflows/trigger-breaking-change-alert.yaml b/.github/workflows/trigger-breaking-change-alert.yaml index 593fcb1086..48bf37afc4 100644 --- a/.github/workflows/trigger-breaking-change-alert.yaml +++ b/.github/workflows/trigger-breaking-change-alert.yaml @@ -12,7 +12,7 @@ jobs: trigger-notifier: if: contains(github.event.pull_request.labels.*.name, 'breaking') secrets: inherit - uses: rapidsai/shared-workflows/.github/workflows/breaking-change-alert.yaml@branch-25.08 + uses: rapidsai/shared-workflows/.github/workflows/breaking-change-alert.yaml@branch-25.10 with: sender_login: ${{ github.event.sender.login }} sender_avatar: ${{ github.event.sender.avatar_url }} diff --git a/README.md b/README.md index bef1a8af6e..1604428492 100755 --- a/README.md +++ b/README.md @@ -104,7 +104,7 @@ pip install cuvs-cu12 --extra-index-url=https://pypi.nvidia.com If installing a version that has not yet been released, the `rapidsai` channel can be replaced with `rapidsai-nightly`: ```bash -conda install -c rapidsai-nightly -c conda-forge cuvs=25.08 +conda install -c rapidsai-nightly -c conda-forge cuvs=25.10 ``` cuVS also has `pip` wheel packages that can be installed. 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. @@ -164,7 +164,7 @@ cuvsCagraIndexParamsDestroy(index_params); cuvsResourcesDestroy(res); ``` -For more code examples of the C APIs, including drop-in Cmake project templates, please refer to the [C examples](https://github.com/rapidsai/cuvs/tree/branch-25.08/examples/c) +For more code examples of the C APIs, including drop-in Cmake project templates, please refer to the [C examples](https://github.com/rapidsai/cuvs/tree/branch-25.10/examples/c) ### Rust API @@ -227,7 +227,7 @@ fn cagra_example() -> Result<()> { } ``` -For more code examples of the Rust APIs, including a drop-in project templates, please refer to the [Rust examples](https://github.com/rapidsai/cuvs/tree/branch-25.08/examples/rust). +For more code examples of the Rust APIs, including a drop-in project templates, please refer to the [Rust examples](https://github.com/rapidsai/cuvs/tree/branch-25.10/examples/rust). ## Contributing diff --git a/VERSION b/VERSION index 3af4bda020..296e35288d 100644 --- a/VERSION +++ b/VERSION @@ -1 +1 @@ -25.08.00 +25.10.00 diff --git a/conda/environments/all_cuda-129_arch-aarch64.yaml b/conda/environments/all_cuda-129_arch-aarch64.yaml index 61408b951e..796d166f6c 100644 --- a/conda/environments/all_cuda-129_arch-aarch64.yaml +++ b/conda/environments/all_cuda-129_arch-aarch64.yaml @@ -31,7 +31,7 @@ dependencies: - libcurand-dev - libcusolver-dev - libcusparse-dev -- librmm==25.8.*,>=0.0.0a0 +- librmm==25.10.*,>=0.0.0a0 - make - nccl>=2.19 - ninja @@ -39,7 +39,7 @@ dependencies: - numpydoc - openblas - pre-commit -- pylibraft==25.8.*,>=0.0.0a0 +- pylibraft==25.10.*,>=0.0.0a0 - pytest - pytest-cov - rapids-build-backend>=0.3.0,<0.4.0.dev0 diff --git a/conda/environments/all_cuda-129_arch-x86_64.yaml b/conda/environments/all_cuda-129_arch-x86_64.yaml index 64f9c75fd0..57675489b2 100644 --- a/conda/environments/all_cuda-129_arch-x86_64.yaml +++ b/conda/environments/all_cuda-129_arch-x86_64.yaml @@ -31,7 +31,7 @@ dependencies: - libcurand-dev - libcusolver-dev - libcusparse-dev -- librmm==25.8.*,>=0.0.0a0 +- librmm==25.10.*,>=0.0.0a0 - make - nccl>=2.19 - ninja @@ -39,7 +39,7 @@ dependencies: - numpydoc - openblas - pre-commit -- pylibraft==25.8.*,>=0.0.0a0 +- pylibraft==25.10.*,>=0.0.0a0 - pytest - pytest-cov - rapids-build-backend>=0.3.0,<0.4.0.dev0 diff --git a/conda/environments/bench_ann_cuda-129_arch-aarch64.yaml b/conda/environments/bench_ann_cuda-129_arch-aarch64.yaml index d3c0fae8c9..cf9a064f32 100644 --- a/conda/environments/bench_ann_cuda-129_arch-aarch64.yaml +++ b/conda/environments/bench_ann_cuda-129_arch-aarch64.yaml @@ -18,7 +18,7 @@ dependencies: - cuda-python>=12.6.2,<13.0a0 - cuda-version=12.9 - cupy>=12.0.0 -- cuvs==25.8.*,>=0.0.0a0 +- cuvs==25.10.*,>=0.0.0a0 - cxx-compiler - cython>=3.0.0 - dlpack>=0.8,<1.0 @@ -30,15 +30,15 @@ dependencies: - libcurand-dev - libcusolver-dev - libcusparse-dev -- libcuvs==25.8.*,>=0.0.0a0 -- librmm==25.8.*,>=0.0.0a0 +- libcuvs==25.10.*,>=0.0.0a0 +- librmm==25.10.*,>=0.0.0a0 - matplotlib-base>=3.9 - nccl>=2.19 - ninja - nlohmann_json>=3.11.2 - openblas - pandas -- pylibraft==25.8.*,>=0.0.0a0 +- pylibraft==25.10.*,>=0.0.0a0 - pyyaml - rapids-build-backend>=0.3.0,<0.4.0.dev0 - requests diff --git a/conda/environments/bench_ann_cuda-129_arch-x86_64.yaml b/conda/environments/bench_ann_cuda-129_arch-x86_64.yaml index 4686f6b2bb..8b8bf10a58 100644 --- a/conda/environments/bench_ann_cuda-129_arch-x86_64.yaml +++ b/conda/environments/bench_ann_cuda-129_arch-x86_64.yaml @@ -18,7 +18,7 @@ dependencies: - cuda-python>=12.6.2,<13.0a0 - cuda-version=12.9 - cupy>=12.0.0 -- cuvs==25.8.*,>=0.0.0a0 +- cuvs==25.10.*,>=0.0.0a0 - cxx-compiler - cython>=3.0.0 - dlpack>=0.8,<1.0 @@ -32,8 +32,8 @@ dependencies: - libcurand-dev - libcusolver-dev - libcusparse-dev -- libcuvs==25.8.*,>=0.0.0a0 -- librmm==25.8.*,>=0.0.0a0 +- libcuvs==25.10.*,>=0.0.0a0 +- librmm==25.10.*,>=0.0.0a0 - matplotlib-base>=3.9 - mkl-devel=2023 - nccl>=2.19 @@ -41,7 +41,7 @@ dependencies: - nlohmann_json>=3.11.2 - openblas - pandas -- pylibraft==25.8.*,>=0.0.0a0 +- pylibraft==25.10.*,>=0.0.0a0 - pyyaml - rapids-build-backend>=0.3.0,<0.4.0.dev0 - requests diff --git a/conda/environments/go_cuda-129_arch-aarch64.yaml b/conda/environments/go_cuda-129_arch-aarch64.yaml index ed5e82f6cf..1a0bbb0003 100644 --- a/conda/environments/go_cuda-129_arch-aarch64.yaml +++ b/conda/environments/go_cuda-129_arch-aarch64.yaml @@ -24,8 +24,8 @@ dependencies: - libcurand-dev - libcusolver-dev - libcusparse-dev -- libcuvs==25.8.*,>=0.0.0a0 -- libraft==25.8.*,>=0.0.0a0 +- libcuvs==25.10.*,>=0.0.0a0 +- libraft==25.10.*,>=0.0.0a0 - nccl>=2.19 - ninja - sysroot_linux-aarch64==2.28 diff --git a/conda/environments/go_cuda-129_arch-x86_64.yaml b/conda/environments/go_cuda-129_arch-x86_64.yaml index 8a2f40d01a..c73052a292 100644 --- a/conda/environments/go_cuda-129_arch-x86_64.yaml +++ b/conda/environments/go_cuda-129_arch-x86_64.yaml @@ -24,8 +24,8 @@ dependencies: - libcurand-dev - libcusolver-dev - libcusparse-dev -- libcuvs==25.8.*,>=0.0.0a0 -- libraft==25.8.*,>=0.0.0a0 +- libcuvs==25.10.*,>=0.0.0a0 +- libraft==25.10.*,>=0.0.0a0 - nccl>=2.19 - ninja - sysroot_linux-64==2.28 diff --git a/conda/environments/rust_cuda-129_arch-aarch64.yaml b/conda/environments/rust_cuda-129_arch-aarch64.yaml index f25258a1af..a0e6b868d7 100644 --- a/conda/environments/rust_cuda-129_arch-aarch64.yaml +++ b/conda/environments/rust_cuda-129_arch-aarch64.yaml @@ -21,8 +21,8 @@ dependencies: - libcurand-dev - libcusolver-dev - libcusparse-dev -- libcuvs==25.8.*,>=0.0.0a0 -- libraft==25.8.*,>=0.0.0a0 +- libcuvs==25.10.*,>=0.0.0a0 +- libraft==25.10.*,>=0.0.0a0 - make - nccl>=2.19 - ninja diff --git a/conda/environments/rust_cuda-129_arch-x86_64.yaml b/conda/environments/rust_cuda-129_arch-x86_64.yaml index c6f36470a1..2e845864db 100644 --- a/conda/environments/rust_cuda-129_arch-x86_64.yaml +++ b/conda/environments/rust_cuda-129_arch-x86_64.yaml @@ -21,8 +21,8 @@ dependencies: - libcurand-dev - libcusolver-dev - libcusparse-dev -- libcuvs==25.8.*,>=0.0.0a0 -- libraft==25.8.*,>=0.0.0a0 +- libcuvs==25.10.*,>=0.0.0a0 +- libraft==25.10.*,>=0.0.0a0 - make - nccl>=2.19 - ninja diff --git a/dependencies.yaml b/dependencies.yaml index 0f77bcf3aa..88469b62c3 100644 --- a/dependencies.yaml +++ b/dependencies.yaml @@ -489,7 +489,7 @@ dependencies: - output_types: [conda, pyproject, requirements] packages: - click - - cuvs==25.8.*,>=0.0.0a0 + - cuvs==25.10.*,>=0.0.0a0 - pandas - pyyaml - requests @@ -504,17 +504,17 @@ dependencies: common: - output_types: conda packages: - - cuvs==25.8.*,>=0.0.0a0 + - cuvs==25.10.*,>=0.0.0a0 depends_on_cuvs_bench: common: - output_types: conda packages: - - cuvs-bench==25.8.*,>=0.0.0a0 + - cuvs-bench==25.10.*,>=0.0.0a0 depends_on_libcuvs: common: - output_types: conda packages: - - &libcuvs_unsuffixed libcuvs==25.8.*,>=0.0.0a0 + - &libcuvs_unsuffixed libcuvs==25.10.*,>=0.0.0a0 - output_types: requirements packages: # pip recognizes the index as a global option for the requirements.txt file @@ -527,18 +527,18 @@ dependencies: cuda: "12.*" cuda_suffixed: "true" packages: - - libcuvs-cu12==25.8.*,>=0.0.0a0 + - libcuvs-cu12==25.10.*,>=0.0.0a0 - {matrix: null, packages: [*libcuvs_unsuffixed]} depends_on_libcuvs_tests: common: - output_types: conda packages: - - libcuvs-tests==25.8.*,>=0.0.0a0 + - libcuvs-tests==25.10.*,>=0.0.0a0 depends_on_libraft: common: - output_types: conda packages: - - &libraft_unsuffixed libraft==25.8.*,>=0.0.0a0 + - &libraft_unsuffixed libraft==25.10.*,>=0.0.0a0 - output_types: requirements packages: # pip recognizes the index as a global option for the requirements.txt file @@ -551,13 +551,13 @@ dependencies: cuda: "12.*" cuda_suffixed: "true" packages: - - libraft-cu12==25.8.*,>=0.0.0a0 + - libraft-cu12==25.10.*,>=0.0.0a0 - {matrix: null, packages: [*libraft_unsuffixed]} depends_on_librmm: common: - output_types: conda packages: - - &librmm_unsuffixed librmm==25.8.*,>=0.0.0a0 + - &librmm_unsuffixed librmm==25.10.*,>=0.0.0a0 - output_types: requirements packages: # pip recognizes the index as a global option for the requirements.txt file @@ -570,13 +570,13 @@ dependencies: cuda: "12.*" cuda_suffixed: "true" packages: - - librmm-cu12==25.8.*,>=0.0.0a0 + - librmm-cu12==25.10.*,>=0.0.0a0 - {matrix: null, packages: [*librmm_unsuffixed]} depends_on_pylibraft: common: - output_types: conda packages: - - &pylibraft_unsuffixed pylibraft==25.8.*,>=0.0.0a0 + - &pylibraft_unsuffixed pylibraft==25.10.*,>=0.0.0a0 - output_types: requirements packages: # pip recognizes the index as a global option for the requirements.txt file @@ -589,7 +589,7 @@ dependencies: cuda: "12.*" cuda_suffixed: "true" packages: - - pylibraft-cu12==25.8.*,>=0.0.0a0 + - pylibraft-cu12==25.10.*,>=0.0.0a0 - {matrix: null, packages: [*pylibraft_unsuffixed]} depends_on_nccl: common: diff --git a/docs/source/developer_guide.md b/docs/source/developer_guide.md index 8e8cc28017..b4bc021e3d 100644 --- a/docs/source/developer_guide.md +++ b/docs/source/developer_guide.md @@ -187,7 +187,7 @@ RAFT relies on `clang-format` to enforce code style across all C++ and CUDA sour 1. Do not split empty functions/records/namespaces. 2. Two-space indentation everywhere, including the line continuations. 3. Disable reflowing of comments. - The reasons behind these deviations from the Google style guide are given in comments [here](https://github.com/rapidsai/raft/blob/branch-25.08/cpp/.clang-format). + The reasons behind these deviations from the Google style guide are given in comments [here](https://github.com/rapidsai/raft/blob/branch-25.10/cpp/.clang-format). [`doxygen`](https://doxygen.nl/) is used as documentation generator and also as a documentation linter. In order to run doxygen as a linter on C++/CUDA code, run @@ -205,7 +205,7 @@ you can run `codespell -i 3 -w .` from the repository root directory. This will bring up an interactive prompt to select which spelling fixes to apply. ### #include style -[include_checker.py](https://github.com/rapidsai/raft/blob/branch-25.08/cpp/scripts/include_checker.py) is used to enforce the include style as follows: +[include_checker.py](https://github.com/rapidsai/raft/blob/branch-25.10/cpp/scripts/include_checker.py) is used to enforce the include style as follows: 1. `#include "..."` should be used for referencing local files only. It is acceptable to be used for referencing files in a sub-folder/parent-folder of the same algorithm, but should never be used to include files in other algorithms or between algorithms and the primitives or other dependencies. 2. `#include <...>` should be used for referencing everything else diff --git a/examples/go/README.md b/examples/go/README.md index ebeb012308..bcbea09642 100644 --- a/examples/go/README.md +++ b/examples/go/README.md @@ -24,7 +24,7 @@ export CC=clang 2. Install the Go module: ```bash -go get github.com/rapidsai/cuvs/go@v25.08.00 # 25.02.00 being your desired version, selected from https://github.com/rapidsai/cuvs/tags +go get github.com/rapidsai/cuvs/go@v25.10.00 # 25.02.00 being your desired version, selected from https://github.com/rapidsai/cuvs/tags ``` Then you can build your project with the usual `go build`. diff --git a/java/benchmarks/pom.xml b/java/benchmarks/pom.xml index f5c77a14bc..f80e197c70 100644 --- a/java/benchmarks/pom.xml +++ b/java/benchmarks/pom.xml @@ -21,7 +21,7 @@ com.nvidia.cuvs benchmarks - 25.08.0 + 25.10.0 jar cuvs-java-benchmarks @@ -41,7 +41,7 @@ com.nvidia.cuvs cuvs-java - 25.08.0 + 25.10.0 jar diff --git a/java/build.sh b/java/build.sh index 89c2b4bef2..7d2628d7e5 100755 --- a/java/build.sh +++ b/java/build.sh @@ -8,7 +8,7 @@ ARGS="$*" NUMARGS=$# CURDIR=$(cd "$(dirname "$0")"; pwd) -VERSION="25.08.0" # Note: The version is updated automatically when ci/release/update-version.sh is invoked +VERSION="25.10.0" # Note: The version is updated automatically when ci/release/update-version.sh is invoked GROUP_ID="com.nvidia.cuvs" if [ -z "${CMAKE_PREFIX_PATH:=}" ]; then diff --git a/java/cuvs-java/pom.xml b/java/cuvs-java/pom.xml index deaeea5755..b3b862d291 100644 --- a/java/cuvs-java/pom.xml +++ b/java/cuvs-java/pom.xml @@ -22,7 +22,7 @@ com.nvidia.cuvs cuvs-java - 25.08.0 + 25.10.0 cuvs-java This project provides Java bindings for cuVS, enabling approximate nearest neighbors search and clustering diff --git a/java/examples/README.md b/java/examples/README.md index 8477a3859f..29826e96a6 100644 --- a/java/examples/README.md +++ b/java/examples/README.md @@ -11,17 +11,17 @@ This maven project contains examples for CAGRA, HNSW, and Bruteforce algorithms. ### CAGRA Example In the current directory do: ``` -mvn package && java --enable-native-access=ALL-UNNAMED -cp target/cuvs-java-examples-25.08.0.jar:$HOME/.m2/repository/com/nvidia/cuvs/cuvs-java/25.08.0/cuvs-java-25.08.0.jar com.nvidia.cuvs.examples.CagraExample +mvn package && java --enable-native-access=ALL-UNNAMED -cp target/cuvs-java-examples-25.10.0.jar:$HOME/.m2/repository/com/nvidia/cuvs/cuvs-java/25.10.0/cuvs-java-25.10.0.jar com.nvidia.cuvs.examples.CagraExample ``` ### HNSW Example In the current directory do: ``` -mvn package && java --enable-native-access=ALL-UNNAMED -cp target/cuvs-java-examples-25.08.0.jar:$HOME/.m2/repository/com/nvidia/cuvs/cuvs-java/25.08.0/cuvs-java-25.08.0.jar com.nvidia.cuvs.examples.HnswExample +mvn package && java --enable-native-access=ALL-UNNAMED -cp target/cuvs-java-examples-25.10.0.jar:$HOME/.m2/repository/com/nvidia/cuvs/cuvs-java/25.10.0/cuvs-java-25.10.0.jar com.nvidia.cuvs.examples.HnswExample ``` ### Bruteforce Example In the current directory do: ``` -mvn package && java --enable-native-access=ALL-UNNAMED -cp target/cuvs-java-examples-25.08.0.jar:$HOME/.m2/repository/com/nvidia/cuvs/cuvs-java/25.08.0/cuvs-java-25.08.0.jar com.nvidia.cuvs.examples.BruteForceExample +mvn package && java --enable-native-access=ALL-UNNAMED -cp target/cuvs-java-examples-25.10.0.jar:$HOME/.m2/repository/com/nvidia/cuvs/cuvs-java/25.10.0/cuvs-java-25.10.0.jar com.nvidia.cuvs.examples.BruteForceExample ``` diff --git a/java/examples/pom.xml b/java/examples/pom.xml index 4c18fd5a28..0c52685672 100644 --- a/java/examples/pom.xml +++ b/java/examples/pom.xml @@ -5,7 +5,7 @@ com.nvidia.cuvs.examples cuvs-java-examples - 25.08.0 + 25.10.0 cuvs-java-examples @@ -18,7 +18,7 @@ com.nvidia.cuvs cuvs-java - 25.08.0 + 25.10.0 diff --git a/python/cuvs/pyproject.toml b/python/cuvs/pyproject.toml index 7a00b824df..373e313492 100644 --- a/python/cuvs/pyproject.toml +++ b/python/cuvs/pyproject.toml @@ -32,9 +32,9 @@ license = { text = "Apache-2.0" } requires-python = ">=3.10" dependencies = [ "cuda-python", - "libcuvs==25.8.*,>=0.0.0a0", + "libcuvs==25.10.*,>=0.0.0a0", "numpy>=1.23,<3.0a0", - "pylibraft==25.8.*,>=0.0.0a0", + "pylibraft==25.10.*,>=0.0.0a0", ] # This list was generated by `rapids-dependency-file-generator`. To make changes, edit ../../dependencies.yaml and run `rapids-dependency-file-generator`. classifiers = [ "Intended Audience :: Developers", @@ -119,9 +119,9 @@ requires = [ "cmake>=3.30.4", "cuda-python", "cython>=3.0.0", - "libcuvs==25.8.*,>=0.0.0a0", - "libraft==25.8.*,>=0.0.0a0", - "librmm==25.8.*,>=0.0.0a0", + "libcuvs==25.10.*,>=0.0.0a0", + "libraft==25.10.*,>=0.0.0a0", + "librmm==25.10.*,>=0.0.0a0", "ninja", ] # This list was generated by `rapids-dependency-file-generator`. To make changes, edit ../../dependencies.yaml and run `rapids-dependency-file-generator`. build-backend = "scikit_build_core.build" diff --git a/python/cuvs_bench/cuvs_bench/plot/__main__.py b/python/cuvs_bench/cuvs_bench/plot/__main__.py index e0dbc27837..1c8a3005f6 100644 --- a/python/cuvs_bench/cuvs_bench/plot/__main__.py +++ b/python/cuvs_bench/cuvs_bench/plot/__main__.py @@ -17,7 +17,7 @@ # 1: https://github.com/erikbern/ann-benchmarks/blob/main/plot.py # 2: https://github.com/erikbern/ann-benchmarks/blob/main/ann_benchmarks/plotting/utils.py # noqa: E501 # 3: https://github.com/erikbern/ann-benchmarks/blob/main/ann_benchmarks/plotting/metrics.py # noqa: E501 -# License: https://github.com/rapidsai/cuvs/blob/branch-25.08/thirdparty/LICENSES/LICENSE.ann-benchmark # noqa: E501 +# License: https://github.com/rapidsai/cuvs/blob/branch-25.10/thirdparty/LICENSES/LICENSE.ann-benchmark # noqa: E501 import itertools import os diff --git a/python/cuvs_bench/pyproject.toml b/python/cuvs_bench/pyproject.toml index 59c3829317..901a37e74d 100644 --- a/python/cuvs_bench/pyproject.toml +++ b/python/cuvs_bench/pyproject.toml @@ -19,7 +19,7 @@ license = { text = "Apache-2.0" } requires-python = ">=3.10" dependencies = [ "click", - "cuvs==25.8.*,>=0.0.0a0", + "cuvs==25.10.*,>=0.0.0a0", "matplotlib>=3.9", "pandas", "pyyaml", diff --git a/python/libcuvs/pyproject.toml b/python/libcuvs/pyproject.toml index 9d357e6928..00e49155da 100644 --- a/python/libcuvs/pyproject.toml +++ b/python/libcuvs/pyproject.toml @@ -30,8 +30,8 @@ authors = [ license = { text = "Apache-2.0" } requires-python = ">=3.10" dependencies = [ - "libraft==25.8.*,>=0.0.0a0", - "librmm==25.8.*,>=0.0.0a0", + "libraft==25.10.*,>=0.0.0a0", + "librmm==25.10.*,>=0.0.0a0", "nvidia-cublas", "nvidia-curand", "nvidia-cusolver", @@ -93,8 +93,8 @@ regex = "(?P.*)" build-backend = "scikit_build_core.build" requires = [ "cmake>=3.30.4", - "libraft==25.8.*,>=0.0.0a0", - "librmm==25.8.*,>=0.0.0a0", + "libraft==25.10.*,>=0.0.0a0", + "librmm==25.10.*,>=0.0.0a0", "ninja", ] # This list was generated by `rapids-dependency-file-generator`. To make changes, edit ../../dependencies.yaml and run `rapids-dependency-file-generator`. dependencies-file = "../../dependencies.yaml" diff --git a/rust/Cargo.toml b/rust/Cargo.toml index 98c2c94018..d8a98b521f 100644 --- a/rust/Cargo.toml +++ b/rust/Cargo.toml @@ -6,7 +6,7 @@ members = [ resolver = "2" [workspace.package] -version = "25.8.0" +version = "25.10.0" edition = "2021" repository = "https://github.com/rapidsai/cuvs" homepage = "https://github.com/rapidsai/cuvs" diff --git a/rust/cuvs/Cargo.toml b/rust/cuvs/Cargo.toml index d2774e9140..c7b713f696 100644 --- a/rust/cuvs/Cargo.toml +++ b/rust/cuvs/Cargo.toml @@ -9,7 +9,7 @@ authors.workspace = true license.workspace = true [dependencies] -ffi = { package = "cuvs-sys", path = "../cuvs-sys", version = "25.8.0" } +ffi = { package = "cuvs-sys", path = "../cuvs-sys", version = "25.10.0" } ndarray = "0.15" [dev-dependencies] From a8249e3b7d7d72d5d93bedb5f511847de148a0d2 Mon Sep 17 00:00:00 2001 From: Jake Awe Date: Fri, 18 Jul 2025 00:07:36 -0500 Subject: [PATCH 02/93] 25.08 -> 25.10 version updates --- .github/workflows/build.yaml | 8 ++++---- .github/workflows/pr.yaml | 8 ++++---- .github/workflows/publish-rust.yaml | 2 +- .github/workflows/test.yaml | 2 +- docs/source/cuvs_bench/index.rst | 8 ++++---- 5 files changed, 14 insertions(+), 14 deletions(-) diff --git a/.github/workflows/build.yaml b/.github/workflows/build.yaml index 3ebc825f6f..210ade79ca 100644 --- a/.github/workflows/build.yaml +++ b/.github/workflows/build.yaml @@ -50,7 +50,7 @@ jobs: branch: ${{ inputs.branch }} arch: "amd64" date: ${{ inputs.date }} - container_image: "rapidsai/ci-conda:25.08-latest" + container_image: "rapidsai/ci-conda:25.10-latest" node_type: "gpu-l4-latest-1" script: "ci/build_rust.sh" sha: ${{ inputs.sha }} @@ -63,7 +63,7 @@ jobs: branch: ${{ inputs.branch }} arch: "amd64" date: ${{ inputs.date }} - container_image: "rapidsai/ci-conda:25.08-latest" + container_image: "rapidsai/ci-conda:25.10-latest" node_type: "gpu-l4-latest-1" script: "ci/build_go.sh" sha: ${{ inputs.sha }} @@ -76,7 +76,7 @@ jobs: branch: ${{ inputs.branch }} arch: "amd64" date: ${{ inputs.date }} - container_image: "rapidsai/ci-conda:25.08-latest" + container_image: "rapidsai/ci-conda:25.10-latest" script: "ci/build_java.sh" file_to_upload: "java/cuvs-java/target/" sha: ${{ inputs.sha }} @@ -113,7 +113,7 @@ jobs: arch: "amd64" branch: ${{ inputs.branch }} build_type: ${{ inputs.build_type || 'branch' }} - container_image: "rapidsai/ci-conda:25.08-latest" + container_image: "rapidsai/ci-conda:25.10-latest" date: ${{ inputs.date }} node_type: "gpu-l4-latest-1" script: "ci/build_docs.sh" diff --git a/.github/workflows/pr.yaml b/.github/workflows/pr.yaml index 2083429482..eb4bd6b094 100644 --- a/.github/workflows/pr.yaml +++ b/.github/workflows/pr.yaml @@ -160,7 +160,7 @@ jobs: build_type: pull-request node_type: "gpu-l4-latest-1" arch: "amd64" - container_image: "rapidsai/ci-conda:25.08-latest" + container_image: "rapidsai/ci-conda:25.10-latest" script: "ci/test_java.sh" file_to_upload: "java/cuvs-java/target/" docs-build: @@ -171,7 +171,7 @@ jobs: build_type: pull-request node_type: "gpu-l4-latest-1" arch: "amd64" - container_image: "rapidsai/ci-conda:25.08-latest" + container_image: "rapidsai/ci-conda:25.10-latest" script: "ci/build_docs.sh" rust-build: needs: conda-cpp-build @@ -181,7 +181,7 @@ jobs: build_type: pull-request node_type: "gpu-l4-latest-1" arch: "amd64" - container_image: "rapidsai/ci-conda:25.08-latest" + container_image: "rapidsai/ci-conda:25.10-latest" script: "ci/build_rust.sh" go-build: needs: conda-cpp-build @@ -191,7 +191,7 @@ jobs: build_type: pull-request node_type: "gpu-l4-latest-1" arch: "amd64" - container_image: "rapidsai/ci-conda:25.08-latest" + container_image: "rapidsai/ci-conda:25.10-latest" script: "ci/build_go.sh" wheel-build-libcuvs: needs: checks diff --git a/.github/workflows/publish-rust.yaml b/.github/workflows/publish-rust.yaml index 6b9846ec24..7834e6bc96 100644 --- a/.github/workflows/publish-rust.yaml +++ b/.github/workflows/publish-rust.yaml @@ -11,7 +11,7 @@ jobs: rust-publish: runs-on: ubuntu-latest container: - image: rapidsai/ci-conda:25.08-latest + image: rapidsai/ci-conda:25.10-latest steps: - uses: actions/checkout@v4 - name: Check if release build diff --git a/.github/workflows/test.yaml b/.github/workflows/test.yaml index 643b5744d2..8479ce38a4 100644 --- a/.github/workflows/test.yaml +++ b/.github/workflows/test.yaml @@ -60,7 +60,7 @@ jobs: sha: ${{ inputs.sha }} node_type: "gpu-l4-latest-1" arch: "amd64" - container_image: "rapidsai/ci-conda:25.08-latest" + container_image: "rapidsai/ci-conda:25.10-latest" script: "ci/test_java.sh" wheel-tests-cuvs: secrets: inherit diff --git a/docs/source/cuvs_bench/index.rst b/docs/source/cuvs_bench/index.rst index 1c231c6e56..26f862c868 100644 --- a/docs/source/cuvs_bench/index.rst +++ b/docs/source/cuvs_bench/index.rst @@ -89,7 +89,7 @@ The following command pulls the nightly container for Python version 3.10, CUDA .. code-block:: bash - docker pull rapidsai/cuvs-bench:24.12a-cuda12.5-py3.10 # substitute cuvs-bench for the exact desired container. + docker pull rapidsai/cuvs-bench:25.10a-cuda12.5-py3.10 # substitute cuvs-bench for the exact desired container. The CUDA and python versions can be changed for the supported values: - Supported CUDA versions: 11.8 and 12.5 @@ -236,7 +236,7 @@ For GPU-enabled systems, the `DATA_FOLDER` variable should be a local folder whe export DATA_FOLDER=path/to/store/datasets/and/results docker run --gpus all --rm -it -u $(id -u) \ -v $DATA_FOLDER:/data/benchmarks \ - rapidsai/cuvs-bench:25.08-cuda12.9-py3.13 \ + rapidsai/cuvs-bench:25.10-cuda12.9-py3.13 \ "--dataset deep-image-96-angular" \ "--normalize" \ "--algorithms cuvs_cagra,cuvs_ivf_pq --batch-size 10 -k 10" \ @@ -249,7 +249,7 @@ Usage of the above command is as follows: * - Argument - Description - * - `rapidsai/cuvs-bench:25.08-cuda12.9-py3.13` + * - `rapidsai/cuvs-bench:25.10-cuda12.9-py3.13` - Image to use. Can be either `cuvs-bench` or `cuvs-bench-datasets` * - `"--dataset deep-image-96-angular"` @@ -296,7 +296,7 @@ All of the `cuvs-bench` images contain the Conda packages, so they can be used d --entrypoint /bin/bash \ --workdir /data/benchmarks \ -v $DATA_FOLDER:/data/benchmarks \ - rapidsai/cuvs-bench:25.08-cuda12.9-py3.13 + rapidsai/cuvs-bench:25.10-cuda12.9-py3.13 This will drop you into a command line in the container, with the `cuvs-bench` python package ready to use, as described in the [Running the benchmarks](#running-the-benchmarks) section above: From 09193514a820256b0db8ca92c8a628b08f97ce00 Mon Sep 17 00:00:00 2001 From: Bradley Dice Date: Thu, 24 Jul 2025 10:40:50 -0500 Subject: [PATCH 03/93] Allow latest OS in devcontainers (#1169) This PR removes the OS suffix from devcontainers, allowing the upstream devcontainer images to determine the OS version. Contributes to https://github.com/rapidsai/build-planning/issues/200. Authors: - Bradley Dice (https://github.com/bdice) Approvers: - Kyle Edwards (https://github.com/KyleFromNVIDIA) URL: https://github.com/rapidsai/cuvs/pull/1169 --- .devcontainer/cuda12.9-conda/devcontainer.json | 2 +- .devcontainer/cuda12.9-pip/devcontainer.json | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/.devcontainer/cuda12.9-conda/devcontainer.json b/.devcontainer/cuda12.9-conda/devcontainer.json index ef5383cd80..3c1da8fe73 100644 --- a/.devcontainer/cuda12.9-conda/devcontainer.json +++ b/.devcontainer/cuda12.9-conda/devcontainer.json @@ -5,7 +5,7 @@ "args": { "CUDA": "12.9", "PYTHON_PACKAGE_MANAGER": "conda", - "BASE": "rapidsai/devcontainers:25.10-cpp-mambaforge-ubuntu22.04" + "BASE": "rapidsai/devcontainers:25.10-cpp-mambaforge" } }, "runArgs": [ diff --git a/.devcontainer/cuda12.9-pip/devcontainer.json b/.devcontainer/cuda12.9-pip/devcontainer.json index 5902c948f3..0305705fb9 100644 --- a/.devcontainer/cuda12.9-pip/devcontainer.json +++ b/.devcontainer/cuda12.9-pip/devcontainer.json @@ -5,7 +5,7 @@ "args": { "CUDA": "12.9", "PYTHON_PACKAGE_MANAGER": "pip", - "BASE": "rapidsai/devcontainers:25.10-cpp-cuda12.9-ucx1.18.0-openmpi-ubuntu22.04" + "BASE": "rapidsai/devcontainers:25.10-cpp-cuda12.9-ucx1.18.0-openmpi5.0.7" } }, "runArgs": [ From 42068e57eded36d3c5b57463b7c6acfaec2e190d Mon Sep 17 00:00:00 2001 From: Vyas Ramasubramani Date: Fri, 25 Jul 2025 08:43:55 -0700 Subject: [PATCH 04/93] Use GCC 14 in conda builds. (#1030) conda-forge is migrating to gcc 14, so this PR is updating for alignment. See https://github.com/rapidsai/build-planning/issues/188 Authors: - Vyas Ramasubramani (https://github.com/vyasr) Approvers: - Jake Awe (https://github.com/AyodeAwe) - Ben Frederickson (https://github.com/benfred) URL: https://github.com/rapidsai/cuvs/pull/1030 --- ci/build_cpp.sh | 6 +++++- conda/environments/all_cuda-129_arch-aarch64.yaml | 2 +- conda/environments/all_cuda-129_arch-x86_64.yaml | 2 +- .../environments/bench_ann_cuda-129_arch-aarch64.yaml | 2 +- conda/environments/bench_ann_cuda-129_arch-x86_64.yaml | 2 +- conda/environments/go_cuda-129_arch-aarch64.yaml | 2 +- conda/environments/go_cuda-129_arch-x86_64.yaml | 2 +- conda/environments/rust_cuda-129_arch-aarch64.yaml | 2 +- conda/environments/rust_cuda-129_arch-x86_64.yaml | 2 +- conda/recipes/cuvs-bench-cpu/conda_build_config.yaml | 4 ++-- conda/recipes/cuvs/conda_build_config.yaml | 4 ++-- conda/recipes/libcuvs/conda_build_config.yaml | 4 ++-- cpp/bench/ann/src/cuvs/cuvs_ann_bench_param_parser.h | 10 +++++++--- cpp/bench/ann/src/faiss/faiss_gpu_benchmark.cu | 4 ++-- dependencies.yaml | 4 ++-- 15 files changed, 30 insertions(+), 22 deletions(-) diff --git a/ci/build_cpp.sh b/ci/build_cpp.sh index e55ba04fdb..f8fb90dbb9 100755 --- a/ci/build_cpp.sh +++ b/ci/build_cpp.sh @@ -1,5 +1,5 @@ #!/bin/bash -# Copyright (c) 2022-2024, NVIDIA CORPORATION. +# Copyright (c) 2022-2025, NVIDIA CORPORATION. set -euo pipefail @@ -19,6 +19,10 @@ sccache --zero-stats 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 + # populates `RATTLER_CHANNELS` array and `RATTLER_ARGS` array source rapids-rattler-channel-string diff --git a/conda/environments/all_cuda-129_arch-aarch64.yaml b/conda/environments/all_cuda-129_arch-aarch64.yaml index 796d166f6c..04fe7a5c21 100644 --- a/conda/environments/all_cuda-129_arch-aarch64.yaml +++ b/conda/environments/all_cuda-129_arch-aarch64.yaml @@ -22,7 +22,7 @@ dependencies: - cython>=3.0.0 - dlpack>=0.8,<1.0 - doxygen>=1.8.20 -- gcc_linux-aarch64=13.* +- gcc_linux-aarch64=14.* - go - graphviz - ipython diff --git a/conda/environments/all_cuda-129_arch-x86_64.yaml b/conda/environments/all_cuda-129_arch-x86_64.yaml index 57675489b2..643dbfe04c 100644 --- a/conda/environments/all_cuda-129_arch-x86_64.yaml +++ b/conda/environments/all_cuda-129_arch-x86_64.yaml @@ -22,7 +22,7 @@ dependencies: - cython>=3.0.0 - dlpack>=0.8,<1.0 - doxygen>=1.8.20 -- gcc_linux-64=13.* +- gcc_linux-64=14.* - go - graphviz - ipython diff --git a/conda/environments/bench_ann_cuda-129_arch-aarch64.yaml b/conda/environments/bench_ann_cuda-129_arch-aarch64.yaml index cf9a064f32..7e5efd0693 100644 --- a/conda/environments/bench_ann_cuda-129_arch-aarch64.yaml +++ b/conda/environments/bench_ann_cuda-129_arch-aarch64.yaml @@ -22,7 +22,7 @@ dependencies: - cxx-compiler - cython>=3.0.0 - dlpack>=0.8,<1.0 -- gcc_linux-aarch64=13.* +- gcc_linux-aarch64=14.* - glog>=0.6.0 - h5py>=3.8.0 - libclang==20.1.4 diff --git a/conda/environments/bench_ann_cuda-129_arch-x86_64.yaml b/conda/environments/bench_ann_cuda-129_arch-x86_64.yaml index 8b8bf10a58..5604f56252 100644 --- a/conda/environments/bench_ann_cuda-129_arch-x86_64.yaml +++ b/conda/environments/bench_ann_cuda-129_arch-x86_64.yaml @@ -22,7 +22,7 @@ dependencies: - cxx-compiler - cython>=3.0.0 - dlpack>=0.8,<1.0 -- gcc_linux-64=13.* +- gcc_linux-64=14.* - glog>=0.6.0 - h5py>=3.8.0 - libaio diff --git a/conda/environments/go_cuda-129_arch-aarch64.yaml b/conda/environments/go_cuda-129_arch-aarch64.yaml index 1a0bbb0003..a23a3e636b 100644 --- a/conda/environments/go_cuda-129_arch-aarch64.yaml +++ b/conda/environments/go_cuda-129_arch-aarch64.yaml @@ -17,7 +17,7 @@ dependencies: - cuda-version=12.9 - cxx-compiler - dlpack>=0.8,<1.0 -- gcc_linux-aarch64=13.* +- gcc_linux-aarch64=14.* - go - libclang==20.1.4 - libcublas-dev diff --git a/conda/environments/go_cuda-129_arch-x86_64.yaml b/conda/environments/go_cuda-129_arch-x86_64.yaml index c73052a292..a6bdbe0159 100644 --- a/conda/environments/go_cuda-129_arch-x86_64.yaml +++ b/conda/environments/go_cuda-129_arch-x86_64.yaml @@ -17,7 +17,7 @@ dependencies: - cuda-version=12.9 - cxx-compiler - dlpack>=0.8,<1.0 -- gcc_linux-64=13.* +- gcc_linux-64=14.* - go - libclang==20.1.4 - libcublas-dev diff --git a/conda/environments/rust_cuda-129_arch-aarch64.yaml b/conda/environments/rust_cuda-129_arch-aarch64.yaml index a0e6b868d7..3a94cc2488 100644 --- a/conda/environments/rust_cuda-129_arch-aarch64.yaml +++ b/conda/environments/rust_cuda-129_arch-aarch64.yaml @@ -15,7 +15,7 @@ dependencies: - cuda-profiler-api - cuda-version=12.9 - cxx-compiler -- gcc_linux-aarch64=13.* +- gcc_linux-aarch64=14.* - libclang==20.1.4 - libcublas-dev - libcurand-dev diff --git a/conda/environments/rust_cuda-129_arch-x86_64.yaml b/conda/environments/rust_cuda-129_arch-x86_64.yaml index 2e845864db..98545cc195 100644 --- a/conda/environments/rust_cuda-129_arch-x86_64.yaml +++ b/conda/environments/rust_cuda-129_arch-x86_64.yaml @@ -15,7 +15,7 @@ dependencies: - cuda-profiler-api - cuda-version=12.9 - cxx-compiler -- gcc_linux-64=13.* +- gcc_linux-64=14.* - libclang==20.1.4 - libcublas-dev - libcurand-dev diff --git a/conda/recipes/cuvs-bench-cpu/conda_build_config.yaml b/conda/recipes/cuvs-bench-cpu/conda_build_config.yaml index 8bad5c58e9..b4e099cde8 100644 --- a/conda/recipes/cuvs-bench-cpu/conda_build_config.yaml +++ b/conda/recipes/cuvs-bench-cpu/conda_build_config.yaml @@ -1,8 +1,8 @@ c_compiler_version: - - 13 + - 14 cxx_compiler_version: - - 13 + - 14 c_stdlib: - sysroot diff --git a/conda/recipes/cuvs/conda_build_config.yaml b/conda/recipes/cuvs/conda_build_config.yaml index caf5489e1d..1082f0d214 100644 --- a/conda/recipes/cuvs/conda_build_config.yaml +++ b/conda/recipes/cuvs/conda_build_config.yaml @@ -1,8 +1,8 @@ c_compiler_version: - - 13 + - 14 cxx_compiler_version: - - 13 + - 14 cuda_compiler: - cuda-nvcc diff --git a/conda/recipes/libcuvs/conda_build_config.yaml b/conda/recipes/libcuvs/conda_build_config.yaml index 65b9e3e8da..21989d9803 100644 --- a/conda/recipes/libcuvs/conda_build_config.yaml +++ b/conda/recipes/libcuvs/conda_build_config.yaml @@ -1,8 +1,8 @@ c_compiler_version: - - 13 + - 14 cxx_compiler_version: - - 13 + - 14 cuda_compiler: - cuda-nvcc diff --git a/cpp/bench/ann/src/cuvs/cuvs_ann_bench_param_parser.h b/cpp/bench/ann/src/cuvs/cuvs_ann_bench_param_parser.h index 7ad8ef91b5..fa462252e3 100644 --- a/cpp/bench/ann/src/cuvs/cuvs_ann_bench_param_parser.h +++ b/cpp/bench/ann/src/cuvs/cuvs_ann_bench_param_parser.h @@ -1,5 +1,5 @@ /* - * Copyright (c) 2023-2024, NVIDIA CORPORATION. + * Copyright (c) 2023-2025, NVIDIA CORPORATION. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -83,7 +83,9 @@ void parse_build_param(const nlohmann::json& conf, { param.n_lists = conf.at("nlist"); if (conf.contains("niter")) { param.kmeans_n_iters = conf.at("niter"); } - if (conf.contains("ratio")) { param.kmeans_trainset_fraction = 1.0 / (double)conf.at("ratio"); } + if (conf.contains("ratio")) { + param.kmeans_trainset_fraction = 1.0 / static_cast(conf.at("ratio")); + } } template @@ -102,7 +104,9 @@ void parse_build_param(const nlohmann::json& conf, { if (conf.contains("nlist")) { param.n_lists = conf.at("nlist"); } if (conf.contains("niter")) { param.kmeans_n_iters = conf.at("niter"); } - if (conf.contains("ratio")) { param.kmeans_trainset_fraction = 1.0 / (double)conf.at("ratio"); } + if (conf.contains("ratio")) { + param.kmeans_trainset_fraction = 1.0 / static_cast(conf.at("ratio")); + } if (conf.contains("pq_bits")) { param.pq_bits = conf.at("pq_bits"); } if (conf.contains("pq_dim")) { param.pq_dim = conf.at("pq_dim"); } if (conf.contains("codebook_kind")) { diff --git a/cpp/bench/ann/src/faiss/faiss_gpu_benchmark.cu b/cpp/bench/ann/src/faiss/faiss_gpu_benchmark.cu index f9ad76dc59..fbe042979d 100644 --- a/cpp/bench/ann/src/faiss/faiss_gpu_benchmark.cu +++ b/cpp/bench/ann/src/faiss/faiss_gpu_benchmark.cu @@ -1,5 +1,5 @@ /* - * Copyright (c) 2023-2024, NVIDIA CORPORATION. + * Copyright (c) 2023-2025, NVIDIA CORPORATION. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -132,7 +132,7 @@ void parse_build_param(const nlohmann::json& conf, ivf_pq_build_p.kmeans_n_iters = ivf_pq_build_conf.at("niter"); } if (ivf_pq_build_conf.contains("ratio")) { - ivf_pq_build_p.kmeans_trainset_fraction = 1.0 / (double)conf.at("ratio"); + ivf_pq_build_p.kmeans_trainset_fraction = 1.0 / static_cast(conf.at("ratio")); } if (ivf_pq_build_conf.contains("pq_bits")) { ivf_pq_build_p.pq_bits = ivf_pq_build_conf.at("pq_bits"); diff --git a/dependencies.yaml b/dependencies.yaml index 88469b62c3..960614a891 100644 --- a/dependencies.yaml +++ b/dependencies.yaml @@ -247,13 +247,13 @@ dependencies: arch: x86_64 cuda: "12.*" packages: - - gcc_linux-64=13.* + - gcc_linux-64=14.* - sysroot_linux-64==2.28 - matrix: arch: aarch64 cuda: "12.*" packages: - - gcc_linux-aarch64=13.* + - gcc_linux-aarch64=14.* - sysroot_linux-aarch64==2.28 - output_types: conda matrices: From 5bd425dc2a3d2a7f0871030b4716a4fb35a70e49 Mon Sep 17 00:00:00 2001 From: Robert Maynard Date: Mon, 28 Jul 2025 10:11:15 -0400 Subject: [PATCH 05/93] Update build infra to support new branching strategy (#1143) rapids_config will use `RAPIDS_BRANCH` contents to determine what branch to use Authors: - Robert Maynard (https://github.com/robertmaynard) Approvers: - Bradley Dice (https://github.com/bdice) - Ben Frederickson (https://github.com/benfred) URL: https://github.com/rapidsai/cuvs/pull/1143 --- RAPIDS_BRANCH | 1 + cmake/RAPIDS.cmake | 6 +++--- cmake/rapids_config.cmake | 10 ++++++++++ rust/cuvs-sys/RAPIDS_BRANCH | 1 + 4 files changed, 15 insertions(+), 3 deletions(-) create mode 100644 RAPIDS_BRANCH create mode 120000 rust/cuvs-sys/RAPIDS_BRANCH diff --git a/RAPIDS_BRANCH b/RAPIDS_BRANCH new file mode 100644 index 0000000000..9b1c52d941 --- /dev/null +++ b/RAPIDS_BRANCH @@ -0,0 +1 @@ +branch-25.10 diff --git a/cmake/RAPIDS.cmake b/cmake/RAPIDS.cmake index d112951d3c..40de7cefcd 100644 --- a/cmake/RAPIDS.cmake +++ b/cmake/RAPIDS.cmake @@ -18,9 +18,9 @@ cmake_minimum_required(VERSION 3.30.4 FATAL_ERROR) # Allow users to control which version is used -if(NOT rapids-cmake-version OR NOT rapids-cmake-version MATCHES [[^([0-9][0-9])\.([0-9][0-9])$]]) +if(NOT rapids-cmake-branch OR NOT rapids-cmake-version) message( - FATAL_ERROR "The CMake variable rapids-cmake-version must be defined in the format MAJOR.MINOR." + FATAL_ERROR "The CMake variable `rapids-cmake-branch` or `rapids-cmake-version` must be defined" ) endif() @@ -33,7 +33,7 @@ endif() # Allow users to control which branch is fetched if(NOT rapids-cmake-branch) # Define a default branch if the user doesn't set one - set(rapids-cmake-branch "branch-${rapids-cmake-version}") + set(rapids-cmake-branch "release/${rapids-cmake-version}") endif() # Allow users to control the exact URL passed to FetchContent diff --git a/cmake/rapids_config.cmake b/cmake/rapids_config.cmake index abe468dce8..b706c926e7 100644 --- a/cmake/rapids_config.cmake +++ b/cmake/rapids_config.cmake @@ -26,5 +26,15 @@ else() ) endif() +# Use STRINGS to trim whitespace/newlines +file(STRINGS "${CMAKE_CURRENT_LIST_DIR}/../RAPIDS_BRANCH" _rapids_branch) +if(NOT _rapids_branch) + message( + FATAL_ERROR + "Could not determine branch name to use for checking out rapids-cmake. The file \"${CMAKE_CURRENT_LIST_DIR}/../RAPIDS_BRANCH\" is missing." + ) +endif() + set(rapids-cmake-version "${RAPIDS_VERSION_MAJOR_MINOR}") +set(rapids-cmake-branch "${_rapids_branch}") include("${CMAKE_CURRENT_LIST_DIR}/RAPIDS.cmake") diff --git a/rust/cuvs-sys/RAPIDS_BRANCH b/rust/cuvs-sys/RAPIDS_BRANCH new file mode 120000 index 0000000000..15b9ac7c23 --- /dev/null +++ b/rust/cuvs-sys/RAPIDS_BRANCH @@ -0,0 +1 @@ +../../RAPIDS_BRANCH \ No newline at end of file From 4f7910614ea9175890fa0adf01ce9d1a8434d170 Mon Sep 17 00:00:00 2001 From: rhdong Date: Tue, 29 Jul 2025 22:11:39 -0700 Subject: [PATCH 06/93] [Feat] Add Dockerfile for reproducible installation (#1195) Authors: - rhdong (https://github.com/rhdong) Approvers: - Corey J. Nolet (https://github.com/cjnolet) URL: https://github.com/rapidsai/cuvs/pull/1195 --- Dockerfile | 88 ++++++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 88 insertions(+) create mode 100644 Dockerfile diff --git a/Dockerfile b/Dockerfile new file mode 100644 index 0000000000..2dbdecf9d8 --- /dev/null +++ b/Dockerfile @@ -0,0 +1,88 @@ +# syntax=docker/dockerfile:1.5 + +# ============================================================================== +# RAPIDS cuVS Docker Image - Vector Search and Clustering on GPU +# +# Usage: +# docker build -t cuvs:latest . +# docker run --gpus all -it cuvs:latest +# +# Customizable Build Arguments: +# docker build --build-arg CUDA_VER=12.8 --build-arg RAPIDS_VER=25.06 -t cuvs:custom . +# ============================================================================== + +# Configurable build arguments - override with --build-arg during build +ARG CUDA_VER=12.9.1 + +# Python version for the conda environment (supported: 3.10, 3.11, 3.12) +ARG PYTHON_VER=3.12 + +# RAPIDS/cuVS version +ARG RAPIDS_VER=25.06 + +FROM nvidia/cuda:${CUDA_VER}-devel-ubuntu24.04 + +# Display build configuration for verification +RUN echo " Building cuVS Docker image with:" && \ + echo " CUDA Version: ${CUDA_VER}" && \ + echo " Python Version: ${PYTHON_VER}" && \ + echo " RAPIDS Version: ${RAPIDS_VER}" + +# Container metadata +LABEL maintainer="RAPIDS cuVS Team" +LABEL description="RAPIDS cuVS - Vector Search and Clustering on GPU" +LABEL org.opencontainers.image.source="https://github.com/rapidsai/cuvs" +LABEL org.opencontainers.image.usage="docker run --gpus all -it " + +# Environment setup +ENV DEBIAN_FRONTEND=noninteractive +ENV PYTHONUNBUFFERED=1 +ENV CONDA_DIR=/opt/conda +ENV PATH=${CONDA_DIR}/bin:$PATH + +# Install system dependencies +RUN apt-get update && apt-get install -y --no-install-recommends \ + wget \ + ca-certificates \ + git \ + && rm -rf /var/lib/apt/lists/* + +# Install Miniforge (lightweight conda with conda-forge defaults) +RUN wget --quiet https://github.com/conda-forge/miniforge/releases/latest/download/Miniforge3-Linux-x86_64.sh -O /tmp/miniforge.sh && \ + /bin/bash /tmp/miniforge.sh -b -p ${CONDA_DIR} && \ + rm /tmp/miniforge.sh && \ + ${CONDA_DIR}/bin/conda clean -afy + +# Create conda environment and install cuVS packages with pinned dependencies +RUN conda create -n cuvs python=${PYTHON_VER} -y && \ + conda run -n cuvs conda install -c rapidsai -c conda-forge -c nvidia \ + cuvs=${RAPIDS_VER} \ + libcuvs=${RAPIDS_VER} \ + cuda-version=${CUDA_VER} \ + "numpy>=1.23,<3.0a0" \ + "cupy>=12.0.0" \ + jupyter \ + ipython \ + -y && \ + conda clean -afy + +# Create non-root user +RUN useradd -m -s /bin/bash cuvs +USER cuvs +WORKDIR /home/cuvs + +# Initialize conda for the user and configure environment activation +RUN ${CONDA_DIR}/bin/conda init bash && \ + echo "conda activate cuvs" >> ~/.bashrc && \ + echo "conda activate cuvs" >> ~/.bash_profile + +# Set environment variables for automatic cuvs environment activation +ENV PATH=${CONDA_DIR}/envs/cuvs/bin:$PATH +ENV CONDA_DEFAULT_ENV=cuvs +ENV CONDA_PREFIX=${CONDA_DIR}/envs/cuvs + +# Verify installation and show version info +RUN python -c "import cuvs; print('cuVS version:', cuvs.__version__)" + +# Default command +CMD ["bash"] From 7ecac66f2bf358b361148ffe3f95394d38077ead Mon Sep 17 00:00:00 2001 From: Robert Maynard Date: Wed, 30 Jul 2025 11:04:55 -0400 Subject: [PATCH 07/93] cuvs RAPIDS dependency now use the new rapids-cmake branch info (#1189) This is required to handle when we switch RAPIDS branching strategy Authors: - Robert Maynard (https://github.com/robertmaynard) Approvers: - Bradley Dice (https://github.com/bdice) URL: https://github.com/rapidsai/cuvs/pull/1189 --- cpp/cmake/thirdparty/get_cuvs.cmake | 6 +++--- cpp/cmake/thirdparty/get_raft.cmake | 6 +++--- examples/cmake/thirdparty/get_cuvs.cmake | 4 ++-- 3 files changed, 8 insertions(+), 8 deletions(-) diff --git a/cpp/cmake/thirdparty/get_cuvs.cmake b/cpp/cmake/thirdparty/get_cuvs.cmake index b37249745b..6ef4791474 100644 --- a/cpp/cmake/thirdparty/get_cuvs.cmake +++ b/cpp/cmake/thirdparty/get_cuvs.cmake @@ -1,5 +1,5 @@ # ============================================================================= -# Copyright (c) 2023-2024, NVIDIA CORPORATION. +# Copyright (c) 2023-2025, NVIDIA CORPORATION. # # Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except # in compliance with the License. You may obtain a copy of the License at @@ -14,14 +14,14 @@ # Use RAPIDS_VERSION_MAJOR_MINOR from cmake/rapids_config.cmake set(CUVS_VERSION "${RAPIDS_VERSION_MAJOR_MINOR}") set(CUVS_FORK "rapidsai") -set(CUVS_PINNED_TAG "branch-${RAPIDS_VERSION_MAJOR_MINOR}") +set(CUVS_PINNED_TAG "${rapids-cmake-checkout-tag}") function(find_and_configure_cuvs) set(oneValueArgs VERSION FORK PINNED_TAG ENABLE_NVTX CLONE_ON_PIN BUILD_CPU_ONLY BUILD_SHARED_LIBS) cmake_parse_arguments(PKG "${options}" "${oneValueArgs}" "${multiValueArgs}" ${ARGN} ) - if(PKG_CLONE_ON_PIN AND NOT PKG_PINNED_TAG STREQUAL "branch-${CUVS_VERSION}") + if(PKG_CLONE_ON_PIN AND NOT PKG_PINNED_TAG STREQUAL "${rapids-cmake-checkout-tag}") message(STATUS "cuVS: pinned tag found: ${PKG_PINNED_TAG}. Cloning cuVS locally.") set(CPM_DOWNLOAD_cuvs ON) endif() diff --git a/cpp/cmake/thirdparty/get_raft.cmake b/cpp/cmake/thirdparty/get_raft.cmake index a71dc4cb43..424e829cb6 100644 --- a/cpp/cmake/thirdparty/get_raft.cmake +++ b/cpp/cmake/thirdparty/get_raft.cmake @@ -1,5 +1,5 @@ # ============================================================================= -# Copyright (c) 2023-2024, NVIDIA CORPORATION. +# Copyright (c) 2023-2025, NVIDIA CORPORATION. # # Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except # in compliance with the License. You may obtain a copy of the License at @@ -14,14 +14,14 @@ # Use RAPIDS_VERSION_MAJOR_MINOR from rapids_config.cmake set(RAFT_VERSION "${RAPIDS_VERSION_MAJOR_MINOR}") set(RAFT_FORK "rapidsai") -set(RAFT_PINNED_TAG "branch-${RAPIDS_VERSION_MAJOR_MINOR}") +set(RAFT_PINNED_TAG "${rapids-cmake-checkout-tag}") function(find_and_configure_raft) set(oneValueArgs VERSION FORK PINNED_TAG USE_RAFT_STATIC ENABLE_NVTX ENABLE_MNMG_DEPENDENCIES CLONE_ON_PIN) cmake_parse_arguments(PKG "${options}" "${oneValueArgs}" "${multiValueArgs}" ${ARGN} ) - if(PKG_CLONE_ON_PIN AND NOT PKG_PINNED_TAG STREQUAL "branch-${RAFT_VERSION}") + if(PKG_CLONE_ON_PIN AND NOT PKG_PINNED_TAG STREQUAL "${rapids-cmake-checkout-tag}") message(STATUS "cuVS: RAFT pinned tag found: ${PKG_PINNED_TAG}. Cloning raft locally.") set(CPM_DOWNLOAD_raft ON) elseif(PKG_USE_RAFT_STATIC AND (NOT CPM_raft_SOURCE)) diff --git a/examples/cmake/thirdparty/get_cuvs.cmake b/examples/cmake/thirdparty/get_cuvs.cmake index 49f0ab7017..911adbfa4a 100644 --- a/examples/cmake/thirdparty/get_cuvs.cmake +++ b/examples/cmake/thirdparty/get_cuvs.cmake @@ -1,5 +1,5 @@ # ============================================================================= -# Copyright (c) 2023-2024, NVIDIA CORPORATION. +# Copyright (c) 2023-2025, NVIDIA CORPORATION. # # Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except # in compliance with the License. You may obtain a copy of the License at @@ -14,7 +14,7 @@ # Use RAPIDS_VERSION_MAJOR_MINOR from cmake/rapids_config.cmake set(CUVS_VERSION "${RAPIDS_VERSION_MAJOR_MINOR}") set(CUVS_FORK "rapidsai") -set(CUVS_PINNED_TAG "branch-${RAPIDS_VERSION_MAJOR_MINOR}") +set(CUVS_PINNED_TAG "${rapids-cmake-branch}") function(find_and_configure_cuvs) set(oneValueArgs VERSION FORK PINNED_TAG ENABLE_NVTX BUILD_CUVS_C_LIBRARY) From 02d834623d11a4cda6f971043539711aaf53e281 Mon Sep 17 00:00:00 2001 From: Anupam <54245698+aamijar@users.noreply.github.com> Date: Thu, 31 Jul 2025 08:31:50 -0700 Subject: [PATCH 08/93] Spectral Embedding precomputed graph api (#1197) Resolves #1184 This API is needed so that we have the option to pass in a precomputed connectivity graph. This allows users to build a connectivity graph using their preferred method whether it is through brute_force or other means and then pass the input to the spectral embedding computation. Authors: - Anupam (https://github.com/aamijar) Approvers: - Corey J. Nolet (https://github.com/cjnolet) URL: https://github.com/rapidsai/cuvs/pull/1197 --- .../cuvs/preprocessing/spectral_embedding.hpp | 5 +++ .../spectral/spectral_embedding.cu | 42 +++++++++++++------ 2 files changed, 35 insertions(+), 12 deletions(-) diff --git a/cpp/include/cuvs/preprocessing/spectral_embedding.hpp b/cpp/include/cuvs/preprocessing/spectral_embedding.hpp index 90197e1bd5..3d935bfd9a 100644 --- a/cpp/include/cuvs/preprocessing/spectral_embedding.hpp +++ b/cpp/include/cuvs/preprocessing/spectral_embedding.hpp @@ -48,4 +48,9 @@ void transform(raft::resources const& handle, raft::device_matrix_view dataset, raft::device_matrix_view embedding); +void transform(raft::resources const& handle, + params config, + raft::device_coo_matrix_view connectivity_graph, + raft::device_matrix_view embedding); + } // namespace cuvs::preprocessing::spectral_embedding diff --git a/cpp/src/preprocessing/spectral/spectral_embedding.cu b/cpp/src/preprocessing/spectral/spectral_embedding.cu index d3466802d1..7565ae63b1 100644 --- a/cpp/src/preprocessing/spectral/spectral_embedding.cu +++ b/cpp/src/preprocessing/spectral/spectral_embedding.cu @@ -112,35 +112,35 @@ raft::device_csr_matrix_view coo_to_csr_matrix( raft::resources const& handle, const int n_samples, raft::device_vector_view sym_coo_row_ind, - raft::device_coo_matrix& sym_coo_matrix) + raft::device_coo_matrix_view sym_coo_matrix_view) { auto stream = raft::resource::get_cuda_stream(handle); raft::sparse::op::coo_sort(n_samples, n_samples, - sym_coo_matrix.structure_view().get_nnz(), - sym_coo_matrix.structure_view().get_rows().data(), - sym_coo_matrix.structure_view().get_cols().data(), - sym_coo_matrix.get_elements().data(), + sym_coo_matrix_view.structure_view().get_nnz(), + sym_coo_matrix_view.structure_view().get_rows().data(), + sym_coo_matrix_view.structure_view().get_cols().data(), + sym_coo_matrix_view.get_elements().data(), stream); - raft::sparse::convert::sorted_coo_to_csr(sym_coo_matrix.structure_view().get_rows().data(), - sym_coo_matrix.structure_view().get_nnz(), + raft::sparse::convert::sorted_coo_to_csr(sym_coo_matrix_view.structure_view().get_rows().data(), + sym_coo_matrix_view.structure_view().get_nnz(), sym_coo_row_ind.data_handle(), n_samples, stream); - auto sym_coo_nnz = sym_coo_matrix.structure_view().get_nnz(); + auto sym_coo_nnz = sym_coo_matrix_view.structure_view().get_nnz(); raft::copy(sym_coo_row_ind.data_handle() + sym_coo_row_ind.size() - 1, &sym_coo_nnz, 1, stream); auto csr_matrix_view = raft::make_device_csr_matrix_view( - const_cast(sym_coo_matrix.get_elements().data()), + const_cast(sym_coo_matrix_view.get_elements().data()), raft::make_device_compressed_structure_view( const_cast(sym_coo_row_ind.data_handle()), - const_cast(sym_coo_matrix.structure_view().get_cols().data()), + const_cast(sym_coo_matrix_view.structure_view().get_cols().data()), n_samples, n_samples, - sym_coo_matrix.structure_view().get_nnz())); + sym_coo_matrix_view.structure_view().get_nnz())); return csr_matrix_view; } @@ -250,7 +250,25 @@ void transform(raft::resources const& handle, create_connectivity_graph(handle, spectral_embedding_config, dataset, embedding, sym_coo_matrix); auto csr_matrix_view = - coo_to_csr_matrix(handle, n_samples, sym_coo_row_ind.view(), sym_coo_matrix); + coo_to_csr_matrix(handle, n_samples, sym_coo_row_ind.view(), sym_coo_matrix.view()); + auto laplacian = + create_laplacian(handle, spectral_embedding_config, csr_matrix_view, diagonal.view()); + compute_eigenpairs( + handle, spectral_embedding_config, n_samples, laplacian, diagonal.view(), embedding); +} + +void transform(raft::resources const& handle, + params spectral_embedding_config, + raft::device_coo_matrix_view connectivity_graph, + raft::device_matrix_view embedding) +{ + const int n_samples = connectivity_graph.structure_view().get_n_rows(); + + auto sym_coo_row_ind = raft::make_device_vector(handle, n_samples + 1); + auto diagonal = raft::make_device_vector(handle, n_samples); + + auto csr_matrix_view = + coo_to_csr_matrix(handle, n_samples, sym_coo_row_ind.view(), connectivity_graph); auto laplacian = create_laplacian(handle, spectral_embedding_config, csr_matrix_view, diagonal.view()); compute_eigenpairs( From 99a6b753c9db848fecf0e69cdd0788c6cdd6f40a Mon Sep 17 00:00:00 2001 From: Micka Date: Thu, 31 Jul 2025 21:00:11 +0200 Subject: [PATCH 09/93] Add error handling for CPU-only bench (#1203) Authors: - Micka (https://github.com/lowener) - Corey J. Nolet (https://github.com/cjnolet) Approvers: - Corey J. Nolet (https://github.com/cjnolet) URL: https://github.com/rapidsai/cuvs/pull/1203 --- cpp/bench/ann/src/common/util.hpp | 17 +++++++++++++++-- python/cuvs_bench/cuvs_bench/run/run.py | 5 ++++- 2 files changed, 19 insertions(+), 3 deletions(-) diff --git a/cpp/bench/ann/src/common/util.hpp b/cpp/bench/ann/src/common/util.hpp index ebcae4f73b..b3bcc77950 100644 --- a/cpp/bench/ann/src/common/util.hpp +++ b/cpp/bench/ann/src/common/util.hpp @@ -431,8 +431,21 @@ inline auto cuda_info() cudaRuntimeGetVersion(&runtime); cudaDeviceProp device_prop; - cudaGetDevice(&dev); - cudaGetDeviceProperties(&device_prop, dev); + auto err_code = cudaGetDevice(&dev); + if (err_code == cudaErrorNoDevice || err_code == cudaErrorInvalidDevice) { + throw std::runtime_error{"cuda_info: call to cudaGetDevice failed with code " + + std::to_string(err_code) + + ". If you are running on a CPU-only machine, please " + "use the CPU package instead."}; + } else if (err_code != cudaSuccess) { + throw std::runtime_error{"cuda_info: call to cudaGetDevice failed with code " + + std::to_string(err_code)}; + } + err_code = cudaGetDeviceProperties(&device_prop, dev); + if (err_code != cudaSuccess) { + throw std::runtime_error{"cuda_info: call to cudaGetDeviceProperties failed with code " + + std::to_string(err_code)}; + } props.emplace_back("gpu_name", std::string(device_prop.name)); props.emplace_back("gpu_sm_count", std::to_string(device_prop.multiProcessorCount)); props.emplace_back("gpu_sm_freq", std::to_string(device_prop.clockRate * 1e3)); diff --git a/python/cuvs_bench/cuvs_bench/run/run.py b/python/cuvs_bench/cuvs_bench/run/run.py index bec7766793..8821286951 100644 --- a/python/cuvs_bench/cuvs_bench/run/run.py +++ b/python/cuvs_bench/cuvs_bench/run/run.py @@ -140,7 +140,10 @@ def gather_algorithm_configs( algos_conf_fs = [ os.path.join(scripts_path, "../config", "algos", f) for f in algos_conf_fs - if ".json" not in f and "constraint" not in f and ".py" not in f + if ".json" not in f + and "constraint" not in f + and ".py" not in f + and "__pycache__" not in f ] if configuration: From d29a06a31d31bcd393de7875ffbc9b94f6a0f708 Mon Sep 17 00:00:00 2001 From: Victor Lafargue Date: Thu, 31 Jul 2025 23:16:34 +0200 Subject: [PATCH 10/93] Fix single GPU sharded search merge (#1094) https://github.com/rapidsai/cuvs/issues/904 Authors: - Victor Lafargue (https://github.com/viclafargue) - Corey J. Nolet (https://github.com/cjnolet) Approvers: - Corey J. Nolet (https://github.com/cjnolet) URL: https://github.com/rapidsai/cuvs/pull/1094 --- cpp/src/neighbors/mg/snmg.cuh | 28 +++++++++++++++++++++++++--- 1 file changed, 25 insertions(+), 3 deletions(-) diff --git a/cpp/src/neighbors/mg/snmg.cuh b/cpp/src/neighbors/mg/snmg.cuh index 16aec08b37..abff1f82b6 100644 --- a/cpp/src/neighbors/mg/snmg.cuh +++ b/cpp/src/neighbors/mg/snmg.cuh @@ -1,5 +1,5 @@ /* - * Copyright (c) 2024, NVIDIA CORPORATION. + * Copyright (c) 2024-2025, NVIDIA CORPORATION. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -602,7 +602,7 @@ void search(const raft::resources& clique, int64_t n_batches = raft::ceildiv(n_rows, (int64_t)n_rows_per_batch); if (n_batches <= 1) n_rows_per_batch = n_rows; - if (merge_mode == MERGE_ON_ROOT_RANK) { + if (merge_mode == MERGE_ON_ROOT_RANK && index.num_ranks_ > 1) { RAFT_LOG_DEBUG("SHARDED SEARCH WITH MERGE_ON_ROOT_RANK MERGE MODE: %d*%drows", n_batches, n_rows_per_batch); @@ -617,7 +617,7 @@ void search(const raft::resources& clique, n_cols, n_neighbors, n_batches); - } else if (merge_mode == TREE_MERGE) { + } else if (merge_mode == TREE_MERGE && index.num_ranks_ > 1) { RAFT_LOG_DEBUG( "SHARDED SEARCH WITH TREE_MERGE MERGE MODE %d*%drows", n_batches, n_rows_per_batch); sharded_search_with_tree_merge(clique, @@ -631,6 +631,28 @@ void search(const raft::resources& clique, n_cols, n_neighbors, n_batches); + } else { + const int rank = 0; +#pragma omp parallel for + for (int64_t batch_idx = 0; batch_idx < n_batches; batch_idx++) { + int64_t offset = batch_idx * n_rows_per_batch; + int64_t query_offset = offset * n_cols; + int64_t output_offset = offset * n_neighbors; + int64_t n_rows_of_current_batch = std::min(n_rows_per_batch, n_rows - offset); + + run_search_batch(clique, + index, + rank, + search_params, + queries, + neighbors, + distances, + query_offset, + output_offset, + n_rows_of_current_batch, + n_cols, + n_neighbors); + } } } } From 60e8af26cc72a5336f8d9c19dd95a25462913018 Mon Sep 17 00:00:00 2001 From: Vinay Deshpande Date: Thu, 31 Jul 2025 22:30:37 +0100 Subject: [PATCH 11/93] Remove unused templated instances (#1204) The `fusedDistanceNNMinReduce` is not documented or exposed downstream, so removing it. This reduced libcuvs.so size by ~0.5 MB, when I compiled for single arch. Authors: - Vinay Deshpande (https://github.com/vinaydes) Approvers: - Corey J. Nolet (https://github.com/cjnolet) URL: https://github.com/rapidsai/cuvs/pull/1204 --- cpp/CMakeLists.txt | 1 - cpp/src/distance/detail/fused_distance_nn.cu | 53 -------------------- 2 files changed, 54 deletions(-) delete mode 100644 cpp/src/distance/detail/fused_distance_nn.cu diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt index 7fb6141f9b..50ee1a0ce2 100644 --- a/cpp/CMakeLists.txt +++ b/cpp/CMakeLists.txt @@ -377,7 +377,6 @@ if(BUILD_SHARED_LIBS) src/distance/detail/pairwise_matrix/dispatch_rbf.cu src/distance/detail/pairwise_matrix/dispatch_l2_expanded_double_double_double_int64_t.cu src/distance/detail/pairwise_matrix/dispatch_l2_expanded_float_float_float_int64_t.cu - src/distance/detail/fused_distance_nn.cu src/distance/distance.cu src/distance/pairwise_distance.cu src/distance/sparse_distance.cu diff --git a/cpp/src/distance/detail/fused_distance_nn.cu b/cpp/src/distance/detail/fused_distance_nn.cu deleted file mode 100644 index 6d2e0f3cdf..0000000000 --- a/cpp/src/distance/detail/fused_distance_nn.cu +++ /dev/null @@ -1,53 +0,0 @@ -/* - * Copyright (c) 2024, NVIDIA CORPORATION. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -#include "../fused_distance_nn.cuh" -#include // raft::KeyValuePair - -#include // int64_t - -#define instantiate_cuvs_distance_fusedDistanceNNMinReduce(DataT, OutT, IdxT) \ - template void cuvs::distance::fusedDistanceNNMinReduce( \ - OutT * min, \ - const DataT* x, \ - const DataT* y, \ - const DataT* xn, \ - const DataT* yn, \ - IdxT m, \ - IdxT n, \ - IdxT k, \ - void* workspace, \ - bool sqrt, \ - bool initOutBuffer, \ - bool isRowMajor, \ - cuvs::distance::DistanceType metric, \ - float metric_arg, \ - cudaStream_t stream) - -instantiate_cuvs_distance_fusedDistanceNNMinReduce(float, float, int); -instantiate_cuvs_distance_fusedDistanceNNMinReduce(float, float, int64_t); - -// We can't have comma's in the macro expansion, so we use the COMMA macro: -#define COMMA , - -instantiate_cuvs_distance_fusedDistanceNNMinReduce(float, raft::KeyValuePair, int); -instantiate_cuvs_distance_fusedDistanceNNMinReduce(float, - raft::KeyValuePair, - int64_t); - -#undef COMMA - -#undef instantiate_cuvs_distance_fusedDistanceNNMinReduce From 81e22d3fcaf7a1205135efbdf5be6bbb2abe03a2 Mon Sep 17 00:00:00 2001 From: Ishan Chattopadhyaya Date: Fri, 1 Aug 2025 09:30:50 +0530 Subject: [PATCH 12/93] [Java][Fix] Multithreaded querying fails without synchronization (#1082) Fix CuVSResources usage for multithreaded queries * Added resources as a constructor argument for Builder classes in CagraIndex, HNSWIndex, TieredIndex and BruteForceIndex * Added a multithreaded CAGRA stability test (which fails with improper resources usage). This test uses CheckedCuVSResources. Authors: - Ishan Chattopadhyaya (https://github.com/chatman) Approvers: - MithunR (https://github.com/mythrocks) URL: https://github.com/rapidsai/cuvs/pull/1082 --- .../java/com/nvidia/cuvs/BruteForceQuery.java | 36 +++- .../main/java/com/nvidia/cuvs/CagraQuery.java | 36 +++- .../com/nvidia/cuvs/CagraSearchParams.java | 12 +- .../main/java/com/nvidia/cuvs/HnswQuery.java | 38 +++- .../com/nvidia/cuvs/TieredIndexQuery.java | 43 +++- .../cuvs/internal/BruteForceIndexImpl.java | 2 +- .../nvidia/cuvs/internal/CagraIndexImpl.java | 2 +- .../nvidia/cuvs/internal/HnswIndexImpl.java | 2 +- .../nvidia/cuvs/internal/TieredIndexImpl.java | 2 +- .../nvidia/cuvs/BruteForceAndSearchIT.java | 8 +- .../nvidia/cuvs/BruteForceRandomizedIT.java | 2 +- .../nvidia/cuvs/CagraBuildAndSearchIT.java | 18 +- .../cuvs/CagraMultiThreadStabilityIT.java | 199 ++++++++++++++++++ .../com/nvidia/cuvs/CagraRandomizedIT.java | 4 +- .../com/nvidia/cuvs/HnswBuildAndSearchIT.java | 6 +- .../com/nvidia/cuvs/HnswRandomizedIT.java | 2 +- .../java/com/nvidia/cuvs/TieredIndexIT.java | 18 +- 17 files changed, 371 insertions(+), 59 deletions(-) create mode 100644 java/cuvs-java/src/test/java/com/nvidia/cuvs/CagraMultiThreadStabilityIT.java diff --git a/java/cuvs-java/src/main/java/com/nvidia/cuvs/BruteForceQuery.java b/java/cuvs-java/src/main/java/com/nvidia/cuvs/BruteForceQuery.java index 93a5218aa2..20e9b8e576 100644 --- a/java/cuvs-java/src/main/java/com/nvidia/cuvs/BruteForceQuery.java +++ b/java/cuvs-java/src/main/java/com/nvidia/cuvs/BruteForceQuery.java @@ -17,11 +17,16 @@ import java.util.Arrays; import java.util.BitSet; +import java.util.Objects; import java.util.function.LongToIntFunction; /** * BruteForceQuery holds the query vectors to be used while invoking search. * + *

Thread Safety: Each BruteForceQuery instance should use its own + * CuVSResources object that is not shared with other threads. Sharing CuVSResources + * between threads can lead to memory allocation errors or JVM crashes. + * * @since 25.02 */ public class BruteForceQuery { @@ -31,6 +36,7 @@ public class BruteForceQuery { private final BitSet[] prefilters; private int numDocs = -1; private final int topK; + private final CuVSResources resources; /** * Constructs an instance of {@link BruteForceQuery} using queryVectors, @@ -43,18 +49,21 @@ public class BruteForceQuery { * index * @param numDocs Maximum of bits in each prefilter, representing number of documents in this index. * Used only when prefilter(s) is/are passed. + * @param resources CuVSResources instance to use for this query */ public BruteForceQuery( float[][] queryVectors, LongToIntFunction mapping, int topK, BitSet[] prefilters, - int numDocs) { + int numDocs, + CuVSResources resources) { this.queryVectors = queryVectors; this.mapping = mapping; this.topK = topK; this.prefilters = prefilters; this.numDocs = numDocs; + this.resources = resources; } /** @@ -100,6 +109,15 @@ public int getNumDocs() { return numDocs; } + /** + * Gets the CuVSResources instance for this query. + * + * @return the CuVSResources instance + */ + public CuVSResources getResources() { + return resources; + } + @Override public String toString() { return "BruteForceQuery [mapping=" @@ -123,6 +141,20 @@ public static class Builder { private int numDocs; private LongToIntFunction mapping = SearchResults.IDENTITY_MAPPING; private int topK = 2; + private final CuVSResources resources; + + /** + * Constructor that requires CuVSResources. + * + *

Important: The provided CuVSResources instance should not be + * shared with other threads. Each thread performing searches should create its own + * CuVSResources instance to avoid memory allocation conflicts and potential JVM crashes. + * + * @param resources the CuVSResources instance to use for this query (must not be shared between threads) + */ + public Builder(CuVSResources resources) { + this.resources = Objects.requireNonNull(resources, "resources cannot be null"); + } /** * Registers the query vectors to be passed in the search call. @@ -176,7 +208,7 @@ public Builder withPrefilters(BitSet[] prefilters, int numDocs) { * @return an instance of {@link BruteForceQuery} */ public BruteForceQuery build() { - return new BruteForceQuery(queryVectors, mapping, topK, prefilters, numDocs); + return new BruteForceQuery(queryVectors, mapping, topK, prefilters, numDocs, resources); } } } diff --git a/java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraQuery.java b/java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraQuery.java index cb318cc2b5..f81ec3191c 100644 --- a/java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraQuery.java +++ b/java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraQuery.java @@ -17,12 +17,17 @@ import java.util.Arrays; import java.util.BitSet; +import java.util.Objects; import java.util.function.LongToIntFunction; /** * CagraQuery holds the CagraSearchParams and the query vectors to be used while * invoking search. * + *

Thread Safety: Each CagraQuery instance should use its own + * CuVSResources object that is not shared with other threads. Sharing CuVSResources + * between threads can lead to memory allocation errors or JVM crashes. + * * @since 25.02 */ public class CagraQuery { @@ -33,6 +38,7 @@ public class CagraQuery { private final int topK; private final BitSet prefilter; private final int numDocs; + private final CuVSResources resources; /** * Constructs an instance of {@link CagraQuery} using cagraSearchParameters, @@ -45,6 +51,7 @@ public class CagraQuery { * @param topK the top k results to return * @param prefilter A single BitSet to use as filter while searching the CAGRA index * @param numDocs Total number of dataset vectors; used to align the prefilter correctly + * @param resources CuVSResources instance to use for this query */ public CagraQuery( CagraSearchParams cagraSearchParameters, @@ -52,7 +59,8 @@ public CagraQuery( LongToIntFunction mapping, int topK, BitSet prefilter, - int numDocs) { + int numDocs, + CuVSResources resources) { super(); this.cagraSearchParameters = cagraSearchParameters; this.queryVectors = queryVectors; @@ -60,6 +68,7 @@ public CagraQuery( this.topK = topK; this.prefilter = prefilter; this.numDocs = numDocs; + this.resources = resources; } /** @@ -114,6 +123,15 @@ public int getNumDocs() { return numDocs; } + /** + * Gets the CuVSResources instance for this query. + * + * @return the CuVSResources instance + */ + public CuVSResources getResources() { + return resources; + } + @Override public String toString() { return "CuVSQuery [cagraSearchParameters=" @@ -138,11 +156,20 @@ public static class Builder { private int topK = 2; private BitSet prefilter; private int numDocs; + private final CuVSResources resources; /** - * Default constructor. + * Constructor that requires CuVSResources. + * + *

Important: The provided CuVSResources instance should not be + * shared with other threads. Each thread performing searches should create its own + * CuVSResources instance to avoid memory allocation conflicts and potential JVM crashes. + * + * @param resources the CuVSResources instance to use for this query (must not be shared between threads) */ - public Builder() {} + public Builder(CuVSResources resources) { + this.resources = Objects.requireNonNull(resources, "resources cannot be null"); + } /** * Sets the instance of configured CagraSearchParams to be passed for search. @@ -211,7 +238,8 @@ public Builder withPrefilter(BitSet prefilter, int numDocs) { * @return an instance of CuVSQuery */ public CagraQuery build() { - return new CagraQuery(cagraSearchParams, queryVectors, mapping, topK, prefilter, numDocs); + return new CagraQuery( + cagraSearchParams, queryVectors, mapping, topK, prefilter, numDocs, resources); } } } diff --git a/java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraSearchParams.java b/java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraSearchParams.java index 83211d8cf8..15073db468 100644 --- a/java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraSearchParams.java +++ b/java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraSearchParams.java @@ -98,7 +98,6 @@ private HashMapMode(int value) { /** * Constructs an instance of CagraSearchParams with passed search parameters. * - * @param resources the resources instance to use * @param maxQueries the maximum number of queries to search at the same * time (batch size) * @param iTopKSize the number of intermediate search results retained @@ -120,7 +119,6 @@ private HashMapMode(int value) { * selection */ private CagraSearchParams( - CuVSResources resources, int maxQueries, int iTopKSize, int maxIterations, @@ -303,7 +301,6 @@ public String toString() { */ public static class Builder { - private CuVSResources resources; private int maxQueries; private int iTopKSize = 64; private int maxIterations; @@ -319,13 +316,9 @@ public static class Builder { private HashMapMode hashMapMode; /** - * Constructs this Builder with an instance of Arena. - * - * @param resources the {@link CuVSResources} instance to use + * Default constructor. */ - public Builder(CuVSResources resources) { - this.resources = resources; - } + public Builder() {} /** * Sets the maximum number of queries to search at the same time (batch size). @@ -486,7 +479,6 @@ public Builder withRandXorMask(long randXORMask) { */ public CagraSearchParams build() { return new CagraSearchParams( - resources, maxQueries, iTopKSize, maxIterations, diff --git a/java/cuvs-java/src/main/java/com/nvidia/cuvs/HnswQuery.java b/java/cuvs-java/src/main/java/com/nvidia/cuvs/HnswQuery.java index 9da7fe8814..28fe4cb6f2 100644 --- a/java/cuvs-java/src/main/java/com/nvidia/cuvs/HnswQuery.java +++ b/java/cuvs-java/src/main/java/com/nvidia/cuvs/HnswQuery.java @@ -16,12 +16,17 @@ package com.nvidia.cuvs; import java.util.Arrays; +import java.util.Objects; import java.util.function.LongToIntFunction; /** * HnswQuery holds the query vectors to be used while invoking search on the * HNSW index. * + *

Thread Safety: Each HnswQuery instance should use its own + * CuVSResources object that is not shared with other threads. Sharing CuVSResources + * between threads can lead to memory allocation errors or JVM crashes. + * * @since 25.02 */ public class HnswQuery { @@ -30,6 +35,7 @@ public class HnswQuery { private final LongToIntFunction mapping; private final float[][] queryVectors; private final int topK; + private final CuVSResources resources; /** * Constructs an instance of {@link HnswQuery} using queryVectors, mapping, and @@ -39,16 +45,19 @@ public class HnswQuery { * @param queryVectors 2D float query vector array * @param mapping a function mapping ordinals (neighbor IDs) to custom user IDs * @param topK the top k results to return + * @param resources CuVSResources instance to use for this query */ private HnswQuery( HnswSearchParams hnswSearchParams, float[][] queryVectors, LongToIntFunction mapping, - int topK) { + int topK, + CuVSResources resources) { this.hnswSearchParams = hnswSearchParams; this.queryVectors = queryVectors; this.mapping = mapping; this.topK = topK; + this.resources = resources; } /** @@ -85,6 +94,15 @@ public int getTopK() { return topK; } + /** + * Gets the CuVSResources instance for this query. + * + * @return the CuVSResources instance + */ + public CuVSResources getResources() { + return resources; + } + @Override public String toString() { return "HnswQuery [mapping=" @@ -97,7 +115,7 @@ public String toString() { } /** - * Builder helps configure and create an instance of BruteForceQuery. + * Builder helps configure and create an instance of HnswQuery. */ public static class Builder { @@ -105,6 +123,20 @@ public static class Builder { private float[][] queryVectors; private LongToIntFunction mapping = SearchResults.IDENTITY_MAPPING; private int topK = 2; + private final CuVSResources resources; + + /** + * Constructor that requires CuVSResources. + * + *

Important: The provided CuVSResources instance should not be + * shared with other threads. Each thread performing searches should create its own + * CuVSResources instance to avoid memory allocation conflicts and potential JVM crashes. + * + * @param resources the CuVSResources instance to use for this query (must not be shared between threads) + */ + public Builder(CuVSResources resources) { + this.resources = Objects.requireNonNull(resources, "resources cannot be null"); + } /** * Sets the instance of configured HnswSearchParams to be passed for search. @@ -157,7 +189,7 @@ public Builder withTopK(int topK) { * @return an instance of {@link HnswQuery} */ public HnswQuery build() { - return new HnswQuery(hnswSearchParams, queryVectors, mapping, topK); + return new HnswQuery(hnswSearchParams, queryVectors, mapping, topK, resources); } } } diff --git a/java/cuvs-java/src/main/java/com/nvidia/cuvs/TieredIndexQuery.java b/java/cuvs-java/src/main/java/com/nvidia/cuvs/TieredIndexQuery.java index 03bfae070e..7aea2d7139 100644 --- a/java/cuvs-java/src/main/java/com/nvidia/cuvs/TieredIndexQuery.java +++ b/java/cuvs-java/src/main/java/com/nvidia/cuvs/TieredIndexQuery.java @@ -19,11 +19,15 @@ import java.util.Arrays; import java.util.BitSet; import java.util.List; +import java.util.Objects; /** * TieredIndexQuery holds the search parameters and query vectors to be used - * while - * invoking search. Currently only supports CAGRA index type. + * while invoking search. Currently only supports CAGRA index type. + * + *

Thread Safety: Each TieredIndexQuery instance should use its own + * CuVSResources object that is not shared with other threads. Sharing CuVSResources + * between threads can lead to memory allocation errors or JVM crashes. * * @since 25.02 */ @@ -35,6 +39,7 @@ public class TieredIndexQuery { private int topK; private BitSet prefilter; private long numDocs; + private final CuVSResources resources; private TieredIndexQuery( TieredIndexType indexType, @@ -43,7 +48,8 @@ private TieredIndexQuery( float[][] queryVectors, int topK, BitSet prefilter, - long numDocs) { + long numDocs, + CuVSResources resources) { super(); this.indexType = indexType; this.cagraSearchParameters = cagraSearchParameters; @@ -52,6 +58,7 @@ private TieredIndexQuery( this.topK = topK; this.prefilter = prefilter; this.numDocs = numDocs; + this.resources = resources; } /** @@ -117,6 +124,15 @@ public long getNumDocs() { return numDocs; } + /** + * Gets the CuVSResources instance for this query. + * + * @return the CuVSResources instance + */ + public CuVSResources getResources() { + return resources; + } + @Override public String toString() { return "TieredIndexQuery [indexType=" @@ -135,10 +151,11 @@ public String toString() { /** * Creates a new Builder instance. * + * @param resources the CuVSResources instance to use for this query * @return a new Builder instance */ - public static Builder newBuilder() { - return new Builder(); + public static Builder newBuilder(CuVSResources resources) { + return new Builder(resources); } /** @@ -152,6 +169,20 @@ public static class Builder { private int topK = 2; private BitSet prefilter; private long numDocs; + private final CuVSResources resources; + + /** + * Constructor that requires CuVSResources. + * + *

Important: The provided CuVSResources instance should not be + * shared with other threads. Each thread performing searches should create its own + * CuVSResources instance to avoid memory allocation conflicts and potential JVM crashes. + * + * @param resources the CuVSResources instance to use for this query (must not be shared between threads) + */ + public Builder(CuVSResources resources) { + this.resources = Objects.requireNonNull(resources, "resources cannot be null"); + } /** * Sets the index type for this query. @@ -231,7 +262,7 @@ public Builder withPrefilter(BitSet prefilter, int numDocs) { */ public TieredIndexQuery build() { return new TieredIndexQuery( - indexType, cagraSearchParams, mapping, queryVectors, topK, prefilter, numDocs); + indexType, cagraSearchParams, mapping, queryVectors, topK, prefilter, numDocs, resources); } } } diff --git a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/BruteForceIndexImpl.java b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/BruteForceIndexImpl.java index f32e311f77..4dfef90261 100644 --- a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/BruteForceIndexImpl.java +++ b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/BruteForceIndexImpl.java @@ -220,7 +220,7 @@ public SearchResults search(BruteForceQuery cuvsQuery) throws Throwable { MemorySegment querySeg = buildMemorySegment(localArena, cuvsQuery.getQueryVectors()); int topk = cuvsQuery.getTopK(); - try (var resourcesAccessor = resources.access()) { + try (var resourcesAccessor = cuvsQuery.getResources().access()) { long cuvsResources = resourcesAccessor.handle(); long queriesBytes = C_FLOAT_BYTE_SIZE * numQueries * vectorDimension; diff --git a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CagraIndexImpl.java b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CagraIndexImpl.java index 4fc373e8ca..704bc72b34 100644 --- a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CagraIndexImpl.java +++ b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CagraIndexImpl.java @@ -239,7 +239,7 @@ public SearchResults search(CagraQuery query) throws Throwable { long neighborsBytes = C_INT_BYTE_SIZE * numQueries * topK; long distancesBytes = C_FLOAT_BYTE_SIZE * numQueries * topK; - try (var resourcesAccessor = resources.access()) { + try (var resourcesAccessor = query.getResources().access()) { var cuvsRes = resourcesAccessor.handle(); MemorySegment queriesDP = allocateRMMSegment(cuvsRes, queriesBytes); diff --git a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/HnswIndexImpl.java b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/HnswIndexImpl.java index 1c97422d15..ed25b5cc3f 100644 --- a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/HnswIndexImpl.java +++ b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/HnswIndexImpl.java @@ -111,7 +111,7 @@ public SearchResults search(HnswQuery query) throws Throwable { MemorySegment distancesTensor = prepareTensor(localArena, distancesMemorySegment, distancesShape, 2, 32, 1, 1); - try (var resourcesAccessor = resources.access()) { + try (var resourcesAccessor = query.getResources().access()) { var cuvsRes = resourcesAccessor.handle(); int returnValue = cuvsStreamSync(cuvsRes); checkCuVSError(returnValue, "cuvsStreamSync"); diff --git a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/TieredIndexImpl.java b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/TieredIndexImpl.java index fd94287787..ef9d33cd43 100644 --- a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/TieredIndexImpl.java +++ b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/TieredIndexImpl.java @@ -196,7 +196,7 @@ public SearchResults search(TieredIndexQuery query) throws Throwable { // Get host query data MemorySegment hostQueriesSeg = Util.buildMemorySegment(localArena, query.getQueryVectors()); - try (var resourceAccess = resources.access()) { + try (var resourceAccess = query.getResources().access()) { long cuvsRes = resourceAccess.handle(); long queriesBytes = C_FLOAT_BYTE_SIZE * numQueries * vectorDimension; 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 92f215165b..ddd1be0b5d 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 @@ -106,7 +106,7 @@ public void testIndexingAndSearchingFlow() throws Throwable { try (CuVSResources resources = CheckedCuVSResources.create()) { BruteForceQuery cuvsQuery = - new BruteForceQuery.Builder() + new BruteForceQuery.Builder(resources) .withTopK(3) .withQueryVectors(queries) .withMapping(SearchResults.IDENTITY_MAPPING) @@ -134,7 +134,7 @@ public void testIndexingAndSearchingWithFiltering() throws Throwable { for (int j = 0; j < 10; j++) { try (CuVSResources resources = CheckedCuVSResources.create()) { BruteForceQuery cuvsQuery = - new BruteForceQuery.Builder() + new BruteForceQuery.Builder(resources) .withTopK(3) .withQueryVectors(queries) .withPrefilters( @@ -160,7 +160,7 @@ public void testIndexingAndSearchingWithFunctionMapping() throws Throwable { for (int j = 0; j < 10; j++) { try (CuVSResources resources = CheckedCuVSResources.create()) { BruteForceQuery cuvsQuery = - new BruteForceQuery.Builder() + new BruteForceQuery.Builder(resources) .withTopK(3) .withQueryVectors(queries) .withMapping(rotate) @@ -185,7 +185,7 @@ public void testIndexingAndSearchingWithListMapping() throws Throwable { for (int j = 0; j < 10; j++) { try (CuVSResources resources = CheckedCuVSResources.create()) { BruteForceQuery cuvsQuery = - new BruteForceQuery.Builder() + new BruteForceQuery.Builder(resources) .withTopK(3) .withQueryVectors(queries) .withMapping(rotate) 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 fb338106f2..4a446c249c 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 @@ -109,7 +109,7 @@ private void tmpResultsTopKWithRandomValues(boolean useNativeMemoryDataset) thro try (CuVSResources resources = CheckedCuVSResources.create()) { BruteForceQuery query = - new BruteForceQuery.Builder() + new BruteForceQuery.Builder(resources) .withTopK(topK) .withQueryVectors(queries) .withPrefilters(prefilters, vectors.length) 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 10ea10bcf5..39e8226250 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 @@ -369,9 +369,9 @@ public void testPrefilteringReducesResults() throws Throwable { .build(); // No prefilter (all points allowed) - CagraSearchParams searchParams = new CagraSearchParams.Builder(resources).build(); + CagraSearchParams searchParams = new CagraSearchParams.Builder().build(); CagraQuery fullQuery = - new CagraQuery.Builder() + new CagraQuery.Builder(resources) .withTopK(2) .withSearchParams(searchParams) .withQueryVectors(queries) @@ -387,7 +387,7 @@ public void testPrefilteringReducesResults() throws Throwable { prefilter.set(2); CagraQuery filteredQuery = - new CagraQuery.Builder() + new CagraQuery.Builder(resources) .withTopK(2) .withSearchParams(searchParams) .withQueryVectors(queries) @@ -446,11 +446,11 @@ private void queryAndCompare( CuVSResources resources) throws Throwable { // Configure search parameters - CagraSearchParams searchParams = new CagraSearchParams.Builder(resources).build(); + CagraSearchParams searchParams = new CagraSearchParams.Builder().build(); // Create a query object with the query vectors CagraQuery cuvsQuery = - new CagraQuery.Builder() + new CagraQuery.Builder(resources) .withTopK(3) .withSearchParams(searchParams) .withQueryVectors(queries) @@ -572,10 +572,10 @@ public void testMergingIndexes() throws Throwable { CagraIndex mergedIndex = CagraIndex.merge(new CagraIndex[] {index1, index2}); log.info("Merge completed successfully"); - CagraSearchParams searchParams = new CagraSearchParams.Builder(resources).build(); + CagraSearchParams searchParams = new CagraSearchParams.Builder().build(); CagraQuery query = - new CagraQuery.Builder() + new CagraQuery.Builder(resources) .withTopK(3) .withSearchParams(searchParams) .withQueryVectors(queries) @@ -678,10 +678,10 @@ public void testMergeStrategies() throws Throwable { CagraIndex.merge(new CagraIndex[] {index1, index2}, physicalMergeParams); log.info("Physical merge completed successfully"); - CagraSearchParams searchParams = new CagraSearchParams.Builder(resources).build(); + CagraSearchParams searchParams = new CagraSearchParams.Builder().build(); CagraQuery query = - new CagraQuery.Builder() + new CagraQuery.Builder(resources) .withTopK(3) .withSearchParams(searchParams) .withQueryVectors(queries) 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 new file mode 100644 index 0000000000..50e8f12a5a --- /dev/null +++ b/java/cuvs-java/src/test/java/com/nvidia/cuvs/CagraMultiThreadStabilityIT.java @@ -0,0 +1,199 @@ +/* + * Copyright (c) 2025, NVIDIA CORPORATION. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.nvidia.cuvs; + +import static com.carrotsearch.randomizedtesting.RandomizedTest.assumeTrue; +import static org.junit.Assert.assertNotNull; +import static org.junit.Assert.assertTrue; +import static org.junit.Assert.fail; + +import com.carrotsearch.randomizedtesting.RandomizedRunner; +import com.nvidia.cuvs.CagraIndexParams.CagraGraphBuildAlgo; +import com.nvidia.cuvs.CagraIndexParams.CuvsDistanceType; +import java.lang.invoke.MethodHandles; +import java.util.ArrayList; +import java.util.List; +import java.util.Random; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.Future; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicReference; +import org.junit.Before; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +@RunWith(RandomizedRunner.class) +public class CagraMultiThreadStabilityIT extends CuVSTestCase { + + private static final Logger log = LoggerFactory.getLogger(MethodHandles.lookup().lookupClass()); + + @Before + public void setup() { + assumeTrue("not supported on " + System.getProperty("os.name"), isLinuxAmd64()); + initializeRandom(); + log.info("Multi-threaded stability test initialized"); + } + + @Test + public void testQueryingUsingMultipleThreads() throws Throwable { + final int dataSize = 10000; + final int dimensions = 256; + final int numThreads = 16; // High thread count to increase contention + final int queriesPerThread = 500; + final int queryBatchSize = 1; // Small batch size to increase frequency of calls + final int topK = 10; + + log.info(" Dataset: {}x{}", dataSize, dimensions); + log.info(" Threads: {}, Queries per thread: {}", numThreads, queriesPerThread); + + float[][] dataset = generateRandomDataset(dataSize, dimensions); + + try (CuVSResources resources = CheckedCuVSResources.create()) { + log.info("Creating CAGRA index for MultiThreaded stability test..."); + + CagraIndexParams indexParams = + new CagraIndexParams.Builder() + .withCagraGraphBuildAlgo(CagraGraphBuildAlgo.NN_DESCENT) + .withGraphDegree(64) + .withIntermediateGraphDegree(128) + .withNumWriterThreads(8) + .withMetric(CuvsDistanceType.L2Expanded) + .build(); + + CagraIndex index = + CagraIndex.newBuilder(resources) + .withDataset(dataset) + .withIndexParams(indexParams) + .build(); + + log.info("CAGRA index created, starting high-contention multi-threaded search..."); + + // Create high contention scenario that would fail without using separate resources in every thread + ExecutorService executor = Executors.newFixedThreadPool(numThreads); + List> futures = new ArrayList<>(); + CountDownLatch startLatch = new CountDownLatch(1); + AtomicInteger successfulQueries = new AtomicInteger(0); + AtomicReference firstError = new AtomicReference<>(); + + for (int threadId = 0; threadId < numThreads; threadId++) { + final int finalThreadId = threadId; + + Future future = + executor.submit( + () -> { + try { + // Wait for all threads to start simultaneously + startLatch.await(); + + for (int queryId = 0; queryId < queriesPerThread; queryId++) { + float[][] queries = generateRandomDataset(queryBatchSize, dimensions); + + try (CuVSResources threadResources = CheckedCuVSResources.create()) { + CagraSearchParams searchParams = new CagraSearchParams.Builder().build(); + CagraQuery query = + new CagraQuery.Builder(threadResources) + .withTopK(topK) + .withSearchParams(searchParams) + .withQueryVectors(queries) + .build(); + + // This call should now work with per-thread resources + SearchResults results = index.search(query); + assertNotNull("Query should return results", results); + assertTrue( + "Query should return some results", !results.getResults().isEmpty()); + + successfulQueries.incrementAndGet(); + } + + // No Thread.yield() - maximize contention + } + + log.info("Thread {} completed successfully", finalThreadId); + + } catch (Throwable t) { + log.error("Thread {} failed: {}", finalThreadId, t.getMessage(), t); + firstError.compareAndSet(null, t); + throw new RuntimeException("Thread failed", t); + } + }); + + futures.add(future); + } + + // Start all threads simultaneously to maximize contention + log.info("Starting all {} threads simultaneously...", numThreads); + startLatch.countDown(); + + // Wait for all threads to complete + boolean allCompleted = true; + for (Future future : futures) { + try { + future.get(120, TimeUnit.SECONDS); // Longer timeout for stress test + } catch (Exception e) { + allCompleted = false; + log.error("Thread failed: {}", e.getMessage(), e); + if (firstError.get() == null) { + firstError.set(e); + } + } + } + + executor.shutdown(); + if (!executor.awaitTermination(10, TimeUnit.SECONDS)) { + executor.shutdownNow(); + } + + // Verify results + int expectedTotalQueries = numThreads * queriesPerThread; + int actualSuccessfulQueries = successfulQueries.get(); + + log.info(" Successful queries: {} / {}", actualSuccessfulQueries, expectedTotalQueries); + + if (firstError.get() != null) { + fail( + "MultiThreaded stablity test failed:" + + " " + + firstError.get().getMessage()); + } + + assertTrue("All threads should complete successfully", allCompleted); + assertTrue( + "All queries should complete successfully", + actualSuccessfulQueries == expectedTotalQueries); + + index.destroyIndex(); + } + } + + private float[][] generateRandomDataset(int size, int dimensions) { + Random random = new Random(42 + System.nanoTime()); + float[][] data = new float[size][dimensions]; + + for (int i = 0; i < size; i++) { + for (int j = 0; j < dimensions; j++) { + data[i][j] = random.nextFloat() * 2.0f - 1.0f; + } + } + + return data; + } +} 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 af4965159b..02b5a6526a 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 @@ -143,10 +143,10 @@ private void tmpResultsTopKWithRandomValues(boolean useNativeMemoryDataset) thro try { // Execute search and retrieve results CagraQuery.Builder queryBuilder = - new CagraQuery.Builder() + new CagraQuery.Builder(resources) .withQueryVectors(queries) .withTopK(topK) - .withSearchParams(new CagraSearchParams.Builder(resources).build()); + .withSearchParams(new CagraSearchParams.Builder().build()); if (sharedPrefilter != null) { queryBuilder.withPrefilter(sharedPrefilter, datasetSize); 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 dd0926d200..ab041ec24b 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 @@ -121,7 +121,7 @@ public void testIndexingAndSearchingFlow() throws Throwable { for (int j = 0; j < 10; j++) { try (CuVSResources resources = CheckedCuVSResources.create()) { HnswQuery hnswQuery = - new HnswQuery.Builder() + new HnswQuery.Builder(resources) .withMapping(SearchResults.IDENTITY_MAPPING) .withQueryVectors(queries) .withSearchParams(new HnswSearchParams.Builder().build()) @@ -146,7 +146,7 @@ public void testIndexingAndSearchingWithFunctionMapping() throws Throwable { for (int j = 0; j < 10; j++) { try (CuVSResources resources = CheckedCuVSResources.create()) { HnswQuery hnswQuery = - new HnswQuery.Builder() + new HnswQuery.Builder(resources) .withQueryVectors(queries) .withMapping(rotate) .withSearchParams(new HnswSearchParams.Builder().build()) @@ -172,7 +172,7 @@ public void testIndexingAndSearchingWithListMapping() throws Throwable { for (int j = 0; j < 10; j++) { try (CuVSResources resources = CheckedCuVSResources.create()) { HnswQuery hnswQuery = - new HnswQuery.Builder() + new HnswQuery.Builder(resources) .withQueryVectors(queries) .withMapping(rotate) .withSearchParams(new HnswSearchParams.Builder().build()) 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 ded5bda54f..3000c125a3 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 @@ -154,7 +154,7 @@ private void tmpResultsTopKWithRandomValues(boolean useNativeMemoryDataset) thro new HnswSearchParams.Builder().withNumThreads(32).build(); HnswQuery hnswQuery = - new HnswQuery.Builder() + new HnswQuery.Builder(resources) .withQueryVectors(queries) .withSearchParams(hnswSearchParams) .withTopK(topK) 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 f8cb70b563..89465be53d 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 @@ -89,10 +89,10 @@ public void testBasicOperations() throws Throwable { log.debug("Initial TieredIndex built successfully"); CagraSearchParams searchParams = - new CagraSearchParams.Builder(resources).withMaxIterations(20).build(); + new CagraSearchParams.Builder().withMaxIterations(20).build(); TieredIndexQuery query = - new TieredIndexQuery.Builder() + TieredIndexQuery.newBuilder(resources) .withTopK(3) .withQueryVectors(queries) .withSearchParams(searchParams) @@ -159,11 +159,10 @@ public void testDifferentKValues() throws Throwable { log.debug("TieredIndex built for K-value testing"); TieredIndexQuery query1 = - new TieredIndexQuery.Builder() + TieredIndexQuery.newBuilder(resources) .withTopK(1) .withQueryVectors(queries) - .withSearchParams( - new CagraSearchParams.Builder(resources).withMaxIterations(20).build()) + .withSearchParams(new CagraSearchParams.Builder().withMaxIterations(20).build()) .build(); log.debug("Searching with K=1"); @@ -176,11 +175,10 @@ public void testDifferentKValues() throws Throwable { assertEquals("Distance to closest vector should be ~0.02", 0.02f, firstResult.get(0), 0.01f); TieredIndexQuery query3 = - new TieredIndexQuery.Builder() + TieredIndexQuery.newBuilder(resources) .withTopK(3) .withQueryVectors(queries) - .withSearchParams( - new CagraSearchParams.Builder(resources).withMaxIterations(20).build()) + .withSearchParams(new CagraSearchParams.Builder().withMaxIterations(20).build()) .build(); log.debug("Searching with K=3"); @@ -221,7 +219,7 @@ public void testPrefilter() throws Throwable { .build(); log.debug("TieredIndex built for prefilter testing"); - CagraSearchParams searchParams = new CagraSearchParams.Builder(resources).build(); + CagraSearchParams searchParams = new CagraSearchParams.Builder().build(); BitSet prefilter = new BitSet(4); prefilter.set(1, true); @@ -229,7 +227,7 @@ public void testPrefilter() throws Throwable { log.debug("Created prefilter allowing indices 1 and 2, excluding 0 and 3"); TieredIndexQuery queryWithFilter = - new TieredIndexQuery.Builder() + TieredIndexQuery.newBuilder(resources) .withTopK(3) .withQueryVectors(queryVectors) .withSearchParams(searchParams) From 53d483966763bc2e2cf4a1a674ccdb30d0414cf8 Mon Sep 17 00:00:00 2001 From: Ben Karsin Date: Mon, 4 Aug 2025 08:56:25 -1000 Subject: [PATCH 13/93] Improve Vamana index build performance and recall (#1032) Series of optimizations for the Vamana index build implementation including: - Reduce shared memory usage by the main kernels: GreedySearch and RobustPrune - Replace cub::blockSort with manual merge that reduces shared memory and improves perf - Remove other cub Sort calls to separate kernel to improve resource utilization - Other changes to resource allocation and management to improve performance Other changes that significantly improve recall (closing most of the gap with CPU build and giving the option to further improve recall with parameters): - Re-worked the RobustPrune kernel to use a multi-pass "occlusion" approach. This closes 60-80% of the recall gap seen between previous GPU Vamana and CPU DiskANN. - Changed the "vamana_iters" parameter to a float to allow fractional iterations. This lets us precisely control recall vs. performance to close the gap (or surpass CPU recall). Further optimizations are underway, but this first round address many needs of users to both improve performance and recall. Authors: - Ben Karsin (https://github.com/bkarsin) - Corey J. Nolet (https://github.com/cjnolet) Approvers: - Tamas Bela Feher (https://github.com/tfeher) - James Xia (https://github.com/jamxia155) - Artem M. Chirkin (https://github.com/achirkin) URL: https://github.com/rapidsai/cuvs/pull/1032 --- cpp/include/cuvs/neighbors/vamana.hpp | 3 +- .../neighbors/detail/vamana/greedy_search.cuh | 35 +- cpp/src/neighbors/detail/vamana/macros.cuh | 58 +- .../neighbors/detail/vamana/robust_prune.cuh | 251 ++++---- .../neighbors/detail/vamana/vamana_build.cuh | 556 ++++++++++++------ .../detail/vamana/vamana_structs.cuh | 51 +- cpp/tests/neighbors/ann_vamana.cuh | 7 + examples/cpp/src/vamana_example.cu | 63 +- 8 files changed, 618 insertions(+), 406 deletions(-) diff --git a/cpp/include/cuvs/neighbors/vamana.hpp b/cpp/include/cuvs/neighbors/vamana.hpp index 05bcadf796..733f9c4e71 100644 --- a/cpp/include/cuvs/neighbors/vamana.hpp +++ b/cpp/include/cuvs/neighbors/vamana.hpp @@ -71,7 +71,7 @@ struct index_params : cuvs::neighbors::index_params { * literature **/ uint32_t visited_size = 64; /** Number of Vamana vector insertion iterations (each iteration inserts all vectors). */ - uint32_t vamana_iters = 1; + float vamana_iters = 1.0; /** Alpha for pruning parameter */ float alpha = 1.2; /** Maximum fraction of dataset inserted per batch. * @@ -83,6 +83,7 @@ struct index_params : cuvs::neighbors::index_params { uint32_t queue_size = 127; /** Max batchsize of reverse edge processing (reduces memory footprint) */ uint32_t reverse_batchsize = 1000000; + /** Codebooks and related parameters */ std::optional> codebooks = std::nullopt; }; diff --git a/cpp/src/neighbors/detail/vamana/greedy_search.cuh b/cpp/src/neighbors/detail/vamana/greedy_search.cuh index 4d94bbaa77..9e884fc3ea 100644 --- a/cpp/src/neighbors/detail/vamana/greedy_search.cuh +++ b/cpp/src/neighbors/detail/vamana/greedy_search.cuh @@ -1,5 +1,5 @@ /* - * Copyright (c) 2024, NVIDIA CORPORATION. + * Copyright (c) 2024-2025, NVIDIA CORPORATION. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -16,10 +16,11 @@ #pragma once +#include + #include "macros.cuh" #include "priority_queue.cuh" #include "vamana_structs.cuh" -#include #include #include @@ -52,7 +53,7 @@ __forceinline__ __device__ void sort_visited( } __syncthreads(); - BlockSortT(*sort_mem).Sort(tmp, CmpDist()); + BlockSortT(*sort_mem).Sort(tmp, CmpDist()); __syncthreads(); for (int i = 0; i < ELTS; i++) { @@ -64,6 +65,23 @@ __forceinline__ __device__ void sort_visited( namespace { +template +__global__ void SortPairsKernel(void* query_list_ptr, int num_queries, int topk) +{ + union ShmemLayout { + typename cub::BlockMergeSort, 32, 1>::TempStorage sort_mem; + }; + extern __shared__ __align__(alignof(ShmemLayout)) char smem[]; + + QueryCandidates* query_list = + static_cast*>(query_list_ptr); + + for (int i = blockIdx.x; i < num_queries; i += gridDim.x) { + __syncthreads(); + SEARCH_SELECT_SORT(topk); + } +} + /******************************************************************************************** GPU kernel to perform a batched GreedySearch on a graph. Since this is used for Vamana construction, the entire visited list is kept and stored within the query_list. @@ -87,7 +105,7 @@ __global__ void GreedySearchKernel( int topk, cuvs::distance::DistanceType metric, int max_queue_size, - int sort_smem_size) + Node* topk_pq_mem) { int n = dataset.extent(0); int dim = dataset.extent(1); @@ -105,9 +123,7 @@ __global__ void GreedySearchKernel( union ShmemLayout { // All blocksort sizes have same alignment (16) - typename cub::BlockMergeSort, 32, 1>::TempStorage sort_mem; T coords; - Node topk_pq; int neighborhood_arr; DistPair candidate_queue; }; @@ -117,13 +133,12 @@ __global__ void GreedySearchKernel( // Dynamic shared memory used for blocksort, temp vector storage, and neighborhood list extern __shared__ __align__(alignof(ShmemLayout)) char smem[]; - size_t smem_offset = sort_smem_size; // temp sorting memory takes first chunk + size_t smem_offset = 0; T* s_coords = reinterpret_cast(&smem[smem_offset]); smem_offset += (dim + align_padding) * sizeof(T); - Node* topk_pq = reinterpret_cast*>(&smem[smem_offset]); - smem_offset += topk * sizeof(Node); + Node* topk_pq = &topk_pq_mem[blockIdx.x * topk]; int* neighbor_array = reinterpret_cast(&smem[smem_offset]); smem_offset += degree * sizeof(int); @@ -272,8 +287,6 @@ __global__ void GreedySearchKernel( __syncthreads(); if (self_found) query_list[i].size--; - - SEARCH_SELECT_SORT(topk); } return; diff --git a/cpp/src/neighbors/detail/vamana/macros.cuh b/cpp/src/neighbors/detail/vamana/macros.cuh index c290413a20..4fddfe4522 100644 --- a/cpp/src/neighbors/detail/vamana/macros.cuh +++ b/cpp/src/neighbors/detail/vamana/macros.cuh @@ -1,5 +1,5 @@ /* - * Copyright (c) 2024, NVIDIA CORPORATION. + * Copyright (c) 2024-2025, NVIDIA CORPORATION. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -19,49 +19,25 @@ namespace cuvs::neighbors::vamana::detail { /* Macros to compute the shared memory requirements for CUB primitives used by search and prune */ -#define COMPUTE_SMEM_SIZES(degree, visited_size, DEG, CANDS) \ - if (degree == DEG && visited_size <= CANDS && visited_size > CANDS / 2) { \ - search_smem_sort_size = static_cast( \ - sizeof(typename cub::BlockMergeSort, 32, CANDS / 32>::TempStorage)); \ - \ - prune_smem_sort_size = static_cast(sizeof( \ - typename cub::BlockMergeSort, 32, (CANDS + DEG) / 32>::TempStorage)); \ +#define COMPUTE_SMEM_SIZE(degree, visited_size, DEG, CANDS) \ + if (degree == DEG && visited_size <= CANDS && visited_size > CANDS / 2) { \ + sort_smem_size = static_cast( \ + sizeof(typename cub::BlockMergeSort, 32, CANDS / 32>::TempStorage)); \ } // Current supported sizes for degree and visited_size. Note that visited_size must be > degree -#define SELECT_SMEM_SIZES(degree, visited_size) \ - COMPUTE_SMEM_SIZES(degree, visited_size, 32, 64); \ - COMPUTE_SMEM_SIZES(degree, visited_size, 32, 128); \ - COMPUTE_SMEM_SIZES(degree, visited_size, 32, 256); \ - COMPUTE_SMEM_SIZES(degree, visited_size, 32, 512); \ - COMPUTE_SMEM_SIZES(degree, visited_size, 64, 128); \ - COMPUTE_SMEM_SIZES(degree, visited_size, 64, 256); \ - COMPUTE_SMEM_SIZES(degree, visited_size, 64, 512); \ - COMPUTE_SMEM_SIZES(degree, visited_size, 128, 256); \ - COMPUTE_SMEM_SIZES(degree, visited_size, 128, 512); \ - COMPUTE_SMEM_SIZES(degree, visited_size, 256, 512); \ - COMPUTE_SMEM_SIZES(degree, visited_size, 256, 1024); - -/* Macros to call the CUB BlockSort primitives for supported sizes for ROBUST_PRUNE*/ -#define PRUNE_CALL_SORT(degree, visited_list, DEG, CANDS) \ - if (degree == DEG && visited_list <= CANDS && visited_list > CANDS / 2) { \ - using BlockSortT = cub::BlockMergeSort, 32, (DEG + CANDS) / 32>; \ - auto& sort_mem = reinterpret_cast(smem); \ - sort_edges_and_cands(new_nbh_list, &query_list[i], &sort_mem); \ - } - -#define PRUNE_SELECT_SORT(degree, visited_list) \ - PRUNE_CALL_SORT(degree, visited_size, 32, 64); \ - PRUNE_CALL_SORT(degree, visited_size, 32, 128); \ - PRUNE_CALL_SORT(degree, visited_size, 32, 256); \ - PRUNE_CALL_SORT(degree, visited_size, 32, 512); \ - PRUNE_CALL_SORT(degree, visited_size, 64, 128); \ - PRUNE_CALL_SORT(degree, visited_size, 64, 256); \ - PRUNE_CALL_SORT(degree, visited_size, 64, 512); \ - PRUNE_CALL_SORT(degree, visited_size, 128, 256); \ - PRUNE_CALL_SORT(degree, visited_size, 128, 512); \ - PRUNE_CALL_SORT(degree, visited_size, 256, 512); \ - PRUNE_CALL_SORT(degree, visited_size, 256, 1024); +#define SELECT_SORT_SMEM_SIZE(degree, visited_size) \ + COMPUTE_SMEM_SIZE(degree, visited_size, 32, 64); \ + COMPUTE_SMEM_SIZE(degree, visited_size, 32, 128); \ + COMPUTE_SMEM_SIZE(degree, visited_size, 32, 256); \ + COMPUTE_SMEM_SIZE(degree, visited_size, 32, 512); \ + COMPUTE_SMEM_SIZE(degree, visited_size, 64, 128); \ + COMPUTE_SMEM_SIZE(degree, visited_size, 64, 256); \ + COMPUTE_SMEM_SIZE(degree, visited_size, 64, 512); \ + COMPUTE_SMEM_SIZE(degree, visited_size, 128, 256); \ + COMPUTE_SMEM_SIZE(degree, visited_size, 128, 512); \ + COMPUTE_SMEM_SIZE(degree, visited_size, 256, 512); \ + COMPUTE_SMEM_SIZE(degree, visited_size, 256, 1024); /* Macros to call the CUB BlockSort primitives for supported sizes for GREEDY SEARCH */ #define SEARCH_CALL_SORT(topk, CANDS) \ diff --git a/cpp/src/neighbors/detail/vamana/robust_prune.cuh b/cpp/src/neighbors/detail/vamana/robust_prune.cuh index 182d20c887..0a76093444 100644 --- a/cpp/src/neighbors/detail/vamana/robust_prune.cuh +++ b/cpp/src/neighbors/detail/vamana/robust_prune.cuh @@ -1,5 +1,5 @@ /* - * Copyright (c) 2024, NVIDIA CORPORATION. + * Copyright (c) 2024-2025, NVIDIA CORPORATION. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -49,62 +49,6 @@ __forceinline__ __device__ void load_to_registers(DistPair* tmp, } } -/* Combines edge and candidate lists, removes duplicates, and sorts by distance - * Uses CUB primitives, so needs to be templated. Called with Macros for supported sizes above */ -template -__forceinline__ __device__ void sort_edges_and_cands( - DistPair* new_nbh_list, - QueryCandidates* query, - typename cub::BlockMergeSort, 32, (DEG + CANDS) / 32>::TempStorage* sort_mem) -{ - const int ELTS = (DEG + CANDS) / 32; - using BlockSortT = cub::BlockMergeSort, 32, ELTS>; - DistPair tmp[ELTS]; - - load_to_registers(tmp, query, new_nbh_list); - - __syncthreads(); - BlockSortT(*sort_mem).Sort(tmp, CmpDist()); - __syncthreads(); - - // Mark duplicates and re-sort - // Copy last element over and shuffle to check for duplicate between threads - new_nbh_list[ELTS * threadIdx.x + (ELTS - 1)] = tmp[ELTS - 1]; - if (tmp[ELTS - 1].idx == tmp[ELTS - 2].idx) { - new_nbh_list[ELTS * threadIdx.x + (ELTS - 1)].idx = raft::upper_bound(); - new_nbh_list[ELTS * threadIdx.x + (ELTS - 1)].dist = raft::upper_bound(); - } - __shfl_up_sync(0xffffffff, tmp[ELTS - 1].idx, 1); - __syncthreads(); - - for (int i = ELTS - 2; i > 0; i--) { - if (tmp[i].idx == tmp[i - 1].idx) { - tmp[i].idx = raft::upper_bound(); - tmp[i].dist = raft::upper_bound(); - } - } - if (threadIdx.x == 0) { - if (tmp[0].idx == tmp[ELTS - 1].idx) { - tmp[0].idx = raft::upper_bound(); - tmp[0].dist = raft::upper_bound(); - } - } - - tmp[ELTS - 1].idx = - new_nbh_list[ELTS * threadIdx.x + (ELTS - 1)].idx; // copy back to tmp for re-shuffling - tmp[ELTS - 1].dist = new_nbh_list[ELTS * threadIdx.x + (ELTS - 1)].dist; - - __syncthreads(); - BlockSortT(*sort_mem).Sort(tmp, CmpDist()); - __syncthreads(); - - for (int i = 0; i < ELTS; i++) { - new_nbh_list[ELTS * threadIdx.x + i].idx = tmp[i].idx; - new_nbh_list[ELTS * threadIdx.x + i].dist = tmp[i].dist; - } - __syncthreads(); -} - namespace { /******************************************************************************************** @@ -129,7 +73,7 @@ __global__ void RobustPruneKernel( int visited_size, cuvs::distance::DistanceType metric, float alpha, - int sort_smem_size) + T* s_coords_mem) { int n = dataset.extent(0); int dim = dataset.extent(1); @@ -139,8 +83,7 @@ __global__ void RobustPruneKernel( union ShmemLayout { // All blocksort sizes have same alignment (16) - typename cub::BlockMergeSort, 32, 3>::TempStorage sort_mem; - T coords; + float occlusion; DistPair nbh_list; }; @@ -149,101 +92,161 @@ __global__ void RobustPruneKernel( int align_padding = raft::alignTo(dim, alignof(ShmemLayout)) - dim; - T* s_coords = reinterpret_cast(&smem[sort_smem_size]); - DistPair* new_nbh_list = reinterpret_cast*>( - &smem[(dim + align_padding) * sizeof(T) + sort_smem_size]); + float* occlusion_list = reinterpret_cast(smem); + DistPair* new_nbh_list = + reinterpret_cast*>(&smem[(degree + visited_size) * sizeof(float)]); static __shared__ Point s_query; - s_query.coords = s_coords; + s_query.coords = &s_coords_mem[blockIdx.x * (dim + align_padding)]; s_query.Dim = dim; + static __shared__ int prev_edges; + static __shared__ accT graphDist; for (int i = blockIdx.x; i < num_queries; i += gridDim.x) { int queryId = query_list[i].queryId; - update_shared_point(&s_query, &dataset(0, 0), query_list[i].queryId, dim); + update_shared_point(&s_query, &dataset(0, 0), queryId, dim, i); - // Load new neighbors to be sorted with candidates - for (int j = threadIdx.x; j < degree; j += blockDim.x) { - new_nbh_list[j].idx = graph(queryId, j); - } + int graphIdx = 0; + int listIdx = 0; + int res_size = degree + visited_size; + + // Count total valid edge candidates __syncthreads(); - for (int j = 0; j < degree; j++) { - if (new_nbh_list[j].idx != raft::upper_bound()) { - new_nbh_list[j].dist = - dist(s_query.coords, &dataset((size_t)new_nbh_list[j].idx, 0), dim, metric); - } else { - new_nbh_list[j].dist = raft::upper_bound(); + if (threadIdx.x == 0) { + prev_edges = degree; + for (int j = 0; j < degree; j++) { + if (graph(queryId, j) == raft::upper_bound()) { + prev_edges = j; + break; + } } } + for (int j = threadIdx.x; j < degree + visited_size; j += blockDim.x) { + occlusion_list[j] = 0.0; + } __syncthreads(); - // combine and sort candidates and existing edges (and removes duplicates) - // Resulting list is stored in new_nbh_list - PRUNE_SELECT_SORT(degree, visited_size); - - __syncthreads(); - - // If less than degree total neighbors, don't need to prune - if (new_nbh_list[degree].idx == raft::upper_bound()) { - if (threadIdx.x == 0) { - int writeId = 0; - for (; new_nbh_list[writeId].idx != raft::upper_bound(); writeId++) { - query_list[i].ids[writeId] = new_nbh_list[writeId].idx; - query_list[i].dists[writeId] = new_nbh_list[writeId].dist; - } - query_list[i].size = writeId; - for (; writeId < degree; writeId++) { - query_list[i].ids[writeId] = raft::upper_bound(); - query_list[i].dists[writeId] = raft::upper_bound(); - } + DistPair next_cand; + // Merge graph and candidate list + for (int outIdx = 0; outIdx < degree + visited_size; outIdx++) { + // Check if no more valid elements from graph or list + if (graphIdx < degree && graph(queryId, graphIdx) == raft::upper_bound()) { + graphIdx = degree; } - } else { - // loop through list, writing nearest to visited_list, - // while nulling out violating neighbors in shared memory - if (threadIdx.x == 0) { - query_list[i].ids[0] = new_nbh_list[0].idx; - query_list[i].dists[0] = new_nbh_list[0].dist; + if (listIdx < visited_size && query_list[i].ids[listIdx] == raft::upper_bound()) { + listIdx = visited_size; } - int writeId = 1; - for (int j = 1; j < degree + query_list[i].size && writeId < degree; j++) { - __syncthreads(); - if (new_nbh_list[j].idx == queryId || new_nbh_list[j].idx == raft::upper_bound()) { + // Get next candidate vector for list + if (graphIdx >= degree) { + if (listIdx >= visited_size) { // Fill remaining list if no candidates + if (res_size > outIdx) res_size = outIdx; // Set result size + new_nbh_list[outIdx].idx = raft::upper_bound(); + new_nbh_list[outIdx].dist = raft::upper_bound(); + __syncthreads(); continue; + } else { + next_cand.idx = query_list[i].ids[listIdx]; + next_cand.dist = query_list[i].dists[listIdx]; + listIdx++; } + } else if (listIdx >= visited_size) { + next_cand.idx = graph(queryId, graphIdx); + accT tempDist = + dist(s_query.coords, &dataset((size_t)graph(queryId, graphIdx), 0), dim, metric); + if (threadIdx.x == 0) graphDist = tempDist; __syncthreads(); - if (threadIdx.x == 0) { - query_list[i].ids[writeId] = new_nbh_list[j].idx; - query_list[i].dists[writeId] = new_nbh_list[j].dist; - } - writeId++; + next_cand.dist = graphDist; + graphIdx++; + } else { + accT listDist = query_list[i].dists[listIdx]; + + accT tempDist = + dist(s_query.coords, &dataset((size_t)graph(queryId, graphIdx), 0), dim, metric); + if (threadIdx.x == 0) graphDist = tempDist; __syncthreads(); - update_shared_point(&s_query, &dataset(0, 0), new_nbh_list[j].idx, dim); + if (listDist <= graphDist) { + next_cand.idx = query_list[i].ids[listIdx]; + next_cand.dist = listDist; + + if (graph(queryId, graphIdx) == query_list[i].ids[listIdx]) { // Duplicate found! + graphIdx++; // Skip the duplicate + } + listIdx++; + } else { + next_cand.idx = graph(queryId, graphIdx); + next_cand.dist = graphDist; + graphIdx++; + } + } + + new_nbh_list[outIdx].idx = next_cand.idx; + new_nbh_list[outIdx].dist = next_cand.dist; + } + + // If we need to prune at all... + if (res_size > degree) { + int accept_count = 0; + + // Go through different alpha values. These constants are hard-coded in the MSFT DiskANN code + for (float cur_alpha = 1.0; cur_alpha <= alpha && accept_count < degree; cur_alpha *= 1.2) { + for (int pass_start = 0; pass_start < res_size && accept_count < degree; pass_start++) { + // pick next non-occluded element + if (occlusion_list[pass_start] == raft::lower_bound() || + occlusion_list[pass_start] > cur_alpha) { + continue; // Skip over elements already pruned or already accepted + } + + if (new_nbh_list[pass_start].idx == queryId) { continue; } + + T* cand_ptr = const_cast(&dataset((size_t)(new_nbh_list[pass_start].idx), 0)); - int tot_size = degree + query_list[i].size; - for (int k = j + 1; k < tot_size; k++) { - T* mem_ptr = const_cast(&dataset((size_t)new_nbh_list[k].idx, 0)); - if (new_nbh_list[k].idx != raft::upper_bound()) { - accT dist_starprime = dist(s_query.coords, mem_ptr, dim, metric); - // TODO - create cosine and selector fcn + occlusion_list[pass_start] = raft::lower_bound(); // Mark as "accepted" + accept_count++; - if (threadIdx.x == 0 && alpha * dist_starprime <= new_nbh_list[k].dist) { - new_nbh_list[k].idx = raft::upper_bound(); + // Update rest of the occlusion list + for (int occId = pass_start + 1; occId < res_size; occId++) { + if (occlusion_list[occId] <= alpha && + occlusion_list[occId] != raft::lower_bound()) { + T* k_ptr = const_cast(&dataset((size_t)(new_nbh_list[occId].idx), 0)); + accT djk = dist(cand_ptr, k_ptr, dim, metric); + accT new_occ = (float)(new_nbh_list[occId].dist / djk); + + occlusion_list[occId] = std::max(occlusion_list[occId], new_occ); } } } } - __syncthreads(); - if (threadIdx.x == 0) { query_list[i].size = writeId; } + // Move all "accepted" candidates to front of list and zero out the rest + if (threadIdx.x == 0) { + int out_idx = 1; + for (int read_idx = 1; out_idx < accept_count; read_idx++) { + if (occlusion_list[read_idx] == raft::lower_bound()) { // If it is "accepted" + new_nbh_list[out_idx].idx = new_nbh_list[read_idx].idx; + new_nbh_list[out_idx].dist = new_nbh_list[read_idx].dist; + out_idx++; + } + } + } __syncthreads(); - for (int j = writeId + threadIdx.x; j < degree; - j += blockDim.x) { // Zero out any unfilled neighbors - query_list[i].ids[j] = raft::upper_bound(); - query_list[i].dists[j] = raft::upper_bound(); + for (int out_idx = accept_count + threadIdx.x; out_idx < degree; out_idx++) { + new_nbh_list[out_idx].idx = raft::upper_bound(); + new_nbh_list[out_idx].dist = raft::upper_bound(); } + + if (threadIdx.x == 0) { res_size = accept_count; } + __syncthreads(); + } + + // Copy results out to graph + for (int j = threadIdx.x; j < degree; j += blockDim.x) { + query_list[i].ids[j] = new_nbh_list[j].idx; + query_list[i].dists[j] = new_nbh_list[j].dist; } + if (threadIdx.x == 0) { query_list[i].size = res_size; } } } diff --git a/cpp/src/neighbors/detail/vamana/vamana_build.cuh b/cpp/src/neighbors/detail/vamana/vamana_build.cuh index 9586bf63a2..25c842947f 100644 --- a/cpp/src/neighbors/detail/vamana/vamana_build.cuh +++ b/cpp/src/neighbors/detail/vamana/vamana_build.cuh @@ -75,6 +75,29 @@ void create_insert_permutation(std::vector& insert_order, uint32_t N) } } +template +__global__ void print_mtx(raft::device_vector_view vec) +{ + printf("extents:%ld\n", vec.extent(0)); + for (int i = 0; i < vec.extent(0); i++) { + printf("%d, ", vec(i)); + } + printf("\n"); +} + +template +__global__ void print_queryIds(void* query_list_ptr) +{ + QueryCandidates* query_list = + static_cast*>(query_list_ptr); + + for (int i = 0; i < 50; i++) { + printf("queryId:%d\n", query_list[i].queryId); + } +} + +#define KERNEL_TIMING (RAFT_LOG_ACTIVE_LEVEL <= RAPIDS_LOGGER_LOG_LEVEL_DEBUG) + /******************************************************************************************** * Main Vamana building function - insert vectors into empty graph in batches * Pre - dataset contains the vector data, host matrix allocated to store the graph @@ -92,7 +115,6 @@ void batched_insert_vamana( raft::host_matrix_view graph, IdxT* medoid_id, cuvs::distance::DistanceType metric) -// int dim) { auto stream = raft::resource::get_cuda_stream(res); int N = dataset.extent(0); @@ -101,18 +123,19 @@ void batched_insert_vamana( // Algorithm params int max_batchsize = (int)(params.max_fraction * (float)N); + if (max_batchsize > (int)dataset.extent(0)) { RAFT_LOG_WARN( "Max fraction is the fraction of the total dataset, so it cannot be larger 1.0, reducing it " "to 1.0"); max_batchsize = (int)dataset.extent(0); } - int insert_iters = (int)(params.vamana_iters); - double base = (double)(params.batch_base); - float alpha = (float)(params.alpha); - int visited_size = params.visited_size; - int queue_size = params.queue_size; - int reverse_batch = params.reverse_batchsize; + float insert_iters = (float)(params.vamana_iters); + double base = (double)(params.batch_base); + float alpha = (float)(params.alpha); + int visited_size = params.visited_size; + int queue_size = params.queue_size; + int reverse_batch = params.reverse_batchsize; if ((visited_size & (visited_size - 1)) != 0) { RAFT_LOG_WARN("visited_size must be a power of 2, rounding up."); @@ -122,7 +145,11 @@ void batched_insert_vamana( visited_size = power; } - // create gpu graph and set to all -1s +#if KERNEL_TIMING + auto start_t = std::chrono::system_clock::now(); +#endif + + // Initialize graph with invalid neighbor indices (raft::upper_bound()). auto d_graph = raft::make_device_matrix(res, graph.extent(0), graph.extent(1)); raft::linalg::map(res, d_graph.view(), raft::const_op{raft::upper_bound()}); @@ -150,36 +177,56 @@ void batched_insert_vamana( visited_ids.data_handle(), visited_dists.data_handle(), (int)max_batchsize, - visited_size); + visited_size, + 1); + auto topk_pq_mem = + raft::make_device_mdarray>(res, + raft::resource::get_large_workspace_resource(res), + raft::make_extents(max_batchsize, visited_size)); + + int align_padding = raft::alignTo(dim, 16) - dim; + + auto s_coords_mem = raft::make_device_mdarray( + res, + raft::resource::get_large_workspace_resource(res), + raft::make_extents(min(maxBlocks, max(max_batchsize, reverse_batch)), + dim + align_padding)); // Create random permutation for order of node inserts into graph std::vector insert_order; create_insert_permutation(insert_order, (uint32_t)N); // Calculate the shared memory sizes of each kernel - int search_smem_sort_size = 0; - int prune_smem_sort_size = 0; - SELECT_SMEM_SIZES(degree, visited_size); // Sets above 2 variables to appropriate sizes + int sort_smem_size = 0; + SELECT_SORT_SMEM_SIZE(degree, visited_size); // Sets sort_smem_size based on dataset // Total dynamic shared memory used by GreedySearch - int align_padding = raft::alignTo(dim, 16) - dim; - int search_smem_total_size = static_cast( - search_smem_sort_size + (dim + align_padding) * sizeof(T) + visited_size * sizeof(Node) + - degree * sizeof(int) + queue_size * sizeof(DistPair)); + int search_smem_total_size = + static_cast((dim + align_padding) * sizeof(T) + // visited_size * sizeof(Node) + + degree * sizeof(int) + queue_size * sizeof(DistPair)); // Total dynamic shared memory size needed by both RobustPrune calls - int prune_smem_total_size = prune_smem_sort_size + (dim + align_padding) * sizeof(T) + + int prune_smem_total_size = (degree + visited_size) * sizeof(float) + // Occlusion list (degree + visited_size) * sizeof(DistPair); - RAFT_LOG_DEBUG("Dynamic shared memory usage (bytes): GreedySearch: %d, RobustPrune: %d", - search_smem_total_size, - prune_smem_total_size); - - if (prune_smem_sort_size == 0) { // If sizes not supported, smem sizes will be 0 - RAFT_FAIL("Vamana graph parameters not supported: graph_degree=%d, visited_size:%d\n", - degree, - visited_size); - } + RAFT_LOG_DEBUG( + "Dynamic shared memory usage (bytes): GreedySearch: %d, Segment Sort: %d, Robust Prune: %d", + search_smem_total_size, + sort_smem_size, + prune_smem_total_size); + +#if KERNEL_TIMING + auto end_t = std::chrono::system_clock::now(); + std::chrono::duration elapsed_seconds = end_t - start_t; + double alloc_time = elapsed_seconds.count(); + + double search_time = 0.0; + double segment_sort_time = 0.0; + double prune1_time = 0.0; + double write1_time = 0.0; + double rev_time = 0.0; + double batch_prune = 0.0; +#endif // Random medoid has minor impact on recall // TODO: use heuristic for better medoid selection, issue: @@ -188,88 +235,146 @@ void batched_insert_vamana( // size of current batch of inserts, increases logarithmically until max_batchsize int step_size = 1; - // Number of passes over dataset (default 1) - for (int iter = 0; iter < insert_iters; iter++) { - // Loop through batches and call the insert and prune kernels - for (int start = 0; start < N;) { - if (start + step_size > N) { - int new_size = N - start; - step_size = new_size; - } - RAFT_LOG_DEBUG("Starting batch of inserts indices_start:%d, batch_size:%d", start, step_size); - - int num_blocks = min(maxBlocks, step_size); - - // Copy ids to be inserted for this batch - raft::copy(query_ids.data_handle(), &insert_order.data()[start], step_size, stream); - set_query_ids<<>>( - query_list_ptr.data_handle(), query_ids.data_handle(), step_size); - - // Call greedy search to get candidates for every vector being inserted - GreedySearchKernel - <<>>(d_graph.view(), - dataset, - query_list_ptr.data_handle(), - step_size, - *medoid_id, - visited_size, - metric, - queue_size, - search_smem_sort_size); - // Run on candidates of vectors being inserted - RobustPruneKernel - <<>>(d_graph.view(), - dataset, - query_list_ptr.data_handle(), - step_size, - visited_size, - metric, - alpha, - prune_smem_sort_size); - - // Write results from first prune to graph edge list - write_graph_edges_kernel<<>>( - d_graph.view(), query_list_ptr.data_handle(), degree, step_size); - - // compute prefix sums of query_list sizes - TODO parallelize prefix sums - auto d_total_edges = raft::make_device_mdarray( - res, raft::resource::get_workspace_resource(res), raft::make_extents(1)); - prefix_sums_sizes - <<<1, 1, 0, stream>>>(query_list, step_size, d_total_edges.data_handle()); - - int total_edges; - raft::copy(&total_edges, d_total_edges.data_handle(), 1, stream); - RAFT_CUDA_TRY(cudaStreamSynchronize(stream)); - - auto edge_dest = - raft::make_device_mdarray(res, - raft::resource::get_large_workspace_resource(res), - raft::make_extents(total_edges)); - auto edge_src = - raft::make_device_mdarray(res, - raft::resource::get_large_workspace_resource(res), - raft::make_extents(total_edges)); - - // Create reverse edge list - create_reverse_edge_list - <<>>(query_list_ptr.data_handle(), - step_size, - degree, - edge_src.data_handle(), - edge_dest.data_handle()); - + // Loop through batches and call the insert and prune kernels - can insert > N times based on + // iters parameter + for (int start = 0; start < (int)(insert_iters * (float)N);) { +#if KERNEL_TIMING + start_t = std::chrono::system_clock::now(); +#endif + + if (start + step_size > (int)(insert_iters * (float)N)) { + step_size = (int)(insert_iters * (float)N) - start; + } + if (start + step_size > N) { step_size = N - start; } + RAFT_LOG_DEBUG("Starting batch of inserts indices_start:%d, batch_size:%d", start, step_size); + + int num_blocks = min(maxBlocks, step_size); + + // Copy ids to be inserted for this batch + raft::copy(query_ids.data_handle(), &insert_order.data()[start], step_size, stream); + set_query_ids<<>>( + query_list_ptr.data_handle(), query_ids.data_handle(), step_size); + + // Call greedy search to get candidates for every vector being inserted + GreedySearchKernel + <<>>(d_graph.view(), + dataset, + query_list_ptr.data_handle(), + step_size, + *medoid_id, + visited_size, + metric, + queue_size, + topk_pq_mem.data_handle()); + RAFT_CUDA_TRY(cudaPeekAtLastError()); + +#if KERNEL_TIMING + RAFT_CUDA_TRY(cudaDeviceSynchronize()); + end_t = std::chrono::system_clock::now(); + elapsed_seconds = end_t - start_t; + search_time += elapsed_seconds.count(); + start_t = std::chrono::system_clock::now(); +#endif + + // Segmented sort on query list + SortPairsKernel<<>>( + query_list_ptr.data_handle(), step_size, visited_size); + RAFT_CUDA_TRY(cudaPeekAtLastError()); + +#if KERNEL_TIMING + RAFT_CUDA_TRY(cudaDeviceSynchronize()); + end_t = std::chrono::system_clock::now(); + elapsed_seconds = end_t - start_t; + segment_sort_time += elapsed_seconds.count(); + start_t = std::chrono::system_clock::now(); +#endif + + // Run on candidates of vectors being inserted + RobustPruneKernel + <<>>(d_graph.view(), + dataset, + query_list_ptr.data_handle(), + step_size, + visited_size, + metric, + alpha, + s_coords_mem.data_handle()); + RAFT_CUDA_TRY(cudaPeekAtLastError()); + + // Segmented sort on query list + SortPairsKernel<<>>( + query_list_ptr.data_handle(), step_size, degree); + RAFT_CUDA_TRY(cudaPeekAtLastError()); + +#if KERNEL_TIMING + RAFT_CUDA_TRY(cudaDeviceSynchronize()); + end_t = std::chrono::system_clock::now(); + elapsed_seconds = end_t - start_t; + prune1_time += elapsed_seconds.count(); + start_t = std::chrono::system_clock::now(); +#endif + + // Write results from first prune to graph edge list + write_graph_edges_kernel<<>>( + d_graph.view(), query_list_ptr.data_handle(), degree, step_size); + RAFT_CUDA_TRY(cudaPeekAtLastError()); + +#if KERNEL_TIMING + RAFT_CUDA_TRY(cudaDeviceSynchronize()); + end_t = std::chrono::system_clock::now(); + elapsed_seconds = end_t - start_t; + write1_time += elapsed_seconds.count(); + start_t = std::chrono::system_clock::now(); +#endif + + // compute prefix sums of query_list sizes - TODO parallelize prefix sums + // auto d_total_edges = raft::make_device_mdarray( + // res, raft::resource::get_workspace_resource(res), raft::make_extents(1)); + rmm::device_scalar d_total_edges(stream); + prefix_sums_sizes<<<1, 1, 0, stream>>>(query_list, step_size, d_total_edges.data()); + RAFT_CUDA_TRY(cudaPeekAtLastError()); + + int total_edges = d_total_edges.value(stream); + // raft::copy(&total_edges, d_total_edges.data_handle(), 1, stream); + // RAFT_CUDA_TRY(cudaStreamSynchronize(stream)); + + auto edge_dist_pair = raft::make_device_mdarray>( + res, + raft::resource::get_large_workspace_resource(res), + raft::make_extents(total_edges)); + + auto edge_dest = + raft::make_device_mdarray(res, + raft::resource::get_large_workspace_resource(res), + raft::make_extents(total_edges)); + auto edge_src = + raft::make_device_mdarray(res, + raft::resource::get_large_workspace_resource(res), + raft::make_extents(total_edges)); + + // Create reverse edge list + create_reverse_edge_list + <<>>(query_list_ptr.data_handle(), + step_size, + degree, + edge_src.data_handle(), + edge_dist_pair.data_handle()); + RAFT_CUDA_TRY(cudaPeekAtLastError()); + + { + // Sort by dists first so final edge lists are each sorted by dist void* d_temp_storage = nullptr; size_t temp_storage_bytes = 0; cub::DeviceMergeSort::SortPairs(d_temp_storage, temp_storage_bytes, - edge_dest.data_handle(), + edge_dist_pair.data_handle(), edge_src.data_handle(), total_edges, - CmpEdge(), + CmpDist(), stream); - RAFT_LOG_DEBUG("Temp storage needed for sorting (bytes): %lu", temp_storage_bytes); + RAFT_LOG_DEBUG("Temp storage needed for sorting dist (bytes): %lu", temp_storage_bytes); auto temp_sort_storage = raft::make_device_mdarray( res, @@ -279,100 +384,177 @@ void batched_insert_vamana( // Sort to group reverse edges by destination cub::DeviceMergeSort::SortPairs(temp_sort_storage.data_handle(), temp_storage_bytes, - edge_dest.data_handle(), + edge_dist_pair.data_handle(), edge_src.data_handle(), total_edges, - CmpEdge(), + CmpDist(), stream); + } - // Get number of unique node destinations - IdxT unique_dests = - cuvs::sparse::neighbors::get_n_components(edge_dest.data_handle(), total_edges, stream); - - // Find which node IDs have reverse edges and their indices in the reverse edge list - thrust::device_vector edge_dest_vec(edge_dest.data_handle(), - edge_dest.data_handle() + total_edges); - auto unique_indices = raft::make_device_vector(res, total_edges); - raft::linalg::map_offset(res, unique_indices.view(), raft::identity_op{}); - - thrust::unique_by_key( - edge_dest_vec.begin(), edge_dest_vec.end(), unique_indices.data_handle()); - - edge_dest_vec.clear(); - edge_dest_vec.shrink_to_fit(); - - // Batch execution of reverse edge creation/application - reverse_batch = params.reverse_batchsize; - for (int rev_start = 0; rev_start < (int)unique_dests; rev_start += reverse_batch) { - if (rev_start + reverse_batch > (int)unique_dests) { - reverse_batch = (int)unique_dests - rev_start; - } - - // Allocate reverse QueryCandidate list based on number of unique destinations - auto reverse_list_ptr = raft::make_device_mdarray>( - res, - raft::resource::get_large_workspace_resource(res), - raft::make_extents(reverse_batch)); - auto rev_ids = - raft::make_device_mdarray(res, - raft::resource::get_large_workspace_resource(res), - raft::make_extents(reverse_batch, visited_size)); - auto rev_dists = - raft::make_device_mdarray(res, - raft::resource::get_large_workspace_resource(res), - raft::make_extents(reverse_batch, visited_size)); - - QueryCandidates* reverse_list = - static_cast*>(reverse_list_ptr.data_handle()); - - init_query_candidate_list<<<256, blockD, 0, stream>>>(reverse_list, - rev_ids.data_handle(), - rev_dists.data_handle(), - (int)reverse_batch, - visited_size); - - // May need more blocks for reverse list - num_blocks = min(maxBlocks, reverse_batch); - - // Populate reverse list ids and candidate lists from edge_src and edge_dest - populate_reverse_list_struct - <<>>(reverse_list, - edge_src.data_handle(), - edge_dest.data_handle(), - unique_indices.data_handle(), - unique_dests, - total_edges, - dataset.extent(0), - rev_start, - reverse_batch); - - // Recompute distances (avoided keeping it during sorting) - recompute_reverse_dists - <<>>(reverse_list, dataset, reverse_batch, metric); - - // Call 2nd RobustPrune on reverse query_list - RobustPruneKernel - <<>>(d_graph.view(), - raft::make_const_mdspan(dataset), - reverse_list_ptr.data_handle(), - reverse_batch, - visited_size, - metric, - alpha, - prune_smem_sort_size); - - // Write new edge lists to graph - write_graph_edges_kernel<<>>( - d_graph.view(), reverse_list_ptr.data_handle(), degree, reverse_batch); + /* + DistPair* temp_ptr = edge_dist_pair.data_handle(); + raft::linalg::map_offset( + res, edge_dest.view(), [temp_ptr] __device__(size_t i) { return temp_ptr[i].idx; }); + */ + raft::linalg::map( + res, + edge_dest.view(), + [] __device__(auto x) { return x.idx; }, + raft::make_const_mdspan(edge_dist_pair.view())); + + void* d_temp_storage = nullptr; + size_t temp_storage_bytes = 0; + + cub::DeviceMergeSort::SortPairs(d_temp_storage, + temp_storage_bytes, + edge_dest.data_handle(), + edge_src.data_handle(), + total_edges, + CmpEdge(), + stream); + + RAFT_LOG_DEBUG("Temp storage needed for sorting (bytes): %lu", temp_storage_bytes); + + auto temp_sort_storage = raft::make_device_mdarray( + res, + raft::resource::get_large_workspace_resource(res), + raft::make_extents(temp_storage_bytes / sizeof(IdxT))); + + // Sort to group reverse edges by destination + cub::DeviceMergeSort::SortPairs(temp_sort_storage.data_handle(), + temp_storage_bytes, + edge_dest.data_handle(), + edge_src.data_handle(), + total_edges, + CmpEdge(), + stream); + + // Get number of unique node destinations + IdxT unique_dests = + cuvs::sparse::neighbors::get_n_components(edge_dest.data_handle(), total_edges, stream); + + // Find which node IDs have reverse edges and their indices in the reverse edge list + thrust::device_vector edge_dest_vec(edge_dest.data_handle(), + edge_dest.data_handle() + total_edges); + auto unique_indices = raft::make_device_vector(res, total_edges); + raft::linalg::map_offset(res, unique_indices.view(), raft::identity_op{}); + + thrust::unique_by_key(edge_dest_vec.begin(), edge_dest_vec.end(), unique_indices.data_handle()); + + edge_dest_vec.clear(); + edge_dest_vec.shrink_to_fit(); + +#if KERNEL_TIMING + RAFT_CUDA_TRY(cudaDeviceSynchronize()); + end_t = std::chrono::system_clock::now(); + elapsed_seconds = end_t - start_t; + rev_time += elapsed_seconds.count(); + start_t = std::chrono::system_clock::now(); +#endif + + // Batch execution of reverse edge creation/application + reverse_batch = params.reverse_batchsize; + for (int rev_start = 0; rev_start < (int)unique_dests; rev_start += reverse_batch) { + if (rev_start + reverse_batch > (int)unique_dests) { + reverse_batch = (int)unique_dests - rev_start; } - start += step_size; - step_size *= base; - if (step_size > max_batchsize) step_size = max_batchsize; + // Allocate reverse QueryCandidate list based on number of unique destinations + auto reverse_list_ptr = raft::make_device_mdarray>( + res, + raft::resource::get_large_workspace_resource(res), + raft::make_extents(reverse_batch)); + auto rev_ids = + raft::make_device_mdarray(res, + raft::resource::get_large_workspace_resource(res), + raft::make_extents(reverse_batch, visited_size)); + + auto rev_dists = + raft::make_device_mdarray(res, + raft::resource::get_large_workspace_resource(res), + raft::make_extents(reverse_batch, visited_size)); + + QueryCandidates* reverse_list = + static_cast*>(reverse_list_ptr.data_handle()); + + init_query_candidate_list<<<256, blockD, 0, stream>>>(reverse_list, + rev_ids.data_handle(), + rev_dists.data_handle(), + (int)reverse_batch, + visited_size); + + // May need more blocks for reverse list + num_blocks = min(maxBlocks, reverse_batch); + + // Populate reverse list ids and candidate lists from edge_src and edge_dest + populate_reverse_list_struct + <<>>(reverse_list, + edge_src.data_handle(), + edge_dest.data_handle(), + unique_indices.data_handle(), + unique_dests, + total_edges, + dataset.extent(0), + rev_start, + reverse_batch); + RAFT_CUDA_TRY(cudaPeekAtLastError()); + + // Recompute distances (avoided keeping it during sorting) + recompute_reverse_dists + <<>>(reverse_list, dataset, reverse_batch, metric); + RAFT_CUDA_TRY(cudaPeekAtLastError()); + + // Call 2nd RobustPrune on reverse query_list + RobustPruneKernel + <<>>(d_graph.view(), + raft::make_const_mdspan(dataset), + reverse_list_ptr.data_handle(), + reverse_batch, + visited_size, + metric, + alpha, + s_coords_mem.data_handle()); + RAFT_CUDA_TRY(cudaPeekAtLastError()); + + // Segmented sort on reverse_list + SortPairsKernel<<>>( + reverse_list_ptr.data_handle(), reverse_batch, degree); + RAFT_CUDA_TRY(cudaPeekAtLastError()); - } // Batch of inserts + // Write new edge lists to graph + write_graph_edges_kernel<<>>( + d_graph.view(), reverse_list_ptr.data_handle(), degree, reverse_batch); + RAFT_CUDA_TRY(cudaPeekAtLastError()); + } - } // insert iterations +#if KERNEL_TIMING + RAFT_CUDA_TRY(cudaDeviceSynchronize()); + end_t = std::chrono::system_clock::now(); + elapsed_seconds = end_t - start_t; + batch_prune += elapsed_seconds.count(); +#endif + + start += step_size; + if (start >= N) { + start = 0; + insert_iters -= 1.0; + step_size = max_batchsize; + } + step_size *= base; + step_size = min(step_size, max_batchsize); + + } // Batch of inserts + +#if KERNEL_TIMING + printf("intro:%lf\ngreedy:%lf\nseg_sort:%lf\nprune1:%lf\nwrite1:%lf\nrev:%lf\nbatch_prune:%lf\n", + alloc_time, + search_time, + segment_sort_time, + prune1_time, + write1_time, + rev_time, + batch_prune); +#endif raft::copy(graph.data_handle(), d_graph.data_handle(), d_graph.size(), stream); @@ -413,6 +595,9 @@ index build( RAFT_EXPECTS(params.visited_size > graph_degree, "visited_size must be > graph_degree"); + RAFT_EXPECTS(params.vamana_iters >= 1.0, + "vamana_iters must be at least 1.0 to insert the entire input dataset"); + int dim = dataset.extent(1); RAFT_LOG_DEBUG("Creating empty graph structure"); @@ -421,7 +606,6 @@ index build( RAFT_LOG_DEBUG("Running Vamana batched insert algorithm"); cuvs::distance::DistanceType metric = cuvs::distance::DistanceType::L2Expanded; - IdxT medoid_id; batched_insert_vamana( res, params, dataset, vamana_graph.view(), &medoid_id, metric); diff --git a/cpp/src/neighbors/detail/vamana/vamana_structs.cuh b/cpp/src/neighbors/detail/vamana/vamana_structs.cuh index 0cf29a5b42..d5812e1ff6 100644 --- a/cpp/src/neighbors/detail/vamana/vamana_structs.cuh +++ b/cpp/src/neighbors/detail/vamana/vamana_structs.cuh @@ -1,5 +1,5 @@ /* - * Copyright (c) 2024, NVIDIA CORPORATION. + * Copyright (c) 2024-2025, NVIDIA CORPORATION. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -50,20 +50,6 @@ template struct __align__(16) DistPair { accT dist; IdxT idx; - - __device__ __host__ DistPair& operator=(const DistPair& other) - { - dist = other.dist; - idx = other.idx; - return *this; - } - - __device__ __host__ DistPair& operator=(const volatile DistPair& other) - { - dist = other.dist; - idx = other.idx; - return *this; - } }; // Swap the values of two DistPair objects @@ -80,8 +66,8 @@ __device__ __host__ void swap(DistPair* a, DistPair* b) } // Structure to sort by distance +template struct CmpDist { - template __device__ bool operator()(const DistPair& lhs, const DistPair& rhs) { return lhs.dist < rhs.dist; @@ -309,19 +295,20 @@ __global__ void init_query_candidate_list(QueryCandidates* query_lis IdxT* visited_id_ptr, accT* visited_dist_ptr, int num_queries, - int maxSize) + int maxSize, + int extra_queries_in_list = 0) { IdxT* ids_ptr = static_cast(visited_id_ptr); accT* dist_ptr = static_cast(visited_dist_ptr); for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < num_queries * maxSize; - i += blockDim.x + gridDim.x) { + i += blockDim.x * gridDim.x) { ids_ptr[i] = raft::upper_bound(); dist_ptr[i] = raft::upper_bound(); } - for (size_t i = blockIdx.x * blockDim.x + threadIdx.x; i < num_queries; - i += blockDim.x + gridDim.x) { + for (size_t i = blockIdx.x * blockDim.x + threadIdx.x; i < num_queries + extra_queries_in_list; + i += blockDim.x * gridDim.x) { query_list[i].maxSize = maxSize; query_list[i].size = 0; query_list[i].ids = &ids_ptr[i * (size_t)(maxSize)]; @@ -359,6 +346,18 @@ __global__ void prefix_sums_sizes(QueryCandidates* query_list, } } +// Device fcn to have a threadblock copy coordinates into shared memory +template +__device__ void update_shared_point( + Point* shared_point, const T* data_ptr, int id, int dim, int idx) +{ + shared_point->id = id; + shared_point->Dim = dim; + for (size_t i = threadIdx.x; i < dim; i += blockDim.x) { + shared_point->coords[i] = data_ptr[(size_t)(id) * (size_t)(dim) + i]; + } +} + // Device fcn to have a threadblock copy coordinates into shared memory template __device__ void update_shared_point(Point* shared_point, @@ -392,8 +391,11 @@ __global__ void write_graph_edges_kernel(raft::device_matrix_view // Create src and dest edge lists used to sort and create reverse edges template -__global__ void create_reverse_edge_list( - void* query_list_ptr, int num_queries, int degree, IdxT* edge_src, IdxT* edge_dest) +__global__ void create_reverse_edge_list(void* query_list_ptr, + int num_queries, + int degree, + IdxT* edge_src, + DistPair* edge_dest) { QueryCandidates* query_list = static_cast*>(query_list_ptr); @@ -404,8 +406,9 @@ __global__ void create_reverse_edge_list( int cand_count = query_list[i + 1].size - query_list[i].size; for (int j = 0; j < cand_count; j++) { - edge_src[query_list[i].size + j] = query_list[i].queryId; - edge_dest[query_list[i].size + j] = query_list[i].ids[j]; + edge_src[query_list[i].size + j] = query_list[i].queryId; + edge_dest[query_list[i].size + j].idx = query_list[i].ids[j]; + edge_dest[query_list[i].size + j].dist = query_list[i].dists[j]; } } } diff --git a/cpp/tests/neighbors/ann_vamana.cuh b/cpp/tests/neighbors/ann_vamana.cuh index 37a9447e95..bb8cd9bd72 100644 --- a/cpp/tests/neighbors/ann_vamana.cuh +++ b/cpp/tests/neighbors/ann_vamana.cuh @@ -66,6 +66,7 @@ struct AnnVamanaInputs { cuvs::distance::DistanceType metric; bool host_dataset; int reverse_batchsize; + double insert_iters; // cagra search params int n_queries; @@ -146,6 +147,8 @@ class AnnVamanaTest : public ::testing::TestWithParam { index_params.visited_size = ps.visited_size; index_params.max_fraction = ps.max_fraction; index_params.reverse_batchsize = ps.reverse_batchsize; + index_params.vamana_iters = ps.insert_iters; + // use randomized codebooks to test serialization & quantization code path if (ps.dim == 384 && std::is_same_v) index_params.codebooks = vamana::deserialize_codebooks( @@ -299,6 +302,7 @@ inline std::vector generate_inputs() {cuvs::distance::DistanceType::L2Expanded}, {false}, {100, 1000000}, + {1.0, 1.5}, {100}, {10}, {cagra::search_algo::AUTO}, @@ -316,6 +320,7 @@ inline std::vector generate_inputs() {cuvs::distance::DistanceType::L2Expanded}, {false}, {1000000}, + {1.0}, {100}, {10}, {cagra::search_algo::AUTO}, @@ -334,6 +339,7 @@ inline std::vector generate_inputs() {cuvs::distance::DistanceType::L2Expanded}, {false}, {1000000}, + {1.0}, {100}, {10}, {cagra::search_algo::AUTO}, @@ -352,6 +358,7 @@ inline std::vector generate_inputs() {cuvs::distance::DistanceType::L2Expanded}, {false}, {1000000}, + {1.0}, {100}, {10}, {cagra::search_algo::AUTO}, diff --git a/examples/cpp/src/vamana_example.cu b/examples/cpp/src/vamana_example.cu index bb8085b27f..9715d0dc38 100644 --- a/examples/cpp/src/vamana_example.cu +++ b/examples/cpp/src/vamana_example.cu @@ -34,7 +34,7 @@ void vamana_build_and_write(raft::device_resources const& dev_resources, int degree, int visited_size, float max_fraction, - int iters, + float iters, std::string codebook_prefix) { using namespace cuvs::neighbors; @@ -45,7 +45,10 @@ void vamana_build_and_write(raft::device_resources const& dev_resources, index_params.visited_size = visited_size; index_params.graph_degree = degree; index_params.vamana_iters = iters; - index_params.codebooks = vamana::deserialize_codebooks(codebook_prefix, dataset.extent(1)); + + if (codebook_prefix != "") { + index_params.codebooks = vamana::deserialize_codebooks(codebook_prefix, dataset.extent(1)); + } std::cout << "Building Vamana index (search graph)" << std::endl; @@ -64,20 +67,24 @@ void vamana_build_and_write(raft::device_resources const& dev_resources, serialize(dev_resources, out_fname, index); // Output index to file (disk sector-aligned format) - serialize(dev_resources, out_fname + ".sector_aligned", index, false, true); + if (codebook_prefix != "") { + serialize(dev_resources, out_fname + ".sector_aligned", index, false, true); + } } void usage() { printf( - "Usage: ./vamana_example <(optional) " + "Usage: ./vamana_example " + " <(optional) " "codebook prefix>\n"); + printf("Input file expected to be binary file of fp32 vectors.\n"); printf("Graph degree sizes supported: 32, 64, 128, 256\n"); + printf("Datatype of input dataset (int8 or float)\n"); printf("Visited_size must be > degree and a power of 2.\n"); printf("max_fraction > 0 and <= 1. Typical values are 0.06 or 0.1.\n"); - printf("Default iterations = 1, increase for better quality graph.\n"); + printf("Default iterations = 1.0, increase for better quality graph.\n"); printf( "Optional path prefix to pq pivots and rotation matrix files. Expects pq pivots file at " "${codebook_prefix}_pq_pivots.bin and rotation matrix file at " @@ -102,23 +109,38 @@ int main(int argc, char* argv[]) // limit. raft::resource::set_workspace_to_pool_resource(dev_resources, 2 * // 1024 * 1024 * 1024ull); - if (argc != 7 && argc != 8) usage(); + if (argc != 8 && argc != 9) usage(); std::string data_fname = (std::string)(argv[1]); // Input filename std::string out_fname = (std::string)argv[2]; // Output index filename - int degree = atoi(argv[3]); - int max_visited = atoi(argv[4]); - float max_fraction = atof(argv[5]); - int iters = atoi(argv[6]); + std::string dtype = (std::string)argv[3]; + int degree = atoi(argv[4]); + int max_visited = atoi(argv[5]); + float max_fraction = atof(argv[6]); + float iters = atof(argv[7]); std::string codebook_prefix = ""; - if (argc >= 8) - codebook_prefix = (std::string)argv[7]; // Path prefix to pq pivots and rotation matrix files - - // Read in binary dataset file - auto dataset = read_bin_dataset(dev_resources, data_fname, INT_MAX); - - // Simple build example to create graph and write to a file - vamana_build_and_write(dev_resources, + if (argc >= 9) + codebook_prefix = (std::string)argv[8]; // Path prefix to pq pivots and rotation matrix files + + if (dtype == "int8") { + // Read in binary dataset file + auto dataset = read_bin_dataset(dev_resources, data_fname, INT_MAX); + + // Simple build example to create graph and write to a file + vamana_build_and_write(dev_resources, + raft::make_const_mdspan(dataset.view()), + out_fname, + degree, + max_visited, + max_fraction, + iters, + codebook_prefix); + } else if (dtype == "float") { + // Read in binary dataset file + auto dataset = read_bin_dataset(dev_resources, data_fname, INT_MAX); + + // Simple build example to create graph and write to a file + vamana_build_and_write(dev_resources, raft::make_const_mdspan(dataset.view()), out_fname, degree, @@ -126,4 +148,7 @@ int main(int argc, char* argv[]) max_fraction, iters, codebook_prefix); + } else { + usage(); + } } From 0a6d5b7feb7ed03b85b2bc45bfedd887ebe804c7 Mon Sep 17 00:00:00 2001 From: Gil Forsyth Date: Mon, 4 Aug 2025 17:21:43 -0400 Subject: [PATCH 14/93] ci(labeler): update labeler action to @v5 (#1210) Bumps the version of `actions/labeler` to `@v5` and updates the syntax in the `labeler.yml` file to account for breaking changes in that version bump. xref: rapidsai/ops#2968 Authors: - Gil Forsyth (https://github.com/gforsyth) Approvers: - Bradley Dice (https://github.com/bdice) - James Lamb (https://github.com/jameslamb) URL: https://github.com/rapidsai/cuvs/pull/1210 --- .github/labeler.yml | 21 +++++++++++++-------- .github/workflows/labeler.yml | 2 +- 2 files changed, 14 insertions(+), 9 deletions(-) diff --git a/.github/labeler.yml b/.github/labeler.yml index 56f77e69c0..bebde421e8 100644 --- a/.github/labeler.yml +++ b/.github/labeler.yml @@ -3,14 +3,19 @@ # Labels culled from https://github.com/rapidsai/raft/labels python: - - 'python/**' - + - changed-files: + any-glob-to-any-file: + - 'python/**' cpp: - - 'cpp/**' - + - changed-files: + any-glob-to-any-file: + - 'cpp/**' CMake: - - '**/CMakeLists.txt' - - '**/cmake/**' - + - changed-files: + any-glob-to-any-file: + - '**/CMakeLists.txt' + - '**/cmake/**' ci: - - 'ci/**' + - changed-files: + any-glob-to-any-file: + - 'ci/**' diff --git a/.github/workflows/labeler.yml b/.github/workflows/labeler.yml index 1ddd5b5cc3..65e68275d3 100644 --- a/.github/workflows/labeler.yml +++ b/.github/workflows/labeler.yml @@ -6,6 +6,6 @@ jobs: triage: runs-on: ubuntu-latest steps: - - uses: actions/labeler@v4 + - uses: actions/labeler@v5 with: repo-token: "${{ secrets.GITHUB_TOKEN }}" From ad4dba2fa33bdee4b607b9fdb68cfa847d9c812c Mon Sep 17 00:00:00 2001 From: Kyle Edwards Date: Tue, 5 Aug 2025 09:49:35 -0400 Subject: [PATCH 15/93] Update rapids-build-backend to 0.4.1 (#1214) Issue: https://github.com/rapidsai/build-planning/issues/207 Authors: - Kyle Edwards (https://github.com/KyleFromNVIDIA) Approvers: - Bradley Dice (https://github.com/bdice) URL: https://github.com/rapidsai/cuvs/pull/1214 --- conda/environments/all_cuda-129_arch-aarch64.yaml | 2 +- conda/environments/all_cuda-129_arch-x86_64.yaml | 2 +- conda/recipes/cuvs-bench-cpu/recipe.yaml | 2 +- conda/recipes/cuvs-bench/recipe.yaml | 2 +- conda/recipes/cuvs/recipe.yaml | 2 +- dependencies.yaml | 2 +- python/cuvs/pyproject.toml | 2 +- python/libcuvs/pyproject.toml | 2 +- 8 files changed, 8 insertions(+), 8 deletions(-) diff --git a/conda/environments/all_cuda-129_arch-aarch64.yaml b/conda/environments/all_cuda-129_arch-aarch64.yaml index 04fe7a5c21..45bb00e8df 100644 --- a/conda/environments/all_cuda-129_arch-aarch64.yaml +++ b/conda/environments/all_cuda-129_arch-aarch64.yaml @@ -42,7 +42,7 @@ dependencies: - pylibraft==25.10.*,>=0.0.0a0 - pytest - pytest-cov -- rapids-build-backend>=0.3.0,<0.4.0.dev0 +- rapids-build-backend>=0.4.0,<0.5.0.dev0 - recommonmark - rust - scikit-build-core>=0.10.0 diff --git a/conda/environments/all_cuda-129_arch-x86_64.yaml b/conda/environments/all_cuda-129_arch-x86_64.yaml index 643dbfe04c..46a0f234c7 100644 --- a/conda/environments/all_cuda-129_arch-x86_64.yaml +++ b/conda/environments/all_cuda-129_arch-x86_64.yaml @@ -42,7 +42,7 @@ dependencies: - pylibraft==25.10.*,>=0.0.0a0 - pytest - pytest-cov -- rapids-build-backend>=0.3.0,<0.4.0.dev0 +- rapids-build-backend>=0.4.0,<0.5.0.dev0 - recommonmark - rust - scikit-build-core>=0.10.0 diff --git a/conda/recipes/cuvs-bench-cpu/recipe.yaml b/conda/recipes/cuvs-bench-cpu/recipe.yaml index fe9d26618e..e856382c68 100644 --- a/conda/recipes/cuvs-bench-cpu/recipe.yaml +++ b/conda/recipes/cuvs-bench-cpu/recipe.yaml @@ -63,7 +63,7 @@ requirements: - openblas - pip - python =${{ py_version }} - - rapids-build-backend>=0.3.0,<0.4.0.dev0 + - rapids-build-backend>=0.4.0,<0.5.0.dev0 - setuptools >=64.0.0 - if: linux64 then: diff --git a/conda/recipes/cuvs-bench/recipe.yaml b/conda/recipes/cuvs-bench/recipe.yaml index b5e001b62d..a823877926 100644 --- a/conda/recipes/cuvs-bench/recipe.yaml +++ b/conda/recipes/cuvs-bench/recipe.yaml @@ -32,7 +32,7 @@ requirements: - libcuvs-bench-ann =${{ version }} - python =${{ py_version }} - pip - - rapids-build-backend >=0.3.0,<0.4.0.dev0 + - rapids-build-backend >=0.4.0,<0.5.0.dev0 - rmm =${{ minor_version }} - setuptools >=64.0.0 run: diff --git a/conda/recipes/cuvs/recipe.yaml b/conda/recipes/cuvs/recipe.yaml index 0c3610771a..e7b9f6de99 100644 --- a/conda/recipes/cuvs/recipe.yaml +++ b/conda/recipes/cuvs/recipe.yaml @@ -56,7 +56,7 @@ requirements: - pip - pylibraft =${{ minor_version }} - python =${{ py_version }} - - rapids-build-backend >=0.3.0,<0.4.0.dev0 + - rapids-build-backend >=0.4.0,<0.5.0.dev0 - scikit-build-core >=0.10.0 - cuda-python >=12.6.2,<13.0a0 - cuda-cudart-dev diff --git a/dependencies.yaml b/dependencies.yaml index 960614a891..dd25ae9afb 100644 --- a/dependencies.yaml +++ b/dependencies.yaml @@ -218,7 +218,7 @@ dependencies: common: - output_types: [conda, requirements, pyproject] packages: - - rapids-build-backend>=0.3.0,<0.4.0.dev0 + - rapids-build-backend>=0.4.0,<0.5.0.dev0 - output_types: [conda] packages: - scikit-build-core>=0.10.0 diff --git a/python/cuvs/pyproject.toml b/python/cuvs/pyproject.toml index 373e313492..99e31fc620 100644 --- a/python/cuvs/pyproject.toml +++ b/python/cuvs/pyproject.toml @@ -15,7 +15,7 @@ [build-system] requires = [ - "rapids-build-backend>=0.3.0,<0.4.0.dev0", + "rapids-build-backend>=0.4.0,<0.5.0.dev0", "scikit-build-core[pyproject]>=0.10.0", ] # This list was generated by `rapids-dependency-file-generator`. To make changes, edit ../../dependencies.yaml and run `rapids-dependency-file-generator`. build-backend = "rapids_build_backend.build" diff --git a/python/libcuvs/pyproject.toml b/python/libcuvs/pyproject.toml index 00e49155da..51bb46316b 100644 --- a/python/libcuvs/pyproject.toml +++ b/python/libcuvs/pyproject.toml @@ -14,7 +14,7 @@ [build-system] requires = [ - "rapids-build-backend>=0.3.0,<0.4.0.dev0", + "rapids-build-backend>=0.4.0,<0.5.0.dev0", "scikit-build-core[pyproject]>=0.10.0", ] # This list was generated by `rapids-dependency-file-generator`. To make changes, edit ../../dependencies.yaml and run `rapids-dependency-file-generator`. build-backend = "rapids_build_backend.build" From c22b7774954d69464cacbd2d647dfe1e84344c7e Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Lorenzo=20Dematt=C3=A9?= Date: Tue, 5 Aug 2025 16:37:12 +0200 Subject: [PATCH 16/93] [Java] Add decorator for synchronized CuVSResource access (#1209) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit As a follow up to #1089 and #1082, this PR introduces a decorator to ensure that access to a shared CuVSResource is synchronized. This could be useful in cases where application logic is complex and it's not easy to efficiently and cleanly guarantee that access to a CuVSResources is done by at most one thread at each given time. Authors: - Lorenzo Dematté (https://github.com/ldematte) - MithunR (https://github.com/mythrocks) Approvers: - MithunR (https://github.com/mythrocks) URL: https://github.com/rapidsai/cuvs/pull/1209 --- .../cuvs/SynchronizedCuVSResources.java | 63 +++++ .../cuvs/CagraMultiThreadStabilityIT.java | 219 ++++++++++-------- 2 files changed, 191 insertions(+), 91 deletions(-) create mode 100644 java/cuvs-java/src/main/java/com/nvidia/cuvs/SynchronizedCuVSResources.java diff --git a/java/cuvs-java/src/main/java/com/nvidia/cuvs/SynchronizedCuVSResources.java b/java/cuvs-java/src/main/java/com/nvidia/cuvs/SynchronizedCuVSResources.java new file mode 100644 index 0000000000..bbba262bb0 --- /dev/null +++ b/java/cuvs-java/src/main/java/com/nvidia/cuvs/SynchronizedCuVSResources.java @@ -0,0 +1,63 @@ +/* + * Copyright (c) 2025, NVIDIA CORPORATION. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.nvidia.cuvs; + +import java.nio.file.Path; +import java.util.concurrent.locks.ReentrantLock; + +/** + * A decorator for CuVSResources that guarantees synchronized (blocking) access to the wrapped CuVSResource + */ +public class SynchronizedCuVSResources implements CuVSResources { + + private final CuVSResources inner; + private final ReentrantLock lock; + + private SynchronizedCuVSResources(CuVSResources inner) { + this.inner = inner; + this.lock = new ReentrantLock(); + } + + static CuVSResources create() throws Throwable { + return new SynchronizedCuVSResources(CuVSResources.create()); + } + + @Override + public ScopedAccess access() { + lock.lock(); + return new ScopedAccess() { + @Override + public long handle() { + return inner.access().handle(); + } + + @Override + public void close() { + lock.unlock(); + } + }; + } + + @Override + public void close() { + inner.close(); + } + + @Override + public Path tempDirectory() { + return inner.tempDirectory(); + } +} 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 50e8f12a5a..3bcef6bb94 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 @@ -16,9 +16,7 @@ package com.nvidia.cuvs; import static com.carrotsearch.randomizedtesting.RandomizedTest.assumeTrue; -import static org.junit.Assert.assertNotNull; -import static org.junit.Assert.assertTrue; -import static org.junit.Assert.fail; +import static org.junit.Assert.*; import com.carrotsearch.randomizedtesting.RandomizedRunner; import com.nvidia.cuvs.CagraIndexParams.CagraGraphBuildAlgo; @@ -45,6 +43,16 @@ public class CagraMultiThreadStabilityIT extends CuVSTestCase { private static final Logger log = LoggerFactory.getLogger(MethodHandles.lookup().lookupClass()); + private final int dimensions = 256; + private final int queriesPerThread = 500; + private final int queryBatchSize = 1; // Small batch size to increase frequency of calls + private final int topK = 10; + + @FunctionalInterface + private interface QueryAction { + void run(CagraIndex index) throws Throwable; + } + @Before public void setup() { assumeTrue("not supported on " + System.getProperty("os.name"), isLinuxAmd64()); @@ -53,18 +61,27 @@ public void setup() { } @Test - public void testQueryingUsingMultipleThreads() throws Throwable { + public void testQueryingUsingMultipleThreadsWithSharedSynchronizedResources() throws Throwable { + try (CuVSResources sharedResources = SynchronizedCuVSResources.create()) { + testQueryingUsingMultipleThreads( + index -> performQueryWithSharedSynchronizedResource(sharedResources, index)); + } + } + + @Test + public void testQueryingUsingMultipleThreadsWithPrivateResources() throws Throwable { + testQueryingUsingMultipleThreads(this::performQueryWithPrivateResource); + } + + private void testQueryingUsingMultipleThreads(QueryAction queryAction) throws Throwable { final int dataSize = 10000; - final int dimensions = 256; - final int numThreads = 16; // High thread count to increase contention - final int queriesPerThread = 500; - final int queryBatchSize = 1; // Small batch size to increase frequency of calls - final int topK = 10; + final int numThreads = 16; log.info(" Dataset: {}x{}", dataSize, dimensions); + // High thread count to increase contention log.info(" Threads: {}, Queries per thread: {}", numThreads, queriesPerThread); - float[][] dataset = generateRandomDataset(dataSize, dimensions); + float[][] dataset = generateRandomDataset(dataSize); try (CuVSResources resources = CheckedCuVSResources.create()) { log.info("Creating CAGRA index for MultiThreaded stability test..."); @@ -86,105 +103,125 @@ public void testQueryingUsingMultipleThreads() throws Throwable { log.info("CAGRA index created, starting high-contention multi-threaded search..."); - // Create high contention scenario that would fail without using separate resources in every thread - ExecutorService executor = Executors.newFixedThreadPool(numThreads); - List> futures = new ArrayList<>(); - CountDownLatch startLatch = new CountDownLatch(1); - AtomicInteger successfulQueries = new AtomicInteger(0); - AtomicReference firstError = new AtomicReference<>(); - - for (int threadId = 0; threadId < numThreads; threadId++) { - final int finalThreadId = threadId; - - Future future = - executor.submit( - () -> { - try { - // Wait for all threads to start simultaneously - startLatch.await(); - - for (int queryId = 0; queryId < queriesPerThread; queryId++) { - float[][] queries = generateRandomDataset(queryBatchSize, dimensions); - - try (CuVSResources threadResources = CheckedCuVSResources.create()) { - CagraSearchParams searchParams = new CagraSearchParams.Builder().build(); - CagraQuery query = - new CagraQuery.Builder(threadResources) - .withTopK(topK) - .withSearchParams(searchParams) - .withQueryVectors(queries) - .build(); - - // This call should now work with per-thread resources - SearchResults results = index.search(query); - assertNotNull("Query should return results", results); - assertTrue( - "Query should return some results", !results.getResults().isEmpty()); - + // Create high contention scenario that would fail without using separate resources in every + // thread + try (ExecutorService executor = Executors.newFixedThreadPool(numThreads)) { + List> futures = new ArrayList<>(); + CountDownLatch startLatch = new CountDownLatch(1); + AtomicInteger successfulQueries = new AtomicInteger(0); + AtomicReference firstError = new AtomicReference<>(); + + for (int threadId = 0; threadId < numThreads; threadId++) { + final int finalThreadId = threadId; + + Future future = + executor.submit( + () -> { + try { + // Wait for all threads to start simultaneously + startLatch.await(); + + for (int queryId = 0; queryId < queriesPerThread; queryId++) { + queryAction.run(index); successfulQueries.incrementAndGet(); - } - // No Thread.yield() - maximize contention - } + // No Thread.yield() - maximize contention + } - log.info("Thread {} completed successfully", finalThreadId); + log.info("Thread {} completed successfully", finalThreadId); - } catch (Throwable t) { - log.error("Thread {} failed: {}", finalThreadId, t.getMessage(), t); - firstError.compareAndSet(null, t); - throw new RuntimeException("Thread failed", t); - } - }); + } catch (Throwable t) { + log.error("Thread {} failed: {}", finalThreadId, t.getMessage(), t); + firstError.compareAndSet(null, t); + throw new RuntimeException("Thread failed", t); + } + }); - futures.add(future); - } + futures.add(future); + } - // Start all threads simultaneously to maximize contention - log.info("Starting all {} threads simultaneously...", numThreads); - startLatch.countDown(); - - // Wait for all threads to complete - boolean allCompleted = true; - for (Future future : futures) { - try { - future.get(120, TimeUnit.SECONDS); // Longer timeout for stress test - } catch (Exception e) { - allCompleted = false; - log.error("Thread failed: {}", e.getMessage(), e); - if (firstError.get() == null) { - firstError.set(e); + // Start all threads simultaneously to maximize contention + log.info("Starting all {} threads simultaneously...", numThreads); + startLatch.countDown(); + + // Wait for all threads to complete + boolean allCompleted = true; + for (Future future : futures) { + try { + future.get(120, TimeUnit.SECONDS); // Longer timeout for stress test + } catch (Exception e) { + allCompleted = false; + log.error("Thread failed: {}", e.getMessage(), e); + if (firstError.get() == null) { + firstError.set(e); + } } } - } - executor.shutdown(); - if (!executor.awaitTermination(10, TimeUnit.SECONDS)) { - executor.shutdownNow(); - } + executor.shutdown(); + if (!executor.awaitTermination(10, TimeUnit.SECONDS)) { + executor.shutdownNow(); + } - // Verify results - int expectedTotalQueries = numThreads * queriesPerThread; - int actualSuccessfulQueries = successfulQueries.get(); + // Verify results + int expectedTotalQueries = numThreads * queriesPerThread; + int actualSuccessfulQueries = successfulQueries.get(); - log.info(" Successful queries: {} / {}", actualSuccessfulQueries, expectedTotalQueries); + log.info(" Successful queries: {} / {}", actualSuccessfulQueries, expectedTotalQueries); - if (firstError.get() != null) { - fail( - "MultiThreaded stablity test failed:" - + " " - + firstError.get().getMessage()); - } + if (firstError.get() != null) { + fail("MultiThreaded stablity test failed:" + " " + firstError.get().getMessage()); + } - assertTrue("All threads should complete successfully", allCompleted); - assertTrue( - "All queries should complete successfully", - actualSuccessfulQueries == expectedTotalQueries); + assertTrue("All threads should complete successfully", allCompleted); + assertEquals( + "All queries should complete successfully", + expectedTotalQueries, + actualSuccessfulQueries); + } index.destroyIndex(); } } - private float[][] generateRandomDataset(int size, int dimensions) { + private void performQueryWithPrivateResource(CagraIndex index) throws Throwable { + float[][] queries = generateRandomDataset(queryBatchSize); + + try (CuVSResources threadResources = CheckedCuVSResources.create()) { + CagraSearchParams searchParams = new CagraSearchParams.Builder().build(); + CagraQuery query = + new CagraQuery.Builder(threadResources) + .withTopK(topK) + .withSearchParams(searchParams) + .withQueryVectors(queries) + .build(); + + // This call should now work with per-thread resources + SearchResults results = index.search(query); + assertNotNull("Query should return results", results); + assertFalse("Query should return some results", results.getResults().isEmpty()); + } + } + + private void performQueryWithSharedSynchronizedResource( + CuVSResources threadResources, CagraIndex index) throws Throwable { + float[][] queries = generateRandomDataset(queryBatchSize); + + CagraSearchParams searchParams = new CagraSearchParams.Builder().build(); + CagraQuery query = + new CagraQuery.Builder(threadResources) + .withTopK(topK) + .withSearchParams(searchParams) + .withQueryVectors(queries) + .build(); + + // This call should now work with per-thread resources + SearchResults results = index.search(query); + assertNotNull("Query should return results", results); + assertFalse("Query should return some results", results.getResults().isEmpty()); + } + + private float[][] generateRandomDataset(int size) { Random random = new Random(42 + System.nanoTime()); float[][] data = new float[size][dimensions]; From 6e734d5eef728c101762233c7b85a7ee8552ef4c Mon Sep 17 00:00:00 2001 From: Micka Date: Tue, 5 Aug 2025 22:26:22 +0200 Subject: [PATCH 17/93] Add NN.extend and kmeans to Python doc (#1217) All the nearest neighbors `extend` functions were missing. Authors: - Micka (https://github.com/lowener) Approvers: - Corey J. Nolet (https://github.com/cjnolet) URL: https://github.com/rapidsai/cuvs/pull/1217 --- docs/source/python_api.rst | 1 + docs/source/python_api/cluster.rst | 12 +++++++++ docs/source/python_api/cluster_kmeans.rst | 27 +++++++++++++++++++ docs/source/python_api/neighbors_cagra.rst | 5 ++++ docs/source/python_api/neighbors_hnsw.rst | 5 ++++ docs/source/python_api/neighbors_ivf_flat.rst | 5 ++++ docs/source/python_api/neighbors_ivf_pq.rst | 5 ++++ 7 files changed, 60 insertions(+) create mode 100644 docs/source/python_api/cluster.rst create mode 100644 docs/source/python_api/cluster_kmeans.rst diff --git a/docs/source/python_api.rst b/docs/source/python_api.rst index aa75160c1f..4c8fc47820 100644 --- a/docs/source/python_api.rst +++ b/docs/source/python_api.rst @@ -7,6 +7,7 @@ Python API Documentation .. toctree:: :maxdepth: 4 + python_api/cluster.rst python_api/distance.rst python_api/neighbors.rst python_api/preprocessing.rst diff --git a/docs/source/python_api/cluster.rst b/docs/source/python_api/cluster.rst new file mode 100644 index 0000000000..b5c0ab957c --- /dev/null +++ b/docs/source/python_api/cluster.rst @@ -0,0 +1,12 @@ +Cluster +======== + +.. role:: py(code) + :language: python + :class: highlight + +.. toctree:: + :maxdepth: 1 + :caption: Contents: + + cluster_kmeans.rst diff --git a/docs/source/python_api/cluster_kmeans.rst b/docs/source/python_api/cluster_kmeans.rst new file mode 100644 index 0000000000..8fda17f80d --- /dev/null +++ b/docs/source/python_api/cluster_kmeans.rst @@ -0,0 +1,27 @@ +K-Means +======= + +.. role:: py(code) + :language: python + :class: highlight + +K-Means Parameters +################## + +.. autoclass:: cuvs.cluster.kmeans.KMeansParams + :members: + +K-Means Fit +########### + +.. autofunction:: cuvs.cluster.kmeans.fit + +K-Means Predict +############### + +.. autofunction:: cuvs.cluster.kmeans.predict + +K-Means Cluster Cost +#################### + +.. autofunction:: cuvs.cluster.kmeans.cluster_cost diff --git a/docs/source/python_api/neighbors_cagra.rst b/docs/source/python_api/neighbors_cagra.rst index e7155efb83..76564ce5e7 100644 --- a/docs/source/python_api/neighbors_cagra.rst +++ b/docs/source/python_api/neighbors_cagra.rst @@ -44,3 +44,8 @@ Index load ########## .. autofunction:: cuvs.neighbors.cagra.load + +Index extend +########### + +.. autofunction:: cuvs.neighbors.cagra.extend diff --git a/docs/source/python_api/neighbors_hnsw.rst b/docs/source/python_api/neighbors_hnsw.rst index 64fe5493ba..fe46a8e4b6 100644 --- a/docs/source/python_api/neighbors_hnsw.rst +++ b/docs/source/python_api/neighbors_hnsw.rst @@ -38,3 +38,8 @@ Index load ########## .. autofunction:: cuvs.neighbors.hnsw.load + +Index extend +########### + +.. autofunction:: cuvs.neighbors.hnsw.extend diff --git a/docs/source/python_api/neighbors_ivf_flat.rst b/docs/source/python_api/neighbors_ivf_flat.rst index f2c21e68a1..d0846b0d67 100644 --- a/docs/source/python_api/neighbors_ivf_flat.rst +++ b/docs/source/python_api/neighbors_ivf_flat.rst @@ -42,3 +42,8 @@ Index load ########## .. autofunction:: cuvs.neighbors.ivf_flat.load + +Index extend +############ + +.. autofunction:: cuvs.neighbors.ivf_flat.extend diff --git a/docs/source/python_api/neighbors_ivf_pq.rst b/docs/source/python_api/neighbors_ivf_pq.rst index 57668fbc33..ec4cfdff6a 100644 --- a/docs/source/python_api/neighbors_ivf_pq.rst +++ b/docs/source/python_api/neighbors_ivf_pq.rst @@ -42,3 +42,8 @@ Index load ########## .. autofunction:: cuvs.neighbors.ivf_pq.load + +Index extend +############ + +.. autofunction:: cuvs.neighbors.ivf_pq.extend From 222362f5f2630f445e783f48aab1529f63e547ac Mon Sep 17 00:00:00 2001 From: Jinsol Park Date: Wed, 6 Aug 2025 12:43:59 -0700 Subject: [PATCH 18/93] Processing mutual reachability dist in connect_knn_graph for host data (#1093) Closes https://github.com/rapidsai/cuvs/issues/1096 Existing `connect_knn_graph` function for host data does not consider the reduction op. This PR adds a fix to calculate mutual reachability distances for selected random cross-component neighbors instead of simple pairwise distances when `reduction_op == cuvs::sparse::neighbors::MutualReachabilityFixConnectivitiesRedOp`. Authors: - Jinsol Park (https://github.com/jinsolp) Approvers: - Tarang Jain (https://github.com/tarang-jain) URL: https://github.com/rapidsai/cuvs/pull/1093 --- cpp/src/cluster/detail/mst.cuh | 84 +++++++++++++++++++++++--------- cpp/tests/cluster/connect_knn.cu | 79 ++++++++++++++++++++++++------ 2 files changed, 127 insertions(+), 36 deletions(-) diff --git a/cpp/src/cluster/detail/mst.cuh b/cpp/src/cluster/detail/mst.cuh index f637426ee6..e73c54ac3c 100644 --- a/cpp/src/cluster/detail/mst.cuh +++ b/cpp/src/cluster/detail/mst.cuh @@ -17,10 +17,13 @@ #pragma once #include "../../sparse/neighbors/cross_component_nn.cuh" +#include "raft/core/device_mdspan.hpp" +#include "raft/core/operators.hpp" #include #include #include #include +#include #include #include #include @@ -74,7 +77,7 @@ void merge_msts(raft::sparse::solver::Graph_COO& template void connect_knn_graph( raft::resources const& handle, - const value_t* X, + raft::device_matrix_view X, raft::sparse::solver::Graph_COO& msf, size_t m, size_t n, @@ -93,7 +96,7 @@ void connect_knn_graph( cuvs::sparse::neighbors::cross_component_nn(handle, connected_edges, - X, + X.data_handle(), color, m, n, @@ -132,22 +135,30 @@ void connect_knn_graph( * @param[inout] msf edge list containing the mst result * @param[in] m number of rows in X * @param[in] n number of columns in X - * @param[in] n_components number of components in color * @param[inout] color the color labels array returned from the mst invocation * @return updated MST edge list */ -template +template void connect_knn_graph( raft::resources const& handle, - const value_t* X, + raft::host_matrix_view X, raft::sparse::solver::Graph_COO& msf, size_t m, size_t n, - int n_components, value_idx* color, + red_op reduction_op = red_op{}, cuvs::distance::DistanceType metric = cuvs::distance::DistanceType::L2SqrtExpanded) { - auto stream = raft::resource::get_cuda_stream(handle); + using namespace cuvs::sparse::neighbors; + static_assert( + std::is_same_v || + std::is_same_v> || + std::is_same_v>, + "reduction_op must be identity_op, MutualReachabilityFixConnectivitiesRedOp, or " + "FixConnectivitiesRedOp"); + + auto stream = raft::resource::get_cuda_stream(handle); + int n_components = get_n_components(color, m, stream); rmm::device_uvector d_color_remapped(m, stream); raft::label::make_monotonic(d_color_remapped.data(), color, m, stream, true); @@ -190,34 +201,48 @@ void connect_knn_graph( host_v_indices.push_back(v); } - auto device_u_indices = raft::make_device_vector(handle, n_components - 1); - auto device_v_indices = raft::make_device_vector(handle, n_components - 1); + size_t new_nnz = n_components - 1; + + auto device_u_indices = raft::make_device_vector(handle, new_nnz); + auto device_v_indices = raft::make_device_vector(handle, new_nnz); - raft::copy(device_u_indices.data_handle(), host_u_indices.data(), n_components - 1, stream); - raft::copy(device_v_indices.data_handle(), host_v_indices.data(), n_components - 1, stream); + raft::copy(device_u_indices.data_handle(), host_u_indices.data(), new_nnz, stream); + raft::copy(device_v_indices.data_handle(), host_v_indices.data(), new_nnz, stream); - auto X_view = raft::make_host_matrix_view(X, m, n); - auto data_u = raft::make_device_matrix(handle, n_components - 1, n); - auto data_v = raft::make_device_matrix(handle, n_components - 1, n); + auto data_u = raft::make_device_matrix(handle, new_nnz, n); + auto data_v = raft::make_device_matrix(handle, new_nnz, n); raft::matrix::detail::gather( - handle, X_view, raft::make_const_mdspan(device_u_indices.view()), data_u.view()); + handle, X, raft::make_const_mdspan(device_u_indices.view()), data_u.view()); raft::matrix::detail::gather( - handle, X_view, raft::make_const_mdspan(device_v_indices.view()), data_v.view()); + handle, X, raft::make_const_mdspan(device_v_indices.view()), data_v.view()); - auto pairwise_dist = - raft::make_device_matrix(handle, n_components - 1, n_components - 1); + auto pairwise_dist = raft::make_device_matrix(handle, new_nnz, new_nnz); cuvs::distance::pairwise_distance(handle, raft::make_const_mdspan(data_u.view()), raft::make_const_mdspan(data_v.view()), pairwise_dist.view(), metric); - auto pairwise_dist_vec = raft::make_device_vector(handle, n_components - 1); + auto pairwise_dist_vec = raft::make_device_vector(handle, n_components - 1); raft::matrix::get_diagonal( handle, raft::make_const_mdspan(pairwise_dist.view()), pairwise_dist_vec.view()); - size_t new_nnz = n_components - 1; + if constexpr (std::is_same>::value) { + raft::linalg::map_offset( + handle, + pairwise_dist_vec.view(), + [pairwise_dist_ptr = pairwise_dist_vec.data_handle(), + core_dist_ptr = reduction_op.core_dists, + device_u_indices_ptr = device_u_indices.data_handle(), + device_v_indices_ptr = device_v_indices.data_handle()] __device__(auto i) { + float dist = pairwise_dist_ptr[i]; + float u_core_dist = core_dist_ptr[device_u_indices_ptr[i]]; + float v_core_dist = core_dist_ptr[device_v_indices_ptr[i]]; + return fmaxf(dist, fmaxf(u_core_dist, v_core_dist)); + }); + } // sort in order of rows to run sorted_coo_to_csr auto rows_begin = thrust::device_pointer_cast(device_u_indices.data_handle()); @@ -304,9 +329,24 @@ void build_sorted_mst( while (n_components > 1 && iters < max_iter) { if (data_on_device) { - connect_knn_graph(handle, X, mst_coo, m, n, color, reduction_op); + connect_knn_graph( + handle, + raft::make_device_matrix_view(X, m, n), + mst_coo, + m, + n, + color, + reduction_op); } else { - connect_knn_graph(handle, X, mst_coo, m, n, n_components, color, metric); + connect_knn_graph( + handle, + raft::make_host_matrix_view(X, m, n), + mst_coo, + m, + n, + color, + reduction_op, + metric); } iters++; diff --git a/cpp/tests/cluster/connect_knn.cu b/cpp/tests/cluster/connect_knn.cu index 70e999406d..b4fd55e5b6 100644 --- a/cpp/tests/cluster/connect_knn.cu +++ b/cpp/tests/cluster/connect_knn.cu @@ -25,7 +25,12 @@ #include #include "../../src/cluster/detail/mst.cuh" +#include "../../src/neighbors/detail/knn_brute_force.cuh" +#include "../../src/neighbors/detail/reachability.cuh" +#include "../../src/sparse/neighbors/cross_component_nn.cuh" #include "../neighbors/naive_knn.cuh" +#include +#include #include #include @@ -41,6 +46,7 @@ struct ConnectKNNInputs { int n_clusters; int k; cuvs::distance::DistanceType metric; + bool mutual_reach; }; template @@ -69,6 +75,36 @@ class ConnectKNNTest : public ::testing::TestWithParam { ps.k, ps.metric); + rmm::device_uvector core_dists(ps.n_rows, stream); + if (ps.mutual_reach) { + cuvs::neighbors::detail::reachability::core_distances( + dists.data(), ps.k, ps.k, (size_t)ps.n_rows, core_dists.data(), stream); + + auto epilogue = cuvs::neighbors::detail::reachability::ReachabilityPostProcess{ + core_dists.data(), 1.0}; + cuvs::neighbors::detail::tiled_brute_force_knn< + T, + int64_t, + T, + cuvs::neighbors::detail::reachability::ReachabilityPostProcess>(handle, + database.data(), + database.data(), + ps.n_rows, + ps.n_rows, + ps.dim, + ps.k, + dists.data(), + inds.data(), + ps.metric, + 2.0, + 0, + 0, + nullptr, + nullptr, + nullptr, + epilogue); + } + rmm::device_uvector coo_rows(queries_size, stream); raft::sparse::COO knn_coo(stream, queries_size * 2); @@ -116,14 +152,29 @@ class ConnectKNNTest : public ::testing::TestWithParam { auto database_h = raft::make_host_matrix(ps.n_rows, ps.dim); raft::copy(database_h.data_handle(), database.data(), ps.n_rows * ps.dim, stream); - cuvs::cluster::agglomerative::detail::connect_knn_graph(handle, - database_h.data_handle(), - mst_coo, - ps.n_rows, - ps.dim, - n_components, - color.data(), - ps.metric); + if (ps.mutual_reach) { + cuvs::sparse::neighbors::MutualReachabilityFixConnectivitiesRedOp reduction_op( + core_dists.data(), (int64_t)ps.n_rows); + cuvs::cluster::agglomerative::detail::connect_knn_graph( + handle, + raft::make_const_mdspan(database_h.view()), + mst_coo, + ps.n_rows, + ps.dim, + color.data(), + reduction_op, + ps.metric); + } else { + cuvs::cluster::agglomerative::detail::connect_knn_graph( + handle, + raft::make_const_mdspan(database_h.view()), + mst_coo, + ps.n_rows, + ps.dim, + color.data(), + raft::identity_op{}, + ps.metric); + } n_components = cuvs::sparse::neighbors::get_n_components(color.data(), ps.n_rows, stream); ASSERT_TRUE(n_components == 1); @@ -147,12 +198,12 @@ class ConnectKNNTest : public ::testing::TestWithParam { }; const std::vector inputs = raft::util::itertools::product( - {5000, 7151}, // n_rows - {64, 137}, // dim - {5, 10}, // n_clusters of make_blobs data - {16}, // k - {cuvs::distance::DistanceType::L2SqrtExpanded} // metric -); + {5000, 7151}, // n_rows + {64, 137}, // dim + {5, 10}, // n_clusters of make_blobs data + {16}, // k + {cuvs::distance::DistanceType::L2SqrtExpanded}, // metric + {true, false}); // mutual_reach typedef ConnectKNNTest ConnectKNNTestF; TEST_P(ConnectKNNTestF, ConnectKNN) { this->basicTest(); } From 1b29b0408f9e2772195d22a48d2f179440f97f2d Mon Sep 17 00:00:00 2001 From: Kyle Edwards Date: Wed, 6 Aug 2025 17:30:19 -0400 Subject: [PATCH 19/93] Update another instance of rapids-build-backend (#1225) Issue: https://github.com/rapidsai/build-planning/issues/207 Authors: - Kyle Edwards (https://github.com/KyleFromNVIDIA) Approvers: - James Lamb (https://github.com/jameslamb) URL: https://github.com/rapidsai/cuvs/pull/1225 --- conda/environments/bench_ann_cuda-129_arch-aarch64.yaml | 2 +- conda/environments/bench_ann_cuda-129_arch-x86_64.yaml | 2 +- dependencies.yaml | 4 ++-- python/cuvs_bench/pyproject.toml | 2 +- 4 files changed, 5 insertions(+), 5 deletions(-) diff --git a/conda/environments/bench_ann_cuda-129_arch-aarch64.yaml b/conda/environments/bench_ann_cuda-129_arch-aarch64.yaml index 7e5efd0693..7ffb09777b 100644 --- a/conda/environments/bench_ann_cuda-129_arch-aarch64.yaml +++ b/conda/environments/bench_ann_cuda-129_arch-aarch64.yaml @@ -40,7 +40,7 @@ dependencies: - pandas - pylibraft==25.10.*,>=0.0.0a0 - pyyaml -- rapids-build-backend>=0.3.0,<0.4.0.dev0 +- rapids-build-backend>=0.4.0,<0.5.0.dev0 - requests - scikit-learn - setuptools diff --git a/conda/environments/bench_ann_cuda-129_arch-x86_64.yaml b/conda/environments/bench_ann_cuda-129_arch-x86_64.yaml index 5604f56252..fcd0f89053 100644 --- a/conda/environments/bench_ann_cuda-129_arch-x86_64.yaml +++ b/conda/environments/bench_ann_cuda-129_arch-x86_64.yaml @@ -43,7 +43,7 @@ dependencies: - pandas - pylibraft==25.10.*,>=0.0.0a0 - pyyaml -- rapids-build-backend>=0.3.0,<0.4.0.dev0 +- rapids-build-backend>=0.4.0,<0.5.0.dev0 - requests - scikit-learn - setuptools diff --git a/dependencies.yaml b/dependencies.yaml index dd25ae9afb..7570ce1973 100644 --- a/dependencies.yaml +++ b/dependencies.yaml @@ -218,7 +218,7 @@ dependencies: common: - output_types: [conda, requirements, pyproject] packages: - - rapids-build-backend>=0.4.0,<0.5.0.dev0 + - &rapids_build_backend rapids-build-backend>=0.4.0,<0.5.0.dev0 - output_types: [conda] packages: - scikit-build-core>=0.10.0 @@ -263,7 +263,7 @@ dependencies: common: - output_types: [conda, requirements, pyproject] packages: - - &rapids_build_backend rapids-build-backend>=0.3.0,<0.4.0.dev0 + - *rapids_build_backend - setuptools - wheel build_py_cuvs: diff --git a/python/cuvs_bench/pyproject.toml b/python/cuvs_bench/pyproject.toml index 901a37e74d..0e052c46e1 100644 --- a/python/cuvs_bench/pyproject.toml +++ b/python/cuvs_bench/pyproject.toml @@ -3,7 +3,7 @@ [build-system] build-backend = "rapids_build_backend.build" requires = [ - "rapids-build-backend>=0.3.0,<0.4.0.dev0", + "rapids-build-backend>=0.4.0,<0.5.0.dev0", "setuptools", "wheel", ] # This list was generated by `rapids-dependency-file-generator`. To make changes, edit ../../dependencies.yaml and run `rapids-dependency-file-generator`. From 2cdb85d82956bcdefc594c1f835c7a497f1a9234 Mon Sep 17 00:00:00 2001 From: Robert Maynard Date: Thu, 7 Aug 2025 09:49:46 -0400 Subject: [PATCH 20/93] Updates needed when building with CUDA 13 (#1219) cuvs also requires https://github.com/rapidsai/raft/pull/2774 to be merged Authors: - Robert Maynard (https://github.com/robertmaynard) - Corey J. Nolet (https://github.com/cjnolet) Approvers: - Corey J. Nolet (https://github.com/cjnolet) URL: https://github.com/rapidsai/cuvs/pull/1219 --- cpp/include/cuvs/neighbors/common.hpp | 2 +- cpp/tests/CMakeLists.txt | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/cpp/include/cuvs/neighbors/common.hpp b/cpp/include/cuvs/neighbors/common.hpp index 8b2edaa424..2e91f3c15f 100644 --- a/cpp/include/cuvs/neighbors/common.hpp +++ b/cpp/include/cuvs/neighbors/common.hpp @@ -485,7 +485,7 @@ namespace filtering { enum class FilterType { None, Bitmap, Bitset }; struct base_filter { - virtual ~base_filter() = default; + ~base_filter() = default; virtual FilterType get_filter_type() const = 0; }; diff --git a/cpp/tests/CMakeLists.txt b/cpp/tests/CMakeLists.txt index a9ec6dfd00..a1b037f706 100644 --- a/cpp/tests/CMakeLists.txt +++ b/cpp/tests/CMakeLists.txt @@ -212,7 +212,7 @@ if(BUILD_TESTS) ConfigureTest( NAME NEIGHBORS_ALL_NEIGHBORS_TEST PATH neighbors/all_neighbors/test_float.cu GPUS 1 PERCENT 100 - ADDITIONAL_DEP nccl + ADDITIONAL_DEP NCCL::NCCL ) ConfigureTest( From 3cd48dc5a24999a6def8fcf79cde81160fc7d061 Mon Sep 17 00:00:00 2001 From: rmaschal Date: Thu, 7 Aug 2025 10:01:45 -0700 Subject: [PATCH 21/93] Fix missing sync_stream in ScaNN build (#1224) As title. For larger number of clusters (or faster PCI-E), the work required to find cluster assignments would take longer than the HtoD copy for each batch. This would lead to overwriting of one of the batch_load_iterator device buffers before kmeans predict completed, and thus lower recall. This adds a raft::resource::sync_stream at the end of the kmeans predict loop to avoid this. Although not technically necessary (as the prior DtoH copy into pageable memory is synchronous), sync_stream is also added to the end of the quantization loop for consistency/correctness. Also removed one unused comment. Authors: - https://github.com/rmaschal Approvers: - Corey J. Nolet (https://github.com/cjnolet) - Tarang Jain (https://github.com/tarang-jain) URL: https://github.com/rapidsai/cuvs/pull/1224 --- cpp/src/neighbors/scann/detail/scann_build.cuh | 15 ++++++--------- 1 file changed, 6 insertions(+), 9 deletions(-) diff --git a/cpp/src/neighbors/scann/detail/scann_build.cuh b/cpp/src/neighbors/scann/detail/scann_build.cuh index 5520684f4c..da723d56e9 100644 --- a/cpp/src/neighbors/scann/detail/scann_build.cuh +++ b/cpp/src/neighbors/scann/detail/scann_build.cuh @@ -173,6 +173,9 @@ index build( res, kmeans_params, batch_view, raft::make_const_mdspan(centroids_view), batch_labels_view); dataset_vec_batches.prefetch_next_batch(); + + // Make sure work on device is finished before swapping buffers + raft::resource::sync_stream(res); } // AVQ update of KMeans centroids @@ -304,18 +307,9 @@ index build( // quantize dataset to bfloat16, if enabled. Similar to SOAR, quantization // is performed in this loop to improve locality // TODO (rmaschal): Might be more efficient to do on CPU, to avoid DtoH copy - auto bf16_dataset = raft::make_device_matrix(res, batch_view.extent(0), dim); if (params.bf16_enabled) { - // raft::linalg::map_offset(res, bf16_dataset.view(), [batch_view, dim] __device__(size_t i) { - // int64_t row_idx = i / dim; - // int64_t col_idx = i % dim; - - // nv_bfloat16 val = __float2bfloat16(batch_view(row_idx, col_idx)); - - // return reinterpret_cast(val); - //}); raft::linalg::unaryOp( bf16_dataset.data_handle(), batch_view.data_handle(), @@ -348,6 +342,9 @@ index build( bf16_dataset.size(), stream); } + + // Make sure work on device is finished before swapping buffers + raft::resource::sync_stream(res); } // Codebooks from VPQ have the shape [subspace idx, subspace dim, code] From e2ec0c80399d1559eaa66071bee45485b1bb20f0 Mon Sep 17 00:00:00 2001 From: Robert Maynard Date: Thu, 7 Aug 2025 17:36:10 -0400 Subject: [PATCH 22/93] Update rapids_config to handle user defined branch name (#1227) rapids_config will use a user defined branch over `RAPIDS_BRANCH` contents Authors: - Robert Maynard (https://github.com/robertmaynard) Approvers: - Kyle Edwards (https://github.com/KyleFromNVIDIA) URL: https://github.com/rapidsai/cuvs/pull/1227 --- cmake/RAPIDS.cmake | 2 +- cmake/rapids_config.cmake | 8 ++++++-- 2 files changed, 7 insertions(+), 3 deletions(-) diff --git a/cmake/RAPIDS.cmake b/cmake/RAPIDS.cmake index 40de7cefcd..ddef819498 100644 --- a/cmake/RAPIDS.cmake +++ b/cmake/RAPIDS.cmake @@ -18,7 +18,7 @@ cmake_minimum_required(VERSION 3.30.4 FATAL_ERROR) # Allow users to control which version is used -if(NOT rapids-cmake-branch OR NOT rapids-cmake-version) +if(NOT (rapids-cmake-branch OR rapids-cmake-version)) message( FATAL_ERROR "The CMake variable `rapids-cmake-branch` or `rapids-cmake-version` must be defined" ) diff --git a/cmake/rapids_config.cmake b/cmake/rapids_config.cmake index b706c926e7..b2c54a3f27 100644 --- a/cmake/rapids_config.cmake +++ b/cmake/rapids_config.cmake @@ -35,6 +35,10 @@ if(NOT _rapids_branch) ) endif() -set(rapids-cmake-version "${RAPIDS_VERSION_MAJOR_MINOR}") -set(rapids-cmake-branch "${_rapids_branch}") +if(NOT rapids-cmake-version) + set(rapids-cmake-version "${RAPIDS_VERSION_MAJOR_MINOR}") +endif() +if(NOT rapids-cmake-branch) + set(rapids-cmake-branch "${_rapids_branch}") +endif() include("${CMAKE_CURRENT_LIST_DIR}/RAPIDS.cmake") From afc24ee8302a57a5c7192a38fc2d74a1a47aa95f Mon Sep 17 00:00:00 2001 From: Jinsol Park Date: Thu, 7 Aug 2025 17:58:21 -0700 Subject: [PATCH 23/93] Make duplicate removal in all neighbors robust to distance drift across batches (#1185) This is to fix an edge case that happens and the root cause is in issue: https://github.com/rapidsai/cuvs/issues/1056, which is about different distance results from `raft::linalg::gemm` based on the input sizes. Right now, when merging two knn graphs from different batches, we sort by distances (i.e. keys), and if the distances are same we sort by indices (i.e. values). After doing so, we compare indices right next to each other to check for duplicates under assumption that same vectors end up with same distances. However, due to the problem stated in issue 1056, distance for same index can be slightly different based on the size of the input matrix to gemm (or where the vector is in the entire matrix). For example, say we are calculating nearest neighbors for vector 0. we could end up with ``` indices = [1, 2, 3, 2, ....] distances = [0.023, 0.02355981, 0.02355983, 0.02355987] ``` because distance between vector 0 and vector 2 is calculated as 0.02355981 in the first batch, and 0.02355987 in the second batch. This PR fixes this issue by checking 4 neighbors to its left for duplicates, instead of checking the one next to itself. Authors: - Jinsol Park (https://github.com/jinsolp) - Corey J. Nolet (https://github.com/cjnolet) Approvers: - Corey J. Nolet (https://github.com/cjnolet) URL: https://github.com/rapidsai/cuvs/pull/1185 --- .../all_neighbors/all_neighbors_merge.cuh | 77 ++++++++++++++----- 1 file changed, 56 insertions(+), 21 deletions(-) diff --git a/cpp/src/neighbors/all_neighbors/all_neighbors_merge.cuh b/cpp/src/neighbors/all_neighbors/all_neighbors_merge.cuh index 1ce8576f5a..a766a7fcad 100644 --- a/cpp/src/neighbors/all_neighbors/all_neighbors_merge.cuh +++ b/cpp/src/neighbors/all_neighbors/all_neighbors_merge.cuh @@ -30,7 +30,7 @@ namespace cuvs::neighbors::all_neighbors::detail { using namespace cuvs::neighbors; -template +template RAFT_KERNEL merge_subgraphs_kernel(IdxT* cluster_data_indices, size_t graph_degree, size_t num_cluster_in_batch, @@ -102,11 +102,12 @@ RAFT_KERNEL merge_subgraphs_kernel(IdxT* cluster_data_indices, __syncthreads(); + size_t limit = 2 * graph_degree; // load sorted result into shared memory to get unique values idxBase = threadIdx.x * ITEMS_PER_THREAD; for (int i = 0; i < ITEMS_PER_THREAD; i++) { size_t colId = idxBase + i; - if (colId < 2 * graph_degree) { + if (colId < limit) { blockKeys[colId] = threadKeyValuePair[i].key; blockValues[colId] = threadKeyValuePair[i].value; } @@ -118,8 +119,39 @@ RAFT_KERNEL merge_subgraphs_kernel(IdxT* cluster_data_indices, if (threadIdx.x == 0) { uniqueMask[0] = 1; } for (int i = 0; i < ITEMS_PER_THREAD; i++) { size_t colId = idxBase + i; - if (colId > 0 && colId < 2 * graph_degree) { - uniqueMask[colId] = static_cast(blockValues[colId] != blockValues[colId - 1]); + if (colId > 0 && colId < limit) { + // this assumes same distance between vector from two different batches. however, currently + // there are subtle differences in the result based on the matrix size used to call gemm. + // This makes it difficult to remove duplicates, because they might no longer be right next + // to each other after sorting by distances. Thus, for now we sweep a neighboring window of + // size 4 or sweep the entire row to check for duplicates, and keep the first occurrence + // only. + // related issue: https://github.com/rapidsai/cuvs/issues/1056 + // uniqueMask[colId] = static_cast(blockValues[colId] != blockValues[colId - 1]); + + int is_unique = 1; + + if constexpr (SweepAll) { // sweep whole row for better deduplication + for (int j = 0; j < limit; ++j) { + if (j < colId && blockValues[j] == blockValues[colId]) { + is_unique = 0; + break; + } + } + } else { // otherwise sweep a small window + IdxT curr_val = blockValues[colId]; +#pragma unroll + for (int offset = -4; offset < 0; offset++) { + int neighbor_idx = static_cast(colId) + offset; + if (neighbor_idx >= 0) { + if (blockValues[neighbor_idx] == curr_val) { + is_unique = 0; + break; + } + } + } + } + uniqueMask[colId] = static_cast(is_unique); } } @@ -127,7 +159,7 @@ RAFT_KERNEL merge_subgraphs_kernel(IdxT* cluster_data_indices, // prefix sum if (threadIdx.x == 0) { - for (int i = 1; i < 2 * graph_degree; i++) { + for (int i = 1; i < limit; i++) { uniqueMask[i] += uniqueMask[i - 1]; } } @@ -141,7 +173,7 @@ RAFT_KERNEL merge_subgraphs_kernel(IdxT* cluster_data_indices, for (int i = 0; i < ITEMS_PER_THREAD; i++) { size_t colId = idxBase + i; - if (colId > 0 && colId < 2 * graph_degree) { + if (colId > 0 && colId < limit) { bool is_unique = uniqueMask[colId] != uniqueMask[colId - 1]; int16_t global_colId = uniqueMask[colId] - 1; if (is_unique && static_cast(global_colId) < graph_degree) { @@ -153,7 +185,7 @@ RAFT_KERNEL merge_subgraphs_kernel(IdxT* cluster_data_indices, } } -template +template void merge_subgraphs(raft::resources const& res, size_t k, size_t num_data_in_cluster, @@ -170,7 +202,7 @@ void merge_subgraphs(raft::resources const& res, #pragma omp critical // for omp-using multi-gpu purposes { if (num_elems <= 128) { - merge_subgraphs_kernel + merge_subgraphs_kernel <<>>( inverted_indices_d, k, @@ -181,7 +213,7 @@ void merge_subgraphs(raft::resources const& res, batch_neighbors_d, select_min); } else if (num_elems <= 512) { - merge_subgraphs_kernel + merge_subgraphs_kernel <<>>( inverted_indices_d, k, @@ -192,7 +224,7 @@ void merge_subgraphs(raft::resources const& res, batch_neighbors_d, select_min); } else if (num_elems <= 1024) { - merge_subgraphs_kernel + merge_subgraphs_kernel <<>>( inverted_indices_d, k, @@ -203,7 +235,7 @@ void merge_subgraphs(raft::resources const& res, batch_neighbors_d, select_min); } else if (num_elems <= 2048) { - merge_subgraphs_kernel + merge_subgraphs_kernel <<>>( inverted_indices_d, k, @@ -221,7 +253,10 @@ void merge_subgraphs(raft::resources const& res, } } -template +template void remap_and_merge_subgraphs(raft::resources const& res, raft::device_vector_view inverted_indices_d, raft::host_vector_view inverted_indices, @@ -254,15 +289,15 @@ void remap_and_merge_subgraphs(raft::resources const& res, num_data_in_cluster * k, raft::resource::get_cuda_stream(res)); - merge_subgraphs(res, - k, - num_data_in_cluster, - inverted_indices_d.data_handle(), - global_distances.data_handle(), - batch_distances_d.data_handle(), - global_neighbors.data_handle(), - batch_neighbors_d.data_handle(), - select_min); + merge_subgraphs(res, + k, + num_data_in_cluster, + inverted_indices_d.data_handle(), + global_distances.data_handle(), + batch_distances_d.data_handle(), + global_neighbors.data_handle(), + batch_neighbors_d.data_handle(), + select_min); } } // namespace cuvs::neighbors::all_neighbors::detail From 5e3a5a6391c5ac4366279430d3ac778e60b900db Mon Sep 17 00:00:00 2001 From: MithunR Date: Fri, 8 Aug 2025 10:28:40 -0700 Subject: [PATCH 24/93] [Java] Exception-safe RMM Allocations (#1215) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit This commit introduces exception safety for RMM allocations. Previously, device memory allocated through `cuvsRmmAlloc()` was freed manually using `cuvsRmmFree()`, in all the index impl classes. The problem there is that if an exception is thrown in the intervening time between alloc and free, it would lead to a leak of device memory. This commit extends the `CloseableHandle` class to encapsulate the allocation of device memory. This new class is used in try-with-resources blocks, to make device memory allocations exception-safe. Authors: - MithunR (https://github.com/mythrocks) Approvers: - Lorenzo Dematté (https://github.com/ldematte) - Corey J. Nolet (https://github.com/cjnolet) URL: https://github.com/rapidsai/cuvs/pull/1215 --- .../cuvs/internal/BruteForceIndexImpl.java | 206 ++++++------ .../nvidia/cuvs/internal/CagraIndexImpl.java | 177 +++++------ .../nvidia/cuvs/internal/TieredIndexImpl.java | 296 +++++++++--------- .../common/CloseableRMMAllocation.java | 86 +++++ .../com/nvidia/cuvs/internal/common/Util.java | 13 - 5 files changed, 424 insertions(+), 354 deletions(-) create mode 100644 java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/common/CloseableRMMAllocation.java diff --git a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/BruteForceIndexImpl.java b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/BruteForceIndexImpl.java index 4dfef90261..83c08fedc6 100644 --- a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/BruteForceIndexImpl.java +++ b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/BruteForceIndexImpl.java @@ -15,6 +15,7 @@ */ package com.nvidia.cuvs.internal; +import static com.nvidia.cuvs.internal.common.CloseableRMMAllocation.allocateRMMSegment; import static com.nvidia.cuvs.internal.common.LinkerHelper.C_FLOAT; import static com.nvidia.cuvs.internal.common.LinkerHelper.C_FLOAT_BYTE_SIZE; import static com.nvidia.cuvs.internal.common.LinkerHelper.C_INT_BYTE_SIZE; @@ -22,7 +23,6 @@ import static com.nvidia.cuvs.internal.common.LinkerHelper.C_LONG_BYTE_SIZE; import static com.nvidia.cuvs.internal.common.Util.CudaMemcpyKind.HOST_TO_DEVICE; import static com.nvidia.cuvs.internal.common.Util.CudaMemcpyKind.INFER_DIRECTION; -import static com.nvidia.cuvs.internal.common.Util.allocateRMMSegment; import static com.nvidia.cuvs.internal.common.Util.buildMemorySegment; import static com.nvidia.cuvs.internal.common.Util.checkCuVSError; import static com.nvidia.cuvs.internal.common.Util.concatenate; @@ -35,7 +35,6 @@ import static com.nvidia.cuvs.internal.panama.headers_h.cuvsBruteForceIndex_t; import static com.nvidia.cuvs.internal.panama.headers_h.cuvsBruteForceSearch; import static com.nvidia.cuvs.internal.panama.headers_h.cuvsBruteForceSerialize; -import static com.nvidia.cuvs.internal.panama.headers_h.cuvsRMMFree; import static com.nvidia.cuvs.internal.panama.headers_h.cuvsStreamSync; import static com.nvidia.cuvs.internal.panama.headers_h.omp_set_num_threads; @@ -45,6 +44,7 @@ import com.nvidia.cuvs.CuVSMatrix; import com.nvidia.cuvs.CuVSResources; import com.nvidia.cuvs.SearchResults; +import com.nvidia.cuvs.internal.common.CloseableRMMAllocation; import com.nvidia.cuvs.internal.panama.cuvsFilter; import java.io.InputStream; import java.io.OutputStream; @@ -118,20 +118,7 @@ public void destroyIndex() { try { int returnValue = cuvsBruteForceIndexDestroy(bruteForceIndexReference.indexPtr); checkCuVSError(returnValue, "cuvsBruteForceIndexDestroy"); - - if (bruteForceIndexReference.datasetBytes > 0) { - try (var resourcesAccessor = resources.access()) { - checkCuVSError( - cuvsRMMFree( - resourcesAccessor.handle(), - bruteForceIndexReference.datasetPtr, - bruteForceIndexReference.datasetBytes), - "cuvsRMMFree"); - } - } - if (bruteForceIndexReference.tensorDataArena != null) { - bruteForceIndexReference.tensorDataArena.close(); - } + bruteForceIndexReference.close(resources); } finally { destroyed = true; } @@ -158,25 +145,31 @@ private IndexReference build( try (var resourcesAccessor = resources.access()) { long cuvsResources = resourcesAccessor.handle(); - MemorySegment datasetMemorySegmentP = allocateRMMSegment(cuvsResources, datasetBytes); + try (var closeableDataMemorySegmentP = allocateRMMSegment(cuvsResources, datasetBytes)) { + MemorySegment datasetMemorySegmentP = closeableDataMemorySegmentP.handle(); - cudaMemcpy(datasetMemorySegmentP, datasetMemSegment, datasetBytes, INFER_DIRECTION); + cudaMemcpy(datasetMemorySegmentP, datasetMemSegment, datasetBytes, INFER_DIRECTION); - long[] datasetShape = {rows, cols}; - var tensorDataArena = Arena.ofShared(); - MemorySegment datasetTensor = - prepareTensor(tensorDataArena, datasetMemorySegmentP, datasetShape, 2, 32, 2, 1); + long[] datasetShape = {rows, cols}; + var tensorDataArena = Arena.ofShared(); + MemorySegment datasetTensor = + prepareTensor(tensorDataArena, datasetMemorySegmentP, datasetShape, 2, 32, 2, 1); - var returnValue = cuvsStreamSync(cuvsResources); - checkCuVSError(returnValue, "cuvsStreamSync"); + var returnValue = cuvsStreamSync(cuvsResources); + checkCuVSError(returnValue, "cuvsStreamSync"); - returnValue = cuvsBruteForceBuild(cuvsResources, datasetTensor, 0, 0.0f, index); - checkCuVSError(returnValue, "cuvsBruteForceBuild"); + returnValue = cuvsBruteForceBuild(cuvsResources, datasetTensor, 0, 0.0f, index); + checkCuVSError(returnValue, "cuvsBruteForceBuild"); - returnValue = cuvsStreamSync(cuvsResources); - checkCuVSError(returnValue, "cuvsStreamSync"); + returnValue = cuvsStreamSync(cuvsResources); + checkCuVSError(returnValue, "cuvsStreamSync"); - return new IndexReference(datasetMemorySegmentP, datasetBytes, tensorDataArena, index); + return new IndexReference( + new CloseableRMMAllocation(closeableDataMemorySegmentP), + datasetBytes, + tensorDataArena, + index); + } } finally { omp_set_num_threads(1); } @@ -205,6 +198,7 @@ public SearchResults search(BruteForceQuery cuvsQuery) throws Throwable { // prepare the prefiltering data final long prefilterDataLength; + final long prefilterBytes; final MemorySegment prefilterDataMemorySegment; BitSet[] prefilters = cuvsQuery.getPrefilters(); if (prefilters != null && prefilters.length > 0) { @@ -212,8 +206,11 @@ public SearchResults search(BruteForceQuery cuvsQuery) throws Throwable { long[] filters = concatenatedFilters.toLongArray(); prefilterDataMemorySegment = buildMemorySegment(localArena, filters); prefilterDataLength = (long) cuvsQuery.getNumDocs() * prefilters.length; + long[] prefilterShape = {(prefilterDataLength + 31) / 32}; + prefilterBytes = C_INT_BYTE_SIZE * prefilterShape[0]; } else { prefilterDataLength = 0; + prefilterBytes = 0; prefilterDataMemorySegment = MemorySegment.NULL; } @@ -223,77 +220,66 @@ public SearchResults search(BruteForceQuery cuvsQuery) throws Throwable { try (var resourcesAccessor = cuvsQuery.getResources().access()) { long cuvsResources = resourcesAccessor.handle(); - long queriesBytes = C_FLOAT_BYTE_SIZE * numQueries * vectorDimension; - long neighborsBytes = C_LONG_BYTE_SIZE * numQueries * topk; - long distanceBytes = C_FLOAT_BYTE_SIZE * numQueries * topk; - long prefilterBytes = 0; // size assigned later - - MemorySegment queriesDP = allocateRMMSegment(cuvsResources, queriesBytes); - MemorySegment neighborsDP = allocateRMMSegment(cuvsResources, neighborsBytes); - MemorySegment distancesDP = allocateRMMSegment(cuvsResources, distanceBytes); - MemorySegment prefilterDP = MemorySegment.NULL; - - cudaMemcpy(queriesDP, querySeg, queriesBytes, INFER_DIRECTION); - - long[] queriesShape = {numQueries, vectorDimension}; - MemorySegment queriesTensor = - prepareTensor(localArena, queriesDP, queriesShape, 2, 32, 2, 1); - long[] neighborsShape = {numQueries, topk}; - MemorySegment neighborsTensor = - prepareTensor(localArena, neighborsDP, neighborsShape, 0, 64, 2, 1); - long[] distancesShape = {numQueries, topk}; - MemorySegment distancesTensor = - prepareTensor(localArena, distancesDP, distancesShape, 2, 32, 2, 1); - - MemorySegment prefilter = cuvsFilter.allocate(localArena); - MemorySegment prefilterTensor; - - if (prefilterDataMemorySegment == MemorySegment.NULL) { - cuvsFilter.type(prefilter, 0); // NO_FILTER - cuvsFilter.addr(prefilter, 0); - } else { - long[] prefilterShape = {(prefilterDataLength + 31) / 32}; - long prefilterLen = prefilterShape[0]; - prefilterBytes = C_INT_BYTE_SIZE * prefilterLen; - - prefilterDP = allocateRMMSegment(cuvsResources, prefilterBytes); - - cudaMemcpy(prefilterDP, prefilterDataMemorySegment, prefilterBytes, HOST_TO_DEVICE); - - prefilterTensor = prepareTensor(localArena, prefilterDP, prefilterShape, 1, 32, 2, 1); - - cuvsFilter.type(prefilter, 2); - cuvsFilter.addr(prefilter, prefilterTensor.address()); - } - - var returnValue = cuvsStreamSync(cuvsResources); - checkCuVSError(returnValue, "cuvsStreamSync"); - - returnValue = - cuvsBruteForceSearch( - cuvsResources, - bruteForceIndexReference.indexPtr, - queriesTensor, - neighborsTensor, - distancesTensor, - prefilter); - checkCuVSError(returnValue, "cuvsBruteForceSearch"); - - returnValue = cuvsStreamSync(cuvsResources); - checkCuVSError(returnValue, "cuvsStreamSync"); - - cudaMemcpy(neighborsMemorySegment, neighborsDP, neighborsBytes, INFER_DIRECTION); - cudaMemcpy(distancesMemorySegment, distancesDP, distanceBytes, INFER_DIRECTION); - - returnValue = cuvsRMMFree(cuvsResources, neighborsDP, neighborsBytes); - checkCuVSError(returnValue, "cuvsRMMFree"); - returnValue = cuvsRMMFree(cuvsResources, distancesDP, distanceBytes); - checkCuVSError(returnValue, "cuvsRMMFree"); - returnValue = cuvsRMMFree(cuvsResources, queriesDP, queriesBytes); - checkCuVSError(returnValue, "cuvsRMMFree"); - if (prefilterBytes > 0) { - returnValue = cuvsRMMFree(cuvsResources, prefilterDP, prefilterBytes); - checkCuVSError(returnValue, "cuvsRMMFree"); + final long queriesBytes = C_FLOAT_BYTE_SIZE * numQueries * vectorDimension; + final long neighborsBytes = C_LONG_BYTE_SIZE * numQueries * topk; + final long distanceBytes = C_FLOAT_BYTE_SIZE * numQueries * topk; + + try (var queriesDP = allocateRMMSegment(cuvsResources, queriesBytes); + var neighborsDP = allocateRMMSegment(cuvsResources, neighborsBytes); + var distancesDP = allocateRMMSegment(cuvsResources, distanceBytes); + var prefilterDP = + prefilterBytes > 0 + ? allocateRMMSegment(cuvsResources, prefilterBytes) + : CloseableRMMAllocation.EMPTY) { + + cudaMemcpy(queriesDP.handle(), querySeg, queriesBytes, INFER_DIRECTION); + + long[] queriesShape = {numQueries, vectorDimension}; + MemorySegment queriesTensor = + prepareTensor(localArena, queriesDP.handle(), queriesShape, 2, 32, 2, 1); + long[] neighborsShape = {numQueries, topk}; + MemorySegment neighborsTensor = + prepareTensor(localArena, neighborsDP.handle(), neighborsShape, 0, 64, 2, 1); + long[] distancesShape = {numQueries, topk}; + MemorySegment distancesTensor = + prepareTensor(localArena, distancesDP.handle(), distancesShape, 2, 32, 2, 1); + + MemorySegment prefilter = cuvsFilter.allocate(localArena); + MemorySegment prefilterTensor; + + if (prefilterDataMemorySegment == MemorySegment.NULL) { + cuvsFilter.type(prefilter, 0); // NO_FILTER + cuvsFilter.addr(prefilter, 0); + } else { + long[] prefilterShape = {(prefilterDataLength + 31) / 32}; + cudaMemcpy( + prefilterDP.handle(), prefilterDataMemorySegment, prefilterBytes, HOST_TO_DEVICE); + + prefilterTensor = + prepareTensor(localArena, prefilterDP.handle(), prefilterShape, 1, 32, 2, 1); + + cuvsFilter.type(prefilter, 2); + cuvsFilter.addr(prefilter, prefilterTensor.address()); + } + + var returnValue = cuvsStreamSync(cuvsResources); + checkCuVSError(returnValue, "cuvsStreamSync"); + + returnValue = + cuvsBruteForceSearch( + cuvsResources, + bruteForceIndexReference.indexPtr, + queriesTensor, + neighborsTensor, + distancesTensor, + prefilter); + checkCuVSError(returnValue, "cuvsBruteForceSearch"); + + returnValue = cuvsStreamSync(cuvsResources); + checkCuVSError(returnValue, "cuvsStreamSync"); + + cudaMemcpy(neighborsMemorySegment, neighborsDP.handle(), neighborsBytes, INFER_DIRECTION); + cudaMemcpy(distancesMemorySegment, distancesDP.handle(), distanceBytes, INFER_DIRECTION); } } return BruteForceSearchResults.create( @@ -479,27 +465,39 @@ public BruteForceIndexImpl build() throws Throwable { */ private static class IndexReference { - private final MemorySegment datasetPtr; + private final CloseableRMMAllocation datasetAllocationHandle; private final long datasetBytes; private final Arena tensorDataArena; private final MemorySegment indexPtr; private IndexReference( - MemorySegment datasetPtr, + CloseableRMMAllocation datasetAllocationHandle, long datasetBytes, Arena tensorDataArena, MemorySegment indexPtr) { - this.datasetPtr = datasetPtr; + this.datasetAllocationHandle = datasetAllocationHandle; this.datasetBytes = datasetBytes; this.tensorDataArena = tensorDataArena; this.indexPtr = indexPtr; } private IndexReference(MemorySegment indexPtr) { - this.datasetPtr = MemorySegment.NULL; + this.datasetAllocationHandle = CloseableRMMAllocation.EMPTY; this.datasetBytes = 0; this.tensorDataArena = null; this.indexPtr = indexPtr; } + + /** + * Free up the memory used for dataset, tensor-data. + */ + private void close(CuVSResources resources) { + try (var resourcesAccessor = resources.access()) { + datasetAllocationHandle.close(); + } + if (tensorDataArena != null) { + tensorDataArena.close(); + } + } } } diff --git a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CagraIndexImpl.java b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CagraIndexImpl.java index 704bc72b34..ac3e9ba43a 100644 --- a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CagraIndexImpl.java +++ b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CagraIndexImpl.java @@ -16,13 +16,13 @@ package com.nvidia.cuvs.internal; import static com.nvidia.cuvs.internal.CuVSParamsHelper.*; +import static com.nvidia.cuvs.internal.common.CloseableRMMAllocation.allocateRMMSegment; import static com.nvidia.cuvs.internal.common.LinkerHelper.C_FLOAT; import static com.nvidia.cuvs.internal.common.LinkerHelper.C_FLOAT_BYTE_SIZE; import static com.nvidia.cuvs.internal.common.LinkerHelper.C_INT; import static com.nvidia.cuvs.internal.common.LinkerHelper.C_INT_BYTE_SIZE; import static com.nvidia.cuvs.internal.common.Util.CudaMemcpyKind.HOST_TO_DEVICE; import static com.nvidia.cuvs.internal.common.Util.CudaMemcpyKind.INFER_DIRECTION; -import static com.nvidia.cuvs.internal.common.Util.allocateRMMSegment; import static com.nvidia.cuvs.internal.common.Util.buildMemorySegment; import static com.nvidia.cuvs.internal.common.Util.checkCuVSError; import static com.nvidia.cuvs.internal.common.Util.concatenate; @@ -43,6 +43,7 @@ import com.nvidia.cuvs.CuVSResources; import com.nvidia.cuvs.SearchResults; import com.nvidia.cuvs.internal.common.CloseableHandle; +import com.nvidia.cuvs.internal.common.CloseableRMMAllocation; import com.nvidia.cuvs.internal.common.CompositeCloseableHandle; import com.nvidia.cuvs.internal.panama.cuvsCagraCompressionParams; import com.nvidia.cuvs.internal.panama.cuvsCagraIndexParams; @@ -212,6 +213,8 @@ private static MemorySegment createCagraIndex() { } } + private static final BitSet[] EMPTY_PREFILTER_BITSET = new BitSet[0]; + /** * Invokes the native search_cagra_index via the Panama API for searching a * CAGRA index. @@ -235,100 +238,92 @@ public SearchResults search(CagraQuery query) throws Throwable { MemorySegment distancesMemorySegment = localArena.allocate(distancesSequenceLayout); MemorySegment floatsSeg = buildMemorySegment(localArena, query.getQueryVectors()); - long queriesBytes = C_FLOAT_BYTE_SIZE * numQueries * vectorDimension; - long neighborsBytes = C_INT_BYTE_SIZE * numQueries * topK; - long distancesBytes = C_FLOAT_BYTE_SIZE * numQueries * topK; + final long queriesBytes = C_FLOAT_BYTE_SIZE * numQueries * vectorDimension; + final long neighborsBytes = C_INT_BYTE_SIZE * numQueries * topK; + final long distancesBytes = C_FLOAT_BYTE_SIZE * numQueries * topK; + final boolean hasPreFilter = query.getPrefilter() != null; + final BitSet[] prefilters = + hasPreFilter ? new BitSet[] {query.getPrefilter()} : EMPTY_PREFILTER_BITSET; + final long prefilterDataLength = hasPreFilter ? query.getNumDocs() * prefilters.length : 0; + final long prefilterLen = hasPreFilter ? (prefilterDataLength + 31) / 32 : 0; + final long prefilterBytes = C_INT_BYTE_SIZE * prefilterLen; try (var resourcesAccessor = query.getResources().access()) { var cuvsRes = resourcesAccessor.handle(); - MemorySegment queriesDP = allocateRMMSegment(cuvsRes, queriesBytes); - MemorySegment neighborsDP = allocateRMMSegment(cuvsRes, neighborsBytes); - MemorySegment distancesDP = allocateRMMSegment(cuvsRes, distancesBytes); - MemorySegment prefilterDP = MemorySegment.NULL; - long prefilterLen = 0; - - cudaMemcpy(queriesDP, floatsSeg, queriesBytes, INFER_DIRECTION); - - long[] queriesShape = {numQueries, vectorDimension}; - MemorySegment queriesTensor = - prepareTensor(localArena, queriesDP, queriesShape, 2, 32, 2, 1); - long[] neighborsShape = {numQueries, topK}; - MemorySegment neighborsTensor = - prepareTensor(localArena, neighborsDP, neighborsShape, 1, 32, 2, 1); - long[] distancesShape = {numQueries, topK}; - MemorySegment distancesTensor = - prepareTensor(localArena, distancesDP, distancesShape, 2, 32, 2, 1); - - var returnValue = cuvsStreamSync(cuvsRes); - checkCuVSError(returnValue, "cuvsStreamSync"); - - // prepare the prefiltering data - long prefilterDataLength = 0; - MemorySegment prefilterDataMemorySegment = MemorySegment.NULL; - BitSet[] prefilters; - if (query.getPrefilter() != null) { - prefilters = new BitSet[] {query.getPrefilter()}; - BitSet concatenatedFilters = concatenate(prefilters, query.getNumDocs()); - long[] filters = concatenatedFilters.toLongArray(); - prefilterDataMemorySegment = buildMemorySegment(localArena, filters); - prefilterDataLength = query.getNumDocs() * prefilters.length; - } - - MemorySegment prefilter = cuvsFilter.allocate(localArena); - MemorySegment prefilterTensor; - - final long prefilterBytes; - - if (prefilterDataMemorySegment == MemorySegment.NULL) { - cuvsFilter.type(prefilter, 0); // NO_FILTER - cuvsFilter.addr(prefilter, 0); - prefilterBytes = 0; - } else { - long[] prefilterShape = {(prefilterDataLength + 31) / 32}; - prefilterLen = prefilterShape[0]; - prefilterBytes = C_INT_BYTE_SIZE * prefilterLen; - - prefilterDP = allocateRMMSegment(cuvsRes, prefilterBytes); - - cudaMemcpy(prefilterDP, prefilterDataMemorySegment, prefilterBytes, HOST_TO_DEVICE); - - prefilterTensor = prepareTensor(localArena, prefilterDP, prefilterShape, 1, 32, 2, 1); - - cuvsFilter.type(prefilter, 1); - cuvsFilter.addr(prefilter, prefilterTensor.address()); - } - - returnValue = cuvsStreamSync(cuvsRes); - checkCuVSError(returnValue, "cuvsStreamSync"); - - returnValue = - cuvsCagraSearch( - cuvsRes, - segmentFromSearchParams(localArena, query.getCagraSearchParameters()), - cagraIndexReference.getMemorySegment(), - queriesTensor, - neighborsTensor, - distancesTensor, - prefilter); - checkCuVSError(returnValue, "cuvsCagraSearch"); - - returnValue = cuvsStreamSync(cuvsRes); - checkCuVSError(returnValue, "cuvsStreamSync"); - - cudaMemcpy(neighborsMemorySegment, neighborsDP, neighborsBytes, INFER_DIRECTION); - cudaMemcpy(distancesMemorySegment, distancesDP, distancesBytes, INFER_DIRECTION); - - returnValue = cuvsRMMFree(cuvsRes, distancesDP, distancesBytes); - checkCuVSError(returnValue, "cuvsRMMFree"); - returnValue = cuvsRMMFree(cuvsRes, neighborsDP, neighborsBytes); - checkCuVSError(returnValue, "cuvsRMMFree"); - returnValue = cuvsRMMFree(cuvsRes, queriesDP, queriesBytes); - checkCuVSError(returnValue, "cuvsRMMFree"); - - if (prefilterLen > 0) { - returnValue = cuvsRMMFree(cuvsRes, prefilterDP, C_INT_BYTE_SIZE * prefilterBytes); - checkCuVSError(returnValue, "cuvsRMMFree"); + try (var queriesDP = allocateRMMSegment(cuvsRes, queriesBytes); + var neighborsDP = allocateRMMSegment(cuvsRes, neighborsBytes); + var distancesDP = allocateRMMSegment(cuvsRes, distancesBytes); + var prefilterDP = + hasPreFilter + ? allocateRMMSegment(cuvsRes, prefilterBytes) + : CloseableRMMAllocation.EMPTY) { + + cudaMemcpy(queriesDP.handle(), floatsSeg, queriesBytes, INFER_DIRECTION); + + long[] queriesShape = {numQueries, vectorDimension}; + MemorySegment queriesTensor = + prepareTensor( + localArena, queriesDP.handle(), queriesShape, kDLFloat(), 32, kDLCUDA(), 1); + long[] neighborsShape = {numQueries, topK}; + MemorySegment neighborsTensor = + prepareTensor( + localArena, neighborsDP.handle(), neighborsShape, kDLUInt(), 32, kDLCUDA(), 1); + long[] distancesShape = {numQueries, topK}; + MemorySegment distancesTensor = + prepareTensor( + localArena, distancesDP.handle(), distancesShape, kDLFloat(), 32, kDLCUDA(), 1); + + var returnValue = cuvsStreamSync(cuvsRes); + checkCuVSError(returnValue, "cuvsStreamSync"); + + // prepare the prefiltering data + MemorySegment prefilterDataMemorySegment = MemorySegment.NULL; + if (hasPreFilter) { + BitSet concatenatedFilters = concatenate(prefilters, query.getNumDocs()); + long[] filters = concatenatedFilters.toLongArray(); + prefilterDataMemorySegment = buildMemorySegment(localArena, filters); + } + + MemorySegment prefilter = cuvsFilter.allocate(localArena); + MemorySegment prefilterTensor; + + if (!hasPreFilter) { + cuvsFilter.type(prefilter, 0); // NO_FILTER + cuvsFilter.addr(prefilter, 0); + } else { + long[] prefilterShape = {prefilterLen}; + + cudaMemcpy( + prefilterDP.handle(), prefilterDataMemorySegment, prefilterBytes, HOST_TO_DEVICE); + + prefilterTensor = + prepareTensor( + localArena, prefilterDP.handle(), prefilterShape, kDLUInt(), 32, kDLCUDA(), 1); + + cuvsFilter.type(prefilter, 1); + cuvsFilter.addr(prefilter, prefilterTensor.address()); + } + + returnValue = cuvsStreamSync(cuvsRes); + checkCuVSError(returnValue, "cuvsStreamSync"); + + returnValue = + cuvsCagraSearch( + cuvsRes, + segmentFromSearchParams(localArena, query.getCagraSearchParameters()), + cagraIndexReference.getMemorySegment(), + queriesTensor, + neighborsTensor, + distancesTensor, + prefilter); + checkCuVSError(returnValue, "cuvsCagraSearch"); + + returnValue = cuvsStreamSync(cuvsRes); + checkCuVSError(returnValue, "cuvsStreamSync"); + + cudaMemcpy(neighborsMemorySegment, neighborsDP.handle(), neighborsBytes, INFER_DIRECTION); + cudaMemcpy(distancesMemorySegment, distancesDP.handle(), distancesBytes, INFER_DIRECTION); } } diff --git a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/TieredIndexImpl.java b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/TieredIndexImpl.java index ef9d33cd43..31d8cd4bec 100644 --- a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/TieredIndexImpl.java +++ b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/TieredIndexImpl.java @@ -16,13 +16,13 @@ package com.nvidia.cuvs.internal; import static com.nvidia.cuvs.internal.CuVSParamsHelper.createTieredIndexParams; +import static com.nvidia.cuvs.internal.common.CloseableRMMAllocation.allocateRMMSegment; import static com.nvidia.cuvs.internal.common.LinkerHelper.C_FLOAT; import static com.nvidia.cuvs.internal.common.LinkerHelper.C_FLOAT_BYTE_SIZE; import static com.nvidia.cuvs.internal.common.LinkerHelper.C_INT_BYTE_SIZE; import static com.nvidia.cuvs.internal.common.LinkerHelper.C_LONG; import static com.nvidia.cuvs.internal.common.LinkerHelper.C_LONG_BYTE_SIZE; import static com.nvidia.cuvs.internal.common.LinkerHelper.C_POINTER; -import static com.nvidia.cuvs.internal.common.Util.*; import static com.nvidia.cuvs.internal.common.Util.CudaMemcpyKind.*; import static com.nvidia.cuvs.internal.common.Util.buildMemorySegment; import static com.nvidia.cuvs.internal.common.Util.checkCuVSError; @@ -34,6 +34,7 @@ import com.nvidia.cuvs.*; import com.nvidia.cuvs.CuVSMatrix; import com.nvidia.cuvs.internal.common.CloseableHandle; +import com.nvidia.cuvs.internal.common.CloseableRMMAllocation; import com.nvidia.cuvs.internal.common.Util; import com.nvidia.cuvs.internal.panama.cuvsCagraIndexParams; import com.nvidia.cuvs.internal.panama.cuvsCagraSearchParams; @@ -48,7 +49,7 @@ import java.util.Objects; /** - * {@link TieredIndex} encapscaps a Tiered index, along with methods to interact + * {@link TieredIndex} encapsulates a Tiered index, along with methods to interact * with it. *

* TieredIndex is a hybrid index that combines brute force search for small datasets @@ -134,39 +135,39 @@ private IndexReference build() { // TieredIndex REQUIRES device memory - allocate it long datasetSize = C_FLOAT_BYTE_SIZE * rows * cols; - MemorySegment datasetDP = allocateRMMSegment(cuvsRes, datasetSize); + try (var datasetDP = allocateRMMSegment(cuvsRes, datasetSize)) { + // Copy host to device + Util.cudaMemcpy(datasetDP.handle(), hostDataSeg, datasetSize, HOST_TO_DEVICE); - // Copy host to device - Util.cudaMemcpy(datasetDP, hostDataSeg, datasetSize, HOST_TO_DEVICE); + // Create tensor from device memory + long[] datasetShape = {rows, cols}; + MemorySegment datasetTensor = + prepareTensor( + localArena, datasetDP.handle(), datasetShape, kDLFloat(), 32, kDLCUDA(), 1); - // Create tensor from device memory - long[] datasetShape = {rows, cols}; - MemorySegment datasetTensor = - prepareTensor(localArena, datasetDP, datasetShape, kDLFloat(), 32, kDLCUDA(), 1); + MemorySegment index = localArena.allocate(cuvsTieredIndex_t); + var returnValue = cuvsTieredIndexCreate(index); + checkCuVSError(returnValue, "cuvsTieredIndexCreate"); - MemorySegment index = localArena.allocate(cuvsTieredIndex_t); - var returnValue = cuvsTieredIndexCreate(index); - checkCuVSError(returnValue, "cuvsTieredIndexCreate"); + returnValue = cuvsStreamSync(cuvsRes); + checkCuVSError(returnValue, "cuvsStreamSync"); - returnValue = cuvsStreamSync(cuvsRes); - checkCuVSError(returnValue, "cuvsStreamSync"); + // Extract the actual index pointer that was written by Create + MemorySegment actualIndexPtr = index.get(C_POINTER, 0); - // Extract the actual index pointer that was written by Create - MemorySegment actualIndexPtr = index.get(C_POINTER, 0); + returnValue = + cuvsTieredIndexBuild( + cuvsRes, indexParamsMemorySegment, datasetTensor, actualIndexPtr); + checkCuVSError(returnValue, "cuvsTieredIndexBuild"); - returnValue = - cuvsTieredIndexBuild(cuvsRes, indexParamsMemorySegment, datasetTensor, actualIndexPtr); - checkCuVSError(returnValue, "cuvsTieredIndexBuild"); - - // Clean up device memory after build - returnValue = cuvsRMMFree(cuvsRes, datasetDP, datasetSize); - checkCuVSError(returnValue, "cuvsRMMFree"); - - return new IndexReference(actualIndexPtr); + return new IndexReference(actualIndexPtr); + } } } } + private static final BitSet[] EMPTY_PREFILTER_BITSET = new BitSet[0]; + /** * Translates C search_tiered_index function to Java * Invokes the native search_tiered_index via the Panama API for searching a @@ -196,116 +197,119 @@ public SearchResults search(TieredIndexQuery query) throws Throwable { // Get host query data MemorySegment hostQueriesSeg = Util.buildMemorySegment(localArena, query.getQueryVectors()); + final long queriesBytes = C_FLOAT_BYTE_SIZE * numQueries * vectorDimension; + final long neighborsBytes = C_LONG_BYTE_SIZE * numQueries * topK; // 64-bit for tiered index + final long distancesBytes = C_FLOAT_BYTE_SIZE * numQueries * topK; + final boolean hasPreFilter = query.getPrefilter() != null; + final BitSet[] prefilters = + hasPreFilter ? new BitSet[] {query.getPrefilter()} : EMPTY_PREFILTER_BITSET; + final long prefilterDataLength = hasPreFilter ? query.getNumDocs() * prefilters.length : 0; + final long prefilterLen = hasPreFilter ? (prefilterDataLength + 31) / 32 : 0; + final long prefilterBytes = C_INT_BYTE_SIZE * prefilterLen; + try (var resourceAccess = query.getResources().access()) { long cuvsRes = resourceAccess.handle(); - long queriesBytes = C_FLOAT_BYTE_SIZE * numQueries * vectorDimension; - long neighborsBytes = C_LONG_BYTE_SIZE * numQueries * topK; // 64-bit for tiered index - long distancesBytes = C_FLOAT_BYTE_SIZE * numQueries * topK; - // Allocate DEVICE memory for all data - MemorySegment queriesDP = allocateRMMSegment(cuvsRes, queriesBytes); - MemorySegment neighborsDP = allocateRMMSegment(cuvsRes, neighborsBytes); - MemorySegment distancesDP = allocateRMMSegment(cuvsRes, distancesBytes); - - // Copy queries from host to device - var returnValue = - cudaMemcpy(queriesDP, hostQueriesSeg, queriesBytes, cudaMemcpyHostToDevice()); - checkCudaError(returnValue, "cudaMemcpy"); - - // Create tensors from device memory - long[] queriesShape = {numQueries, vectorDimension}; - MemorySegment queriesTensor = - prepareTensor(localArena, queriesDP, queriesShape, kDLFloat(), 32, kDLCUDA(), 1); - long[] neighborsShape = {numQueries, topK}; - MemorySegment neighborsTensor = - prepareTensor( - localArena, neighborsDP, neighborsShape, kDLInt(), 64, kDLCUDA(), 1); // 64-bit int - long[] distancesShape = {numQueries, topK}; - MemorySegment distancesTensor = - prepareTensor(localArena, distancesDP, distancesShape, kDLFloat(), 32, kDLCUDA(), 1); - - // Sync before prefilter setup - returnValue = cuvsStreamSync(cuvsRes); - checkCuVSError(returnValue, "cuvsStreamSync"); - - // Handle prefilter - MemorySegment prefilter = cuvsFilter.allocate(localArena); - final MemorySegment prefilterDP; - final long prefilterBytes; - - if (query.getPrefilter() != null) { - BitSet[] prefilters = new BitSet[] {query.getPrefilter()}; - BitSet concatenatedFilters = concatenate(prefilters, (int) query.getNumDocs()); - long[] filters = concatenatedFilters.toLongArray(); - MemorySegment hostPrefilterSeg = buildMemorySegment(localArena, filters); - - long prefilterDataLength = query.getNumDocs() * prefilters.length; - long[] prefilterShape = {(prefilterDataLength + 31) / 32}; - long prefilterLen = prefilterShape[0]; - prefilterBytes = C_INT_BYTE_SIZE * prefilterLen; - - // Allocate device memory for prefilter - prefilterDP = allocateRMMSegment(cuvsRes, prefilterBytes); - - // Copy prefilter to device - checkCudaError( - cudaMemcpy(prefilterDP, hostPrefilterSeg, prefilterBytes, cudaMemcpyHostToDevice()), - "cudaMemcpy"); - - MemorySegment prefilterTensor = - prepareTensor(localArena, prefilterDP, prefilterShape, kDLUInt(), 32, kDLCUDA(), 1); - - cuvsFilter.type(prefilter, 1); // BITSET - cuvsFilter.addr(prefilter, prefilterTensor.address()); - } else { - prefilterDP = MemorySegment.NULL; - prefilterBytes = 0; - cuvsFilter.type(prefilter, 0); // NO_FILTER - cuvsFilter.addr(prefilter, 0); - } - - // Perform search - returnValue = - cuvsTieredIndexSearch( - cuvsRes, - segmentFromSearchParams(query.getCagraSearchParameters(), localArena), - tieredIndexReference.getMemorySegment(), - queriesTensor, - neighborsTensor, - distancesTensor, - prefilter); - checkCuVSError(returnValue, "cuvsTieredIndexSearch"); - - // Copy results from device to host - returnValue = - cudaMemcpy(neighborsSeg, neighborsDP, neighborsBytes, cudaMemcpyDeviceToHost()); - checkCudaError(returnValue, "cudaMemcpy"); - returnValue = - cudaMemcpy(distancesSeg, distancesDP, distancesBytes, cudaMemcpyDeviceToHost()); - checkCudaError(returnValue, "cudaMemcpy"); - - // Clean up device memory - returnValue = cuvsRMMFree(cuvsRes, queriesDP, queriesBytes); - checkCuVSError(returnValue, "cuvsRMMFree"); - returnValue = cuvsRMMFree(cuvsRes, neighborsDP, neighborsBytes); - checkCuVSError(returnValue, "cuvsRMMFree"); - returnValue = cuvsRMMFree(cuvsRes, distancesDP, distancesBytes); - checkCuVSError(returnValue, "cuvsRMMFree"); - - if (prefilterDP != MemorySegment.NULL) { - returnValue = cuvsRMMFree(cuvsRes, prefilterDP, prefilterBytes); - checkCuVSError(returnValue, "cuvsRMMFree"); + try (var queriesDP = allocateRMMSegment(cuvsRes, queriesBytes); + var neighborsDP = allocateRMMSegment(cuvsRes, neighborsBytes); + var distancesDP = allocateRMMSegment(cuvsRes, distancesBytes); + var prefilterDP = + hasPreFilter + ? allocateRMMSegment(cuvsRes, prefilterBytes) + : CloseableRMMAllocation.EMPTY) { + + // Copy queries from host to device + var returnValue = + cudaMemcpy( + queriesDP.handle(), hostQueriesSeg, queriesBytes, cudaMemcpyHostToDevice()); + checkCudaError(returnValue, "cudaMemcpy"); + + // Create tensors from device memory + long[] queriesShape = {numQueries, vectorDimension}; + MemorySegment queriesTensor = + prepareTensor( + localArena, queriesDP.handle(), queriesShape, kDLFloat(), 32, kDLCUDA(), 1); + long[] neighborsShape = {numQueries, topK}; + MemorySegment neighborsTensor = + prepareTensor( + localArena, + neighborsDP.handle(), + neighborsShape, + kDLInt(), + 64, + kDLCUDA(), + 1); // 64-bit int + long[] distancesShape = {numQueries, topK}; + MemorySegment distancesTensor = + prepareTensor( + localArena, distancesDP.handle(), distancesShape, kDLFloat(), 32, kDLCUDA(), 1); + + // Sync before prefilter setup + returnValue = cuvsStreamSync(cuvsRes); + checkCuVSError(returnValue, "cuvsStreamSync"); + + // Handle prefilter + MemorySegment prefilter = cuvsFilter.allocate(localArena); + + if (hasPreFilter) { + BitSet concatenatedFilters = concatenate(prefilters, (int) query.getNumDocs()); + long[] filters = concatenatedFilters.toLongArray(); + MemorySegment hostPrefilterSeg = buildMemorySegment(localArena, filters); + + long[] prefilterShape = {prefilterLen}; + + // Copy prefilter to device + checkCudaError( + cudaMemcpy( + prefilterDP.handle(), + hostPrefilterSeg, + prefilterBytes, + cudaMemcpyHostToDevice()), + "cudaMemcpy"); + + MemorySegment prefilterTensor = + prepareTensor( + localArena, prefilterDP.handle(), prefilterShape, kDLUInt(), 32, kDLCUDA(), 1); + + cuvsFilter.type(prefilter, 1); // BITSET + cuvsFilter.addr(prefilter, prefilterTensor.address()); + } else { + cuvsFilter.type(prefilter, 0); // NO_FILTER + cuvsFilter.addr(prefilter, 0); + } + + // Perform search + returnValue = + cuvsTieredIndexSearch( + cuvsRes, + segmentFromSearchParams(query.getCagraSearchParameters(), localArena), + tieredIndexReference.getMemorySegment(), + queriesTensor, + neighborsTensor, + distancesTensor, + prefilter); + checkCuVSError(returnValue, "cuvsTieredIndexSearch"); + + // Copy results from device to host + returnValue = + cudaMemcpy( + neighborsSeg, neighborsDP.handle(), neighborsBytes, cudaMemcpyDeviceToHost()); + checkCudaError(returnValue, "cudaMemcpy"); + returnValue = + cudaMemcpy( + distancesSeg, distancesDP.handle(), distancesBytes, cudaMemcpyDeviceToHost()); + checkCudaError(returnValue, "cudaMemcpy"); + + return TieredSearchResultsImpl.create( + neighborsLayout, + distancesLayout, + neighborsSeg, + distancesSeg, + topK, + query.getMapping(), + numQueries); } - - return TieredSearchResultsImpl.create( - neighborsLayout, - distancesLayout, - neighborsSeg, - distancesSeg, - topK, - query.getMapping(), - numQueries); } } } @@ -333,25 +337,25 @@ private void performExtend(CuVSMatrix extendDataset) { // Allocate device memory for extend data long dataSize = C_FLOAT_BYTE_SIZE * rows * cols; - MemorySegment datasetDP = allocateRMMSegment(cuvsRes, dataSize); - - // Copy host to device - checkCudaError( - cudaMemcpy(datasetDP, hostDataSeg, dataSize, cudaMemcpyHostToDevice()), "cudaMemcpy"); - - // Create tensor from device memory - long[] datasetShape = {rows, cols}; - MemorySegment datasetTensor = - prepareTensor(localArena, datasetDP, datasetShape, kDLFloat(), 32, kDLCUDA(), 1); + try (var datasetDP = allocateRMMSegment(cuvsRes, dataSize)) { + // Copy host to device + checkCudaError( + cudaMemcpy(datasetDP.handle(), hostDataSeg, dataSize, cudaMemcpyHostToDevice()), + "cudaMemcpy"); - checkCuVSError(cuvsStreamSync(cuvsRes), "cuvsStreamSync"); + // Create tensor from device memory + long[] datasetShape = {rows, cols}; + MemorySegment datasetTensor = + prepareTensor( + localArena, datasetDP.handle(), datasetShape, kDLFloat(), 32, kDLCUDA(), 1); - checkCuVSError( - cuvsTieredIndexExtend(cuvsRes, datasetTensor, tieredIndexReference.getMemorySegment()), - "cuvsTieredIndexExtend"); + checkCuVSError(cuvsStreamSync(cuvsRes), "cuvsStreamSync"); - // Clean up device memory - checkCuVSError(cuvsRMMFree(cuvsRes, datasetDP, dataSize), "cuvsRMMFree"); + checkCuVSError( + cuvsTieredIndexExtend( + cuvsRes, datasetTensor, tieredIndexReference.getMemorySegment()), + "cuvsTieredIndexExtend"); + } } } } diff --git a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/common/CloseableRMMAllocation.java b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/common/CloseableRMMAllocation.java new file mode 100644 index 0000000000..6dcc4c892f --- /dev/null +++ b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/common/CloseableRMMAllocation.java @@ -0,0 +1,86 @@ +/* + * Copyright (c) 2025, NVIDIA CORPORATION. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.nvidia.cuvs.internal.common; + +import static com.nvidia.cuvs.internal.common.LinkerHelper.C_POINTER; +import static com.nvidia.cuvs.internal.common.Util.checkCuVSError; +import static com.nvidia.cuvs.internal.panama.headers_h.cuvsRMMAlloc; +import static com.nvidia.cuvs.internal.panama.headers_h.cuvsRMMFree; + +import java.lang.foreign.Arena; +import java.lang.foreign.MemorySegment; + +/** + * A closeable handle for RMM allocations that can be used with try-with-resources. + * This class manages the lifecycle of a memory allocation made with RMM (RAPIDS Memory Manager). + * It ensures that the allocated memory is properly released when no longer needed. + */ +public class CloseableRMMAllocation implements CloseableHandle { + + private final long cuvsResourceHandle; + private long numBytes; + private MemorySegment pointer; + + private CloseableRMMAllocation(long cuvsResourceHandle, long numBytes, MemorySegment pointer) { + this.cuvsResourceHandle = cuvsResourceHandle; + this.numBytes = numBytes; + this.pointer = pointer; + } + + /** + * Copy constructor transfers the ownership of the argument's MemorySegment to this object. + */ + public CloseableRMMAllocation(CloseableRMMAllocation other) { + this.cuvsResourceHandle = other.cuvsResourceHandle; + this.numBytes = other.numBytes; + this.pointer = other.release(); + } + + public static CloseableRMMAllocation allocateRMMSegment(long cuvsResourceHandle, long numBytes) { + try (var localArena = Arena.ofConfined()) { + MemorySegment datasetMemorySegment = localArena.allocate(C_POINTER); + checkCuVSError( + cuvsRMMAlloc(cuvsResourceHandle, datasetMemorySegment, numBytes), "cuvsRMMAlloc"); + return new CloseableRMMAllocation( + cuvsResourceHandle, numBytes, datasetMemorySegment.get(C_POINTER, 0)); + } + } + + @Override + public MemorySegment handle() { + return pointer; + } + + private MemorySegment release() { + var oldPointer = pointer; + pointer = MemorySegment.NULL; + numBytes = 0; + return oldPointer; + } + + private boolean mustClose() { + return pointer != MemorySegment.NULL; + } + + @Override + public void close() { + if (mustClose()) { + checkCuVSError(cuvsRMMFree(cuvsResourceHandle, pointer, numBytes), "cuvsRMMFree"); + } + } + + public static CloseableRMMAllocation EMPTY = new CloseableRMMAllocation(0, 0, MemorySegment.NULL); +} diff --git a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/common/Util.java b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/common/Util.java index 6d08878650..d5972ecca7 100644 --- a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/common/Util.java +++ b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/common/Util.java @@ -19,12 +19,10 @@ import static com.nvidia.cuvs.internal.common.LinkerHelper.C_FLOAT; import static com.nvidia.cuvs.internal.common.LinkerHelper.C_INT; import static com.nvidia.cuvs.internal.common.LinkerHelper.C_LONG; -import static com.nvidia.cuvs.internal.common.LinkerHelper.C_POINTER; import static com.nvidia.cuvs.internal.panama.headers_h.cudaGetDeviceCount; import static com.nvidia.cuvs.internal.panama.headers_h.cudaGetDeviceProperties_v2; import static com.nvidia.cuvs.internal.panama.headers_h.cudaMemGetInfo; import static com.nvidia.cuvs.internal.panama.headers_h.cudaSetDevice; -import static com.nvidia.cuvs.internal.panama.headers_h.cuvsRMMAlloc; import static com.nvidia.cuvs.internal.panama.headers_h.size_t; import com.nvidia.cuvs.GPUInfo; @@ -351,15 +349,4 @@ public static MemorySegment prepareTensor( return tensor; } - - public static MemorySegment allocateRMMSegment(long resourceHandle, long datasetBytes) { - try (var localArena = Arena.ofConfined()) { - MemorySegment datasetMemorySegment = localArena.allocate(C_POINTER); - - var returnValue = cuvsRMMAlloc(resourceHandle, datasetMemorySegment, datasetBytes); - checkCuVSError(returnValue, "cuvsRMMAlloc"); - - return datasetMemorySegment.get(C_POINTER, 0); - } - } } From 6947ac41166a5511b28e023999aa093f0b54e424 Mon Sep 17 00:00:00 2001 From: MithunR Date: Fri, 8 Aug 2025 15:16:17 -0700 Subject: [PATCH 25/93] Drop log level for CAGRA trace messages (#1229) This is a follow-up to #1226, which aims to reduce the amount of noise in the Java test logs. In the course of making the changes in #1226, one observes that there exist some `TRACE` messages from `libcuvs` that are being logged at the `INFO` level. Thousands of these fill up the Java logs, obscuring some legitimate warnings. For instance: 1. `Graph optimized, creating index` and `optimizing graph` appear 4574 times. 2. `Saving CAGRA index with dataset` appears 1025 times. There isn't any further information in the messages, beyond indicating that position in code. This commit changes the level of the first message to `TRACE`, and the second to `DEBUG` (to match the logs around it). AFAICT, this shouldn't have an effect on the tests, or have impact on user-logs, beyond reducing log volume. Authors: - MithunR (https://github.com/mythrocks) Approvers: - Jinsol Park (https://github.com/jinsolp) - Corey J. Nolet (https://github.com/cjnolet) URL: https://github.com/rapidsai/cuvs/pull/1229 --- cpp/src/neighbors/detail/cagra/cagra_build.cuh | 4 ++-- cpp/src/neighbors/detail/cagra/cagra_serialize.cuh | 4 ++-- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/cpp/src/neighbors/detail/cagra/cagra_build.cuh b/cpp/src/neighbors/detail/cagra/cagra_build.cuh index 651e6bbb85..07c9157b72 100644 --- a/cpp/src/neighbors/detail/cagra/cagra_build.cuh +++ b/cpp/src/neighbors/detail/cagra/cagra_build.cuh @@ -769,14 +769,14 @@ index build( cagra_graph = raft::make_host_matrix(dataset.extent(0), graph_degree); - RAFT_LOG_INFO("optimizing graph"); + RAFT_LOG_TRACE("optimizing graph"); optimize(res, knn_graph->view(), cagra_graph.view(), params.guarantee_connectivity); // free intermediate graph before trying to create the index knn_graph.reset(); } - RAFT_LOG_INFO("Graph optimized, creating index"); + RAFT_LOG_TRACE("Graph optimized, creating index"); // Construct an index from dataset and optimized knn graph. if (params.compression.has_value()) { diff --git a/cpp/src/neighbors/detail/cagra/cagra_serialize.cuh b/cpp/src/neighbors/detail/cagra/cagra_serialize.cuh index 4bd761dc60..8b0990ac6f 100644 --- a/cpp/src/neighbors/detail/cagra/cagra_serialize.cuh +++ b/cpp/src/neighbors/detail/cagra/cagra_serialize.cuh @@ -1,5 +1,5 @@ /* - * Copyright (c) 2023-2024, NVIDIA CORPORATION. + * Copyright (c) 2023-2025, NVIDIA CORPORATION. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -74,7 +74,7 @@ void serialize(raft::resources const& res, raft::serialize_scalar(res, os, include_dataset); if (include_dataset) { - RAFT_LOG_INFO("Saving CAGRA index with dataset"); + RAFT_LOG_DEBUG("Saving CAGRA index with dataset"); neighbors::detail::serialize(res, os, index_.data()); } else { RAFT_LOG_DEBUG("Saving CAGRA index WITHOUT dataset"); From 811bc0068e5023d9988afbf3b6fd424f8c1423f3 Mon Sep 17 00:00:00 2001 From: MithunR Date: Mon, 11 Aug 2025 07:03:04 -0700 Subject: [PATCH 26/93] Reduce noise in Java tests (#1226) When the Java tests run currently, they produce a large amount of logs to the console. This includes contents of the test datasets, and various messages that ought to be trace/debug messages. This commit changes a lot of these logs to trace/debug messages, to reduce the console log noise. Note that this doesn't completely eliminate unnecessary noise. There will need to be follow-up changes to `libcuvs` itself. We might also need to investigate some of the warnings that were previously obscured by the log volume. Authors: - MithunR (https://github.com/mythrocks) Approvers: - Corey J. Nolet (https://github.com/cjnolet) URL: https://github.com/rapidsai/cuvs/pull/1226 --- .../nvidia/cuvs/BruteForceRandomizedIT.java | 12 +++---- .../nvidia/cuvs/CagraBuildAndSearchIT.java | 32 +++++++++---------- .../cuvs/CagraMultiThreadStabilityIT.java | 16 +++++----- .../com/nvidia/cuvs/CagraRandomizedIT.java | 14 ++++---- .../java/com/nvidia/cuvs/CuVSResourcesIT.java | 12 +++---- .../java/com/nvidia/cuvs/CuVSTestCase.java | 10 +++--- .../com/nvidia/cuvs/HnswBuildAndSearchIT.java | 2 +- .../com/nvidia/cuvs/HnswRandomizedIT.java | 12 +++---- 8 files changed, 55 insertions(+), 55 deletions(-) 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 4a446c249c..6bfbf4d2c7 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 @@ -36,7 +36,7 @@ public class BruteForceRandomizedIT extends CuVSTestCase { public void setup() { assumeTrue(isLinuxAmd64()); initializeRandom(); - log.info("Random context initialized for test."); + log.trace("Random context initialized for test."); } @Test @@ -80,10 +80,10 @@ private void tmpResultsTopKWithRandomValues(boolean useNativeMemoryDataset) thro // Generate random query vectors float[][] queries = generateData(random, numQueries, dimensions); - log.info("Dataset size: {}x{}", datasetSize, dimensions); - log.info("Query size: {}x{}", numQueries, dimensions); - log.info("TopK: {}", topK); - log.info("Use native memory dataset? " + useNativeMemoryDataset); + log.debug("Dataset size: {}x{}", datasetSize, dimensions); + log.debug("Query size: {}x{}", numQueries, dimensions); + log.debug("TopK: {}", topK); + log.debug("Use native memory dataset? " + useNativeMemoryDataset); // Debugging: Log dataset and queries if (log.isDebugEnabled()) { @@ -138,7 +138,7 @@ private void tmpResultsTopKWithRandomValues(boolean useNativeMemoryDataset) thro .build(); } - log.info("Index built successfully. Executing search..."); + log.trace("Index built successfully. Executing search..."); SearchResults results = index.search(query); compareResults(results, expected, topK, datasetSize, numQueries); 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 39e8226250..cc7879a43f 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 @@ -62,7 +62,7 @@ public class CagraBuildAndSearchIT extends CuVSTestCase { public void setup() { assumeTrue("not supported on " + System.getProperty("os.name"), isLinuxAmd64()); initializeRandom(); - log.info("Random context initialized for test."); + log.trace("Random context initialized for test."); } private static void runConcurrently(int nThreads, Supplier runnableSupplier) @@ -379,7 +379,7 @@ public void testPrefilteringReducesResults() throws Throwable { SearchResults fullSearchResults = index.search(fullQuery); List> fullResults = fullSearchResults.getResults(); - log.info("Full results: {}", fullResults); + log.debug("Full results: {}", fullResults); // Apply prefilter: only allow ids 0 and 2 (bitset: 1100) BitSet prefilter = new BitSet(4); @@ -396,7 +396,7 @@ public void testPrefilteringReducesResults() throws Throwable { SearchResults filteredSearchResults = index.search(filteredQuery); List> filteredResults = filteredSearchResults.getResults(); - log.info("Filtered results: {}", filteredResults); + log.debug("Filtered results: {}", filteredResults); assertEquals(expectedResults, fullResults); assertEquals(expectedFilteredResults, filteredResults); @@ -461,14 +461,14 @@ private void queryAndCompare( SearchResults results = index1.search(cuvsQuery); // Check results - log.info(results.getResults().toString()); + log.debug(results.getResults().toString()); checkResults(expectedResults, results.getResults()); // Search from the second index results = index2.search(cuvsQuery); // Check results - log.info(results.getResults().toString()); + log.debug(results.getResults().toString()); checkResults(expectedResults, results.getResults()); } @@ -554,23 +554,23 @@ public void testMergingIndexes() throws Throwable { .withMetric(CuvsDistanceType.L2Expanded) .build(); - log.info("Building first index..."); + log.trace("Building first index..."); CagraIndex index1 = CagraIndex.newBuilder(resources) .withDataset(vector1) .withIndexParams(indexParams) .build(); - log.info("Building second index..."); + log.trace("Building second index..."); CagraIndex index2 = CagraIndex.newBuilder(resources) .withDataset(vector2) .withIndexParams(indexParams) .build(); - log.info("Merging indexes..."); + log.trace("Merging indexes..."); CagraIndex mergedIndex = CagraIndex.merge(new CagraIndex[] {index1, index2}); - log.info("Merge completed successfully"); + log.trace("Merge completed successfully"); CagraSearchParams searchParams = new CagraSearchParams.Builder().build(); @@ -582,9 +582,9 @@ public void testMergingIndexes() throws Throwable { .withMapping(SearchResults.IDENTITY_MAPPING) .build(); - log.info("Searching merged index..."); + log.trace("Searching merged index..."); SearchResults results = mergedIndex.search(query); - log.info("Search results: " + results.getResults().toString()); + log.debug("Search results: " + results.getResults().toString()); assertEquals(expectedResults, results.getResults()); @@ -644,14 +644,14 @@ public void testMergeStrategies() throws Throwable { .withMetric(CuvsDistanceType.L2Expanded) .build(); - log.info("Building first index..."); + log.trace("Building first index..."); CagraIndex index1 = CagraIndex.newBuilder(resources) .withDataset(vector1) .withIndexParams(indexParams) .build(); - log.info("Building second index..."); + log.trace("Building second index..."); CagraIndex index2 = CagraIndex.newBuilder(resources) .withDataset(vector2) @@ -673,10 +673,10 @@ public void testMergeStrategies() throws Throwable { .withStrategy(MergeStrategy.PHYSICAL) .build(); - log.info("Merging indexes with PHYSICAL strategy..."); + log.trace("Merging indexes with PHYSICAL strategy..."); CagraIndex physicalMergedIndex = CagraIndex.merge(new CagraIndex[] {index1, index2}, physicalMergeParams); - log.info("Physical merge completed successfully"); + log.trace("Physical merge completed successfully"); CagraSearchParams searchParams = new CagraSearchParams.Builder().build(); @@ -688,7 +688,7 @@ public void testMergeStrategies() throws Throwable { .withMapping(SearchResults.IDENTITY_MAPPING) .build(); - log.info("Searching physically merged index..."); + log.trace("Searching physically merged index..."); SearchResults physicalResults = physicalMergedIndex.search(query); assertNotNull("Physical merge search results should not be null", physicalResults); assertEquals( 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 3bcef6bb94..1df57d1e0b 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 @@ -57,7 +57,7 @@ private interface QueryAction { public void setup() { assumeTrue("not supported on " + System.getProperty("os.name"), isLinuxAmd64()); initializeRandom(); - log.info("Multi-threaded stability test initialized"); + log.trace("Multi-threaded stability test initialized"); } @Test @@ -77,14 +77,14 @@ private void testQueryingUsingMultipleThreads(QueryAction queryAction) throws Th final int dataSize = 10000; final int numThreads = 16; - log.info(" Dataset: {}x{}", dataSize, dimensions); + log.debug(" Dataset: {}x{}", dataSize, dimensions); // High thread count to increase contention - log.info(" Threads: {}, Queries per thread: {}", numThreads, queriesPerThread); + log.debug(" Threads: {}, Queries per thread: {}", numThreads, queriesPerThread); float[][] dataset = generateRandomDataset(dataSize); try (CuVSResources resources = CheckedCuVSResources.create()) { - log.info("Creating CAGRA index for MultiThreaded stability test..."); + log.trace("Creating CAGRA index for MultiThreaded stability test..."); CagraIndexParams indexParams = new CagraIndexParams.Builder() @@ -101,7 +101,7 @@ private void testQueryingUsingMultipleThreads(QueryAction queryAction) throws Th .withIndexParams(indexParams) .build(); - log.info("CAGRA index created, starting high-contention multi-threaded search..."); + log.trace("CAGRA index created, starting high-contention multi-threaded search..."); // Create high contention scenario that would fail without using separate resources in every // thread @@ -128,7 +128,7 @@ private void testQueryingUsingMultipleThreads(QueryAction queryAction) throws Th // No Thread.yield() - maximize contention } - log.info("Thread {} completed successfully", finalThreadId); + log.trace("Thread {} completed successfully", finalThreadId); } catch (Throwable t) { log.error("Thread {} failed: {}", finalThreadId, t.getMessage(), t); @@ -141,7 +141,7 @@ private void testQueryingUsingMultipleThreads(QueryAction queryAction) throws Th } // Start all threads simultaneously to maximize contention - log.info("Starting all {} threads simultaneously...", numThreads); + log.debug("Starting all {} threads simultaneously...", numThreads); startLatch.countDown(); // Wait for all threads to complete @@ -167,7 +167,7 @@ private void testQueryingUsingMultipleThreads(QueryAction queryAction) throws Th int expectedTotalQueries = numThreads * queriesPerThread; int actualSuccessfulQueries = successfulQueries.get(); - log.info(" Successful queries: {} / {}", actualSuccessfulQueries, expectedTotalQueries); + log.debug(" Successful queries: {} / {}", actualSuccessfulQueries, expectedTotalQueries); if (firstError.get() != null) { fail("MultiThreaded stablity test failed:" + " " + firstError.get().getMessage()); 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 02b5a6526a..69e28d20aa 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 @@ -36,7 +36,7 @@ public class CagraRandomizedIT extends CuVSTestCase { public void setup() { assumeTrue("not supported on " + System.getProperty("os.name"), isLinuxAmd64()); initializeRandom(); - log.info("Random context initialized for test."); + log.trace("Random context initialized for test."); } @Test @@ -87,10 +87,10 @@ private void tmpResultsTopKWithRandomValues(boolean useNativeMemoryDataset) thro // Generate random query vectors float[][] queries = generateData(random, numQueries, dimensions); - log.info("Dataset size: {}x{}", datasetSize, dimensions); - log.info("Query size: {}x{}", numQueries, dimensions); - log.info("TopK: {}", topK); - log.info("Use native memory dataset? " + useNativeMemoryDataset); + log.debug("Dataset size: {}x{}", datasetSize, dimensions); + log.debug("Query size: {}x{}", numQueries, dimensions); + log.debug("TopK: {}", topK); + log.debug("Use native memory dataset? " + useNativeMemoryDataset); // Debugging: Log dataset and queries if (log.isDebugEnabled()) { @@ -138,7 +138,7 @@ private void tmpResultsTopKWithRandomValues(boolean useNativeMemoryDataset) thro .withIndexParams(indexParams) .build(); } - log.info("Index built successfully."); + log.trace("Index built successfully."); try { // Execute search and retrieve results @@ -153,7 +153,7 @@ private void tmpResultsTopKWithRandomValues(boolean useNativeMemoryDataset) thro } CagraQuery query = queryBuilder.build(); - log.info("Query built successfully. Executing search..."); + log.trace("Query built successfully. Executing search..."); SearchResults results = index.search(query); compareResults(results, expected, topK, datasetSize, numQueries); 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 a599089dce..e479a6ebda 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 @@ -43,7 +43,7 @@ public void testConcurrentAccessViaCheckedCuVSResourcesIsForbidden() throws Thro var executor = Executors.newFixedThreadPool(1)) { try (var access1 = resources.access()) { - log.info( + log.debug( "Outer access to resource {} from {}", access1.handle(), Thread.currentThread().threadId()); @@ -56,18 +56,18 @@ public void testConcurrentAccessViaCheckedCuVSResourcesIsForbidden() throws Thro executor.submit( () -> { try (var access2 = resources.access()) { - log.info( + log.debug( "Nested access to resource {} from {}", access2.handle(), Thread.currentThread().threadId()); - log.info("Nested access finished"); + log.debug("Nested access finished"); } }); future.get(); }); assertEquals(IllegalStateException.class, exception.getCause().getClass()); assertEquals(expectedError, exception.getCause().getMessage()); - log.info("Outer access finished"); + log.debug("Outer access finished"); } } } @@ -78,7 +78,7 @@ public void testSequentialAccessViaCheckedCuVSResourcesIsAllowed() throws Throwa var executor = Executors.newFixedThreadPool(1)) { try (var access1 = resources.access()) { - log.info( + log.debug( "Access 1 to resource {} from {}", access1.handle(), Thread.currentThread().threadId()); } @@ -86,7 +86,7 @@ public void testSequentialAccessViaCheckedCuVSResourcesIsAllowed() throws Throwa executor.submit( () -> { try (var access2 = resources.access()) { - log.info( + log.debug( "Access 2 to resource {} from {}", access2.handle(), Thread.currentThread().threadId()); 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 17fe7d5206..fee3c87cca 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 @@ -35,7 +35,7 @@ public abstract class CuVSTestCase { protected void initializeRandom() { random = RandomizedContext.current().getRandom(); - log.info("Test seed: " + RandomizedContext.current().getRunnerSeedAsString()); + log.debug("Test seed: " + RandomizedContext.current().getRunnerSeedAsString()); } protected float[][] generateData(Random random, int rows, int cols) { @@ -77,7 +77,7 @@ protected List> generateExpectedResults( neighborsResult.add(neighbors.subList(0, Math.min(topK * 2 + 10, dataset.length))); } - log.info("Expected results generated successfully."); + log.trace("Expected results generated successfully."); return neighborsResult; } @@ -89,8 +89,8 @@ protected void compareResults( int numQueries) { for (int i = 0; i < numQueries; i++) { - log.info("Results returned for query " + i + ": " + results.getResults().get(i).keySet()); - log.info( + log.debug("Results returned for query " + i + ": " + results.getResults().get(i).keySet()); + log.debug( "Expected results for query " + i + ": " @@ -109,7 +109,7 @@ protected void compareResults( .toList(); // just make sure that the first 5 results are in the expected list (which - // comprises of 2*topK results) + // consists of 2*topK results) for (int j = 0; j < Math.min(5, sortedResultKeys.size()); j++) { assertTrue( "Not found in expected list: " + sortedResultKeys.get(j), 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 ab041ec24b..b2b86c37ad 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 @@ -93,7 +93,7 @@ private static void indexAndQueryOnce( SearchResults results = hnswIndex.search(hnswQuery); // Check results - log.info(results.getResults().toString()); + log.debug(results.getResults().toString()); checkResults(expectedResults, results.getResults()); // Cleanup 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 3000c125a3..8cda810c6c 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 @@ -40,7 +40,7 @@ public class HnswRandomizedIT extends CuVSTestCase { public void setup() { assumeTrue("not supported on " + System.getProperty("os.name"), isLinuxAmd64()); initializeRandom(); - log.info("Random context initialized for test."); + log.trace("Random context initialized for test."); } @Test @@ -72,10 +72,10 @@ private void tmpResultsTopKWithRandomValues(boolean useNativeMemoryDataset) thro // Generate random query vectors float[][] queries = generateData(random, numQueries, dimensions); - log.info("Dataset size: {}x{}", datasetSize, dimensions); - log.info("Query size: {}x{}", numQueries, dimensions); - log.info("TopK: {}", topK); - log.info("Use native memory dataset? " + useNativeMemoryDataset); + log.debug("Dataset size: {}x{}", datasetSize, dimensions); + log.debug("Query size: {}x{}", numQueries, dimensions); + log.debug("TopK: {}", topK); + log.debug("Use native memory dataset? " + useNativeMemoryDataset); // Debugging: Log dataset and queries if (log.isDebugEnabled()) { @@ -160,7 +160,7 @@ private void tmpResultsTopKWithRandomValues(boolean useNativeMemoryDataset) thro .withTopK(topK) .build(); - log.info("Index built successfully. Executing search..."); + log.trace("Index built successfully. Executing search..."); SearchResults results = hnswIndex.search(hnswQuery); compareResults(results, expected, topK, datasetSize, numQueries); From aff1e3b4291b735fa1838dccd57b5ec491141ec4 Mon Sep 17 00:00:00 2001 From: Robert Maynard Date: Mon, 11 Aug 2025 20:20:33 -0400 Subject: [PATCH 27/93] Use rapids_cuda_enable_fatbin_compression (#1239) Standardize compression flags via rapids_cuda_enable_fatbin_compression Authors: - Robert Maynard (https://github.com/robertmaynard) Approvers: - Bradley Dice (https://github.com/bdice) URL: https://github.com/rapidsai/cuvs/pull/1239 --- cpp/cmake/modules/ConfigureCUDA.cmake | 9 ++------- 1 file changed, 2 insertions(+), 7 deletions(-) diff --git a/cpp/cmake/modules/ConfigureCUDA.cmake b/cpp/cmake/modules/ConfigureCUDA.cmake index 0b6ebbaad2..5ce9b918b4 100644 --- a/cpp/cmake/modules/ConfigureCUDA.cmake +++ b/cpp/cmake/modules/ConfigureCUDA.cmake @@ -48,13 +48,8 @@ list(APPEND CUVS_CUDA_FLAGS --expt-extended-lambda --expt-relaxed-constexpr) list(APPEND CUVS_CXX_FLAGS "-DCUDA_API_PER_THREAD_DEFAULT_STREAM") list(APPEND CUVS_CUDA_FLAGS "-DCUDA_API_PER_THREAD_DEFAULT_STREAM") # make sure we produce smallest binary size -list(APPEND CUVS_CUDA_FLAGS -Xfatbin=-compress-all) -if(CMAKE_CUDA_COMPILER_ID STREQUAL "NVIDIA" - AND (CMAKE_CUDA_COMPILER_VERSION VERSION_GREATER_EQUAL 12.9 AND CMAKE_CUDA_COMPILER_VERSION - VERSION_LESS 13.0) -) - list(APPEND CUVS_CUDA_FLAGS -Xfatbin=--compress-level=3) -endif() +include(${rapids-cmake-dir}/cuda/enable_fatbin_compression.cmake) +rapids_cuda_enable_fatbin_compression(VARIABLE CUVS_CUDA_FLAGS TUNE_FOR rapids) # Option to enable line info in CUDA device compilation to allow introspection when profiling / # memchecking From 34032a813c9c38573a6cc091d98669870ee2f053 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Lorenzo=20Dematt=C3=A9?= Date: Tue, 12 Aug 2025 07:50:12 +0200 Subject: [PATCH 28/93] [Java] Add CAGRA index graph accessor/build from graph (host memory) (#1216) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit This PR leverages the functions introduced by #1086 and the data structures introduced by #1111 to access, copy, and re-create an index to/from a CAGRA graph. Supersedes #1105 Authors: - Lorenzo Dematté (https://github.com/ldematte) - MithunR (https://github.com/mythrocks) Approvers: - MithunR (https://github.com/mythrocks) URL: https://github.com/rapidsai/cuvs/pull/1216 --- .../main/java/com/nvidia/cuvs/CagraIndex.java | 13 ++ .../main/java/com/nvidia/cuvs/CuVSMatrix.java | 1 + .../nvidia/cuvs/internal/CagraIndexImpl.java | 125 +++++++++++++++++- .../cuvs/internal/CuVSHostMatrixImpl.java | 15 ++- .../cuvs/internal/CuVSMatrixBaseImpl.java | 20 +++ .../com/nvidia/cuvs/internal/common/Util.java | 13 +- .../nvidia/cuvs/CagraBuildAndSearchIT.java | 42 ++++++ 7 files changed, 217 insertions(+), 12 deletions(-) diff --git a/java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraIndex.java b/java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraIndex.java index 84487857c8..427b5e22d2 100644 --- a/java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraIndex.java +++ b/java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraIndex.java @@ -50,6 +50,13 @@ public interface CagraIndex { */ SearchResults search(CagraQuery query) throws Throwable; + /** Returns the CAGRA graph + * + * @return a {@link CuVSMatrix} encapsulating the native int (uint32_t) array used to represent + * the cagra graph + */ + CuVSMatrix getGraph(); + /** * A method to persist a CAGRA index using an instance of {@link OutputStream} * for writing index bytes. @@ -208,6 +215,12 @@ interface Builder { */ Builder from(InputStream inputStream); + /** + * Sets a CAGRA graph instance to re-create an index from a + * previously built graph. + */ + Builder from(CuVSMatrix graph); + /** * Sets the dataset vectors for building the {@link CagraIndex}. * diff --git a/java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSMatrix.java b/java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSMatrix.java index 0b892fbc5f..f8e60c5555 100644 --- a/java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSMatrix.java +++ b/java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSMatrix.java @@ -30,6 +30,7 @@ public interface CuVSMatrix extends AutoCloseable { enum DataType { FLOAT, INT, + UINT, BYTE } diff --git a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CagraIndexImpl.java b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CagraIndexImpl.java index ac3e9ba43a..489213f4da 100644 --- a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CagraIndexImpl.java +++ b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CagraIndexImpl.java @@ -45,6 +45,7 @@ import com.nvidia.cuvs.internal.common.CloseableHandle; import com.nvidia.cuvs.internal.common.CloseableRMMAllocation; import com.nvidia.cuvs.internal.common.CompositeCloseableHandle; +import com.nvidia.cuvs.internal.common.Util; import com.nvidia.cuvs.internal.panama.cuvsCagraCompressionParams; import com.nvidia.cuvs.internal.panama.cuvsCagraIndexParams; import com.nvidia.cuvs.internal.panama.cuvsCagraMergeParams; @@ -122,6 +123,31 @@ private CagraIndexImpl(IndexReference indexReference, CuVSResources resources) { this.destroyed = false; } + /** + * Constructor for creating an index from a pre-build CAGRA graph + * + * @param metric the distance type used + * @param graph a previously built CAGRA graph + * @param dataset the dataset used for indexing + * @param resources an instance of {@link CuVSResources} + */ + private CagraIndexImpl( + CagraIndexParams.CuvsDistanceType metric, + CuVSMatrix graph, + CuVSMatrix dataset, + CuVSResources resources) { + Objects.requireNonNull(graph); + Objects.requireNonNull(dataset); + + this.resources = resources; + + assert graph instanceof CuVSMatrixBaseImpl; + assert dataset instanceof CuVSMatrixBaseImpl; + + this.cagraIndexReference = + fromGraph(metric, (CuVSMatrixBaseImpl) graph, (CuVSMatrixBaseImpl) dataset); + } + private void checkNotDestroyed() { if (destroyed) { throw new IllegalStateException("destroyed"); @@ -164,9 +190,12 @@ private IndexReference build(CagraIndexParams indexParameters, CuVSMatrixBaseImp omp_set_num_threads(numWriterThreads); MemorySegment dataSeg = dataset.memorySegment(); + // TODO: type kDLCPU()/kDLCUDA() should be aligned with the CuVSMatrixBaseImpl type (host or + // device?) long[] datasetShape = {rows, cols}; - MemorySegment datasetTensor = prepareTensor(localArena, dataSeg, datasetShape, 2, 32, 2, 1); + MemorySegment datasetTensor = + prepareTensor(localArena, dataSeg, datasetShape, kDLFloat(), 32, kDLCPU(), 1); var index = createCagraIndex(); @@ -381,6 +410,81 @@ public void serialize(OutputStream outputStream, Path tempFile, int bufferLength } } + @Override + public CuVSMatrix getGraph() { + try (var localArena = Arena.ofConfined()) { + var outPtr = localArena.allocate(uint32_t); + checkCuVSError( + cuvsCagraIndexGetGraphDegree(cagraIndexReference.getMemorySegment(), outPtr), + "cuvsCagraIndexGetGraphDegree"); + long graphDegree = Util.dereferenceUnsignedInt(outPtr); + + checkCuVSError( + cuvsCagraIndexGetSize(cagraIndexReference.getMemorySegment(), outPtr), + "cuvsCagraIndexGetSize"); + long size = Util.dereferenceUnsignedInt(outPtr); + + // TODO: use a "device" graph + tensor, avoid (defer) copy + var graph = new CuVSHostMatrixArenaImpl(size, graphDegree, CuVSMatrix.DataType.UINT); + var graphHostTensor = graph.toTensor(localArena); + var graphDeviceTensor = + prepareTensor( + localArena, + MemorySegment.NULL, + new long[] {size, graphDegree}, + kDLUInt(), + 32, + kDLCUDA(), + 1); + checkCuVSError( + cuvsCagraIndexGetGraph(cagraIndexReference.getMemorySegment(), graphDeviceTensor), + "cuvsCagraIndexGetGraph"); + + try (var resourceAccess = resources.access()) { + var cuvsRes = resourceAccess.handle(); + checkCuVSError(cuvsStreamSync(cuvsRes), "cuvsStreamSync"); + + checkCuVSError( + cuvsMatrixCopy(cuvsRes, graphDeviceTensor, graphHostTensor), "cuvsMatrixCopy"); + + checkCuVSError(cuvsStreamSync(cuvsRes), "cuvsStreamSync"); + } + + return graph; + } + } + + private IndexReference fromGraph( + CagraIndexParams.CuvsDistanceType metric, + CuVSMatrixBaseImpl graph, + CuVSMatrixBaseImpl dataset) { + try (var localArena = Arena.ofConfined()) { + long rows = dataset.size(); + long cols = dataset.columns(); + + var index = createCagraIndex(); + try (var resourcesAccess = resources.access()) { + long cuvsRes = resourcesAccess.handle(); + + long[] datasetShape = {rows, cols}; + MemorySegment datasetTensor = + prepareTensor( + localArena, dataset.memorySegment(), datasetShape, kDLFloat(), 32, kDLCPU(), 1); + + long[] graphShape = {graph.size(), graph.columns()}; + MemorySegment graphTensor = + prepareTensor( + localArena, graph.memorySegment(), graphShape, kDLUInt(), 32, kDLCPU(), 1); + + checkCuVSError( + cuvsCagraIndexFromArgs(cuvsRes, metric.value, graphTensor, datasetTensor, index), + "cuvsCagraIndexFromArgs"); + } + + return new IndexReference(index, dataset); + } + } + @Override public void serializeToHNSW(OutputStream outputStream) throws Throwable { Path path = @@ -678,6 +782,7 @@ public static class Builder implements CagraIndex.Builder { private CagraIndexParams cagraIndexParams; private final CuVSResources cuvsResources; private InputStream inputStream; + private CuVSMatrix graph; public Builder(CuVSResources cuvsResources) { this.cuvsResources = cuvsResources; @@ -689,6 +794,12 @@ public Builder from(InputStream inputStream) { return this; } + @Override + public Builder from(CuVSMatrix graph) { + this.graph = graph; + return this; + } + @Override public Builder withDataset(float[][] vectors) { this.dataset = CuVSMatrix.ofArray(vectors); @@ -712,7 +823,17 @@ public CagraIndexImpl build() throws Throwable { if (inputStream != null) { return new CagraIndexImpl(inputStream, cuvsResources); } else { - return new CagraIndexImpl(cagraIndexParams, dataset, cuvsResources); + if (graph != null) { + if (cagraIndexParams == null || dataset == null) { + throw new IllegalArgumentException( + "In order to reconstruct a CAGRA index from a graph, " + + "you must specify the original dataset and the metric used."); + } + return new CagraIndexImpl( + cagraIndexParams.getCuvsDistanceType(), graph, dataset, cuvsResources); + } else { + return new CagraIndexImpl(cagraIndexParams, dataset, cuvsResources); + } } } } diff --git a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CuVSHostMatrixImpl.java b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CuVSHostMatrixImpl.java index a7096e207c..c2895282ca 100644 --- a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CuVSHostMatrixImpl.java +++ b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CuVSHostMatrixImpl.java @@ -18,13 +18,12 @@ import static com.nvidia.cuvs.internal.common.LinkerHelper.C_CHAR; import static com.nvidia.cuvs.internal.common.LinkerHelper.C_FLOAT; import static com.nvidia.cuvs.internal.common.LinkerHelper.C_INT; +import static com.nvidia.cuvs.internal.common.Util.prepareTensor; +import static com.nvidia.cuvs.internal.panama.headers_h.*; import com.nvidia.cuvs.CuVSHostMatrix; import com.nvidia.cuvs.RowView; -import java.lang.foreign.MemoryLayout; -import java.lang.foreign.MemorySegment; -import java.lang.foreign.SequenceLayout; -import java.lang.foreign.ValueLayout; +import java.lang.foreign.*; import java.lang.invoke.VarHandle; /** @@ -61,7 +60,7 @@ protected CuVSHostMatrixImpl( protected static ValueLayout valueLayoutFromType(DataType dataType) { return switch (dataType) { case FLOAT -> C_FLOAT; - case INT -> C_INT; + case INT, UINT -> C_INT; case BYTE -> C_CHAR; }; } @@ -131,6 +130,12 @@ public ValueLayout valueLayout() { return valueLayout; } + @Override + public MemorySegment toTensor(Arena arena) { + return prepareTensor( + arena, memorySegment, new long[] {size, columns}, code(), bits(), kDLCPU(), 1); + } + private static class SliceRowView implements RowView { private final MemorySegment memorySegment; private final long size; diff --git a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CuVSMatrixBaseImpl.java b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CuVSMatrixBaseImpl.java index fc9daf9387..120960cfc4 100644 --- a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CuVSMatrixBaseImpl.java +++ b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CuVSMatrixBaseImpl.java @@ -15,7 +15,10 @@ */ package com.nvidia.cuvs.internal; +import static com.nvidia.cuvs.internal.panama.headers_h.*; + import com.nvidia.cuvs.CuVSMatrix; +import java.lang.foreign.Arena; import java.lang.foreign.MemorySegment; public abstract class CuVSMatrixBaseImpl implements CuVSMatrix { @@ -45,4 +48,21 @@ public long columns() { public MemorySegment memorySegment() { return memorySegment; } + + protected int bits() { + return switch (dataType) { + case FLOAT, INT, UINT -> 32; + case BYTE -> 8; + }; + } + + protected int code() { + return switch (dataType) { + case FLOAT -> kDLFloat(); + case INT -> kDLInt(); + case UINT, BYTE -> kDLUInt(); + }; + } + + public abstract MemorySegment toTensor(Arena arena); } diff --git a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/common/Util.java b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/common/Util.java index d5972ecca7..8125497078 100644 --- a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/common/Util.java +++ b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/common/Util.java @@ -19,11 +19,7 @@ import static com.nvidia.cuvs.internal.common.LinkerHelper.C_FLOAT; import static com.nvidia.cuvs.internal.common.LinkerHelper.C_INT; import static com.nvidia.cuvs.internal.common.LinkerHelper.C_LONG; -import static com.nvidia.cuvs.internal.panama.headers_h.cudaGetDeviceCount; -import static com.nvidia.cuvs.internal.panama.headers_h.cudaGetDeviceProperties_v2; -import static com.nvidia.cuvs.internal.panama.headers_h.cudaMemGetInfo; -import static com.nvidia.cuvs.internal.panama.headers_h.cudaSetDevice; -import static com.nvidia.cuvs.internal.panama.headers_h.size_t; +import static com.nvidia.cuvs.internal.panama.headers_h.*; import com.nvidia.cuvs.GPUInfo; import com.nvidia.cuvs.internal.panama.DLDataType; @@ -73,6 +69,13 @@ public static void checkCudaError(int value, String caller) { } } + private static final long UNSIGNED_INT_MASK = 0xFFFFFFFFL; + + public static long dereferenceUnsignedInt(MemorySegment ptr) { + assert ptr.byteSize() == 4; + return ptr.get(uint32_t, 0) & UNSIGNED_INT_MASK; + } + /** * Java analog to CUDA's cudaMemcpyKind, used for cudaMemcpy() calls. * @see CUDA Runtime API 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 cc7879a43f..8223e199f6 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 @@ -290,6 +290,48 @@ private Path createSerializedIndex(CuVSMatrix dataset) throws Throwable { } } + @Test + public void testReconstructIndexFromGraph() throws Throwable { + try (var dataset = CuVSMatrix.ofArray(createSampleData())) { + var queries = createSampleQueries(); + List> expectedResults = getExpectedResults(); + + try (CuVSResources resources = CuVSResources.create()) { + var index = indexOnce(dataset, resources); + var graph = index.getGraph(); + + var reconstructedIndex = + CagraIndex.newBuilder(resources) + .from(graph) + .withDataset(dataset) + .withIndexParams( + new CagraIndexParams.Builder().withMetric(CuvsDistanceType.L2Expanded).build()) + .build(); + queryAndCompare( + index, + reconstructedIndex, + SearchResults.IDENTITY_MAPPING, + queries, + expectedResults, + resources); + + var originalIndexPath = serializeOnce(index); + var reconstructedIndexPath = serializeOnce(reconstructedIndex); + + var originalBytes = Files.readAllBytes(originalIndexPath); + var reconstructedBytes = Files.readAllBytes(reconstructedIndexPath); + + assertArrayEquals(originalBytes, reconstructedBytes); + + index.destroyIndex(); + reconstructedIndex.destroyIndex(); + + Files.deleteIfExists(originalIndexPath); + Files.deleteIfExists(reconstructedIndexPath); + } + } + } + @Test public void testIndexingAndSearchingFlowWithCustomMappingFunction() throws Throwable { var dataset = CuVSMatrix.ofArray(createSampleData()); From 690610ce9b023a491e0c69a4b0549a65bc8daec7 Mon Sep 17 00:00:00 2001 From: Tarang Jain <40517122+tarang-jain@users.noreply.github.com> Date: Wed, 13 Aug 2025 08:10:02 -0700 Subject: [PATCH 29/93] Pin Latest Faiss Version (#1247) Authors: - Tarang Jain (https://github.com/tarang-jain) Approvers: - Corey J. Nolet (https://github.com/cjnolet) URL: https://github.com/rapidsai/cuvs/pull/1247 --- cpp/cmake/patches/faiss.diff | 19 ------------------- cpp/cmake/patches/faiss_override.json | 11 ++--------- 2 files changed, 2 insertions(+), 28 deletions(-) delete mode 100644 cpp/cmake/patches/faiss.diff diff --git a/cpp/cmake/patches/faiss.diff b/cpp/cmake/patches/faiss.diff deleted file mode 100644 index bf92f89d77..0000000000 --- a/cpp/cmake/patches/faiss.diff +++ /dev/null @@ -1,19 +0,0 @@ -diff --git a/faiss/gpu/impl/CuvsIVFFlat.cu b/faiss/gpu/impl/CuvsIVFFlat.cu -index 2cccee860..c4cb13f48 100644 ---- a/faiss/gpu/impl/CuvsIVFFlat.cu -+++ b/faiss/gpu/impl/CuvsIVFFlat.cu -@@ -427,13 +427,11 @@ void CuvsIVFFlat::copyInvertedListsFrom(const InvertedLists* ivf) { - // Precompute the centers vector norms for L2Expanded distance - if (this->metric_ == faiss::METRIC_L2) { - cuvs_index->allocate_center_norms(raft_handle); -- raft::linalg::rowNorm( -+ raft::linalg::rowNorm( - cuvs_index->center_norms().value().data_handle(), - cuvs_index->centers().data_handle(), - cuvs_index->dim(), - (uint32_t)nlist, -- raft::linalg::L2Norm, -- true, - raft_handle.get_stream()); - } - } diff --git a/cpp/cmake/patches/faiss_override.json b/cpp/cmake/patches/faiss_override.json index 0252cf3103..f3a1aaf795 100644 --- a/cpp/cmake/patches/faiss_override.json +++ b/cpp/cmake/patches/faiss_override.json @@ -1,16 +1,9 @@ { "packages" : { "faiss" : { - "version": "1.11.0", + "version": "1.12.0", "git_url": "https://github.com/facebookresearch/faiss.git", - "git_tag": "v1.11.0", - "patches" : [ - { - "file" : "${current_json_dir}/faiss.diff", - "issue" : "Apply RAFT breaking changes", - "fixed_in" : "" - } - ] + "git_tag": "v1.12.0" } } } From 378440e36b023812daf7e741034d67db72f45299 Mon Sep 17 00:00:00 2001 From: Tarang Jain <40517122+tarang-jain@users.noreply.github.com> Date: Wed, 13 Aug 2025 14:35:27 -0700 Subject: [PATCH 30/93] [Improvement] Replace Calls to thrust::transform with raft::linalg (#1238) Authors: - Tarang Jain (https://github.com/tarang-jain) Approvers: - Jinsol Park (https://github.com/jinsolp) - Micka (https://github.com/lowener) URL: https://github.com/rapidsai/cuvs/pull/1238 --- cpp/src/cluster/detail/connectivities.cuh | 62 ++++++++--------- cpp/src/neighbors/ball_cover/ball_cover.cuh | 48 +++++--------- cpp/src/neighbors/ball_cover/registers.cuh | 16 +++-- cpp/src/neighbors/detail/reachability.cuh | 66 +++++++++---------- cpp/tests/cluster/connect_knn.cu | 15 ++--- cpp/tests/neighbors/ann_ivf_flat.cuh | 14 ++-- cpp/tests/neighbors/ann_nn_descent.cuh | 8 +-- cpp/tests/neighbors/ball_cover.cu | 24 ++----- .../neighbors/brute_force_prefiltered.cu | 53 +++++---------- examples/cpp/src/ivf_flat_example.cu | 11 +--- 10 files changed, 131 insertions(+), 186 deletions(-) diff --git a/cpp/src/cluster/detail/connectivities.cuh b/cpp/src/cluster/detail/connectivities.cuh index 39be707bb7..1847da325b 100644 --- a/cpp/src/cluster/detail/connectivities.cuh +++ b/cpp/src/cluster/detail/connectivities.cuh @@ -1,5 +1,5 @@ /* - * Copyright (c) 2021-2024, NVIDIA CORPORATION. + * Copyright (c) 2021-2025, NVIDIA CORPORATION. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -23,6 +23,7 @@ #include #include #include +#include #include #include #include @@ -31,10 +32,6 @@ #include -#include -#include -#include - #include namespace cuvs::cluster::agglomerative::detail { @@ -83,18 +80,25 @@ struct distance_graph_impl { data.resize(knn_graph_coo.nnz, stream); // self-loops get max distance - auto transform_in = thrust::make_zip_iterator( - thrust::make_tuple(knn_graph_coo.rows(), knn_graph_coo.cols(), knn_graph_coo.vals())); - - thrust::transform(thrust_policy, - transform_in, - transform_in + knn_graph_coo.nnz, - knn_graph_coo.vals(), - [=] __device__(const thrust::tuple& tup) { - bool self_loop = thrust::get<0>(tup) == thrust::get<1>(tup); - return (self_loop * std::numeric_limits::max()) + - (!self_loop * thrust::get<2>(tup)); - }); + auto rows_view = raft::make_device_vector_view(knn_graph_coo.rows(), + knn_graph_coo.nnz); + auto cols_view = raft::make_device_vector_view(knn_graph_coo.cols(), + knn_graph_coo.nnz); + auto vals_in_view = raft::make_device_vector_view( + knn_graph_coo.vals(), knn_graph_coo.nnz); + auto vals_out_view = + raft::make_device_vector_view(knn_graph_coo.vals(), knn_graph_coo.nnz); + + raft::linalg::map( + handle, + vals_out_view, + [=] __device__(const value_idx row, const value_idx col, const value_t val) { + bool self_loop = row == col; + return (self_loop * std::numeric_limits::max()) + (!self_loop * val); + }, + rows_view, + cols_view, + vals_in_view); raft::sparse::convert::sorted_coo_to_csr( knn_graph_coo.rows(), knn_graph_coo.nnz, indptr.data(), m + 1, stream); @@ -147,7 +151,9 @@ void pairwise_distances(const raft::resources& handle, value_idx blocks = raft::ceildiv(nnz, (value_idx)256); fill_indices2<<>>(indices, m, nnz); - thrust::sequence(exec_policy, indptr, indptr + m, 0, (int)m); + raft::linalg::map_offset(handle, + raft::make_device_vector_view(indptr, m), + [=] __device__(value_idx idx) { return idx * m; }); raft::update_device(indptr + m, &nnz, 1, stream); @@ -160,19 +166,13 @@ void pairwise_distances(const raft::resources& handle, handle, X_view, X_view, raft::make_device_matrix_view(data, m, m), metric); // self-loops get max distance - auto transform_in = - thrust::make_zip_iterator(thrust::make_tuple(thrust::make_counting_iterator(0), data)); - - thrust::transform(exec_policy, - transform_in, - transform_in + nnz, - data, - [=] __device__(const thrust::tuple& tup) { - value_idx idx = thrust::get<0>(tup); - bool self_loop = idx % m == idx / m; - return (self_loop * std::numeric_limits::max()) + - (!self_loop * thrust::get<1>(tup)); - }); + auto data_view = raft::make_device_vector_view(data, nnz); + + raft::linalg::map_offset(handle, data_view, [=] __device__(value_idx idx) { + value_t val = data[idx]; + bool self_loop = idx % m == idx / m; + return (self_loop * std::numeric_limits::max()) + (!self_loop * val); + }); } /** diff --git a/cpp/src/neighbors/ball_cover/ball_cover.cuh b/cpp/src/neighbors/ball_cover/ball_cover.cuh index b0c44c3526..99c1de882e 100644 --- a/cpp/src/neighbors/ball_cover/ball_cover.cuh +++ b/cpp/src/neighbors/ball_cover/ball_cover.cuh @@ -25,6 +25,7 @@ #include #include #include +#include #include #include #include @@ -34,12 +35,9 @@ #include #include -#include #include -#include #include #include -#include #include #include @@ -68,10 +66,7 @@ void sample_landmarks(raft::resources const& handle, rmm::device_uvector R_indices(index.n_landmarks, raft::resource::get_cuda_stream(handle)); - thrust::sequence(raft::resource::get_thrust_policy(handle), - index.get_R_1nn_cols().data_handle(), - index.get_R_1nn_cols().data_handle() + index.m, - (value_idx)0); + raft::linalg::map_offset(handle, index.get_R_1nn_cols(), raft::identity_op{}); thrust::fill(raft::resource::get_thrust_policy(handle), R_1nn_ones.data(), @@ -124,25 +119,23 @@ void construct_landmark_1nn(raft::resources const& handle, int64_t k, cuvs::neighbors::ball_cover::index& index) { - rmm::device_uvector R_1nn_inds(index.m, raft::resource::get_cuda_stream(handle)); + auto R_1nn_inds = raft::make_device_vector(handle, index.m); thrust::fill(raft::resource::get_thrust_policy(handle), - R_1nn_inds.data(), - R_1nn_inds.data() + index.m, + R_1nn_inds.data_handle(), + R_1nn_inds.data_handle() + index.m, std::numeric_limits::max()); - value_idx* R_1nn_inds_ptr = R_1nn_inds.data(); - value_t* R_1nn_dists_ptr = index.get_R_1nn_dists().data_handle(); - - auto idxs = thrust::make_counting_iterator(0); - thrust::for_each( - raft::resource::get_thrust_policy(handle), idxs, idxs + index.m, [=] __device__(value_idx i) { - R_1nn_inds_ptr[i] = R_knn_inds_ptr[i * k]; - R_1nn_dists_ptr[i] = R_knn_dists_ptr[i * k]; + raft::linalg::map_offset(handle, R_1nn_inds.view(), [R_knn_inds_ptr, k] __device__(value_idx i) { + return R_knn_inds_ptr[i * k]; + }); + raft::linalg::map_offset( + handle, index.get_R_1nn_dists(), [R_knn_dists_ptr, k] __device__(value_idx i) { + return R_knn_dists_ptr[i * k]; }); auto keys = thrust::make_zip_iterator( - thrust::make_tuple(R_1nn_inds.data(), index.get_R_1nn_dists().data_handle())); + thrust::make_tuple(R_1nn_inds.data_handle(), index.get_R_1nn_dists().data_handle())); // group neighborhoods for each reference landmark and sort each group by distance thrust::sort_by_key(raft::resource::get_thrust_policy(handle), @@ -152,7 +145,7 @@ void construct_landmark_1nn(raft::resources const& handle, NNComp()); // convert to CSR for fast lookup - raft::sparse::convert::sorted_coo_to_csr(R_1nn_inds.data(), + raft::sparse::convert::sorted_coo_to_csr(R_1nn_inds.data_handle(), index.m, index.get_R_indptr().data_handle(), index.n_landmarks + 1, @@ -212,18 +205,13 @@ template void compute_landmark_radii(raft::resources const& handle, cuvs::neighbors::ball_cover::index& index) { - auto entries = thrust::make_counting_iterator(0); - const value_idx* R_indptr_ptr = index.get_R_indptr().data_handle(); const value_t* R_1nn_dists_ptr = index.get_R_1nn_dists().data_handle(); - value_t* R_radius_ptr = index.get_R_radius().data_handle(); - thrust::for_each(raft::resource::get_thrust_policy(handle), - entries, - entries + index.n_landmarks, - [=] __device__(value_idx input) { - value_idx last_row_idx = R_indptr_ptr[input + 1] - 1; - R_radius_ptr[input] = R_1nn_dists_ptr[last_row_idx]; - }); + raft::linalg::map_offset( + handle, index.get_R_radius(), [R_indptr_ptr, R_1nn_dists_ptr] __device__(value_idx input) { + value_idx last_row_idx = R_indptr_ptr[input + 1] - 1; + return R_1nn_dists_ptr[last_row_idx]; + }); } /** diff --git a/cpp/src/neighbors/ball_cover/registers.cuh b/cpp/src/neighbors/ball_cover/registers.cuh index 9931275f33..191dfc1fb5 100644 --- a/cpp/src/neighbors/ball_cover/registers.cuh +++ b/cpp/src/neighbors/ball_cover/registers.cuh @@ -24,6 +24,7 @@ #include #include +#include #include #include @@ -1458,13 +1459,14 @@ void rbc_eps_pass(raft::resources const& handle, if (actual_max > max_k_in) { // ceil vd to max_k - thrust::transform(raft::resource::get_thrust_policy(handle), - vd_ptr, - vd_ptr + n_query_rows, - vd_ptr, - [max_k_in] __device__(value_idx vd_count) { - return vd_count > max_k_in ? max_k_in : vd_count; - }); + raft::linalg::unaryOp( + vd_ptr, + vd_ptr, + n_query_rows, + [max_k_in] __device__(value_idx vd_count) { + return vd_count > max_k_in ? max_k_in : vd_count; + }, + raft::resource::get_cuda_stream(handle)); } thrust::exclusive_scan(raft::resource::get_thrust_policy(handle), diff --git a/cpp/src/neighbors/detail/reachability.cuh b/cpp/src/neighbors/detail/reachability.cuh index 0f73e014e5..4a73179b6a 100644 --- a/cpp/src/neighbors/detail/reachability.cuh +++ b/cpp/src/neighbors/detail/reachability.cuh @@ -17,6 +17,9 @@ #pragma once #include "./knn_brute_force.cuh" +#include +#include +#include #include #include #include @@ -24,12 +27,6 @@ #include #include -#include - -#include -#include -#include -#include namespace cuvs::neighbors::detail::reachability { @@ -47,17 +44,19 @@ namespace cuvs::neighbors::detail::reachability { * @param[in] stream stream for which to order cuda operations */ template -void core_distances( - value_t* knn_dists, int min_samples, int n_neighbors, size_t n, value_t* out, cudaStream_t stream) +void core_distances(raft::resources const& handle, + value_t* knn_dists, + int min_samples, + int n_neighbors, + size_t n, + value_t* out) { ASSERT(n_neighbors >= min_samples, "the size of the neighborhood should be greater than or equal to min_samples"); - auto exec_policy = rmm::exec_policy(stream); - - auto indices = thrust::make_counting_iterator(0); + auto out_view = raft::make_device_vector_view(out, n); - thrust::transform(exec_policy, indices, indices + n, out, [=] __device__(value_idx row) { + raft::linalg::map_offset(handle, out_view, [=] __device__(value_idx row) { return knn_dists[row * n_neighbors + (min_samples - 1)]; }); } @@ -118,7 +117,7 @@ void _compute_core_dists(const raft::resources& handle, compute_knn(handle, X, inds.data(), dists.data(), m, n, X, m, min_samples, metric); // Slice core distances (distances to kth nearest neighbor) - core_distances(dists.data(), min_samples, min_samples, m, core_dists, stream); + core_distances(handle, dists.data(), min_samples, min_samples, m, core_dists); } // Functor to post-process distances into reachability space @@ -202,8 +201,7 @@ void mutual_reachability_graph(const raft::resources& handle, RAFT_EXPECTS(metric == cuvs::distance::DistanceType::L2SqrtExpanded, "Currently only L2 expanded distance is supported"); - auto stream = raft::resource::get_cuda_stream(handle); - auto exec_policy = raft::resource::get_thrust_policy(handle); + auto stream = raft::resource::get_cuda_stream(handle); rmm::device_uvector coo_rows(min_samples * m, stream); rmm::device_uvector inds(min_samples * m, stream); @@ -213,7 +211,7 @@ void mutual_reachability_graph(const raft::resources& handle, compute_knn(handle, X, inds.data(), dists.data(), m, n, X, m, min_samples, metric); // Slice core distances (distances to kth nearest neighbor) - core_distances(dists.data(), min_samples, min_samples, m, core_dists, stream); + core_distances(handle, dists.data(), min_samples, min_samples, m, core_dists); /** * Compute L2 norm @@ -222,12 +220,12 @@ void mutual_reachability_graph(const raft::resources& handle, handle, inds.data(), dists.data(), X, m, n, min_samples, core_dists, (value_t)1.0 / alpha); // self-loops get max distance - auto coo_rows_counting_itr = thrust::make_counting_iterator(0); - thrust::transform(exec_policy, - coo_rows_counting_itr, - coo_rows_counting_itr + (m * min_samples), - coo_rows.data(), - [min_samples] __device__(value_idx c) -> value_idx { return c / min_samples; }); + auto coo_rows_view = + raft::make_device_vector_view(coo_rows.data(), m * min_samples); + raft::linalg::map_offset( + handle, coo_rows_view, [min_samples] __device__(value_idx c) -> value_idx { + return c / min_samples; + }); raft::sparse::linalg::symmetrize(handle, coo_rows.data(), @@ -241,18 +239,20 @@ void mutual_reachability_graph(const raft::resources& handle, raft::sparse::convert::sorted_coo_to_csr(out.rows(), out.nnz, indptr, m + 1, stream); // self-loops get max distance - auto transform_in = - thrust::make_zip_iterator(thrust::make_tuple(out.rows(), out.cols(), out.vals())); + auto rows_view = raft::make_device_vector_view(out.rows(), out.nnz); + auto cols_view = raft::make_device_vector_view(out.cols(), out.nnz); + auto vals_in_view = raft::make_device_vector_view(out.vals(), out.nnz); + auto vals_out_view = raft::make_device_vector_view(out.vals(), out.nnz); - thrust::transform(exec_policy, - transform_in, - transform_in + out.nnz, - out.vals(), - [=] __device__(const thrust::tuple& tup) { - return thrust::get<0>(tup) == thrust::get<1>(tup) - ? std::numeric_limits::max() - : thrust::get<2>(tup); - }); + raft::linalg::map( + handle, + vals_out_view, + [=] __device__(const value_idx row, const value_idx col, const value_t val) { + return row == col ? std::numeric_limits::max() : val; + }, + rows_view, + cols_view, + vals_in_view); } } // namespace cuvs::neighbors::detail::reachability diff --git a/cpp/tests/cluster/connect_knn.cu b/cpp/tests/cluster/connect_knn.cu index b4fd55e5b6..ecca492489 100644 --- a/cpp/tests/cluster/connect_knn.cu +++ b/cpp/tests/cluster/connect_knn.cu @@ -16,6 +16,7 @@ #include #include +#include #include #include #include @@ -78,7 +79,7 @@ class ConnectKNNTest : public ::testing::TestWithParam { rmm::device_uvector core_dists(ps.n_rows, stream); if (ps.mutual_reach) { cuvs::neighbors::detail::reachability::core_distances( - dists.data(), ps.k, ps.k, (size_t)ps.n_rows, core_dists.data(), stream); + handle, dists.data(), ps.k, ps.k, (size_t)ps.n_rows, core_dists.data()); auto epilogue = cuvs::neighbors::detail::reachability::ReachabilityPostProcess{ core_dists.data(), 1.0}; @@ -111,13 +112,11 @@ class ConnectKNNTest : public ::testing::TestWithParam { rmm::device_uvector indptr(ps.n_rows + 1, stream); // changing inds and dists to sparse format - int64_t k = ps.k; - auto coo_rows_counting_itr = thrust::make_counting_iterator(0); - thrust::transform(raft::resource::get_thrust_policy(handle), - coo_rows_counting_itr, - coo_rows_counting_itr + (ps.n_rows * ps.k), - coo_rows.data(), - [k] __device__(int64_t c) -> int64_t { return c / k; }); + int64_t k = ps.k; + auto coo_rows_view = + raft::make_device_vector_view(coo_rows.data(), ps.n_rows * ps.k); + raft::linalg::map_offset( + handle, coo_rows_view, [k] __device__(int64_t c) -> int64_t { return c / k; }); raft::sparse::linalg::symmetrize(handle, coo_rows.data(), diff --git a/cpp/tests/neighbors/ann_ivf_flat.cuh b/cpp/tests/neighbors/ann_ivf_flat.cuh index aab1868473..bb8c7d1cd7 100644 --- a/cpp/tests/neighbors/ann_ivf_flat.cuh +++ b/cpp/tests/neighbors/ann_ivf_flat.cuh @@ -25,7 +25,6 @@ #include #include #include -#include #include #include @@ -135,10 +134,8 @@ class AnnIVFFlatTest : public ::testing::TestWithParam> { auto database_view = raft::make_device_matrix_view( (const DataT*)database.data(), ps.num_db_vecs, ps.dim); idx = cuvs::neighbors::ivf_flat::build(handle_, index_params, database_view); - rmm::device_uvector vector_indices(ps.num_db_vecs, stream_); - thrust::sequence(raft::resource::get_thrust_policy(handle_), - thrust::device_pointer_cast(vector_indices.data()), - thrust::device_pointer_cast(vector_indices.data() + ps.num_db_vecs)); + auto vector_indices = raft::make_device_vector(handle_, ps.num_db_vecs); + raft::linalg::map_offset(handle_, vector_indices.view(), raft::identity_op{}); raft::resource::sync_stream(handle_); IdxT half_of_data = ps.num_db_vecs / 2; @@ -153,7 +150,7 @@ class AnnIVFFlatTest : public ::testing::TestWithParam> { database.data() + half_of_data * ps.dim, IdxT(ps.num_db_vecs) - half_of_data, ps.dim); auto new_half_of_data_indices_view = raft::make_device_vector_view( - vector_indices.data() + half_of_data, IdxT(ps.num_db_vecs) - half_of_data); + vector_indices.data_handle() + half_of_data, IdxT(ps.num_db_vecs) - half_of_data); cuvs::neighbors::ivf_flat::extend( handle_, @@ -461,10 +458,7 @@ class AnnIVFFlatTest : public ::testing::TestWithParam> { // Create Bitset filter auto removed_indices = raft::make_device_vector(handle_, test_ivf_sample_filter::offset); - thrust::sequence(raft::resource::get_thrust_policy(handle_), - thrust::device_pointer_cast(removed_indices.data_handle()), - thrust::device_pointer_cast(removed_indices.data_handle() + - test_ivf_sample_filter::offset)); + raft::linalg::map_offset(handle_, removed_indices.view(), raft::identity_op{}); raft::resource::sync_stream(handle_); cuvs::core::bitset removed_indices_bitset( diff --git a/cpp/tests/neighbors/ann_nn_descent.cuh b/cpp/tests/neighbors/ann_nn_descent.cuh index 6edae3d6d5..e1c27dc48c 100644 --- a/cpp/tests/neighbors/ann_nn_descent.cuh +++ b/cpp/tests/neighbors/ann_nn_descent.cuh @@ -236,12 +236,12 @@ class AnnNNDescentDistEpiTest : public ::testing::TestWithParam core_dists_dev(ps.n_rows, stream_); cuvs::neighbors::detail::reachability::core_distances( + handle_, distances_naive_dev.data(), ps.graph_degree, ps.graph_degree, ps.n_rows, - core_dists_dev.data(), - stream_); + core_dists_dev.data()); auto epilogue = cuvs::neighbors::detail::reachability::ReachabilityPostProcess{ core_dists_dev.data(), 1.0, static_cast(ps.n_rows)}; @@ -289,12 +289,12 @@ class AnnNNDescentDistEpiTest : public ::testing::TestWithParam core_dists_dev(ps.n_rows, stream_); cuvs::neighbors::detail::reachability::core_distances( + handle_, index.distances().value().data_handle(), ps.graph_degree, ps.graph_degree, ps.n_rows, - core_dists_dev.data(), - stream_); + core_dists_dev.data()); size_t extended_graph_degree, graph_degree; auto build_config = nn_descent::detail::get_build_config( diff --git a/cpp/tests/neighbors/ball_cover.cu b/cpp/tests/neighbors/ball_cover.cu index 429c939d71..5e7fb51bac 100644 --- a/cpp/tests/neighbors/ball_cover.cu +++ b/cpp/tests/neighbors/ball_cover.cu @@ -23,15 +23,14 @@ #include #include #include +#include #include #include #include #include -#include #include -#include #include @@ -193,16 +192,10 @@ class BallCoverKNNQueryTest : public ::testing::TestWithParam #include +#include #include #include #include @@ -32,10 +33,6 @@ #include #include -#include -#include -#include -#include #include #include @@ -431,28 +428,19 @@ class PrefilteredBruteForceOnBitmapTest raft::copy(dataset_h.data(), blobs_in_val.data_handle(), dataset_size, stream); if constexpr (std::is_same_v) { - thrust::device_ptr d_output_ptr = - thrust::device_pointer_cast(blobs_in_val.data_handle()); - thrust::device_ptr d_value_ptr = thrust::device_pointer_cast(dataset_d.data()); - thrust::transform(thrust::cuda::par.on(stream), - d_output_ptr, - d_output_ptr + dataset_size, - d_value_ptr, - float_to_half()); + raft::linalg::unaryOp( + dataset_d.data(), blobs_in_val.data_handle(), dataset_size, float_to_half(), stream); } else { raft::copy(dataset_d.data(), blobs_in_val.data_handle(), dataset_size, stream); } raft::copy(queries_h.data(), blobs_in_val.data_handle() + dataset_size, queries_size, stream); if constexpr (std::is_same_v) { - thrust::device_ptr d_output_ptr = - thrust::device_pointer_cast(blobs_in_val.data_handle() + dataset_size); - thrust::device_ptr d_value_ptr = thrust::device_pointer_cast(queries_d.data()); - thrust::transform(thrust::cuda::par.on(stream), - d_output_ptr, - d_output_ptr + queries_size, - d_value_ptr, - float_to_half()); + raft::linalg::unaryOp(queries_d.data(), + blobs_in_val.data_handle() + dataset_size, + queries_size, + float_to_half(), + stream); } else { raft::copy(queries_d.data(), blobs_in_val.data_handle() + dataset_size, queries_size, stream); } @@ -861,28 +849,19 @@ class PrefilteredBruteForceOnBitsetTest raft::copy(dataset_h.data(), blobs_in_val.data_handle(), dataset_size, stream); if constexpr (std::is_same_v) { - thrust::device_ptr d_output_ptr = - thrust::device_pointer_cast(blobs_in_val.data_handle()); - thrust::device_ptr d_value_ptr = thrust::device_pointer_cast(dataset_d.data()); - thrust::transform(thrust::cuda::par.on(stream), - d_output_ptr, - d_output_ptr + dataset_size, - d_value_ptr, - float_to_half()); + raft::linalg::unaryOp( + dataset_d.data(), blobs_in_val.data_handle(), dataset_size, float_to_half(), stream); } else { raft::copy(dataset_d.data(), blobs_in_val.data_handle(), dataset_size, stream); } raft::copy(queries_h.data(), blobs_in_val.data_handle() + dataset_size, queries_size, stream); if constexpr (std::is_same_v) { - thrust::device_ptr d_output_ptr = - thrust::device_pointer_cast(blobs_in_val.data_handle() + dataset_size); - thrust::device_ptr d_value_ptr = thrust::device_pointer_cast(queries_d.data()); - thrust::transform(thrust::cuda::par.on(stream), - d_output_ptr, - d_output_ptr + queries_size, - d_value_ptr, - float_to_half()); + raft::linalg::unaryOp(queries_d.data(), + blobs_in_val.data_handle() + dataset_size, + queries_size, + float_to_half(), + stream); } else { raft::copy(queries_d.data(), blobs_in_val.data_handle() + dataset_size, queries_size, stream); } diff --git a/examples/cpp/src/ivf_flat_example.cu b/examples/cpp/src/ivf_flat_example.cu index 4744d3e42d..8f1d8c44f6 100644 --- a/examples/cpp/src/ivf_flat_example.cu +++ b/examples/cpp/src/ivf_flat_example.cu @@ -19,16 +19,12 @@ #include #include #include -#include +#include #include #include #include -#include -#include -#include - #include #include @@ -77,10 +73,7 @@ void ivf_flat_build_extend_search(raft::device_resources const& dev_resources, // Define dataset indices. auto data_indices = raft::make_device_vector(dev_resources, dataset.extent(0)); - thrust::counting_iterator first(0); - thrust::device_ptr ptr(data_indices.data_handle()); - thrust::copy( - raft::resource::get_thrust_policy(dev_resources), first, first + dataset.extent(0), ptr); + raft::linalg::map_offset(dev_resources, data_indices.view(), raft::identity_op{}); // Sub-sample the dataset to create a training set. auto trainset = From e3a61265c9eeeb64a883d9f3d13a21809b552d08 Mon Sep 17 00:00:00 2001 From: Jinsol Park Date: Wed, 13 Aug 2025 17:03:24 -0700 Subject: [PATCH 31/93] Improve NN Descent documentation (#1246) Closing https://github.com/rapidsai/cuvs/issues/1245 Authors: - Jinsol Park (https://github.com/jinsolp) Approvers: - Micka (https://github.com/lowener) URL: https://github.com/rapidsai/cuvs/pull/1246 --- cpp/include/cuvs/neighbors/nn_descent.hpp | 70 +++++++++++++------ cpp/src/neighbors/nn_descent.cuh | 38 +++++++--- .../cuvs/neighbors/nn_descent/nn_descent.pyx | 4 +- 3 files changed, 79 insertions(+), 33 deletions(-) diff --git a/cpp/include/cuvs/neighbors/nn_descent.hpp b/cpp/include/cuvs/neighbors/nn_descent.hpp index 94300e2737..5f6260b0af 100644 --- a/cpp/include/cuvs/neighbors/nn_descent.hpp +++ b/cpp/include/cuvs/neighbors/nn_descent.hpp @@ -1,5 +1,5 @@ /* - * Copyright (c) 2023-2024, NVIDIA CORPORATION. + * Copyright (c) 2023-2025, NVIDIA CORPORATION. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -37,26 +37,25 @@ namespace cuvs::neighbors::nn_descent { /** * @brief Parameters used to build an nn-descent index - * - * `graph_degree`: For an input dataset of dimensions (N, D), + * - `graph_degree`: For an input dataset of dimensions (N, D), * determines the final dimensions of the all-neighbors knn graph * which turns out to be of dimensions (N, graph_degree) - * `intermediate_graph_degree`: Internally, nn-descent builds an + * - `intermediate_graph_degree`: Internally, nn-descent builds an * all-neighbors knn graph of dimensions (N, intermediate_graph_degree) * before selecting the final `graph_degree` neighbors. It's recommended * that `intermediate_graph_degree` >= 1.5 * graph_degree - * `max_iterations`: The number of iterations that nn-descent will refine + * - `max_iterations`: The number of iterations that nn-descent will refine * the graph for. More iterations produce a better quality graph at cost of performance - * `termination_threshold`: The delta at which nn-descent will terminate its iterations - * + * - `termination_threshold`: The delta at which nn-descent will terminate its iterations + * - `return_distances`: Boolean to decide whether to return distances array */ struct index_params : cuvs::neighbors::index_params { - size_t graph_degree = 64; // Degree of output graph. - size_t intermediate_graph_degree = 128; // Degree of input graph for pruning. - size_t max_iterations = 20; // Number of nn-descent iterations. - float termination_threshold = 0.0001; // Termination threshold of nn-descent. - bool return_distances = true; // return distances if true - size_t n_clusters = 1; // defaults to not using any batching + size_t graph_degree = 64; + size_t intermediate_graph_degree = 128; + size_t max_iterations = 20; + float termination_threshold = 0.0001; + bool return_distances = true; + size_t n_clusters = 1; /** @brief Construct NN descent parameters for a specific kNN graph degree * @@ -211,7 +210,10 @@ struct index : cuvs::neighbors::index { * @brief Build nn-descent Index with dataset in device memory * * The following distance metrics are supported: - * - L2 + * - L2Expanded + * - L2SqrtExpanded + * - CosineExpanded + * - InnerProduct * * Usage example: * @code{.cpp} @@ -244,7 +246,10 @@ auto build(raft::resources const& res, * @brief Build nn-descent Index with dataset in host memory * * The following distance metrics are supported: - * - L2 + * - L2Expanded + * - L2SqrtExpanded + * - CosineExpanded + * - InnerProduct * * Usage example: * @code{.cpp} @@ -279,7 +284,10 @@ auto build(raft::resources const& res, * @brief Build nn-descent Index with dataset in device memory * * The following distance metrics are supported: - * - L2 + * - L2Expanded + * - L2SqrtExpanded + * - CosineExpanded + * - InnerProduct * * Usage example: * @code{.cpp} @@ -312,7 +320,10 @@ auto build(raft::resources const& res, * @brief Build nn-descent Index with dataset in host memory * * The following distance metrics are supported: - * - L2 + * - L2Expanded + * - L2SqrtExpanded + * - CosineExpanded + * - InnerProduct * * Usage example: * @code{.cpp} @@ -347,7 +358,11 @@ auto build(raft::resources const& res, * @brief Build nn-descent Index with dataset in device memory * * The following distance metrics are supported: - * - L2 + * - L2Expanded + * - L2SqrtExpanded + * - CosineExpanded + * - InnerProduct + * - BitwiseHamming * * Usage example: * @code{.cpp} @@ -380,7 +395,11 @@ auto build(raft::resources const& res, * @brief Build nn-descent Index with dataset in host memory * * The following distance metrics are supported: - * - L2 + * - L2Expanded + * - L2SqrtExpanded + * - CosineExpanded + * - InnerProduct + * - BitwiseHamming * * Usage example: * @code{.cpp} @@ -415,7 +434,11 @@ auto build(raft::resources const& res, * @brief Build nn-descent Index with dataset in device memory * * The following distance metrics are supported: - * - L2 + * - L2Expanded + * - L2SqrtExpanded + * - CosineExpanded + * - InnerProduct + * - BitwiseHamming * * Usage example: * @code{.cpp} @@ -448,7 +471,11 @@ auto build(raft::resources const& res, * @brief Build nn-descent Index with dataset in host memory * * The following distance metrics are supported: - * - L2 + * - L2Expanded + * - L2SqrtExpanded + * - CosineExpanded + * - InnerProduct + * - BitwiseHamming * * Usage example: * @code{.cpp} @@ -479,6 +506,7 @@ auto build(raft::resources const& res, std::optional> graph = std::nullopt) -> cuvs::neighbors::nn_descent::index; +/** @} */ /** * @brief Test if we have enough GPU memory to run NN descent algorithm. * diff --git a/cpp/src/neighbors/nn_descent.cuh b/cpp/src/neighbors/nn_descent.cuh index c4da01c561..78c63dc31e 100644 --- a/cpp/src/neighbors/nn_descent.cuh +++ b/cpp/src/neighbors/nn_descent.cuh @@ -38,15 +38,19 @@ namespace cuvs::neighbors::nn_descent { * @brief Build nn-descent Index with dataset in device memory * * The following distance metrics are supported: - * - L2 + * - L2Expanded + * - L2SqrtExpanded + * - CosineExpanded + * - InnerProduct + * - BitwiseHamming (when T == int8, uint8) * * Usage example: * @code{.cpp} - * using namespace cuvs::neighbors::experimental; + * using namespace cuvs::neighbors; * // use default index parameters * nn_descent::index_params index_params; * // create and fill the index from a [N, D] raft::device_matrix_view dataset - * auto index = cagra::build(res, index_params, dataset); + * auto index = nn_descent::build(res, index_params, dataset); * // index.graph() provides a raft::host_matrix_view of an * // all-neighbors knn graph of dimensions [N, k] of the input * // dataset @@ -85,17 +89,21 @@ auto build(raft::resources const& res, * @brief Build nn-descent Index with dataset in device memory * * The following distance metrics are supported: - * - L2 + * - L2Expanded + * - L2SqrtExpanded + * - CosineExpanded + * - InnerProduct + * - BitwiseHamming (when T == int8, uint8) * * Usage example: * @code{.cpp} - * using namespace cuvs::neighbors::experimental; + * using namespace cuvs::neighbors; * // use default index parameters * nn_descent::index_params index_params; * // create and fill the index from a [N, D] raft::device_matrix_view dataset * auto knn_graph = raft::make_host_matrix(N, D); * auto index = nn_descent::index{res, knn_graph.view()}; - * cagra::build(res, index_params, dataset, index); + * nn_descent::build(res, index_params, dataset, index); * // index.graph() provides a raft::host_matrix_view of an * // all-neighbors knn graph of dimensions [N, k] of the input * // dataset @@ -136,15 +144,19 @@ void build(raft::resources const& res, * @brief Build nn-descent Index with dataset in host memory * * The following distance metrics are supported: - * - L2 + * - L2Expanded + * - L2SqrtExpanded + * - CosineExpanded + * - InnerProduct + * - BitwiseHamming (when T == int8, uint8) * * Usage example: * @code{.cpp} - * using namespace cuvs::neighbors::experimental; + * using namespace cuvs::neighbors; * // use default index parameters * nn_descent::index_params index_params; * // create and fill the index from a [N, D] raft::host_matrix_view dataset - * auto index = cagra::build(res, index_params, dataset); + * auto index = nn_descent::build(res, index_params, dataset); * // index.graph() provides a raft::host_matrix_view of an * // all-neighbors knn graph of dimensions [N, k] of the input * // dataset @@ -183,7 +195,11 @@ auto build(raft::resources const& res, * @brief Build nn-descent Index with dataset in host memory * * The following distance metrics are supported: - * - L2 + * - L2Expanded + * - L2SqrtExpanded + * - CosineExpanded + * - InnerProduct + * - BitwiseHamming (when T == int8, uint8) * * Usage example: * @code{.cpp} @@ -193,7 +209,7 @@ auto build(raft::resources const& res, * // create and fill the index from a [N, D] raft::host_matrix_view dataset * auto knn_graph = raft::make_host_matrix(N, D); * auto index = nn_descent::index{res, knn_graph.view()}; - * cagra::build(res, index_params, dataset, index); + * nn_descent::build(res, index_params, dataset, index); * // index.graph() provides a raft::host_matrix_view of an * // all-neighbors knn graph of dimensions [N, k] of the input * // dataset diff --git a/python/cuvs/cuvs/neighbors/nn_descent/nn_descent.pyx b/python/cuvs/cuvs/neighbors/nn_descent/nn_descent.pyx index a6277b911d..0f0589eced 100644 --- a/python/cuvs/cuvs/neighbors/nn_descent/nn_descent.pyx +++ b/python/cuvs/cuvs/neighbors/nn_descent/nn_descent.pyx @@ -55,7 +55,9 @@ cdef class IndexParams: ---------- metric : str, default = "sqeuclidean" String denoting the metric type. - distribution of the newly added data. + Supported metrics are `l2`, `euclidean`, `sqeuclidean`, + `inner_product`, `cosine`, and `bitwise_hamming` + (`bitwise_hamming` is for int8 and uint8 data types only) graph_degree : int For an input dataset of dimensions (N, D), determines the final dimensions of the all-neighbors knn graph which turns out to be of From 279afdb041b4e7f201f21345231673dcd833bf60 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Lorenzo=20Dematt=C3=A9?= Date: Thu, 14 Aug 2025 05:45:14 +0200 Subject: [PATCH 32/93] [REVIEW][Java] Rename destroyIndex() to close(), extend AutoCloseable (#1252) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The various `Index` classes in Java all have a `destroyIndex` method that must be called in order to free resources associated with that index, when done. Other classes in the repository use `close()` instead, and extend or implement `AutoCloseable`; this have the advantage of support in language (try-with-resources) and tooling (e.g. IntelliJ warnings). Additionally, having a single pattern is better for ease of use and consistency. This PR renames all instances of `destroyIndex` to `close` and make the declaring interfaces extend `AutoCloseable` Authors: - Lorenzo Dematté (https://github.com/ldematte) - MithunR (https://github.com/mythrocks) Approvers: - MithunR (https://github.com/mythrocks) URL: https://github.com/rapidsai/cuvs/pull/1252 --- .../com/nvidia/cuvs/CagraIndexBenchmarks.java | 18 +-- .../java/com/nvidia/cuvs/BruteForceIndex.java | 5 +- .../main/java/com/nvidia/cuvs/CagraIndex.java | 5 +- .../main/java/com/nvidia/cuvs/HnswIndex.java | 5 +- .../java/com/nvidia/cuvs/TieredIndex.java | 7 +- .../cuvs/internal/BruteForceIndexImpl.java | 2 +- .../nvidia/cuvs/internal/CagraIndexImpl.java | 2 +- .../nvidia/cuvs/internal/HnswIndexImpl.java | 2 +- .../nvidia/cuvs/internal/TieredIndexImpl.java | 2 +- .../nvidia/cuvs/BruteForceAndSearchIT.java | 43 +++--- .../nvidia/cuvs/CagraBuildAndSearchIT.java | 85 ++++++------ .../cuvs/CagraMultiThreadStabilityIT.java | 129 +++++++++--------- .../com/nvidia/cuvs/CagraRandomizedIT.java | 2 +- .../com/nvidia/cuvs/HnswBuildAndSearchIT.java | 4 +- .../com/nvidia/cuvs/HnswRandomizedIT.java | 4 +- 15 files changed, 155 insertions(+), 160 deletions(-) diff --git a/java/benchmarks/src/main/java/com/nvidia/cuvs/CagraIndexBenchmarks.java b/java/benchmarks/src/main/java/com/nvidia/cuvs/CagraIndexBenchmarks.java index 18d8495424..9b3a7dffc3 100644 --- a/java/benchmarks/src/main/java/com/nvidia/cuvs/CagraIndexBenchmarks.java +++ b/java/benchmarks/src/main/java/com/nvidia/cuvs/CagraIndexBenchmarks.java @@ -121,20 +121,20 @@ public void testIndexingAndSerializeToFile() throws Throwable { .build(); // Create the index with the dataset - CagraIndex index = CagraIndex.newBuilder(resources) + try (CagraIndex index = CagraIndex.newBuilder(resources) .withDataset(arrayDataset) .withIndexParams(indexParams) - .build(); + .build()) { - // Saving the index on to the disk. - var indexFilePath = Path.of(UUID.randomUUID() + ".cag"); - try (var outputStream = Files.newOutputStream(indexFilePath)) { + // Saving the index on to the disk. + var indexFilePath = Path.of(UUID.randomUUID() + ".cag"); + try (var outputStream = Files.newOutputStream(indexFilePath)) { index.serialize(outputStream); - } + } - // Cleanup - Files.deleteIfExists(indexFilePath); - index.destroyIndex(); + // Cleanup + Files.deleteIfExists(indexFilePath); + } } } diff --git a/java/cuvs-java/src/main/java/com/nvidia/cuvs/BruteForceIndex.java b/java/cuvs-java/src/main/java/com/nvidia/cuvs/BruteForceIndex.java index 2b50a0b741..47a5f64c2a 100644 --- a/java/cuvs-java/src/main/java/com/nvidia/cuvs/BruteForceIndex.java +++ b/java/cuvs-java/src/main/java/com/nvidia/cuvs/BruteForceIndex.java @@ -28,13 +28,14 @@ * * @since 25.02 */ -public interface BruteForceIndex { +public interface BruteForceIndex extends AutoCloseable { /** * Invokes the native destroy_brute_force_index function to de-allocate * BRUTEFORCE index */ - void destroyIndex() throws Throwable; + @Override + void close() throws Exception; /** * Invokes the native search_brute_force_index via the Panama API for searching diff --git a/java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraIndex.java b/java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraIndex.java index 427b5e22d2..e3b885a333 100644 --- a/java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraIndex.java +++ b/java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraIndex.java @@ -33,12 +33,13 @@ * * @since 25.02 */ -public interface CagraIndex { +public interface CagraIndex extends AutoCloseable { /** * Invokes the native destroy_cagra_index to de-allocate the CAGRA index */ - void destroyIndex() throws Throwable; + @Override + void close() throws Exception; /** * Invokes the native search_cagra_index via the Panama API for searching a diff --git a/java/cuvs-java/src/main/java/com/nvidia/cuvs/HnswIndex.java b/java/cuvs-java/src/main/java/com/nvidia/cuvs/HnswIndex.java index 746dc05fd1..cda73c9ea4 100644 --- a/java/cuvs-java/src/main/java/com/nvidia/cuvs/HnswIndex.java +++ b/java/cuvs-java/src/main/java/com/nvidia/cuvs/HnswIndex.java @@ -25,12 +25,13 @@ * * @since 25.02 */ -public interface HnswIndex { +public interface HnswIndex extends AutoCloseable { /** * Invokes the native destroy_hnsw_index to de-allocate the HNSW index */ - void destroyIndex() throws Throwable; + @Override + void close() throws Exception; /** * Invokes the native search_hnsw_index via the Panama API for searching a HNSW diff --git a/java/cuvs-java/src/main/java/com/nvidia/cuvs/TieredIndex.java b/java/cuvs-java/src/main/java/com/nvidia/cuvs/TieredIndex.java index 2726382f23..3e8d1e0752 100644 --- a/java/cuvs-java/src/main/java/com/nvidia/cuvs/TieredIndex.java +++ b/java/cuvs-java/src/main/java/com/nvidia/cuvs/TieredIndex.java @@ -23,15 +23,16 @@ * {@link TieredIndex} encapsulates a Tiered index, along with methods to * interact with it. */ -public interface TieredIndex { +public interface TieredIndex extends AutoCloseable { /** * Destroys the underlying native TieredIndex object and releases associated * resources. * - * @throws Throwable if an error occurs during index destruction + * @throws Exception if an error occurs during index destruction */ - void destroyIndex() throws Throwable; + @Override + void close() throws Exception; /** * Searches the index with the specified query and search parameters. diff --git a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/BruteForceIndexImpl.java b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/BruteForceIndexImpl.java index 83c08fedc6..e5f98ad2dd 100644 --- a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/BruteForceIndexImpl.java +++ b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/BruteForceIndexImpl.java @@ -113,7 +113,7 @@ private void checkNotDestroyed() { * BRUTEFORCE index */ @Override - public void destroyIndex() { + public void close() { checkNotDestroyed(); try { int returnValue = cuvsBruteForceIndexDestroy(bruteForceIndexReference.indexPtr); diff --git a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CagraIndexImpl.java b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CagraIndexImpl.java index 489213f4da..0d5fc8d90f 100644 --- a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CagraIndexImpl.java +++ b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CagraIndexImpl.java @@ -158,7 +158,7 @@ private void checkNotDestroyed() { * Invokes the native destroy_cagra_index to de-allocate the CAGRA index */ @Override - public void destroyIndex() throws Throwable { + public void close() { checkNotDestroyed(); try { int returnValue = cuvsCagraIndexDestroy(cagraIndexReference.getMemorySegment()); diff --git a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/HnswIndexImpl.java b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/HnswIndexImpl.java index ed25b5cc3f..736b22f567 100644 --- a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/HnswIndexImpl.java +++ b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/HnswIndexImpl.java @@ -74,7 +74,7 @@ private HnswIndexImpl( * Invokes the native destroy_hnsw_index to de-allocate the HNSW index */ @Override - public void destroyIndex() { + public void close() { int returnValue = cuvsHnswIndexDestroy(hnswIndexReference.getMemorySegment()); checkCuVSError(returnValue, "cuvsHnswIndexDestroy"); } diff --git a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/TieredIndexImpl.java b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/TieredIndexImpl.java index 31d8cd4bec..c507045326 100644 --- a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/TieredIndexImpl.java +++ b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/TieredIndexImpl.java @@ -94,7 +94,7 @@ private void checkNotDestroyed() { * Invokes the native destroy_tiered_index to de-allocate the Tiered index */ @Override - public void destroyIndex() { + public void close() { checkNotDestroyed(); try { int returnValue = cuvsTieredIndexDestroy(tieredIndexReference.getMemorySegment()); 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 ddd1be0b5d..900b16ba1a 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 @@ -57,37 +57,34 @@ private static void indexAndQueryOnce( new BruteForceIndexParams.Builder().withNumWriterThreads(32).build(); // Create the index with the dataset - BruteForceIndex index = + try (BruteForceIndex index = BruteForceIndex.newBuilder(resources) .withDataset(dataset) .withIndexParams(indexParams) - .build(); + .build()) { - // Saving the index on to the disk. - String indexFileName = UUID.randomUUID() + ".bf"; - try (var outputStream = new FileOutputStream(indexFileName)) { - index.serialize(outputStream); - } - - // Loading a BRUTEFORCE index from disk. - Path indexFile = Path.of(indexFileName); - try (var inputStream = Files.newInputStream(indexFile)) { - BruteForceIndex loadedIndex = BruteForceIndex.newBuilder(resources).from(inputStream).build(); + // Saving the index on to the disk. + String indexFileName = UUID.randomUUID() + ".bf"; + try (var outputStream = new FileOutputStream(indexFileName)) { + index.serialize(outputStream); + } - // search the loaded index - SearchResults results = loadedIndex.search(cuvsQuery); - checkResults(expectedResults, results.getResults()); + // Loading a BRUTEFORCE index from disk. + Path indexFile = Path.of(indexFileName); + try (var inputStream = Files.newInputStream(indexFile); + BruteForceIndex loadedIndex = + BruteForceIndex.newBuilder(resources).from(inputStream).build()) { - // search the first index - results = index.search(cuvsQuery); - checkResults(expectedResults, results.getResults()); + // search the loaded index + SearchResults results = loadedIndex.search(cuvsQuery); + checkResults(expectedResults, results.getResults()); - // Cleanup - index.destroyIndex(); - loadedIndex.destroyIndex(); + // search the first index + results = index.search(cuvsQuery); + checkResults(expectedResults, results.getResults()); + } + Files.deleteIfExists(indexFile); } - - Files.deleteIfExists(indexFile); } /** 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 8223e199f6..a4133c4dc5 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 @@ -233,7 +233,7 @@ public void testIndexing() throws Throwable { () -> { try (CuVSResources resources = CheckedCuVSResources.create()) { var index = indexOnce(CuVSMatrix.ofArray(dataset), resources); - index.destroyIndex(); + index.close(); } catch (Throwable e) { throw new RuntimeException(e); } @@ -250,11 +250,10 @@ public void testSerialization() throws Throwable { numTestsRuns, () -> () -> { - try (CuVSResources resources = CheckedCuVSResources.create()) { - var index = indexOnce(CuVSMatrix.ofArray(dataset), resources); + try (CuVSResources resources = CheckedCuVSResources.create(); + var index = indexOnce(CuVSMatrix.ofArray(dataset), resources)) { var indexPath = serializeOnce(index); Files.deleteIfExists(indexPath); - index.destroyIndex(); } catch (Throwable e) { throw new RuntimeException(e); } @@ -272,7 +271,7 @@ public void testDeserialization() throws Throwable { () -> () -> { try (CuVSResources resources = CheckedCuVSResources.create()) { - deserializeOnce(indexPath, resources).destroyIndex(); + deserializeOnce(indexPath, resources).close(); } catch (Throwable e) { throw new RuntimeException(e); } @@ -282,11 +281,9 @@ public void testDeserialization() throws Throwable { } private Path createSerializedIndex(CuVSMatrix dataset) throws Throwable { - try (CuVSResources resources = CheckedCuVSResources.create()) { - var index = indexOnce(dataset, resources); - var indexPath = serializeOnce(index); - index.destroyIndex(); - return indexPath; + try (CuVSResources resources = CheckedCuVSResources.create(); + var index = indexOnce(dataset, resources)) { + return serializeOnce(index); } } @@ -296,38 +293,36 @@ public void testReconstructIndexFromGraph() throws Throwable { var queries = createSampleQueries(); List> expectedResults = getExpectedResults(); - try (CuVSResources resources = CuVSResources.create()) { - var index = indexOnce(dataset, resources); + try (CuVSResources resources = CuVSResources.create(); + var index = indexOnce(dataset, resources)) { var graph = index.getGraph(); - var reconstructedIndex = + try (var reconstructedIndex = CagraIndex.newBuilder(resources) .from(graph) .withDataset(dataset) .withIndexParams( new CagraIndexParams.Builder().withMetric(CuvsDistanceType.L2Expanded).build()) - .build(); - queryAndCompare( - index, - reconstructedIndex, - SearchResults.IDENTITY_MAPPING, - queries, - expectedResults, - resources); - - var originalIndexPath = serializeOnce(index); - var reconstructedIndexPath = serializeOnce(reconstructedIndex); - - var originalBytes = Files.readAllBytes(originalIndexPath); - var reconstructedBytes = Files.readAllBytes(reconstructedIndexPath); - - assertArrayEquals(originalBytes, reconstructedBytes); - - index.destroyIndex(); - reconstructedIndex.destroyIndex(); - - Files.deleteIfExists(originalIndexPath); - Files.deleteIfExists(reconstructedIndexPath); + .build()) { + queryAndCompare( + index, + reconstructedIndex, + SearchResults.IDENTITY_MAPPING, + queries, + expectedResults, + resources); + + var originalIndexPath = serializeOnce(index); + var reconstructedIndexPath = serializeOnce(reconstructedIndex); + + var originalBytes = Files.readAllBytes(originalIndexPath); + var reconstructedBytes = Files.readAllBytes(reconstructedIndexPath); + + assertArrayEquals(originalBytes, reconstructedBytes); + + Files.deleteIfExists(originalIndexPath); + Files.deleteIfExists(reconstructedIndexPath); + } } } } @@ -516,8 +511,8 @@ private void queryAndCompare( private void cleanup(CagraIndex index, CagraIndex loadedIndex) throws Throwable { // Cleanup - index.destroyIndex(); - loadedIndex.destroyIndex(); + index.close(); + loadedIndex.close(); } @Test @@ -644,10 +639,10 @@ public void testMergingIndexes() throws Throwable { if (indexFile.exists()) { indexFile.delete(); } - index1.destroyIndex(); - index2.destroyIndex(); - mergedIndex.destroyIndex(); - loadedMergedIndex.destroyIndex(); + index1.close(); + index2.close(); + mergedIndex.close(); + loadedMergedIndex.close(); } } @@ -756,10 +751,10 @@ public void testMergeStrategies() throws Throwable { if (physicalIndexFile.exists()) { physicalIndexFile.delete(); } - index1.destroyIndex(); - index2.destroyIndex(); - physicalMergedIndex.destroyIndex(); - loadedPhysicalIndex.destroyIndex(); + index1.close(); + index2.close(); + physicalMergedIndex.close(); + loadedPhysicalIndex.close(); } } } 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 1df57d1e0b..32f3d86fe0 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 @@ -95,92 +95,91 @@ private void testQueryingUsingMultipleThreads(QueryAction queryAction) throws Th .withMetric(CuvsDistanceType.L2Expanded) .build(); - CagraIndex index = + try (CagraIndex index = CagraIndex.newBuilder(resources) .withDataset(dataset) .withIndexParams(indexParams) - .build(); + .build()) { - log.trace("CAGRA index created, starting high-contention multi-threaded search..."); + log.trace("CAGRA index created, starting high-contention multi-threaded search..."); - // Create high contention scenario that would fail without using separate resources in every - // thread - try (ExecutorService executor = Executors.newFixedThreadPool(numThreads)) { - List> futures = new ArrayList<>(); - CountDownLatch startLatch = new CountDownLatch(1); - AtomicInteger successfulQueries = new AtomicInteger(0); - AtomicReference firstError = new AtomicReference<>(); + // Create high contention scenario that would fail without using separate resources in every + // thread + try (ExecutorService executor = Executors.newFixedThreadPool(numThreads)) { + List> futures = new ArrayList<>(); + CountDownLatch startLatch = new CountDownLatch(1); + AtomicInteger successfulQueries = new AtomicInteger(0); + AtomicReference firstError = new AtomicReference<>(); - for (int threadId = 0; threadId < numThreads; threadId++) { - final int finalThreadId = threadId; + for (int threadId = 0; threadId < numThreads; threadId++) { + final int finalThreadId = threadId; - Future future = - executor.submit( - () -> { - try { - // Wait for all threads to start simultaneously - startLatch.await(); + Future future = + executor.submit( + () -> { + try { + // Wait for all threads to start simultaneously + startLatch.await(); - for (int queryId = 0; queryId < queriesPerThread; queryId++) { - queryAction.run(index); - successfulQueries.incrementAndGet(); + for (int queryId = 0; queryId < queriesPerThread; queryId++) { + queryAction.run(index); + successfulQueries.incrementAndGet(); - // No Thread.yield() - maximize contention - } + // No Thread.yield() - maximize contention + } - log.trace("Thread {} completed successfully", finalThreadId); + log.trace("Thread {} completed successfully", finalThreadId); - } catch (Throwable t) { - log.error("Thread {} failed: {}", finalThreadId, t.getMessage(), t); - firstError.compareAndSet(null, t); - throw new RuntimeException("Thread failed", t); - } - }); + } catch (Throwable t) { + log.error("Thread {} failed: {}", finalThreadId, t.getMessage(), t); + firstError.compareAndSet(null, t); + throw new RuntimeException("Thread failed", t); + } + }); - futures.add(future); - } + futures.add(future); + } - // Start all threads simultaneously to maximize contention - log.debug("Starting all {} threads simultaneously...", numThreads); - startLatch.countDown(); - - // Wait for all threads to complete - boolean allCompleted = true; - for (Future future : futures) { - try { - future.get(120, TimeUnit.SECONDS); // Longer timeout for stress test - } catch (Exception e) { - allCompleted = false; - log.error("Thread failed: {}", e.getMessage(), e); - if (firstError.get() == null) { - firstError.set(e); + // Start all threads simultaneously to maximize contention + log.debug("Starting all {} threads simultaneously...", numThreads); + startLatch.countDown(); + + // Wait for all threads to complete + boolean allCompleted = true; + for (Future future : futures) { + try { + future.get(120, TimeUnit.SECONDS); // Longer timeout for stress test + } catch (Exception e) { + allCompleted = false; + log.error("Thread failed: {}", e.getMessage(), e); + if (firstError.get() == null) { + firstError.set(e); + } } } - } - executor.shutdown(); - if (!executor.awaitTermination(10, TimeUnit.SECONDS)) { - executor.shutdownNow(); - } + executor.shutdown(); + if (!executor.awaitTermination(10, TimeUnit.SECONDS)) { + executor.shutdownNow(); + } - // Verify results - int expectedTotalQueries = numThreads * queriesPerThread; - int actualSuccessfulQueries = successfulQueries.get(); + // Verify results + int expectedTotalQueries = numThreads * queriesPerThread; + int actualSuccessfulQueries = successfulQueries.get(); - log.debug(" Successful queries: {} / {}", actualSuccessfulQueries, expectedTotalQueries); + log.debug(" Successful queries: {} / {}", actualSuccessfulQueries, expectedTotalQueries); - if (firstError.get() != null) { - fail("MultiThreaded stablity test failed:" + " " + firstError.get().getMessage()); - } + if (firstError.get() != null) { + fail("MultiThreaded stablity test failed:" + " " + firstError.get().getMessage()); + } - assertTrue("All threads should complete successfully", allCompleted); - assertEquals( - "All queries should complete successfully", - expectedTotalQueries, - actualSuccessfulQueries); + assertTrue("All threads should complete successfully", allCompleted); + assertEquals( + "All queries should complete successfully", + expectedTotalQueries, + actualSuccessfulQueries); + } } - - index.destroyIndex(); } } 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 69e28d20aa..144090866c 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 @@ -158,7 +158,7 @@ private void tmpResultsTopKWithRandomValues(boolean useNativeMemoryDataset) thro compareResults(results, expected, topK, datasetSize, numQueries); } finally { - index.destroyIndex(); + index.close(); } } } 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 b2b86c37ad..c72ccd8aec 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 @@ -97,10 +97,10 @@ private static void indexAndQueryOnce( checkResults(expectedResults, results.getResults()); // Cleanup - hnswIndex.destroyIndex(); + hnswIndex.close(); } } finally { - index.destroyIndex(); + index.close(); Files.deleteIfExists(hnswIndexPath); } } 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 8cda810c6c..02ac3c3a14 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 @@ -164,10 +164,10 @@ private void tmpResultsTopKWithRandomValues(boolean useNativeMemoryDataset) thro SearchResults results = hnswIndex.search(hnswQuery); compareResults(results, expected, topK, datasetSize, numQueries); - hnswIndex.destroyIndex(); + hnswIndex.close(); } } finally { - index.destroyIndex(); + index.close(); Files.deleteIfExists(hnswIndexPath); } } From 699d32c8be9a344c83456232851430e4c3840484 Mon Sep 17 00:00:00 2001 From: Ben Frederickson Date: Thu, 14 Aug 2025 12:21:49 -0700 Subject: [PATCH 33/93] Merge support for tiered index (#1155) Authors: - Ben Frederickson (https://github.com/benfred) - Corey J. Nolet (https://github.com/cjnolet) Approvers: - Corey J. Nolet (https://github.com/cjnolet) URL: https://github.com/rapidsai/cuvs/pull/1155 --- cpp/include/cuvs/neighbors/cagra.hpp | 3 + cpp/include/cuvs/neighbors/tiered_index.h | 24 ++++ cpp/include/cuvs/neighbors/tiered_index.hpp | 30 +++++ .../neighbors/detail/cagra/cagra_merge.cuh | 6 +- cpp/src/neighbors/detail/tiered_index.cuh | 73 ++++++++++ cpp/src/neighbors/tiered_index.cu | 27 ++++ cpp/src/neighbors/tiered_index_c.cpp | 126 +++++++++++++++--- cpp/tests/neighbors/tiered_index.cu | 79 +++++++---- 8 files changed, 322 insertions(+), 46 deletions(-) diff --git a/cpp/include/cuvs/neighbors/cagra.hpp b/cpp/include/cuvs/neighbors/cagra.hpp index 48ee128271..6c7d1a1979 100644 --- a/cpp/include/cuvs/neighbors/cagra.hpp +++ b/cpp/include/cuvs/neighbors/cagra.hpp @@ -1975,6 +1975,9 @@ void serialize_to_hnswlib( * @note: When device memory is sufficient, the dataset attached to the returned index is allocated * in device memory by default; otherwise, host memory is used automatically. * + * @note: This API only supports physical merge (`merge_strategy = MERGE_STRATEGY_PHYSICAL`), and + * attempting a logical merge here will throw an error. + * * Usage example: * @code{.cpp} * using namespace raft::neighbors; diff --git a/cpp/include/cuvs/neighbors/tiered_index.h b/cpp/include/cuvs/neighbors/tiered_index.h index c72f929d95..6c54d474b6 100644 --- a/cpp/include/cuvs/neighbors/tiered_index.h +++ b/cpp/include/cuvs/neighbors/tiered_index.h @@ -246,6 +246,30 @@ cuvsError_t cuvsTieredIndexSearch(cuvsResources_t res, cuvsError_t cuvsTieredIndexExtend(cuvsResources_t res, DLManagedTensor* new_vectors, cuvsTieredIndex_t index); +/** + * @} + */ + +/** + * @defgroup tiered_c_index_merge Tiered index merge + * @{ + */ +/** + * @brief Merge multiple indices together into a single index + * + * @param[in] res cuvsResources_t opaque C handle + * @param[in] index_params Index parameters to use when merging + * @param[in] indices pointers to indices to merge together + * @param[in] num_indices the number of indices to merge + * @param[out] output_index the merged index + * @return cuvsError_t + */ +cuvsError_t cuvsTieredIndexMerge(cuvsResources_t res, + cuvsTieredIndexParams_t index_params, + cuvsTieredIndex_t* indices, + size_t num_indices, + cuvsTieredIndex_t output_index); + /** * @} */ diff --git a/cpp/include/cuvs/neighbors/tiered_index.hpp b/cpp/include/cuvs/neighbors/tiered_index.hpp index 1037955309..4d62773fc5 100644 --- a/cpp/include/cuvs/neighbors/tiered_index.hpp +++ b/cpp/include/cuvs/neighbors/tiered_index.hpp @@ -197,4 +197,34 @@ void search(raft::resources const& res, raft::device_matrix_view distances, const cuvs::neighbors::filtering::base_filter& sample_filter = cuvs::neighbors::filtering::none_sample_filter{}); + +/** @brief Merge multiple tiered indices into a single index. + * + * This function merges multiple tiered indices into one, combining both the datasets and graph + * structures. + * + * @param[in] res + * @param[in] index_params configure the index building + * @param[in] indices A vector of pointers to the indices to merge. All indices should + * be of the same type, and have datasets with the same dimensionality + * + * @return A new tiered index containing the merged indices + */ +auto merge(raft::resources const& res, + const index_params& index_params, + const std::vector>*>& indices) + -> tiered_index::index>; + +/** @copydoc merge */ +auto merge(raft::resources const& res, + const index_params& index_params, + const std::vector>*>& indices) + -> tiered_index::index>; + +/** @copydoc merge */ +auto merge(raft::resources const& res, + const index_params& index_params, + const std::vector>*>& indices) + -> tiered_index::index>; + } // namespace cuvs::neighbors::tiered_index diff --git a/cpp/src/neighbors/detail/cagra/cagra_merge.cuh b/cpp/src/neighbors/detail/cagra/cagra_merge.cuh index 192f7f1632..86faad02c0 100644 --- a/cpp/src/neighbors/detail/cagra/cagra_merge.cuh +++ b/cpp/src/neighbors/detail/cagra/cagra_merge.cuh @@ -1,5 +1,5 @@ /* - * Copyright (c) 2023-2024, NVIDIA CORPORATION. + * Copyright (c) 2023-2025, NVIDIA CORPORATION. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -44,6 +44,10 @@ index merge(raft::resources const& handle, const cagra::merge_params& params, std::vector*>& indices) { + // we're doing a physical merge here, make sure that this matches the merge_params + RAFT_EXPECTS(params.merge_strategy == cuvs::neighbors::MergeStrategy::MERGE_STRATEGY_PHYSICAL, + "cagra::merge only supports merge_strategy=MERGE_STRATEGY_PHYSICAL"); + using cagra_index_t = cuvs::neighbors::cagra::index; using ds_idx_type = typename cagra_index_t::dataset_index_type; diff --git a/cpp/src/neighbors/detail/tiered_index.cuh b/cpp/src/neighbors/detail/tiered_index.cuh index ef7fd7d015..47f1ff349c 100644 --- a/cpp/src/neighbors/detail/tiered_index.cuh +++ b/cpp/src/neighbors/detail/tiered_index.cuh @@ -297,6 +297,79 @@ auto build( return std::shared_ptr>(ret); } +/** + * @brief merge multiple indices together + */ +template +auto merge(raft::resources const& res, + const index_params& index_params, + const std::vector*>& indices) + -> std::shared_ptr> +{ + using value_type = typename UpstreamT::value_type; + + RAFT_EXPECTS(indices.size() > 0, "Must pass at least one index to merge"); + + for (auto index : indices) { + RAFT_EXPECTS(index != nullptr, + "Null pointer detected in 'indices'. Ensure all elements are valid before usage."); + } + + // handle simple case of only one index being merged + if (indices.size() == 1) { return indices[0]->state; } + + auto dim = indices[0]->state->dim(); + auto include_norms = indices[0]->state->storage->include_norms; + + // validate data and check what needs to be merged + size_t bfknn_rows = 0, ann_rows = 0; + for (auto index : indices) { + RAFT_EXPECTS(dim == index->state->dim(), "Each index must have the same dimensionality"); + bfknn_rows += index->state->bfknn_rows(); + ann_rows += index->state->ann_rows(); + } + + // degenerate case - all indices are empty, just re-use the state from the first index + if (!bfknn_rows && !ann_rows) { return indices[0]->state; } + + // concatenate all the storages together + auto to_allocate = bfknn_rows + ann_rows; + auto new_storage = + std::make_shared>(res, to_allocate, dim, include_norms); + + for (auto index : indices) { + auto storage = index->state->storage; + + // copy over dataset to new storage + raft::copy(res, + raft::make_device_matrix_view( + new_storage->dataset.data() + new_storage->num_rows_used * dim, + storage->num_rows_used, + dim), + raft::make_device_matrix_view( + storage->dataset.data(), storage->num_rows_used, dim)); + + // copy over precalculated norms + if (include_norms) { + raft::copy(res, + raft::make_device_vector_view( + new_storage->norms->data() + new_storage->num_rows_used, storage->num_rows_used), + raft::make_device_vector_view( + storage->norms->data(), storage->num_rows_used)); + } + new_storage->num_rows_used += storage->num_rows_used; + } + + auto next_state = std::make_shared>(*indices[0]->state); + next_state->storage = new_storage; + next_state->build_params = index_params; + + if (next_state->bfknn_rows() > static_cast(next_state->build_params.min_ann_rows)) { + next_state = compact(res, *next_state); + } + return next_state; +} + template auto extend(raft::resources const& res, const index_state& current, diff --git a/cpp/src/neighbors/tiered_index.cu b/cpp/src/neighbors/tiered_index.cu index 1a3a410419..ba109defcb 100644 --- a/cpp/src/neighbors/tiered_index.cu +++ b/cpp/src/neighbors/tiered_index.cu @@ -175,6 +175,33 @@ void search(raft::resources const& res, res, search_params, ivf_pq::typed_search, queries, neighbors, distances, sample_filter); } +auto merge(raft::resources const& res, + const index_params& index_params, + const std::vector>*>& indices) + -> tiered_index::index> +{ + auto state = detail::merge(res, index_params, indices); + return cuvs::neighbors::tiered_index::index>(state); +} + +auto merge(raft::resources const& res, + const index_params& index_params, + const std::vector>*>& indices) + -> tiered_index::index> +{ + auto state = detail::merge(res, index_params, indices); + return cuvs::neighbors::tiered_index::index>(state); +} + +auto merge(raft::resources const& res, + const index_params& index_params, + const std::vector>*>& indices) + -> tiered_index::index> +{ + auto state = detail::merge(res, index_params, indices); + return cuvs::neighbors::tiered_index::index>(state); +} + template int64_t index::size() const noexcept { diff --git a/cpp/src/neighbors/tiered_index_c.cpp b/cpp/src/neighbors/tiered_index_c.cpp index 196b199e10..267c7ba5f0 100644 --- a/cpp/src/neighbors/tiered_index_c.cpp +++ b/cpp/src/neighbors/tiered_index_c.cpp @@ -37,6 +37,33 @@ namespace { using namespace cuvs::neighbors; +template +void convert_c_index_params(cuvsTieredIndexParams params, + int64_t n_rows, + int64_t dim, + tiered_index::index_params* out) +{ + out->min_ann_rows = params.min_ann_rows; + out->create_ann_index_on_extend = params.create_ann_index_on_extend; + out->metric = params.metric; + + if constexpr (std::is_same_v) { + if (params.cagra_params != NULL) { + cagra::convert_c_index_params(*params.cagra_params, n_rows, dim, out); + } + } else if constexpr (std::is_same_v) { + if (params.ivf_flat_params != NULL) { + ivf_flat::convert_c_index_params(*params.ivf_flat_params, out); + } + } else if constexpr (std::is_same_v) { + if (params.ivf_pq_params != NULL) { + ivf_pq::convert_c_index_params(*params.ivf_pq_params, out); + } + } else { + RAFT_FAIL("unhandled index params type"); + } +} + template void* _build(cuvsResources_t res, cuvsTieredIndexParams params, DLManagedTensor* dataset_tensor) { @@ -44,40 +71,26 @@ void* _build(cuvsResources_t res, cuvsTieredIndexParams params, DLManagedTensor* using mdspan_type = raft::device_matrix_view; auto mds = cuvs::core::from_dlpack(dataset_tensor); + auto dataset = dataset_tensor->dl_tensor; + RAFT_EXPECTS(dataset.ndim == 2, "dataset should be a 2-dimensional tensor"); + RAFT_EXPECTS(dataset.shape != nullptr, "dataset should have an initialized shape"); + switch (params.algo) { case CUVS_TIERED_INDEX_ALGO_CAGRA: { auto build_params = tiered_index::index_params(); - if (params.cagra_params != NULL) { - auto dataset = dataset_tensor->dl_tensor; - cagra::convert_c_index_params( - *params.cagra_params, dataset.shape[0], dataset.shape[1], &build_params); - } - build_params.min_ann_rows = params.min_ann_rows; - build_params.create_ann_index_on_extend = params.create_ann_index_on_extend; - build_params.metric = params.metric; + convert_c_index_params(params, dataset.shape[0], dataset.shape[1], &build_params); return new tiered_index::index>( tiered_index::build(*res_ptr, build_params, mds)); } case CUVS_TIERED_INDEX_ALGO_IVF_FLAT: { auto build_params = tiered_index::index_params(); - if (params.ivf_flat_params != NULL) { - ivf_flat::convert_c_index_params(*params.ivf_flat_params, &build_params); - } - build_params.metric = params.metric; - build_params.min_ann_rows = params.min_ann_rows; - build_params.create_ann_index_on_extend = params.create_ann_index_on_extend; + convert_c_index_params(params, dataset.shape[0], dataset.shape[1], &build_params); return new tiered_index::index>( tiered_index::build(*res_ptr, build_params, mds)); } case CUVS_TIERED_INDEX_ALGO_IVF_PQ: { - auto build_params = tiered_index::index_params(); - build_params.metric = params.metric; - if (params.ivf_pq_params != NULL) { - ivf_pq::convert_c_index_params(*params.ivf_pq_params, &build_params); - } - build_params.metric = params.metric; - build_params.min_ann_rows = params.min_ann_rows; - build_params.create_ann_index_on_extend = params.create_ann_index_on_extend; + auto build_params = tiered_index::index_params(); + convert_c_index_params(params, dataset.shape[0], dataset.shape[1], &build_params); return new tiered_index::index>( tiered_index::build(*res_ptr, build_params, mds)); } @@ -157,6 +170,47 @@ void _extend(cuvsResources_t res, DLManagedTensor* new_vectors, cuvsTieredIndex tiered_index::extend(*res_ptr, vectors_mds, index_ptr); } +template +void _merge(cuvsResources_t res, + cuvsTieredIndexParams params, + cuvsTieredIndex_t* indices, + size_t num_indices, + cuvsTieredIndex_t output_index) +{ + auto res_ptr = reinterpret_cast(res); + + std::vector*> cpp_indices; + + int64_t n_rows = 0, dim = 0; + for (size_t i = 0; i < num_indices; ++i) { + RAFT_EXPECTS(indices[i]->dtype.code == indices[0]->dtype.code, + "indices must all have the same dtype"); + RAFT_EXPECTS(indices[i]->dtype.bits == indices[0]->dtype.bits, + "indices must all have the same dtype"); + RAFT_EXPECTS(indices[i]->algo == indices[0]->algo, + "indices must all have the same index algorithm"); + + auto idx_ptr = + reinterpret_cast*>(indices[i]->addr); + n_rows += idx_ptr->size(); + if (dim) { + RAFT_EXPECTS(dim == idx_ptr->dim(), "indices must all have the same dimensionality"); + } else { + dim = idx_ptr->dim(); + } + cpp_indices.push_back(idx_ptr); + } + + auto build_params = tiered_index::index_params(); + convert_c_index_params(params, n_rows, dim, &build_params); + + auto ptr = new cuvs::neighbors::tiered_index::index( + cuvs::neighbors::tiered_index::merge(*res_ptr, build_params, cpp_indices)); + + output_index->addr = reinterpret_cast(ptr); + output_index->dtype = indices[0]->dtype; + output_index->algo = indices[0]->algo; +} } // namespace @@ -305,3 +359,31 @@ extern "C" cuvsError_t cuvsTieredIndexExtend(cuvsResources_t res, } }); } + +extern "C" cuvsError_t cuvsTieredIndexMerge(cuvsResources_t res, + cuvsTieredIndexParams_t params, + cuvsTieredIndex_t* indices, + size_t num_indices, + cuvsTieredIndex_t output_index) +{ + return cuvs::core::translate_exceptions([=] { + RAFT_EXPECTS(num_indices >= 1, "must have at least one index to merge"); + + switch (indices[0]->algo) { + case CUVS_TIERED_INDEX_ALGO_CAGRA: { + _merge>(res, *params, indices, num_indices, output_index); + break; + } + case CUVS_TIERED_INDEX_ALGO_IVF_FLAT: { + _merge>(res, *params, indices, num_indices, output_index); + break; + } + case CUVS_TIERED_INDEX_ALGO_IVF_PQ: { + _merge>( + res, *params, indices, num_indices, output_index); + break; + } + default: RAFT_FAIL("unsupported tiered index algorithm"); + } + }); +} diff --git a/cpp/tests/neighbors/tiered_index.cu b/cpp/tests/neighbors/tiered_index.cu index 4aab6ba6ff..b1ed2b79ea 100644 --- a/cpp/tests/neighbors/tiered_index.cu +++ b/cpp/tests/neighbors/tiered_index.cu @@ -34,18 +34,30 @@ namespace cuvs::neighbors::tiered_index { +enum TieredIndexTestStrategy { TEST_EXTEND, TEST_MERGE }; + struct AnnTieredIndexInputs { int n_rows; int dim; cuvs::distance::DistanceType metric; int k; int n_queries; - double min_recall; + TieredIndexTestStrategy test_strategy; }; +inline ::std::ostream& operator<<(::std::ostream& os, const TieredIndexTestStrategy& p) +{ + switch (p) { + case TEST_EXTEND: os << "TEST_EXTEND"; break; + case TEST_MERGE: os << "TEST_MERGE"; break; + } + return os; +} + inline ::std::ostream& operator<<(::std::ostream& os, const AnnTieredIndexInputs& p) { os << "dataset shape=" << p.n_rows << "x" << p.dim << ", metric=" << print_metric{p.metric} + << ", k=" << p.k << ", n_queries=" << p.n_queries << ", test_strategy=" << p.test_strategy << std::endl; return os; } @@ -95,25 +107,57 @@ class ANNTieredIndexTest : public ::testing::TestWithParam // Calculate tiered index std::vector indices_tiered(ps.n_queries * ps.k); std::vector distances_tiered(ps.n_queries * ps.k); + double min_recall = 0.85; { index_params build_params; build_params.metric = ps.metric; build_params.min_ann_rows = 2000; build_params.create_ann_index_on_extend = true; - // include 50% of the rows in the initial build - and then incrementally add - // the rest row by row + typename UpstreamT::search_params_type search_params; + if constexpr (std::is_same_v) { + // min_recall logic copied from ivf_flat unittest + min_recall = + static_cast(search_params.n_probes) / static_cast(build_params.n_lists); + } + if constexpr (std::is_same_v) { + min_recall = + static_cast(search_params.n_probes) / static_cast(build_params.n_lists); + } + + // include 50% of the rows in the initial build auto initial_rows = ps.n_rows / 2; auto database_view = raft::make_device_matrix_view( (const value_type*)database.data(), initial_rows, ps.dim); - auto index = cuvs::neighbors::tiered_index::build(handle_, build_params, database_view); - for (int i = initial_rows; i < ps.n_rows; ++i) { - cuvs::neighbors::tiered_index::extend( - handle_, - raft::make_device_matrix_view( - (const value_type*)database.data() + i * ps.dim, 1, ps.dim), - &index); + + std::optional> final_index; + + if (ps.test_strategy == TEST_EXTEND) { + // test extend functionality by adding remaining vectors one by one + for (int i = initial_rows; i < ps.n_rows; ++i) { + cuvs::neighbors::tiered_index::extend( + handle_, + raft::make_device_matrix_view( + (const value_type*)database.data() + i * ps.dim, 1, ps.dim), + &index); + } + final_index.emplace(index); + } else if (ps.test_strategy == TEST_MERGE) { + // test merge by creating a new index with the remaining vectors and then merge into the + // main index + auto database_view2 = raft::make_device_matrix_view( + (const value_type*)database.data() + initial_rows * ps.dim, + ps.n_rows - initial_rows, + ps.dim); + auto index2 = cuvs::neighbors::tiered_index::build(handle_, build_params, database_view2); + + std::vector*> indices = {&index, &index2}; + final_index.emplace(cuvs::neighbors::tiered_index::merge(handle_, build_params, indices)); + + } else { + RAFT_FAIL("unknown test_strategy"); } raft::resource::sync_stream(handle_); @@ -122,23 +166,13 @@ class ANNTieredIndexTest : public ::testing::TestWithParam rmm::device_uvector distances_tiered_dev(ps.n_queries * ps.k, stream_); rmm::device_uvector indices_tiered_dev(ps.n_queries * ps.k, stream_); - typename UpstreamT::search_params_type search_params; - - if constexpr (std::is_same_v) { - search_params.n_probes = 128; - } - if constexpr (std::is_same_v) { - search_params.n_probes = 128; - } - auto distances_view = raft::make_device_matrix_view( (value_type*)distances_tiered_dev.data(), ps.n_queries, ps.k); auto indices_view = raft::make_device_matrix_view( (int64_t*)indices_tiered_dev.data(), ps.n_queries, ps.k); cuvs::neighbors::tiered_index::search( - handle_, search_params, index, queries_view, indices_view, distances_view); + handle_, search_params, *final_index, queries_view, indices_view, distances_view); raft::update_host( distances_tiered.data(), distances_tiered_dev.data(), ps.n_queries * ps.k, stream_); @@ -147,7 +181,6 @@ class ANNTieredIndexTest : public ::testing::TestWithParam raft::resource::sync_stream(handle_); } - double min_recall = ps.min_recall; EXPECT_TRUE(eval_neighbours(indices_naive, indices_tiered, distances_naive, @@ -193,7 +226,7 @@ const std::vector inputs = cuvs::distance::DistanceType::InnerProduct}, // metric {10}, // k {10}, // n_queries - {0.85} // min_recall + {TEST_EXTEND, TEST_MERGE} // test_strategy ); typedef ANNTieredIndexTest> CAGRA_F; TEST_P(CAGRA_F, AnnTieredIndex) { this->testTieredIndex(); } From 77d6b6e2eb012c38e547afd2c48204f04a506b4d Mon Sep 17 00:00:00 2001 From: Ben Frederickson Date: Fri, 15 Aug 2025 15:02:02 -0700 Subject: [PATCH 34/93] Run `cargo fmt` in the pre-commit hooks (#1255) Authors: - Ben Frederickson (https://github.com/benfred) Approvers: - Corey J. Nolet (https://github.com/cjnolet) - James Lamb (https://github.com/jameslamb) URL: https://github.com/rapidsai/cuvs/pull/1255 --- .pre-commit-config.yaml | 6 ++++++ rust/cuvs-sys/build.rs | 5 ++--- rust/cuvs/src/brute_force.rs | 4 ++-- rust/cuvs/src/distance/mod.rs | 14 ++++++++++---- rust/cuvs/src/ivf_flat/index.rs | 8 ++++---- rust/cuvs/src/ivf_flat/index_params.rs | 4 ++-- rust/cuvs/src/ivf_pq/index_params.rs | 4 ++-- rust/cuvs/src/lib.rs | 4 ++-- 8 files changed, 30 insertions(+), 19 deletions(-) diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 57674fa1ac..05489ee94d 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -92,6 +92,12 @@ repos: pass_filenames: false language: python additional_dependencies: [gitpython] + - id: cargo-fmt + name: cargo-fmt + entry: cargo fmt --manifest-path rust/Cargo.toml --all + pass_filenames: false + files: rust/.* + language: rust - repo: https://github.com/codespell-project/codespell rev: v2.2.2 hooks: diff --git a/rust/cuvs-sys/build.rs b/rust/cuvs-sys/build.rs index bcb0a88c7d..2b9a98cee4 100644 --- a/rust/cuvs-sys/build.rs +++ b/rust/cuvs-sys/build.rs @@ -1,5 +1,5 @@ /* - * Copyright (c) 2024, NVIDIA CORPORATION. + * Copyright (c) 2024-2025, NVIDIA CORPORATION. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -20,8 +20,7 @@ use std::path::PathBuf; fn main() { // build the cuvs c-api library with cmake, and link it into this crate - let cuvs_build = cmake::Config::new(".") - .build(); + let cuvs_build = cmake::Config::new(".").build(); println!( "cargo:rustc-link-search=native={}/lib", diff --git a/rust/cuvs/src/brute_force.rs b/rust/cuvs/src/brute_force.rs index 63752ee9c1..dd5e93d204 100644 --- a/rust/cuvs/src/brute_force.rs +++ b/rust/cuvs/src/brute_force.rs @@ -1,5 +1,5 @@ /* - * Copyright (c) 2024, NVIDIA CORPORATION. + * Copyright (c) 2024-2025, NVIDIA CORPORATION. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -109,10 +109,10 @@ impl Drop for Index { #[cfg(test)] mod tests { use super::*; + use mark_flaky_tests::flaky; use ndarray::s; use ndarray_rand::rand_distr::Uniform; use ndarray_rand::RandomExt; - use mark_flaky_tests::flaky; fn test_bfknn(metric: DistanceType) { let res = Resources::new().unwrap(); diff --git a/rust/cuvs/src/distance/mod.rs b/rust/cuvs/src/distance/mod.rs index f6911c6836..c29a4dd350 100644 --- a/rust/cuvs/src/distance/mod.rs +++ b/rust/cuvs/src/distance/mod.rs @@ -1,5 +1,5 @@ /* - * Copyright (c) 2024, NVIDIA CORPORATION. + * Copyright (c) 2024-2025, NVIDIA CORPORATION. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -14,7 +14,6 @@ * limitations under the License. */ - use crate::distance_type::DistanceType; use crate::dlpack::ManagedTensor; use crate::error::{check_cuvs, Result}; @@ -72,8 +71,15 @@ mod tests { .to_device(&res) .unwrap(); - pairwise_distance(&res, &dataset_device, &dataset_device, &distances, DistanceType::L2Expanded, - None).unwrap(); + pairwise_distance( + &res, + &dataset_device, + &dataset_device, + &distances, + DistanceType::L2Expanded, + None, + ) + .unwrap(); // Copy back to host memory distances.to_host(&res, &mut distances_host).unwrap(); diff --git a/rust/cuvs/src/ivf_flat/index.rs b/rust/cuvs/src/ivf_flat/index.rs index f66904142b..5892959557 100644 --- a/rust/cuvs/src/ivf_flat/index.rs +++ b/rust/cuvs/src/ivf_flat/index.rs @@ -1,5 +1,5 @@ /* - * Copyright (c) 2024, NVIDIA CORPORATION. + * Copyright (c) 2024-2025, NVIDIA CORPORATION. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -16,9 +16,9 @@ use std::io::{stderr, Write}; -use crate::ivf_flat::{IndexParams, SearchParams}; use crate::dlpack::ManagedTensor; use crate::error::{check_cuvs, Result}; +use crate::ivf_flat::{IndexParams, SearchParams}; use crate::resources::Resources; /// Ivf-Flat ANN Index @@ -127,8 +127,8 @@ mod tests { let dataset_device = ManagedTensor::from(&dataset).to_device(&res).unwrap(); // build the ivf-flat index - let index = - Index::build(&res, &build_params, dataset_device).expect("failed to create ivf-flat index"); + let index = Index::build(&res, &build_params, dataset_device) + .expect("failed to create ivf-flat index"); // use the first 4 points from the dataset as queries : will test that we get them back // as their own nearest neighbor diff --git a/rust/cuvs/src/ivf_flat/index_params.rs b/rust/cuvs/src/ivf_flat/index_params.rs index 6c86ef527f..266de3bdb0 100644 --- a/rust/cuvs/src/ivf_flat/index_params.rs +++ b/rust/cuvs/src/ivf_flat/index_params.rs @@ -1,5 +1,5 @@ /* - * Copyright (c) 2024, NVIDIA CORPORATION. + * Copyright (c) 2024-2025, NVIDIA CORPORATION. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -14,8 +14,8 @@ * limitations under the License. */ -use crate::error::{check_cuvs, Result}; use crate::distance_type::DistanceType; +use crate::error::{check_cuvs, Result}; use std::fmt; use std::io::{stderr, Write}; diff --git a/rust/cuvs/src/ivf_pq/index_params.rs b/rust/cuvs/src/ivf_pq/index_params.rs index 321821bc32..46d1f81d3e 100644 --- a/rust/cuvs/src/ivf_pq/index_params.rs +++ b/rust/cuvs/src/ivf_pq/index_params.rs @@ -1,5 +1,5 @@ /* - * Copyright (c) 2024, NVIDIA CORPORATION. + * Copyright (c) 2024-2025, NVIDIA CORPORATION. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -129,7 +129,7 @@ impl IndexParams { /// time. The parameter is applied to both PQ codebook generation methods, i.e., PER_SUBSPACE and /// PER_CLUSTER. In both cases, we will use `pq_book_size * max_train_points_per_pq_code` training /// points to train each codebook. - pub fn set_max_train_points_per_pq_code(self, max_pq_points: u32)-> IndexParams { + pub fn set_max_train_points_per_pq_code(self, max_pq_points: u32) -> IndexParams { unsafe { (*self.0).max_train_points_per_pq_code = max_pq_points; } diff --git a/rust/cuvs/src/lib.rs b/rust/cuvs/src/lib.rs index 71a387b2c6..c2059b6097 100644 --- a/rust/cuvs/src/lib.rs +++ b/rust/cuvs/src/lib.rs @@ -1,5 +1,5 @@ /* - * Copyright (c) 2024, NVIDIA CORPORATION. + * Copyright (c) 2024-2025, NVIDIA CORPORATION. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -22,9 +22,9 @@ pub mod brute_force; pub mod cagra; pub mod distance; pub mod distance_type; -pub mod ivf_flat; mod dlpack; mod error; +pub mod ivf_flat; pub mod ivf_pq; mod resources; From 2dcf6d8b245669c905c37170dfa004198cb928ee Mon Sep 17 00:00:00 2001 From: Jinsol Park Date: Mon, 18 Aug 2025 12:23:45 -0700 Subject: [PATCH 35/93] Removing deprecated batching code specific to NN Descent (#1249) Closing https://github.com/rapidsai/cuvs/issues/1051 Authors: - Jinsol Park (https://github.com/jinsolp) Approvers: - Corey J. Nolet (https://github.com/cjnolet) URL: https://github.com/rapidsai/cuvs/pull/1249 --- cpp/include/cuvs/neighbors/nn_descent.h | 3 +- cpp/include/cuvs/neighbors/nn_descent.hpp | 1 - cpp/src/neighbors/detail/nn_descent_batch.cuh | 737 ------------------ cpp/src/neighbors/nn_descent.cuh | 59 +- cpp/src/neighbors/nn_descent_c.cpp | 4 +- cpp/tests/neighbors/ann_nn_descent.cuh | 20 - .../ann_nn_descent/test_float_uint32_t.cu | 6 - .../ann_nn_descent/test_uint8_t_uint32_t.cu | 8 - .../cuvs/neighbors/nn_descent/nn_descent.pxd | 1 - .../cuvs/neighbors/nn_descent/nn_descent.pyx | 9 +- 10 files changed, 8 insertions(+), 840 deletions(-) delete mode 100644 cpp/src/neighbors/detail/nn_descent_batch.cuh diff --git a/cpp/include/cuvs/neighbors/nn_descent.h b/cpp/include/cuvs/neighbors/nn_descent.h index 7dd87f2248..118b1c26aa 100644 --- a/cpp/include/cuvs/neighbors/nn_descent.h +++ b/cpp/include/cuvs/neighbors/nn_descent.h @@ -1,5 +1,5 @@ /* - * Copyright (c) 2024, NVIDIA CORPORATION. + * Copyright (c) 2024-2025, NVIDIA CORPORATION. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -54,7 +54,6 @@ struct cuvsNNDescentIndexParams { size_t max_iterations; float termination_threshold; bool return_distances; - size_t n_clusters; }; typedef struct cuvsNNDescentIndexParams* cuvsNNDescentIndexParams_t; diff --git a/cpp/include/cuvs/neighbors/nn_descent.hpp b/cpp/include/cuvs/neighbors/nn_descent.hpp index 5f6260b0af..ba1d417587 100644 --- a/cpp/include/cuvs/neighbors/nn_descent.hpp +++ b/cpp/include/cuvs/neighbors/nn_descent.hpp @@ -55,7 +55,6 @@ struct index_params : cuvs::neighbors::index_params { size_t max_iterations = 20; float termination_threshold = 0.0001; bool return_distances = true; - size_t n_clusters = 1; /** @brief Construct NN descent parameters for a specific kNN graph degree * diff --git a/cpp/src/neighbors/detail/nn_descent_batch.cuh b/cpp/src/neighbors/detail/nn_descent_batch.cuh deleted file mode 100644 index c0726409a2..0000000000 --- a/cpp/src/neighbors/detail/nn_descent_batch.cuh +++ /dev/null @@ -1,737 +0,0 @@ -/* - * Copyright (c) 2024, NVIDIA CORPORATION. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -#include -#include -#include -#undef RAFT_EXPLICIT_INSTANTIATE_ONLY - -#include "nn_descent_gnnd.hpp" -#include -#include - -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include - -#include - -#include - -#include -#include -#include -#include -#include -#include - -namespace cuvs::neighbors::nn_descent::detail::experimental { - -// -// Run balanced kmeans on a subsample of the dataset to get centroids -// -template , memory_type::host>> -void get_balanced_kmeans_centroids( - raft::resources const& res, - cuvs::distance::DistanceType metric, - mdspan, row_major, Accessor> dataset, - raft::device_matrix_view centroids) -{ - size_t num_rows = static_cast(dataset.extent(0)); - size_t num_cols = static_cast(dataset.extent(1)); - size_t n_clusters = centroids.extent(0); - size_t num_subsamples = - std::min(static_cast(num_rows / n_clusters), static_cast(num_rows * 0.1)); - - auto d_subsample_dataset = - raft::make_device_matrix(res, num_subsamples, num_cols); - raft::matrix::sample_rows( - res, raft::random::RngState{0}, dataset, d_subsample_dataset.view()); - - cuvs::cluster::kmeans::balanced_params kmeans_params; - kmeans_params.metric = metric; - - auto d_subsample_dataset_const_view = - raft::make_device_matrix_view( - d_subsample_dataset.data_handle(), num_subsamples, num_cols); - auto centroids_view = raft::make_device_matrix_view( - centroids.data_handle(), n_clusters, num_cols); - cuvs::cluster::kmeans::fit(res, kmeans_params, d_subsample_dataset_const_view, centroids_view); -} - -// -// Get the top k closest centroid indices for each data point -// Loads the data in batches onto device if data is on host for memory efficiency -// -template -void get_global_nearest_k( - raft::resources const& res, - size_t k, - size_t num_rows, - size_t n_clusters, - const T* dataset, - raft::host_matrix_view global_nearest_cluster, - raft::device_matrix_view centroids, - cuvs::distance::DistanceType metric) -{ - size_t num_cols = centroids.extent(1); - auto centroids_view = raft::make_device_matrix_view( - centroids.data_handle(), n_clusters, num_cols); - - cudaPointerAttributes attr; - RAFT_CUDA_TRY(cudaPointerGetAttributes(&attr, dataset)); - float* ptr = reinterpret_cast(attr.devicePointer); - - size_t num_batches = n_clusters; - size_t batch_size = (num_rows + n_clusters) / n_clusters; - if (ptr == nullptr) { // data on host - - auto d_dataset_batch = - raft::make_device_matrix(res, batch_size, num_cols); - - auto nearest_clusters_idx = - raft::make_device_matrix(res, batch_size, k); - auto nearest_clusters_idxt = - raft::make_device_matrix(res, batch_size, k); - auto nearest_clusters_dist = - raft::make_device_matrix(res, batch_size, k); - - for (size_t i = 0; i < num_batches; i++) { - size_t batch_size_ = batch_size; - - if (i == num_batches - 1) { batch_size_ = num_rows - batch_size * i; } - raft::copy(d_dataset_batch.data_handle(), - dataset + i * batch_size * num_cols, - batch_size_ * num_cols, - resource::get_cuda_stream(res)); - - std::optional> norms_view; - cuvs::neighbors::brute_force::index brute_force_index( - res, centroids_view, norms_view, metric); - cuvs::neighbors::brute_force::search(res, - brute_force_index, - raft::make_const_mdspan(d_dataset_batch.view()), - nearest_clusters_idx.view(), - nearest_clusters_dist.view()); - - thrust::copy(raft::resource::get_thrust_policy(res), - nearest_clusters_idx.data_handle(), - nearest_clusters_idx.data_handle() + nearest_clusters_idx.size(), - nearest_clusters_idxt.data_handle()); - raft::copy(global_nearest_cluster.data_handle() + i * batch_size * k, - nearest_clusters_idxt.data_handle(), - batch_size_ * k, - resource::get_cuda_stream(res)); - } - } else { // data on device - auto nearest_clusters_idx = - raft::make_device_matrix(res, num_rows, k); - auto nearest_clusters_dist = - raft::make_device_matrix(res, num_rows, k); - - std::optional> norms_view; - cuvs::neighbors::brute_force::index brute_force_index( - res, centroids_view, norms_view, metric); - auto dataset_view = - raft::make_device_matrix_view(dataset, num_rows, num_cols); - cuvs::neighbors::brute_force::search(res, - brute_force_index, - dataset_view, - nearest_clusters_idx.view(), - nearest_clusters_dist.view()); - - auto nearest_clusters_idxt = - raft::make_device_matrix(res, batch_size, k); - for (size_t i = 0; i < num_batches; i++) { - size_t batch_size_ = batch_size; - - if (i == num_batches - 1) { batch_size_ = num_rows - batch_size * i; } - thrust::copy(raft::resource::get_thrust_policy(res), - nearest_clusters_idx.data_handle() + i * batch_size * k, - nearest_clusters_idx.data_handle() + (i * batch_size + batch_size_) * k, - nearest_clusters_idxt.data_handle()); - raft::copy(global_nearest_cluster.data_handle() + i * batch_size * k, - nearest_clusters_idxt.data_handle(), - batch_size_ * k, - resource::get_cuda_stream(res)); - } - } -} - -// -// global_nearest_cluster [num_rows X k=2] : top 2 closest clusters for each data point -// inverted_indices [num_rows x k vector] : sparse vector for data indices for each cluster -// cluster_size [n_cluster] : cluster size for each cluster -// offset [n_cluster] : offset in inverted_indices for each cluster -// Loads the data in batches onto device if data is on host for memory efficiency -// -template -void get_inverted_indices(raft::resources const& res, - size_t n_clusters, - size_t& max_cluster_size, - size_t& min_cluster_size, - raft::host_matrix_view global_nearest_cluster, - raft::host_vector_view inverted_indices, - raft::host_vector_view cluster_size, - raft::host_vector_view offset) -{ - // build sparse inverted indices and get number of data points for each cluster - size_t num_rows = global_nearest_cluster.extent(0); - size_t k = global_nearest_cluster.extent(1); - - auto local_offset = raft::make_host_vector(n_clusters); - - max_cluster_size = 0; - min_cluster_size = std::numeric_limits::max(); - - std::fill(cluster_size.data_handle(), cluster_size.data_handle() + n_clusters, 0); - std::fill(local_offset.data_handle(), local_offset.data_handle() + n_clusters, 0); - - // TODO: this part isn't really a bottleneck but maybe worth trying omp parallel - // for with atomic add - for (size_t i = 0; i < num_rows; i++) { - for (size_t j = 0; j < k; j++) { - IdxT cluster_id = global_nearest_cluster(i, j); - cluster_size(cluster_id) += 1; - } - } - - offset(0) = 0; - for (size_t i = 1; i < n_clusters; i++) { - offset(i) = offset(i - 1) + cluster_size(i - 1); - } - for (size_t i = 0; i < num_rows; i++) { - for (size_t j = 0; j < k; j++) { - IdxT cluster_id = global_nearest_cluster(i, j); - inverted_indices(offset(cluster_id) + local_offset(cluster_id)) = i; - local_offset(cluster_id) += 1; - } - } - - max_cluster_size = static_cast( - *std::max_element(cluster_size.data_handle(), cluster_size.data_handle() + n_clusters)); - min_cluster_size = static_cast( - *std::min_element(cluster_size.data_handle(), cluster_size.data_handle() + n_clusters)); -} - -template -struct KeyValuePair { - KeyType key; - ValueType value; -}; - -template -struct CustomKeyComparator { - __device__ bool operator()(const KeyValuePair& a, - const KeyValuePair& b) const - { - if (a.key == b.key) { return a.value < b.value; } - return a.key < b.key; - } -}; - -template -RAFT_KERNEL merge_subgraphs(IdxT* cluster_data_indices, - size_t graph_degree, - size_t num_cluster_in_batch, - float* global_distances, - float* batch_distances, - IdxT* global_indices, - IdxT* batch_indices) -{ - size_t batch_row = blockIdx.x; - typedef cub::BlockMergeSort, BLOCK_SIZE, ITEMS_PER_THREAD> - BlockMergeSortType; - __shared__ typename cub::BlockMergeSort, BLOCK_SIZE, ITEMS_PER_THREAD>:: - TempStorage tmpSmem; - - extern __shared__ char sharedMem[]; - float* blockKeys = reinterpret_cast(sharedMem); - IdxT* blockValues = reinterpret_cast(&sharedMem[graph_degree * 2 * sizeof(float)]); - int16_t* uniqueMask = - reinterpret_cast(&sharedMem[graph_degree * 2 * (sizeof(float) + sizeof(IdxT))]); - - if (batch_row < num_cluster_in_batch) { - // load batch or global depending on threadIdx - size_t global_row = cluster_data_indices[batch_row]; - - KeyValuePair threadKeyValuePair[ITEMS_PER_THREAD]; - - size_t halfway = BLOCK_SIZE / 2; - size_t do_global = threadIdx.x < halfway; - - float* distances; - IdxT* indices; - - if (do_global) { - distances = global_distances; - indices = global_indices; - } else { - distances = batch_distances; - indices = batch_indices; - } - - size_t idxBase = (threadIdx.x * do_global + (threadIdx.x - halfway) * (1lu - do_global)) * - static_cast(ITEMS_PER_THREAD); - size_t arrIdxBase = (global_row * do_global + batch_row * (1lu - do_global)) * graph_degree; - for (int i = 0; i < ITEMS_PER_THREAD; i++) { - size_t colId = idxBase + i; - if (colId < graph_degree) { - threadKeyValuePair[i].key = distances[arrIdxBase + colId]; - threadKeyValuePair[i].value = indices[arrIdxBase + colId]; - } else { - threadKeyValuePair[i].key = std::numeric_limits::max(); - threadKeyValuePair[i].value = std::numeric_limits::max(); - } - } - - __syncthreads(); - - BlockMergeSortType(tmpSmem).Sort(threadKeyValuePair, CustomKeyComparator{}); - - // load sorted result into shared memory to get unique values - idxBase = threadIdx.x * ITEMS_PER_THREAD; - for (int i = 0; i < ITEMS_PER_THREAD; i++) { - size_t colId = idxBase + i; - if (colId < 2 * graph_degree) { - blockKeys[colId] = threadKeyValuePair[i].key; - blockValues[colId] = threadKeyValuePair[i].value; - } - } - - __syncthreads(); - - // get unique mask - if (threadIdx.x == 0) { uniqueMask[0] = 1; } - for (int i = 0; i < ITEMS_PER_THREAD; i++) { - size_t colId = idxBase + i; - if (colId > 0 && colId < 2 * graph_degree) { - uniqueMask[colId] = static_cast(blockValues[colId] != blockValues[colId - 1]); - } - } - - __syncthreads(); - - // prefix sum - if (threadIdx.x == 0) { - for (int i = 1; i < 2 * graph_degree; i++) { - uniqueMask[i] += uniqueMask[i - 1]; - } - } - - __syncthreads(); - // load unique values to global memory - if (threadIdx.x == 0) { - global_distances[global_row * graph_degree] = blockKeys[0]; - global_indices[global_row * graph_degree] = blockValues[0]; - } - - for (int i = 0; i < ITEMS_PER_THREAD; i++) { - size_t colId = idxBase + i; - if (colId > 0 && colId < 2 * graph_degree) { - bool is_unique = uniqueMask[colId] != uniqueMask[colId - 1]; - int16_t global_colId = uniqueMask[colId] - 1; - if (is_unique && static_cast(global_colId) < graph_degree) { - global_distances[global_row * graph_degree + global_colId] = blockKeys[colId]; - global_indices[global_row * graph_degree + global_colId] = blockValues[colId]; - } - } - } - } -} - -// -// builds knn graph using NN Descent and merge with global graph -// -template , memory_type::host>> -void build_and_merge(raft::resources const& res, - const index_params& params, - size_t num_data_in_cluster, - size_t graph_degree, - size_t int_graph_node_degree, - T* cluster_data, - IdxT* cluster_data_indices, - int* int_graph, - IdxT* inverted_indices, - IdxT* global_indices_d, - float* global_distances_d, - IdxT* batch_indices_h, - IdxT* batch_indices_d, - float* batch_distances_d, - GNND& nnd) -{ - nnd.build(cluster_data, num_data_in_cluster, int_graph, true, batch_distances_d); - - // remap indices -#pragma omp parallel for - for (size_t i = 0; i < num_data_in_cluster; i++) { - for (size_t j = 0; j < graph_degree; j++) { - size_t local_idx = int_graph[i * int_graph_node_degree + j]; - batch_indices_h[i * graph_degree + j] = inverted_indices[local_idx]; - } - } - - raft::copy(batch_indices_d, - batch_indices_h, - num_data_in_cluster * graph_degree, - raft::resource::get_cuda_stream(res)); - - size_t num_elems = graph_degree * 2; - size_t sharedMemSize = num_elems * (sizeof(float) + sizeof(IdxT) + sizeof(int16_t)); - - if (num_elems <= 128) { - merge_subgraphs - <<>>( - cluster_data_indices, - graph_degree, - num_data_in_cluster, - global_distances_d, - batch_distances_d, - global_indices_d, - batch_indices_d); - } else if (num_elems <= 512) { - merge_subgraphs - <<>>( - cluster_data_indices, - graph_degree, - num_data_in_cluster, - global_distances_d, - batch_distances_d, - global_indices_d, - batch_indices_d); - } else if (num_elems <= 1024) { - merge_subgraphs - <<>>( - cluster_data_indices, - graph_degree, - num_data_in_cluster, - global_distances_d, - batch_distances_d, - global_indices_d, - batch_indices_d); - } else if (num_elems <= 2048) { - merge_subgraphs - <<>>( - cluster_data_indices, - graph_degree, - num_data_in_cluster, - global_distances_d, - batch_distances_d, - global_indices_d, - batch_indices_d); - } else { - // this is as far as we can get due to the shared mem usage of cub::BlockMergeSort - RAFT_FAIL("The degree of knn is too large (%lu). It must be smaller than 1024", graph_degree); - } - raft::resource::sync_stream(res); -} - -// -// For each cluster, gather the data samples that belong to that cluster, and -// call build_and_merge -// -template -void cluster_nnd(raft::resources const& res, - const index_params& params, - size_t graph_degree, - size_t extended_graph_degree, - size_t max_cluster_size, - raft::host_matrix_view dataset, - IdxT* offsets, - IdxT* cluster_size, - IdxT* cluster_data_indices, - int* int_graph, - IdxT* inverted_indices, - IdxT* global_indices_h, - float* global_distances_h, - IdxT* batch_indices_h, - IdxT* batch_indices_d, - float* batch_distances_d, - const BuildConfig& build_config) -{ - size_t num_rows = dataset.extent(0); - size_t num_cols = dataset.extent(1); - - GNND nnd(res, build_config); - - auto cluster_data_matrix = - raft::make_host_matrix(max_cluster_size, num_cols); - - for (size_t cluster_id = 0; cluster_id < params.n_clusters; cluster_id++) { - RAFT_LOG_DEBUG( - "# Data on host. Running clusters: %lu / %lu", cluster_id + 1, params.n_clusters); - size_t num_data_in_cluster = cluster_size[cluster_id]; - size_t offset = offsets[cluster_id]; - -#pragma omp parallel for - for (size_t i = 0; i < num_data_in_cluster; i++) { - for (size_t j = 0; j < num_cols; j++) { - size_t global_row = (inverted_indices + offset)[i]; - cluster_data_matrix(i, j) = dataset(global_row, j); - } - } - - build_and_merge(res, - params, - num_data_in_cluster, - graph_degree, - extended_graph_degree, - cluster_data_matrix.data_handle(), - cluster_data_indices + offset, - int_graph, - inverted_indices + offset, - global_indices_h, - global_distances_h, - batch_indices_h, - batch_indices_d, - batch_distances_d, - nnd); - nnd.reset(res); - } -} - -template -void cluster_nnd(raft::resources const& res, - const index_params& params, - size_t graph_degree, - size_t extended_graph_degree, - size_t max_cluster_size, - raft::device_matrix_view dataset, - IdxT* offsets, - IdxT* cluster_size, - IdxT* cluster_data_indices, - int* int_graph, - IdxT* inverted_indices, - IdxT* global_indices_h, - float* global_distances_h, - IdxT* batch_indices_h, - IdxT* batch_indices_d, - float* batch_distances_d, - const BuildConfig& build_config) -{ - size_t num_rows = dataset.extent(0); - size_t num_cols = dataset.extent(1); - - GNND nnd(res, build_config); - - auto cluster_data_matrix = - raft::make_device_matrix(res, max_cluster_size, num_cols); - - for (size_t cluster_id = 0; cluster_id < params.n_clusters; cluster_id++) { - RAFT_LOG_DEBUG( - "# Data on device. Running clusters: %lu / %lu", cluster_id + 1, params.n_clusters); - size_t num_data_in_cluster = cluster_size[cluster_id]; - size_t offset = offsets[cluster_id]; - - auto cluster_data_view = raft::make_device_matrix_view( - cluster_data_matrix.data_handle(), num_data_in_cluster, num_cols); - auto cluster_data_indices_view = raft::make_device_vector_view( - cluster_data_indices + offset, num_data_in_cluster); - - auto dataset_IdxT = - raft::make_device_matrix_view(dataset.data_handle(), num_rows, num_cols); - raft::matrix::gather(res, dataset_IdxT, cluster_data_indices_view, cluster_data_view); - - build_and_merge(res, - params, - num_data_in_cluster, - graph_degree, - extended_graph_degree, - cluster_data_view.data_handle(), - cluster_data_indices + offset, - int_graph, - inverted_indices + offset, - global_indices_h, - global_distances_h, - batch_indices_h, - batch_indices_d, - batch_distances_d, - nnd); - nnd.reset(res); - } -} - -template , memory_type::host>> -void batch_build(raft::resources const& res, - const index_params& params, - mdspan, row_major, Accessor> dataset, - index& global_idx) -{ - size_t graph_degree = params.graph_degree; - size_t intermediate_degree = params.intermediate_graph_degree; - - size_t num_rows = static_cast(dataset.extent(0)); - size_t num_cols = static_cast(dataset.extent(1)); - - auto centroids = - raft::make_device_matrix(res, params.n_clusters, num_cols); - get_balanced_kmeans_centroids(res, params.metric, dataset, centroids.view()); - - size_t k = 2; - auto global_nearest_cluster = raft::make_host_matrix(num_rows, k); - get_global_nearest_k(res, - k, - num_rows, - params.n_clusters, - dataset.data_handle(), - global_nearest_cluster.view(), - centroids.view(), - params.metric); - - auto inverted_indices = raft::make_host_vector(num_rows * k); - auto cluster_size = raft::make_host_vector(params.n_clusters); - auto offset = raft::make_host_vector(params.n_clusters); - - size_t max_cluster_size, min_cluster_size; - get_inverted_indices(res, - params.n_clusters, - max_cluster_size, - min_cluster_size, - global_nearest_cluster.view(), - inverted_indices.view(), - cluster_size.view(), - offset.view()); - - if (intermediate_degree >= min_cluster_size) { - RAFT_LOG_WARN( - "Intermediate graph degree cannot be larger than minimum cluster size, reducing it to %lu", - dataset.extent(0)); - intermediate_degree = min_cluster_size - 1; - } - if (intermediate_degree < graph_degree) { - RAFT_LOG_WARN( - "Graph degree (%lu) cannot be larger than intermediate graph degree (%lu), reducing " - "graph_degree.", - graph_degree, - intermediate_degree); - graph_degree = intermediate_degree; - } - - size_t extended_graph_degree = - roundUp32(static_cast(graph_degree * (graph_degree <= 32 ? 1.0 : 1.3))); - size_t extended_intermediate_degree = - roundUp32(static_cast(intermediate_degree * (intermediate_degree <= 32 ? 1.0 : 1.3))); - - auto int_graph = raft::make_host_matrix( - max_cluster_size, static_cast(extended_graph_degree)); - - BuildConfig build_config{.max_dataset_size = max_cluster_size, - .dataset_dim = num_cols, - .node_degree = extended_graph_degree, - .internal_node_degree = extended_intermediate_degree, - .max_iterations = params.max_iterations, - .termination_threshold = params.termination_threshold, - .output_graph_degree = graph_degree, - .metric = params.metric}; - - auto global_indices_h = raft::make_managed_matrix(res, num_rows, graph_degree); - auto global_distances_h = raft::make_managed_matrix(res, num_rows, graph_degree); - - std::fill(global_indices_h.data_handle(), - global_indices_h.data_handle() + num_rows * graph_degree, - std::numeric_limits::max()); - std::fill(global_distances_h.data_handle(), - global_distances_h.data_handle() + num_rows * graph_degree, - std::numeric_limits::max()); - - auto batch_indices_h = - raft::make_host_matrix(max_cluster_size, graph_degree); - auto batch_indices_d = - raft::make_device_matrix(res, max_cluster_size, graph_degree); - auto batch_distances_d = - raft::make_device_matrix(res, max_cluster_size, graph_degree); - - auto cluster_data_indices = raft::make_device_vector(res, num_rows * k); - raft::copy(cluster_data_indices.data_handle(), - inverted_indices.data_handle(), - num_rows * k, - resource::get_cuda_stream(res)); - - cluster_nnd(res, - params, - graph_degree, - extended_graph_degree, - max_cluster_size, - dataset, - offset.data_handle(), - cluster_size.data_handle(), - cluster_data_indices.data_handle(), - int_graph.data_handle(), - inverted_indices.data_handle(), - global_indices_h.data_handle(), - global_distances_h.data_handle(), - batch_indices_h.data_handle(), - batch_indices_d.data_handle(), - batch_distances_d.data_handle(), - build_config); - - raft::copy(global_idx.graph().data_handle(), - global_indices_h.data_handle(), - num_rows * graph_degree, - raft::resource::get_cuda_stream(res)); - if (params.return_distances && global_idx.distances().has_value()) { - raft::copy(global_idx.distances().value().data_handle(), - global_distances_h.data_handle(), - num_rows * graph_degree, - raft::resource::get_cuda_stream(res)); - } -} - -template , memory_type::host>> -index batch_build(raft::resources const& res, - const index_params& params, - mdspan, row_major, Accessor> dataset) -{ - size_t intermediate_degree = params.intermediate_graph_degree; - size_t graph_degree = params.graph_degree; - - if (intermediate_degree < graph_degree) { - RAFT_LOG_WARN( - "Graph degree (%lu) cannot be larger than intermediate graph degree (%lu), reducing " - "graph_degree.", - graph_degree, - intermediate_degree); - graph_degree = intermediate_degree; - } - - index idx{ - res, dataset.extent(0), static_cast(graph_degree), params.return_distances}; - - batch_build(res, params, dataset, idx); - - return idx; -} - -} // namespace cuvs::neighbors::nn_descent::detail::experimental diff --git a/cpp/src/neighbors/nn_descent.cuh b/cpp/src/neighbors/nn_descent.cuh index 78c63dc31e..4f0bf01cfa 100644 --- a/cpp/src/neighbors/nn_descent.cuh +++ b/cpp/src/neighbors/nn_descent.cuh @@ -17,14 +17,11 @@ #pragma once #include "detail/nn_descent.cuh" -#include "detail/nn_descent_batch.cuh" -#include #include #include #include -#include #include namespace cuvs::neighbors::nn_descent { @@ -70,19 +67,7 @@ auto build(raft::resources const& res, index_params const& params, raft::device_matrix_view dataset) -> index { - if (params.n_clusters > 1) { - // related issue: https://github.com/rapidsai/cuvs/issues/1051 - RAFT_LOG_WARN( - "NN Descent batch build (using n_clusters > 1) is deprecated and will be removed in a future " - "release. Please use cuvs::all_neighbors::build(...) instead."); - if constexpr (std::is_same_v) { - return detail::experimental::batch_build(res, params, dataset); - } else { - RAFT_FAIL("Batched nn-descent is only supported for float precision"); - } - } else { - return detail::build(res, params, dataset); - } + return detail::build(res, params, dataset); } /** @@ -125,19 +110,7 @@ void build(raft::resources const& res, raft::device_matrix_view dataset, index& idx) { - if (params.n_clusters > 1) { - // related issue: https://github.com/rapidsai/cuvs/issues/1051 - RAFT_LOG_WARN( - "NN Descent batch build (using n_clusters > 1) is deprecated and will be removed in a future " - "release. Please use cuvs::all_neighbors::build(...) instead."); - if constexpr (std::is_same_v) { - detail::experimental::batch_build(res, params, dataset, idx); - } else { - RAFT_FAIL("Batched nn-descent is only supported for float precision"); - } - } else { - detail::build(res, params, dataset, idx); - } + detail::build(res, params, dataset, idx); } /** @@ -176,19 +149,7 @@ auto build(raft::resources const& res, index_params const& params, raft::host_matrix_view dataset) -> index { - if (params.n_clusters > 1) { - // related issue: https://github.com/rapidsai/cuvs/issues/1051 - RAFT_LOG_WARN( - "NN Descent batch build (using n_clusters > 1) is deprecated and will be removed in a future " - "release. Please use cuvs::all_neighbors::build(...) instead."); - if constexpr (std::is_same_v) { - return detail::experimental::batch_build(res, params, dataset); - } else { - RAFT_FAIL("Batched nn-descent is only supported for float precision"); - } - } else { - return detail::build(res, params, dataset); - } + return detail::build(res, params, dataset); } /** @@ -231,19 +192,7 @@ void build(raft::resources const& res, raft::host_matrix_view dataset, index& idx) { - if (params.n_clusters > 1) { - // related issue: https://github.com/rapidsai/cuvs/issues/1051 - RAFT_LOG_WARN( - "NN Descent batch build (using n_clusters > 1) is deprecated and will be removed in a future " - "release. Please use cuvs::all_neighbors::build(...) instead."); - if constexpr (std::is_same_v) { - detail::experimental::batch_build(res, params, dataset, idx); - } else { - RAFT_FAIL("Batched nn-descent is only supported for float precision"); - } - } else { - detail::build(res, params, dataset, idx); - } + detail::build(res, params, dataset, idx); } /** @} */ // end group nn-descent diff --git a/cpp/src/neighbors/nn_descent_c.cpp b/cpp/src/neighbors/nn_descent_c.cpp index ddebb11e73..b0a19879c4 100644 --- a/cpp/src/neighbors/nn_descent_c.cpp +++ b/cpp/src/neighbors/nn_descent_c.cpp @@ -52,7 +52,6 @@ void* _build(cuvsResources_t res, build_params.max_iterations = params.max_iterations; build_params.termination_threshold = params.termination_threshold; build_params.return_distances = params.return_distances; - build_params.n_clusters = params.n_clusters; using graph_type = raft::host_matrix_view; std::optional graph; @@ -187,8 +186,7 @@ extern "C" cuvsError_t cuvsNNDescentIndexParamsCreate(cuvsNNDescentIndexParams_t .intermediate_graph_degree = cpp_params.intermediate_graph_degree, .max_iterations = cpp_params.max_iterations, .termination_threshold = cpp_params.termination_threshold, - .return_distances = cpp_params.return_distances, - .n_clusters = cpp_params.n_clusters}; + .return_distances = cpp_params.return_distances}; }); } diff --git a/cpp/tests/neighbors/ann_nn_descent.cuh b/cpp/tests/neighbors/ann_nn_descent.cuh index e1c27dc48c..ad0bf617d8 100644 --- a/cpp/tests/neighbors/ann_nn_descent.cuh +++ b/cpp/tests/neighbors/ann_nn_descent.cuh @@ -400,7 +400,6 @@ class AnnNNDescentBatchTest : public ::testing::TestWithParam( (const DataT*)database.data(), ps.n_rows, ps.dim); @@ -487,16 +486,6 @@ const std::vector inputs = {false, true}, {0.90}); -// Additional test cases for large datasets to test batching -const std::vector inputsLargeBatch = - raft::util::itertools::product( - {150000}, // n_rows > 100000 to trigger batching - {64}, - {32}, - {cuvs::distance::DistanceType::BitwiseHamming}, - {true}, - {0.90}); - const std::vector inputsDistEpilogue = raft::util::itertools::product( {2000, 4000}, // n_rows @@ -506,13 +495,4 @@ const std::vector inputsDistEpilogue = {true}, // data on host {0.90}); -const std::vector inputsBatch = - raft::util::itertools::product( - {std::make_pair(0.9, 3lu), std::make_pair(0.9, 2lu)}, // min_recall, n_clusters - {4000, 5000}, // n_rows - {192, 512}, // dim - {32, 64}, // graph_degree - {cuvs::distance::DistanceType::L2Expanded}, - {false, true}); - } // namespace cuvs::neighbors::nn_descent diff --git a/cpp/tests/neighbors/ann_nn_descent/test_float_uint32_t.cu b/cpp/tests/neighbors/ann_nn_descent/test_float_uint32_t.cu index 01fa769493..aff8a2e62f 100644 --- a/cpp/tests/neighbors/ann_nn_descent/test_float_uint32_t.cu +++ b/cpp/tests/neighbors/ann_nn_descent/test_float_uint32_t.cu @@ -26,15 +26,9 @@ TEST_P(AnnNNDescentTestF_U32, AnnNNDescent) { this->testNNDescent(); } typedef AnnNNDescentDistEpiTest AnnNNDescentTestDistEpiF_U32; TEST_P(AnnNNDescentTestDistEpiF_U32, AnnNNDescentDistEpi) { this->testNNDescent(); } -typedef AnnNNDescentBatchTest AnnNNDescentBatchTestF_U32; -TEST_P(AnnNNDescentBatchTestF_U32, AnnNNDescentBatch) { this->testNNDescentBatch(); } - INSTANTIATE_TEST_CASE_P(AnnNNDescentTest, AnnNNDescentTestF_U32, ::testing::ValuesIn(inputs)); INSTANTIATE_TEST_CASE_P(AnnNNDescentDistEpi, AnnNNDescentTestDistEpiF_U32, ::testing::ValuesIn(inputsDistEpilogue)); -INSTANTIATE_TEST_CASE_P(AnnNNDescentBatchTest, - AnnNNDescentBatchTestF_U32, - ::testing::ValuesIn(inputsBatch)); } // namespace cuvs::neighbors::nn_descent diff --git a/cpp/tests/neighbors/ann_nn_descent/test_uint8_t_uint32_t.cu b/cpp/tests/neighbors/ann_nn_descent/test_uint8_t_uint32_t.cu index 16159ecc0a..92207b8db7 100644 --- a/cpp/tests/neighbors/ann_nn_descent/test_uint8_t_uint32_t.cu +++ b/cpp/tests/neighbors/ann_nn_descent/test_uint8_t_uint32_t.cu @@ -24,12 +24,4 @@ typedef AnnNNDescentTest AnnNNDescentTestUI8_U32; TEST_P(AnnNNDescentTestUI8_U32, AnnNNDescent) { this->testNNDescent(); } INSTANTIATE_TEST_CASE_P(AnnNNDescentTest, AnnNNDescentTestUI8_U32, ::testing::ValuesIn(inputs)); - -typedef AnnNNDescentTest AnnNNDescentTestUI8_U32_LargeBatch; -TEST_P(AnnNNDescentTestUI8_U32_LargeBatch, AnnNNDescent) { this->testNNDescent(); } - -INSTANTIATE_TEST_CASE_P(AnnNNDescentTestLargeBatch, - AnnNNDescentTestUI8_U32_LargeBatch, - ::testing::ValuesIn(inputsLargeBatch)); - } // namespace cuvs::neighbors::nn_descent diff --git a/python/cuvs/cuvs/neighbors/nn_descent/nn_descent.pxd b/python/cuvs/cuvs/neighbors/nn_descent/nn_descent.pxd index 6753d9f016..5d82b7e7b3 100644 --- a/python/cuvs/cuvs/neighbors/nn_descent/nn_descent.pxd +++ b/python/cuvs/cuvs/neighbors/nn_descent/nn_descent.pxd @@ -33,7 +33,6 @@ cdef extern from "cuvs/neighbors/nn_descent.h" nogil: size_t max_iterations float termination_threshold bool return_distances - size_t n_clusters ctypedef cuvsNNDescentIndexParams* cuvsNNDescentIndexParams_t diff --git a/python/cuvs/cuvs/neighbors/nn_descent/nn_descent.pyx b/python/cuvs/cuvs/neighbors/nn_descent/nn_descent.pyx index 0f0589eced..5c9d0d8155 100644 --- a/python/cuvs/cuvs/neighbors/nn_descent/nn_descent.pyx +++ b/python/cuvs/cuvs/neighbors/nn_descent/nn_descent.pyx @@ -72,6 +72,8 @@ cdef class IndexParams: More iterations produce a better quality graph at cost of performance termination_threshold : float The delta at which nn-descent will terminate its iterations + return_distances : bool + Whether to return distances array """ cdef cuvsNNDescentIndexParams* params @@ -90,7 +92,6 @@ cdef class IndexParams: intermediate_graph_degree=None, max_iterations=None, termination_threshold=None, - n_clusters=None, return_distances=None ): if metric is not None: @@ -103,8 +104,6 @@ cdef class IndexParams: self.params.max_iterations = max_iterations if termination_threshold is not None: self.params.termination_threshold = termination_threshold - if n_clusters is not None: - self.params.n_clusters = n_clusters if return_distances is not None: self.params.return_distances = return_distances @@ -132,10 +131,6 @@ cdef class IndexParams: def termination_threshold(self): return self.params.termination_threshold - @property - def n_clusters(self): - return self.params.n_clusters - cdef class Index: """ NN-Descent index object. This object stores the trained NN-Descent index, From 2922045b7beaf180940ec5e3732fdd70f336dbd0 Mon Sep 17 00:00:00 2001 From: Jinsol Park Date: Mon, 18 Aug 2025 16:41:47 -0700 Subject: [PATCH 36/93] Mutual Reachability in all-neighbors API (#1016) Closing https://github.com/rapidsai/cuvs/issues/856 Making mutual reachability part of all-neighbors. If `core_distances` optional matrix is given, automatically calculates distances and indices in mutual reachability space. Authors: - Jinsol Park (https://github.com/jinsolp) - Corey J. Nolet (https://github.com/cjnolet) Approvers: - Corey J. Nolet (https://github.com/cjnolet) URL: https://github.com/rapidsai/cuvs/pull/1016 --- cpp/include/cuvs/neighbors/all_neighbors.hpp | 16 +- .../cuvs/neighbors/graph_build_types.hpp | 2 +- .../neighbors/all_neighbors/all_neighbors.cu | 40 +-- .../neighbors/all_neighbors/all_neighbors.cuh | 197 +++++++++--- .../all_neighbors/all_neighbors_batched.cuh | 188 +++++++++--- .../all_neighbors/all_neighbors_builder.cuh | 287 +++++++++++++----- cpp/tests/neighbors/all_neighbors.cuh | 93 +++++- .../neighbors/all_neighbors/test_float.cu | 8 + 8 files changed, 655 insertions(+), 176 deletions(-) diff --git a/cpp/include/cuvs/neighbors/all_neighbors.hpp b/cpp/include/cuvs/neighbors/all_neighbors.hpp index a4323b67ba..b52a4ca294 100644 --- a/cpp/include/cuvs/neighbors/all_neighbors.hpp +++ b/cpp/include/cuvs/neighbors/all_neighbors.hpp @@ -126,13 +126,19 @@ struct all_neighbors_params { * in host memory * @param[out] indices nearest neighbor indices of shape [n_row x k] * @param[out] distances nearest neighbor distances [n_row x k] + * @param[out] core_distances array for core distances of size [n_row]. Requires distances matrix to + * compute core_distances. If core_distances is given, the resulting indices and distances will be + * mutual reachability space. + * @param[in] alpha distance scaling parameter as used in robust single linkage. */ void build( const raft::resources& handle, const all_neighbors_params& params, raft::host_matrix_view dataset, raft::device_matrix_view indices, - std::optional> distances = std::nullopt); + std::optional> distances = std::nullopt, + std::optional> core_distances = std::nullopt, + float alpha = 1.0); /** * @brief Builds an approximate all-neighbors knn graph (find nearest neighbors for all the training @@ -156,13 +162,19 @@ void build( * in device memory * @param[out] indices nearest neighbor indices of shape [n_row x k] * @param[out] distances nearest neighbor distances [n_row x k] + * @param[out] core_distances array for core distances of size [n_row]. Requires distances matrix to + * compute core_distances. If core_distances is given, the resulting indices and distances will be + * mutual reachability space. + * @param[in] alpha distance scaling parameter as used in robust single linkage. */ void build( const raft::resources& handle, const all_neighbors_params& params, raft::device_matrix_view dataset, raft::device_matrix_view indices, - std::optional> distances = std::nullopt); + std::optional> distances = std::nullopt, + std::optional> core_distances = std::nullopt, + float alpha = 1.0); /** @} */ } // namespace cuvs::neighbors::all_neighbors diff --git a/cpp/include/cuvs/neighbors/graph_build_types.hpp b/cpp/include/cuvs/neighbors/graph_build_types.hpp index 0711b5c2fb..dc6fdb45bd 100644 --- a/cpp/include/cuvs/neighbors/graph_build_types.hpp +++ b/cpp/include/cuvs/neighbors/graph_build_types.hpp @@ -27,7 +27,7 @@ namespace cuvs::neighbors { * @{ */ -enum GRAPH_BUILD_ALGO { BRUTE_FORCE = 0, IVF_PQ = 1, NN_DESCENT = 1 }; +enum GRAPH_BUILD_ALGO { BRUTE_FORCE = 0, IVF_PQ = 1, NN_DESCENT = 2 }; namespace graph_build_params { diff --git a/cpp/src/neighbors/all_neighbors/all_neighbors.cu b/cpp/src/neighbors/all_neighbors/all_neighbors.cu index 8ba57df398..2e9bf93688 100644 --- a/cpp/src/neighbors/all_neighbors/all_neighbors.cu +++ b/cpp/src/neighbors/all_neighbors/all_neighbors.cu @@ -18,23 +18,29 @@ namespace cuvs::neighbors::all_neighbors { -#define CUVS_INST_ALL_NEIGHBORS(T, IdxT) \ - void build(const raft::resources& handle, \ - const all_neighbors_params& params, \ - raft::host_matrix_view dataset, \ - raft::device_matrix_view indices, \ - std::optional> distances) \ - { \ - return all_neighbors::detail::build(handle, params, dataset, indices, distances); \ - } \ - \ - void build(const raft::resources& handle, \ - const all_neighbors_params& params, \ - raft::device_matrix_view dataset, \ - raft::device_matrix_view indices, \ - std::optional> distances) \ - { \ - return all_neighbors::detail::build(handle, params, dataset, indices, distances); \ +#define CUVS_INST_ALL_NEIGHBORS(T, IdxT) \ + void build(const raft::resources& handle, \ + const all_neighbors_params& params, \ + raft::host_matrix_view dataset, \ + raft::device_matrix_view indices, \ + std::optional> distances, \ + std::optional> core_distances, \ + T alpha) \ + { \ + return all_neighbors::detail::build( \ + handle, params, dataset, indices, distances, core_distances, alpha); \ + } \ + \ + void build(const raft::resources& handle, \ + const all_neighbors_params& params, \ + raft::device_matrix_view dataset, \ + raft::device_matrix_view indices, \ + std::optional> distances, \ + std::optional> core_distances, \ + T alpha) \ + { \ + return all_neighbors::detail::build( \ + handle, params, dataset, indices, distances, core_distances, alpha); \ } CUVS_INST_ALL_NEIGHBORS(float, int64_t); diff --git a/cpp/src/neighbors/all_neighbors/all_neighbors.cuh b/cpp/src/neighbors/all_neighbors/all_neighbors.cuh index 66abd5ac5d..50f0de104c 100644 --- a/cpp/src/neighbors/all_neighbors/all_neighbors.cuh +++ b/cpp/src/neighbors/all_neighbors/all_neighbors.cuh @@ -15,71 +15,104 @@ */ #pragma once +#include "../detail/reachability.cuh" #include "all_neighbors_batched.cuh" #include +#include +#include #include namespace cuvs::neighbors::all_neighbors::detail { using namespace cuvs::neighbors; -void check_metric(const all_neighbors_params& params) +GRAPH_BUILD_ALGO check_params_validity(const all_neighbors_params& params, + bool do_mutual_reachability_dist) { if (std::holds_alternative(params.graph_build_params)) { - auto allowed_metrics_batch = params.metric == cuvs::distance::DistanceType::L2Expanded || - params.metric == cuvs::distance::DistanceType::L2SqrtExpanded; - auto allowed_metrics_single = allowed_metrics_batch || - params.metric == cuvs::distance::DistanceType::CosineExpanded || - params.metric == cuvs::distance::DistanceType::InnerProduct; - // related issue: https://github.com/rapidsai/cuvs/issues/1056 - RAFT_EXPECTS((params.n_clusters <= 1 && allowed_metrics_single) || allowed_metrics_batch, - "Distance metric supported for for all-neighbors build with brute force depends " - "on params.n_clusters. When params.n_clusters <= 1, supported metrics are " - "L2Expanded, L2SqrtExpanded, CosineExpanded, or InnerProduct. When " - "params.n_clusters > 1, supported metrics are L2Expanded, L2SqrtExpanded."); + if (do_mutual_reachability_dist) { + // InnerProduct is not supported for mutual reachability distance, because mutual reachability + // distance takes "max" of core distances and pairwise distance. + auto allowed_metrics = params.metric == cuvs::distance::DistanceType::L2Expanded || + params.metric == cuvs::distance::DistanceType::L2SqrtExpanded || + params.metric == cuvs::distance::DistanceType::CosineExpanded; + RAFT_EXPECTS( + allowed_metrics, + "Distance metric for all-neighbors build with brute force for computing mutual " + "reachability distance should be L2Expanded, L2SqrtExpanded, or CosineExpanded."); + } else { + auto allowed_metrics = params.metric == cuvs::distance::DistanceType::L2Expanded || + params.metric == cuvs::distance::DistanceType::L2SqrtExpanded || + params.metric == cuvs::distance::DistanceType::CosineExpanded || + params.metric == cuvs::distance::DistanceType::InnerProduct; + RAFT_EXPECTS(allowed_metrics, + "Distance metric for all-neighbors build with brute force should be L2Expanded, " + "L2SqrtExpanded, CosineExpanded, or InnerProduct."); + } + return GRAPH_BUILD_ALGO::BRUTE_FORCE; } else if (std::holds_alternative( params.graph_build_params)) { - auto allowed_metrics = params.metric == cuvs::distance::DistanceType::L2Expanded || - params.metric == cuvs::distance::DistanceType::L2SqrtExpanded || - params.metric == cuvs::distance::DistanceType::CosineExpanded || - params.metric == cuvs::distance::DistanceType::InnerProduct; - RAFT_EXPECTS(allowed_metrics, - "Distance metric for all-neighbors build with NN Descent should be L2Expanded, " - "L2SqrtExpanded, CosineExpanded, or InnerProduct"); + if (do_mutual_reachability_dist) { + // InnerProduct is not supported for mutual reachability distance, because mutual reachability + // distance takes "max" of core distances and pairwise distance. + auto allowed_metrics = params.metric == cuvs::distance::DistanceType::L2Expanded || + params.metric == cuvs::distance::DistanceType::L2SqrtExpanded || + params.metric == cuvs::distance::DistanceType::CosineExpanded; + RAFT_EXPECTS( + allowed_metrics, + "Distance metric for all-neighbors build with NN Descent for computing mutual reachability " + "distance should be L2Expanded, L2SqrtExpanded, or CosineExpanded."); + } else { + auto allowed_metrics = params.metric == cuvs::distance::DistanceType::L2Expanded || + params.metric == cuvs::distance::DistanceType::L2SqrtExpanded || + params.metric == cuvs::distance::DistanceType::CosineExpanded || + params.metric == cuvs::distance::DistanceType::InnerProduct; + RAFT_EXPECTS(allowed_metrics, + "Distance metric for all-neighbors build with NN Descent should be L2Expanded, " + "L2SqrtExpanded, CosineExpanded, or InnerProduct."); + } + return GRAPH_BUILD_ALGO::NN_DESCENT; } else if (std::holds_alternative(params.graph_build_params)) { RAFT_EXPECTS(params.metric == cuvs::distance::DistanceType::L2Expanded, "Distance metric for all-neighbors build with IVFPQ should be L2Expanded"); + RAFT_EXPECTS(!do_mutual_reachability_dist, + "mutual reachability distance cannot be calculated using IVFPQ"); + return GRAPH_BUILD_ALGO::IVF_PQ; } else { RAFT_FAIL("Invalid all-neighbors build algo"); } } // Single build (i.e. no batching) supports both host and device datasets -template +template void single_build( const raft::resources& handle, const all_neighbors_params& params, mdspan, row_major, Accessor> dataset, raft::device_matrix_view indices, - std::optional> distances = std::nullopt) + std::optional> distances = std::nullopt, + DistEpilogueT dist_epilogue = DistEpilogueT{}) { size_t num_rows = static_cast(dataset.extent(0)); size_t num_cols = static_cast(dataset.extent(1)); auto knn_builder = get_knn_builder( - handle, params, num_rows, num_rows, indices.extent(1), indices, distances); + handle, params, num_rows, num_rows, indices.extent(1), indices, distances, dist_epilogue); knn_builder->prepare_build(dataset); knn_builder->build_knn(dataset); } template -void build(const raft::resources& handle, - const all_neighbors_params& params, - raft::host_matrix_view dataset, - raft::device_matrix_view indices, - std::optional> distances = std::nullopt) +void build( + const raft::resources& handle, + const all_neighbors_params& params, + raft::host_matrix_view dataset, + raft::device_matrix_view indices, + std::optional> distances = std::nullopt, + std::optional> core_distances = std::nullopt, + T alpha = 1.0) { - check_metric(params); + auto build_algo = check_params_validity(params, core_distances.has_value()); RAFT_EXPECTS(dataset.extent(0) == indices.extent(0), "number of rows in dataset should be the same as number of rows in indices matrix"); @@ -90,21 +123,75 @@ void build(const raft::resources& handle, "indices matrix and distances matrix has to be the same shape."); } + if (core_distances.has_value()) { + RAFT_EXPECTS(distances.has_value(), + "distances matrix should be allocated to get mutual reachability distance."); + } + + std::unique_ptr> aux_vectors; if (params.n_clusters == 1) { single_build(handle, params, dataset, indices, distances); } else { - batch_build(handle, params, dataset, indices, distances); + if (core_distances.has_value()) { + aux_vectors = std::make_unique>( + params.n_clusters, dataset.extent(0), params.overlap_factor); + batch_build(handle, params, dataset, indices, distances, aux_vectors.get()); + } else { + batch_build(handle, params, dataset, indices, distances); + } + } + + // NN Descent doesn't include self loops. Shifted to keep it consistent with brute force and ivfpq + bool need_shift = (build_algo == GRAPH_BUILD_ALGO::NN_DESCENT) && + (params.metric != cuvs::distance::DistanceType::InnerProduct); + + if (need_shift) { + raft::matrix::shift(handle, indices, 1); + if (distances.has_value()) { + raft::matrix::shift(handle, distances.value(), 1, std::make_optional(0.0)); + } + } + + if (core_distances.has_value()) { // calculate mutual reachability distances + size_t k = indices.extent(1); + size_t num_rows = core_distances.value().size(); + cuvs::neighbors::detail::reachability::core_distances( + handle, + distances.value().data_handle(), + k, + k, + num_rows, + core_distances.value().data_handle()); + + using ReachabilityPP = cuvs::neighbors::detail::reachability::ReachabilityPostProcess; + auto dist_epilogue = ReachabilityPP{core_distances.value().data_handle(), alpha, num_rows}; + if (params.n_clusters == 1) { + single_build(handle, params, dataset, indices, distances, dist_epilogue); + } else { + batch_build(handle, params, dataset, indices, distances, aux_vectors.get(), dist_epilogue); + } + + if (need_shift) { + raft::matrix::shift(handle, indices, 1); + raft::matrix::shift(handle, + distances.value(), + raft::make_device_matrix_view( + core_distances.value().data_handle(), num_rows, 1)); + } } } template -void build(const raft::resources& handle, - const all_neighbors_params& params, - raft::device_matrix_view dataset, - raft::device_matrix_view indices, - std::optional> distances = std::nullopt) +void build( + const raft::resources& handle, + const all_neighbors_params& params, + raft::device_matrix_view dataset, + raft::device_matrix_view indices, + std::optional> distances = std::nullopt, + std::optional> core_distances = std::nullopt, + T alpha = 1.0) { - check_metric(params); + auto build_algo = check_params_validity(params, core_distances.has_value()); RAFT_EXPECTS(dataset.extent(0) == indices.extent(0), "number of rows in dataset should be the same as number of rows in indices matrix"); @@ -115,6 +202,11 @@ void build(const raft::resources& handle, "indices matrix and distances matrix has to be the same shape."); } + if (core_distances.has_value()) { + RAFT_EXPECTS(distances.has_value(), + "distances matrix should be allocated to get mutual reachability distance."); + } + if (params.n_clusters > 1) { RAFT_FAIL( "Batched all-neighbors build is not supported with data on device. Put data on host for " @@ -122,5 +214,40 @@ void build(const raft::resources& handle, } else { single_build(handle, params, dataset, indices, distances); } + + // NN Descent doesn't include self loops. Shifted to keep it consistent with brute force and ivfpq + bool need_shift = (build_algo == GRAPH_BUILD_ALGO::NN_DESCENT) && + (params.metric != cuvs::distance::DistanceType::InnerProduct); + + if (need_shift) { + raft::matrix::shift(handle, indices, 1); + if (distances.has_value()) { + raft::matrix::shift(handle, distances.value(), 1, std::make_optional(0.0)); + } + } + + if (core_distances.has_value()) { + size_t k = indices.extent(1); + size_t num_rows = core_distances.value().size(); + cuvs::neighbors::detail::reachability::core_distances( + handle, + distances.value().data_handle(), + k, + k, + num_rows, + core_distances.value().data_handle()); + + using ReachabilityPP = cuvs::neighbors::detail::reachability::ReachabilityPostProcess; + auto dist_epilogue = ReachabilityPP{core_distances.value().data_handle(), alpha, num_rows}; + single_build(handle, params, dataset, indices, distances, dist_epilogue); + + if (need_shift) { + raft::matrix::shift(handle, indices, 1); + raft::matrix::shift(handle, + distances.value(), + raft::make_device_matrix_view( + core_distances.value().data_handle(), num_rows, 1)); + } + } } } // namespace cuvs::neighbors::all_neighbors::detail diff --git a/cpp/src/neighbors/all_neighbors/all_neighbors_batched.cuh b/cpp/src/neighbors/all_neighbors/all_neighbors_batched.cuh index 468d4ee0df..75eec6d50f 100644 --- a/cpp/src/neighbors/all_neighbors/all_neighbors_batched.cuh +++ b/cpp/src/neighbors/all_neighbors/all_neighbors_batched.cuh @@ -15,6 +15,7 @@ */ #pragma once +#include "../detail/reachability.cuh" #include "all_neighbors_builder.cuh" #include "raft/core/logger_macros.hpp" #include @@ -23,11 +24,13 @@ #include #include #include +#include #include #include #include #include #include +#include #include #include #include @@ -35,6 +38,31 @@ namespace cuvs::neighbors::all_neighbors::detail { using namespace cuvs::neighbors; +template +void reset_global_matrices(raft::resources const& res, + cuvs::distance::DistanceType metric, + raft::managed_matrix_view global_neighbors, + raft::managed_matrix_view global_distances) +{ + size_t num_rows = static_cast(global_neighbors.extent(0)); + size_t k = static_cast(global_neighbors.extent(1)); + + bool select_min = cuvs::distance::is_min_close(metric); + IdxT global_neighbors_fill_value = + select_min ? std::numeric_limits::max() : std::numeric_limits::min(); + T global_distances_fill_value = + select_min ? std::numeric_limits::max() : std::numeric_limits::min(); + + raft::matrix::fill( + res, + raft::make_device_matrix_view(global_neighbors.data_handle(), num_rows, k), + global_neighbors_fill_value); + raft::matrix::fill( + res, + raft::make_device_matrix_view(global_distances.data_handle(), num_rows, k), + global_distances_fill_value); +} + /** * Run balanced kmeans on a subsample of the dataset to get centroids. * Arguments: @@ -307,15 +335,16 @@ void single_gpu_batch_build(const raft::resources& handle, } } -template +template void multi_gpu_batch_build(const raft::resources& handle, const all_neighbors_params& params, raft::host_matrix_view dataset, raft::managed_matrix_view global_neighbors, raft::managed_matrix_view global_distances, raft::host_vector_view cluster_sizes, - raft::host_vector_view cluster_offsets, - raft::host_vector_view inverted_indices) + raft::host_vector_view cluster_offsets_c, + raft::host_vector_view inverted_indices, + DistEpilogueT dist_epilogue = DistEpilogueT{}) { size_t num_rows = dataset.extent(0); size_t num_cols = dataset.extent(1); @@ -326,6 +355,23 @@ void multi_gpu_batch_build(const raft::resources& handle, size_t clusters_per_rank = params.n_clusters / num_ranks; size_t rem = params.n_clusters - clusters_per_rank * num_ranks; + auto cluster_offsets = raft::make_host_vector(cluster_offsets_c.size()); + raft::copy(cluster_offsets.data_handle(), + cluster_offsets_c.data_handle(), + cluster_offsets_c.size(), + raft::resource::get_cuda_stream(handle)); + + using ReachabilityPP = cuvs::neighbors::detail::reachability::ReachabilityPostProcess; + const bool mutual_reach_dist = std::is_same_v; + std::optional> core_distances_h; + if constexpr (mutual_reach_dist) { + core_distances_h.emplace(raft::make_host_vector(num_rows)); + raft::copy(core_distances_h.value().data_handle(), + dist_epilogue.core_dists, + num_rows, + raft::resource::get_cuda_stream(handle)); + } + #pragma omp parallel for num_threads(num_ranks) for (int rank = 0; rank < num_ranks; rank++) { auto dev_res = raft::resource::set_current_device_to_rank(handle, rank); @@ -369,8 +415,30 @@ void multi_gpu_batch_build(const raft::resources& handle, size_t max_cluster_size, min_cluster_size; get_min_max_cluster_size(k, max_cluster_size, min_cluster_size, cluster_sizes_for_this_rank); + std::optional> core_distances_d_for_rank; + auto dist_epilgogue_for_rank = [&]() { + if constexpr (mutual_reach_dist) { + core_distances_d_for_rank.emplace(raft::make_device_vector(dev_res, num_rows)); + raft::copy(core_distances_d_for_rank.value().data_handle(), + core_distances_h.value().data_handle(), + num_rows, + raft::resource::get_cuda_stream(dev_res)); + return ReachabilityPP{ + core_distances_d_for_rank.value().data_handle(), dist_epilogue.alpha, num_rows}; + } else { + return dist_epilogue; + } + }(); + std::unique_ptr> knn_builder = - get_knn_builder(dev_res, params, min_cluster_size, max_cluster_size, k); + get_knn_builder(dev_res, + params, + min_cluster_size, + max_cluster_size, + k, + std::nullopt, + std::nullopt, + dist_epilgogue_for_rank); single_gpu_batch_build(dev_res, dataset, @@ -384,13 +452,31 @@ void multi_gpu_batch_build(const raft::resources& handle, } } -template +/* Holds necessary vectors to avoid recomputation when calculating mutual rechability distances */ +template +struct BatchBuildAux { + raft::host_vector cluster_sizes; + raft::host_vector cluster_offsets; + raft::host_vector inverted_indices; + bool is_computed = false; + + BatchBuildAux(IdxT n_clusters, IdxT num_rows, IdxT overlap_factor) + : cluster_sizes(raft::make_host_vector(n_clusters)), + cluster_offsets(raft::make_host_vector(n_clusters)), + inverted_indices(raft::make_host_vector(num_rows * overlap_factor)) + { + } +}; + +template void batch_build( const raft::resources& handle, const all_neighbors_params& params, raft::host_matrix_view dataset, raft::device_matrix_view indices, - std::optional> distances = std::nullopt) + std::optional> distances = std::nullopt, + BatchBuildAux* aux_vectors = nullptr, + DistEpilogueT dist_epilogue = DistEpilogueT{}) { if (raft::resource::is_multi_gpu(handle)) { // For efficient CPU-computation of omp parallel for regions per GPU @@ -405,37 +491,49 @@ void batch_build( "overlap_factor should be smaller than n_clusters. We recommend starting from " "overlap_factor=2 and gradually increase it for better knn graph recall."); - auto centroids = raft::make_device_matrix(handle, params.n_clusters, num_cols); - get_centroids_on_data_subsample(handle, params.metric, dataset, centroids.view()); - - auto global_nearest_cluster = raft::make_host_matrix(num_rows, params.overlap_factor); - assign_clusters( - handle, params, dataset, centroids.view(), global_nearest_cluster.view()); - - auto inverted_indices = - raft::make_host_vector(num_rows * params.overlap_factor); - auto cluster_sizes = raft::make_host_vector(params.n_clusters); - auto cluster_offsets = raft::make_host_vector(params.n_clusters); - get_inverted_indices(handle, - global_nearest_cluster.view(), - inverted_indices.view(), - cluster_sizes.view(), - cluster_offsets.view()); + std::optional> local_inverted_indices; + std::optional> local_cluster_sizes; + std::optional> local_cluster_offsets; + + auto inverted_indices_view = + aux_vectors != nullptr + ? aux_vectors->inverted_indices.view() + : local_inverted_indices + .emplace(raft::make_host_vector(num_rows * params.overlap_factor)) + .view(); + + auto cluster_sizes_view = + aux_vectors != nullptr + ? aux_vectors->cluster_sizes.view() + : local_cluster_sizes.emplace(raft::make_host_vector(params.n_clusters)).view(); + + auto cluster_offsets_view = + aux_vectors != nullptr + ? aux_vectors->cluster_offsets.view() + : local_cluster_offsets.emplace(raft::make_host_vector(params.n_clusters)).view(); + + if (aux_vectors == nullptr || !aux_vectors->is_computed) { + // clustering step is only computed when needed + auto centroids = raft::make_device_matrix(handle, params.n_clusters, num_cols); + get_centroids_on_data_subsample(handle, params.metric, dataset, centroids.view()); + + auto global_nearest_cluster = + raft::make_host_matrix(num_rows, params.overlap_factor); + assign_clusters( + handle, params, dataset, centroids.view(), global_nearest_cluster.view()); + + get_inverted_indices(handle, + global_nearest_cluster.view(), + inverted_indices_view, + cluster_sizes_view, + cluster_offsets_view); + if (aux_vectors != nullptr) { aux_vectors->is_computed = true; } + } auto global_neighbors = raft::make_managed_matrix(handle, num_rows, k); auto global_distances = raft::make_managed_matrix(handle, num_rows, k); - bool select_min = cuvs::distance::is_min_close(params.metric); - IdxT global_neighbors_fill_value = - select_min ? std::numeric_limits::max() : std::numeric_limits::min(); - T global_distances_fill_value = - select_min ? std::numeric_limits::max() : std::numeric_limits::min(); - std::fill(global_neighbors.data_handle(), - global_neighbors.data_handle() + num_rows * k, - global_neighbors_fill_value); - std::fill(global_distances.data_handle(), - global_distances.data_handle() + num_rows * k, - global_distances_fill_value); + reset_global_matrices(handle, params.metric, global_neighbors.view(), global_distances.view()); if (raft::resource::is_multi_gpu(handle)) { multi_gpu_batch_build(handle, @@ -443,23 +541,31 @@ void batch_build( dataset, global_neighbors.view(), global_distances.view(), - cluster_sizes.view(), - cluster_offsets.view(), - inverted_indices.view()); + cluster_sizes_view, + raft::make_const_mdspan(cluster_offsets_view), + inverted_indices_view, + dist_epilogue); } else { size_t max_cluster_size, min_cluster_size; - get_min_max_cluster_size(k, max_cluster_size, min_cluster_size, cluster_sizes.view()); + get_min_max_cluster_size(k, max_cluster_size, min_cluster_size, cluster_sizes_view); std::unique_ptr> knn_builder = - get_knn_builder(handle, params, min_cluster_size, max_cluster_size, k); + get_knn_builder(handle, + params, + min_cluster_size, + max_cluster_size, + k, + std::nullopt, + std::nullopt, + dist_epilogue); single_gpu_batch_build(handle, dataset, *knn_builder, params.n_clusters, global_neighbors.view(), global_distances.view(), - cluster_sizes.view(), - cluster_offsets.view(), - inverted_indices.view()); + cluster_sizes_view, + cluster_offsets_view, + inverted_indices_view); } raft::copy(indices.data_handle(), diff --git a/cpp/src/neighbors/all_neighbors/all_neighbors_builder.cuh b/cpp/src/neighbors/all_neighbors/all_neighbors_builder.cuh index af095baf0f..123d7b5401 100644 --- a/cpp/src/neighbors/all_neighbors/all_neighbors_builder.cuh +++ b/cpp/src/neighbors/all_neighbors/all_neighbors_builder.cuh @@ -15,22 +15,27 @@ */ #pragma once +#include "../detail/knn_brute_force.cuh" #include "../detail/nn_descent_gnnd.hpp" +#include "../detail/reachability.cuh" #include "all_neighbors_merge.cuh" -#include "raft/core/device_mdarray.hpp" -#include "raft/core/mdspan.hpp" -#include "raft/util/cudart_utils.hpp" + #include #include #include #include #include +#include #include #include #include +#include #include +#include #include +#include +#include namespace cuvs::neighbors::all_neighbors::detail { using namespace cuvs::neighbors; @@ -304,7 +309,7 @@ struct all_neighbors_builder_ivfpq : public all_neighbors_builder { std::optional> refined_distances_h; }; -template +template struct all_neighbors_builder_nn_descent : public all_neighbors_builder { all_neighbors_builder_nn_descent( raft::resources const& res, @@ -314,10 +319,12 @@ struct all_neighbors_builder_nn_descent : public all_neighbors_builder size_t k, graph_build_params::nn_descent_params& params, std::optional> indices = std::nullopt, - std::optional> distances = std::nullopt) + std::optional> distances = std::nullopt, + DistEpilogueT dist_epilogue = DistEpilogueT{}) : all_neighbors_builder( res, n_clusters, min_cluster_size, max_cluster_size, k, indices, distances), - nnd_params{params} + nnd_params{params}, + dist_epilogue{dist_epilogue} { } @@ -346,6 +353,13 @@ struct all_neighbors_builder_nn_descent : public all_neighbors_builder nnd_builder.emplace(this->res, build_config); int_graph.emplace(raft::make_host_matrix( this->max_cluster_size, static_cast(extended_graph_degree))); + + if constexpr (std::is_same_v< + DistEpilogueT, + cuvs::neighbors::detail::reachability::ReachabilityPostProcess>) { + batch_core_distances.emplace( + raft::make_device_vector(this->res, this->max_cluster_size)); + } } void prepare_build(raft::device_matrix_view dataset) override @@ -371,37 +385,78 @@ struct all_neighbors_builder_nn_descent : public all_neighbors_builder "need valid inverted_indices, global_neighbors, and global_distances for " "build_knn if doing batching."); + using ReachabilityPP = cuvs::neighbors::detail::reachability::ReachabilityPostProcess; + if (this->n_clusters > 1) { bool return_distances = true; size_t num_data_in_cluster = dataset.extent(0); - nnd_builder.value().build(dataset.data_handle(), - static_cast(num_data_in_cluster), - int_graph.value().data_handle(), - return_distances, - this->batch_distances_d.value().data_handle()); - - remap_and_merge_subgraphs(this->res, - this->inverted_indices_d.value().view(), - inverted_indices.value(), - int_graph.value().view(), - this->batch_neighbors_h.value().view(), - this->batch_neighbors_d.value().view(), - this->batch_distances_d.value().view(), - global_neighbors.value(), - global_distances.value(), - num_data_in_cluster, - this->k, - cuvs::distance::is_min_close(nnd_params.metric)); + if constexpr (std::is_same_v) { + // gather core dists + raft::copy(this->inverted_indices_d.value().data_handle(), + inverted_indices.value().data_handle(), + num_data_in_cluster, + raft::resource::get_cuda_stream(this->res)); + + raft::matrix::gather(this->res, + raft::make_device_matrix_view( + dist_epilogue.core_dists, dist_epilogue.n, 1), + raft::make_device_vector_view( + this->inverted_indices_d.value().data_handle(), num_data_in_cluster), + raft::make_device_matrix_view( + batch_core_distances.value().data_handle(), num_data_in_cluster, 1)); + + nnd_builder.value().build( + dataset.data_handle(), + static_cast(num_data_in_cluster), + int_graph.value().data_handle(), + return_distances, + this->batch_distances_d.value().data_handle(), + cuvs::neighbors::detail::reachability::ReachabilityPostProcess{ + batch_core_distances.value().data_handle(), dist_epilogue.alpha, num_data_in_cluster}); + } else { + nnd_builder.value().build(dataset.data_handle(), + static_cast(num_data_in_cluster), + int_graph.value().data_handle(), + return_distances, + this->batch_distances_d.value().data_handle()); + } + + remap_and_merge_subgraphs>( + this->res, + this->inverted_indices_d.value().view(), + inverted_indices.value(), + int_graph.value().view(), + this->batch_neighbors_h.value().view(), + this->batch_neighbors_d.value().view(), + this->batch_distances_d.value().view(), + global_neighbors.value(), + global_distances.value(), + num_data_in_cluster, + this->k, + cuvs::distance::is_min_close(nnd_params.metric)); } else { size_t num_rows = dataset.extent(0); - nnd_builder.value().build( - dataset.data_handle(), - static_cast(num_rows), - int_graph.value().data_handle(), - this->distances_.has_value(), - this->distances_.value_or(raft::make_device_matrix(this->res, 0, 0).view()) - .data_handle()); + if constexpr (std::is_same_v) { + nnd_builder.value().build( + dataset.data_handle(), + static_cast(num_rows), + int_graph.value().data_handle(), + this->distances_.has_value(), + this->distances_.value_or(raft::make_device_matrix(this->res, 0, 0).view()) + .data_handle(), + cuvs::neighbors::detail::reachability::ReachabilityPostProcess{ + dist_epilogue.core_dists, dist_epilogue.alpha, dist_epilogue.n}); + } else { + nnd_builder.value().build( + dataset.data_handle(), + static_cast(num_rows), + int_graph.value().data_handle(), + this->distances_.has_value(), + this->distances_.value_or(raft::make_device_matrix(this->res, 0, 0).view()) + .data_handle(), + dist_epilogue); + } auto tmp_indices = raft::make_host_matrix(int_graph.value().extent(0), this->k); @@ -443,9 +498,12 @@ struct all_neighbors_builder_nn_descent : public all_neighbors_builder std::optional> nnd_builder; std::optional> int_graph; + + DistEpilogueT dist_epilogue; + std::optional> batch_core_distances; }; -template +template struct all_neighbors_builder_brute_force : public all_neighbors_builder { all_neighbors_builder_brute_force( raft::resources const& res, @@ -455,10 +513,12 @@ struct all_neighbors_builder_brute_force : public all_neighbors_builder size_t k, graph_build_params::brute_force_params& params, std::optional> indices = std::nullopt, - std::optional> distances = std::nullopt) + std::optional> distances = std::nullopt, + DistEpilogueT dist_epilogue = DistEpilogueT{}) : all_neighbors_builder( res, n_clusters, min_cluster_size, max_cluster_size, k, indices, distances), - bf_params{params} + bf_params{params}, + dist_epilogue{dist_epilogue} { } @@ -470,6 +530,12 @@ struct all_neighbors_builder_brute_force : public all_neighbors_builder size_t num_cols = dataset.extent(1); data_d.emplace( raft::make_device_matrix(this->res, this->max_cluster_size, num_cols)); + if constexpr (std::is_same_v< + DistEpilogueT, + cuvs::neighbors::detail::reachability::ReachabilityPostProcess>) { + batch_core_distances.emplace( + raft::make_device_vector(this->res, this->max_cluster_size)); + } } void build_knn_common( @@ -485,26 +551,63 @@ struct all_neighbors_builder_brute_force : public all_neighbors_builder "build_knn if doing batching."); if (this->n_clusters > 1) { - auto idx = cuvs::neighbors::brute_force::build(this->res, bf_params.build_params, dataset); - size_t num_data_in_cluster = dataset.extent(0); + using ReachabilityPP = + cuvs::neighbors::detail::reachability::ReachabilityPostProcess; + + if constexpr (std::is_same_v) { + // gather core dists + raft::copy(this->inverted_indices_d.value().data_handle(), + inverted_indices.value().data_handle(), + num_data_in_cluster, + raft::resource::get_cuda_stream(this->res)); - cuvs::neighbors::brute_force::search( - this->res, - bf_params.search_params, - idx, - dataset, - raft::make_device_matrix_view( - this->batch_neighbors_d.value().data_handle(), num_data_in_cluster, this->k), - raft::make_device_matrix_view( - this->batch_distances_d.value().data_handle(), num_data_in_cluster, this->k)); - + raft::matrix::gather(this->res, + raft::make_device_matrix_view( + dist_epilogue.core_dists, dist_epilogue.n, 1), + raft::make_device_vector_view( + this->inverted_indices_d.value().data_handle(), num_data_in_cluster), + raft::make_device_matrix_view( + batch_core_distances.value().data_handle(), num_data_in_cluster, 1)); + + cuvs::neighbors::detail::tiled_brute_force_knn( + this->res, + dataset.data_handle(), + dataset.data_handle(), + dataset.extent(0), + dataset.extent(0), + dataset.extent(1), + this->k, + this->batch_distances_d.value().data_handle(), + this->batch_neighbors_d.value().data_handle(), + bf_params.build_params.metric, + 2.0, + 0, + 0, + nullptr, + nullptr, + nullptr, + ReachabilityPP{ + batch_core_distances.value().data_handle(), dist_epilogue.alpha, num_data_in_cluster}); + } else { + auto idx = cuvs::neighbors::brute_force::build(this->res, bf_params.build_params, dataset); + + cuvs::neighbors::brute_force::search( + this->res, + bf_params.search_params, + idx, + dataset, + raft::make_device_matrix_view( + this->batch_neighbors_d.value().data_handle(), num_data_in_cluster, this->k), + raft::make_device_matrix_view( + this->batch_distances_d.value().data_handle(), num_data_in_cluster, this->k)); + } raft::copy(this->batch_neighbors_h.value().data_handle(), this->batch_neighbors_d.value().data_handle(), num_data_in_cluster * this->k, raft::resource::get_cuda_stream(this->res)); - remap_and_merge_subgraphs( + remap_and_merge_subgraphs>( this->res, this->inverted_indices_d.value().view(), inverted_indices.value(), @@ -518,17 +621,41 @@ struct all_neighbors_builder_brute_force : public all_neighbors_builder this->k, cuvs::distance::is_min_close(bf_params.build_params.metric)); } else { - auto idx = cuvs::neighbors::brute_force::build(this->res, bf_params.build_params, dataset); - - cuvs::neighbors::brute_force::search( - this->res, - bf_params.search_params, - idx, - dataset, - this->indices_.value(), - this->distances_.has_value() - ? this->distances_.value() - : raft::make_device_matrix(this->res, dataset.extent(0), this->k).view()); + if constexpr (std::is_same_v) { + auto idx = cuvs::neighbors::brute_force::build(this->res, bf_params.build_params, dataset); + + cuvs::neighbors::brute_force::search( + this->res, + bf_params.search_params, + idx, + dataset, + this->indices_.value(), + this->distances_.has_value() + ? this->distances_.value() + : raft::make_device_matrix(this->res, dataset.extent(0), this->k).view()); + } else { + cuvs::neighbors::detail::tiled_brute_force_knn( + this->res, + dataset.data_handle(), + dataset.data_handle(), + dataset.extent(0), + dataset.extent(0), + dataset.extent(1), + this->k, + this->distances_.has_value() + ? this->distances_.value().data_handle() + : raft::make_device_matrix(this->res, dataset.extent(0), this->k) + .data_handle(), + this->indices_.value().data_handle(), + bf_params.build_params.metric, + 2.0, + 0, + 0, + nullptr, + nullptr, + nullptr, + dist_epilogue); + } } } @@ -559,11 +686,13 @@ struct all_neighbors_builder_brute_force : public all_neighbors_builder } graph_build_params::brute_force_params bf_params; + DistEpilogueT dist_epilogue; std::optional> data_d; + std::optional> batch_core_distances; }; -template +template std::unique_ptr> get_knn_builder( const raft::resources& handle, const all_neighbors_params& params, @@ -571,7 +700,8 @@ std::unique_ptr> get_knn_builder( size_t max_cluster_size, size_t k, std::optional> indices = std::nullopt, - std::optional> distances = std::nullopt) + std::optional> distances = std::nullopt, + DistEpilogueT dist_epilogue = DistEpilogueT{}) { if (std::holds_alternative(params.graph_build_params)) { auto brute_force_params = @@ -580,14 +710,17 @@ std::unique_ptr> get_knn_builder( RAFT_LOG_WARN("Setting brute_force_params metric to metric given for batching algorithm"); brute_force_params.build_params.metric = params.metric; } - return std::make_unique>(handle, - params.n_clusters, - min_cluster_size, - max_cluster_size, - k, - brute_force_params, - indices, - distances); + return std::make_unique>( + handle, + params.n_clusters, + min_cluster_size, + max_cluster_size, + k, + brute_force_params, + indices, + distances, + dist_epilogue); + } else if (std::holds_alternative( params.graph_build_params)) { auto nn_descent_params = @@ -596,14 +729,16 @@ std::unique_ptr> get_knn_builder( RAFT_LOG_WARN("Setting nnd_params metric to metric given for batching algorithm"); nn_descent_params.metric = params.metric; } - return std::make_unique>(handle, - params.n_clusters, - min_cluster_size, - max_cluster_size, - k, - nn_descent_params, - indices, - distances); + return std::make_unique>( + handle, + params.n_clusters, + min_cluster_size, + max_cluster_size, + k, + nn_descent_params, + indices, + distances, + dist_epilogue); } else if (std::holds_alternative(params.graph_build_params)) { auto ivf_pq_params = std::get(params.graph_build_params); if (ivf_pq_params.build_params.metric != params.metric) { diff --git a/cpp/tests/neighbors/all_neighbors.cuh b/cpp/tests/neighbors/all_neighbors.cuh index 1edc93ce26..fddfeb61af 100644 --- a/cpp/tests/neighbors/all_neighbors.cuh +++ b/cpp/tests/neighbors/all_neighbors.cuh @@ -15,6 +15,8 @@ */ #pragma once +#include "../../src/neighbors/detail/knn_brute_force.cuh" +#include "../../src/neighbors/detail/reachability.cuh" #include "../test_utils.cuh" #include "ann_utils.cuh" #include "naive_knn.cuh" @@ -52,6 +54,7 @@ struct AllNeighborsInputs { int dim; int k; bool data_on_host; + bool mutual_reach; }; inline ::std::ostream& operator<<(::std::ostream& os, const AllNeighborsInputs& p) @@ -101,6 +104,7 @@ void get_graphs(raft::resources& handle, params.graph_build_params = ivfpq_build_params; } + auto metric = std::get<1>(ps.build_algo_metric_recall); auto cuda_stream = raft::resource::get_cuda_stream(handle); { rmm::device_uvector distances_naive_dev(queries_size, cuda_stream); @@ -115,6 +119,41 @@ void get_graphs(raft::resources& handle, ps.dim, ps.k, std::get<1>(ps.build_algo_metric_recall)); + + if (ps.mutual_reach) { + rmm::device_uvector core_dists_dev(ps.n_rows, cuda_stream); + + cuvs::neighbors::detail::reachability::core_distances( + handle, distances_naive_dev.data(), ps.k, ps.k, ps.n_rows, core_dists_dev.data()); + + auto epilogue = + cuvs::neighbors::detail::reachability::ReachabilityPostProcess{ + core_dists_dev.data(), 1.0, static_cast(ps.n_rows)}; + + cuvs::neighbors::detail::tiled_brute_force_knn< + DataT, + IdxT, + DistanceT, + cuvs::neighbors::detail::reachability::ReachabilityPostProcess>( + handle, + database.data(), + database.data(), + ps.n_rows, + ps.n_rows, + ps.dim, + ps.k, + distances_naive_dev.data(), + indices_naive_dev.data(), + metric, + 2.0, + 0, + 0, + nullptr, + nullptr, + nullptr, + epilogue); + } + raft::update_host(indices_bf.data(), indices_naive_dev.data(), queries_size, cuda_stream); raft::update_host(distances_bf.data(), distances_naive_dev.data(), queries_size, cuda_stream); @@ -134,7 +173,10 @@ void get_graphs(raft::resources& handle, params, raft::make_const_mdspan(database_h.view()), raft::make_device_matrix_view(indices_allNN_dev.data(), ps.n_rows, ps.k), - raft::make_device_matrix_view(distances_allNN_dev.data(), ps.n_rows, ps.k)); + raft::make_device_matrix_view(distances_allNN_dev.data(), ps.n_rows, ps.k), + ps.mutual_reach + ? std::make_optional(raft::make_device_vector(handle, ps.n_rows).view()) + : std::nullopt); } else { all_neighbors::build( @@ -142,7 +184,10 @@ void get_graphs(raft::resources& handle, params, raft::make_device_matrix_view(database.data(), ps.n_rows, ps.dim), raft::make_device_matrix_view(indices_allNN_dev.data(), ps.n_rows, ps.k), - raft::make_device_matrix_view(distances_allNN_dev.data(), ps.n_rows, ps.k)); + raft::make_device_matrix_view(distances_allNN_dev.data(), ps.n_rows, ps.k), + ps.mutual_reach + ? std::make_optional(raft::make_device_vector(handle, ps.n_rows).view()) + : std::nullopt); } raft::copy(indices_allNN.data(), indices_allNN_dev.data(), queries_size, cuda_stream); @@ -223,13 +268,16 @@ const std::vector inputsSingle = {5000, 7151}, // n_rows {64, 137}, // dim {16, 23}, // graph_degree - {false, true} // data on host + {false, true}, // data on host + {false} // mutual_reach ); const std::vector inputsBatch = raft::util::itertools::product( {std::make_tuple(BRUTE_FORCE, cuvs::distance::DistanceType::L2Expanded, 0.9), std::make_tuple(BRUTE_FORCE, cuvs::distance::DistanceType::L2SqrtExpanded, 0.9), + std::make_tuple(BRUTE_FORCE, cuvs::distance::DistanceType::CosineExpanded, 0.9), + std::make_tuple(BRUTE_FORCE, cuvs::distance::DistanceType::InnerProduct, 0.9), std::make_tuple(IVF_PQ, cuvs::distance::DistanceType::L2Expanded, 0.9), std::make_tuple(NN_DESCENT, cuvs::distance::DistanceType::L2Expanded, 0.9), std::make_tuple(NN_DESCENT, cuvs::distance::DistanceType::L2SqrtExpanded, 0.9), @@ -243,7 +291,44 @@ const std::vector inputsBatch = {5000, 7151}, // n_rows {64, 137}, // dim {16, 23}, // graph_degree - {true} // data on host + {true}, // data on host + {false} // mutual_reach + ); + +const std::vector mutualReachSingle = + raft::util::itertools::product( + {std::make_tuple(BRUTE_FORCE, cuvs::distance::DistanceType::L2Expanded, 0.9), + std::make_tuple(BRUTE_FORCE, cuvs::distance::DistanceType::L2SqrtExpanded, 0.9), + std::make_tuple(BRUTE_FORCE, cuvs::distance::DistanceType::CosineExpanded, 0.9), + std::make_tuple(NN_DESCENT, cuvs::distance::DistanceType::L2Expanded, 0.9), + std::make_tuple(NN_DESCENT, cuvs::distance::DistanceType::L2SqrtExpanded, 0.9), + std::make_tuple(NN_DESCENT, cuvs::distance::DistanceType::CosineExpanded, 0.9)}, + {std::make_tuple(1lu, 2lu)}, // n_clusters, overlap_factor + {5000, 7151}, // n_rows + {64, 137}, // dim + {16, 23}, // graph_degree + {false, true}, // data on host + {true} // mutual_reach + ); + +const std::vector mutualReachBatch = + raft::util::itertools::product( + {std::make_tuple(BRUTE_FORCE, cuvs::distance::DistanceType::L2Expanded, 0.9), + std::make_tuple(BRUTE_FORCE, cuvs::distance::DistanceType::L2SqrtExpanded, 0.9), + std::make_tuple(BRUTE_FORCE, cuvs::distance::DistanceType::CosineExpanded, 0.9), + std::make_tuple(NN_DESCENT, cuvs::distance::DistanceType::L2Expanded, 0.9), + std::make_tuple(NN_DESCENT, cuvs::distance::DistanceType::L2SqrtExpanded, 0.9), + std::make_tuple(NN_DESCENT, cuvs::distance::DistanceType::CosineExpanded, 0.9)}, + { + std::make_tuple(4lu, 2lu), + std::make_tuple(7lu, 2lu), + std::make_tuple(10lu, 2lu), + }, // n_clusters, overlap_factor + {5000, 7151}, // n_rows + {64, 137}, // dim + {16, 23}, // graph_degree + {true}, // data on host + {true} // mutual_reach ); } // namespace cuvs::neighbors::all_neighbors diff --git a/cpp/tests/neighbors/all_neighbors/test_float.cu b/cpp/tests/neighbors/all_neighbors/test_float.cu index 3fa41efd5c..88af4c2253 100644 --- a/cpp/tests/neighbors/all_neighbors/test_float.cu +++ b/cpp/tests/neighbors/all_neighbors/test_float.cu @@ -28,4 +28,12 @@ INSTANTIATE_TEST_CASE_P(AllNeighborsSingleTest, ::testing::ValuesIn(inputsSingle)); INSTANTIATE_TEST_CASE_P(AllNeighborsBatchTest, AllNeighborsTestF, ::testing::ValuesIn(inputsBatch)); + +INSTANTIATE_TEST_CASE_P(AllNeighborsSingleMutualTest, + AllNeighborsTestF, + ::testing::ValuesIn(mutualReachSingle)); + +INSTANTIATE_TEST_CASE_P(AllNeighborsBatchMutualTest, + AllNeighborsTestF, + ::testing::ValuesIn(mutualReachBatch)); } // namespace cuvs::neighbors::all_neighbors From e0d413380cb9ca17145b2ca2e4b0c80d463d83be Mon Sep 17 00:00:00 2001 From: Rui Lan Date: Mon, 18 Aug 2025 19:53:11 -0700 Subject: [PATCH 37/93] [REVIEW] Add a public API for CAGRA graph optimize (#1260) This change exposes CAGRA's graph optimization functionality through a new public API in the `cuvs::neighbors::cagra::helpers` namespace. The optimize function allows users to optimize a custom KNN graphs to create CAGRA search graphs. Key Changes: - Added `cpp/include/cuvs/neighbors/cagra_optimize.hpp` with a public API declaration. As cagra_optimize assumes both input and output are host matrices, this PR exposes the API with host inputs only. Given this, users need to explicitly manage device-to-host transfers if KNN graph is not a RAFT host matrix - Modified previously unused `cpp/src/neighbors/cagra_optimize.cu` with implementation that forwards to internal CAGRA optimize. Removed the version using device matrix for `knn_graph`, which generates runtime segfaults - Added unit tests for the new API This PR partially resolves https://github.com/rapidsai/cuvs/issues/456 and completes one of the steps for https://github.com/rapidsai/cuvs/issues/1146 Authors: - Rui Lan (https://github.com/abc99lr) Approvers: - Corey J. Nolet (https://github.com/cjnolet) URL: https://github.com/rapidsai/cuvs/pull/1260 --- cpp/include/cuvs/neighbors/cagra.hpp | 1 + cpp/include/cuvs/neighbors/cagra_optimize.hpp | 49 ++++++++++++++ cpp/src/neighbors/cagra_optimize.cu | 22 ++---- cpp/tests/CMakeLists.txt | 5 ++ .../ann_cagra/test_optimize_uint32_t.cu | 67 +++++++++++++++++++ 5 files changed, 127 insertions(+), 17 deletions(-) create mode 100644 cpp/include/cuvs/neighbors/cagra_optimize.hpp create mode 100644 cpp/tests/neighbors/ann_cagra/test_optimize_uint32_t.cu diff --git a/cpp/include/cuvs/neighbors/cagra.hpp b/cpp/include/cuvs/neighbors/cagra.hpp index 6c7d1a1979..fd1fe72b34 100644 --- a/cpp/include/cuvs/neighbors/cagra.hpp +++ b/cpp/include/cuvs/neighbors/cagra.hpp @@ -2686,3 +2686,4 @@ auto distribute(const raft::resources& clique, const std::string& filename) } // namespace cuvs::neighbors::cagra #include +#include diff --git a/cpp/include/cuvs/neighbors/cagra_optimize.hpp b/cpp/include/cuvs/neighbors/cagra_optimize.hpp new file mode 100644 index 0000000000..ef0ae90061 --- /dev/null +++ b/cpp/include/cuvs/neighbors/cagra_optimize.hpp @@ -0,0 +1,49 @@ +/* + * Copyright (c) 2025, NVIDIA CORPORATION. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include +#include +#include +#include + +namespace cuvs::neighbors::cagra::helpers { + +/** + * @brief Optimize a KNN graph into a CAGRA graph. + * + * This function optimizes a k-NN graph to create a CAGRA graph. + * The input/output graphs must be on host memory. + * + * Usage example: + * @code{.cpp} + * raft::resources res; + * auto h_knn = raft::make_host_matrix(N, K_in); + * // Fill h_knn with KNN graph + * auto h_out = raft::make_host_matrix(N, K_out); + * cuvs::neighbors::cagra::helpers::optimize(res, h_knn.view(), h_out.view()); + * @endcode + * + * @param[in] handle RAFT resources + * @param[in] knn_graph Input KNN graph on host [n_rows, k_in] + * @param[out] new_graph Output CAGRA graph on host [n_rows, k_out] + */ +void optimize(raft::resources const& handle, + raft::host_matrix_view knn_graph, + raft::host_matrix_view new_graph); + +} // namespace cuvs::neighbors::cagra::helpers diff --git a/cpp/src/neighbors/cagra_optimize.cu b/cpp/src/neighbors/cagra_optimize.cu index 436d5d321b..b26239f2d3 100644 --- a/cpp/src/neighbors/cagra_optimize.cu +++ b/cpp/src/neighbors/cagra_optimize.cu @@ -1,5 +1,5 @@ /* - * Copyright (c) 2024, NVIDIA CORPORATION. + * Copyright (c) 2024-2025, NVIDIA CORPORATION. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -15,27 +15,15 @@ */ #include "cagra.cuh" -#include +#include -namespace cuvs::neighbors::cagra { +namespace cuvs::neighbors::cagra::helpers { -void optimize(raft::resources const& handle, - raft::device_matrix_view knn_graph, - raft::host_matrix_view new_graph) -{ - cuvs::neighbors::cagra::optimize< - uint32_t, - raft::host_device_accessor, - raft::memory_type::device>>(handle, knn_graph, new_graph); -} void optimize(raft::resources const& handle, raft::host_matrix_view knn_graph, raft::host_matrix_view new_graph) { - cuvs::neighbors::cagra::optimize< - uint32_t, - raft::host_device_accessor, - raft::memory_type::host>>(handle, knn_graph, new_graph); + cuvs::neighbors::cagra::optimize(handle, knn_graph, new_graph); } -} // namespace cuvs::neighbors::cagra +} // namespace cuvs::neighbors::cagra::helpers diff --git a/cpp/tests/CMakeLists.txt b/cpp/tests/CMakeLists.txt index a1b037f706..a266b577ed 100644 --- a/cpp/tests/CMakeLists.txt +++ b/cpp/tests/CMakeLists.txt @@ -166,6 +166,11 @@ if(BUILD_TESTS) 1 PERCENT 100 ) + ConfigureTest( + NAME NEIGHBORS_ANN_CAGRA_HELPERS_TEST PATH neighbors/ann_cagra/test_optimize_uint32_t.cu GPUS 1 + PERCENT 100 + ) + ConfigureTest( NAME NEIGHBORS_ANN_CAGRA_HALF_UINT32_TEST PATH neighbors/ann_cagra/test_half_uint32_t.cu GPUS 1 PERCENT 100 diff --git a/cpp/tests/neighbors/ann_cagra/test_optimize_uint32_t.cu b/cpp/tests/neighbors/ann_cagra/test_optimize_uint32_t.cu new file mode 100644 index 0000000000..55058affee --- /dev/null +++ b/cpp/tests/neighbors/ann_cagra/test_optimize_uint32_t.cu @@ -0,0 +1,67 @@ +/* + * Copyright (c) 2025, NVIDIA CORPORATION. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include +#include +#include +#include + +// This test targets public API exposure and basic invariants only (shapes, in-range indices). +// Detailed optimization correctness is exercised by CAGRA build tests. + +namespace { + +using IdxT = uint32_t; + +// Helper to create a simple synthetic KNN graph (ring-like neighbors) +auto make_ring_knn_host(int64_t num_rows, int64_t kin) +{ + auto knn_graph = raft::make_host_matrix(num_rows, kin); + for (int64_t i = 0; i < num_rows; ++i) { + for (int64_t j = 0; j < kin; ++j) { + knn_graph(i, j) = static_cast((i + j + 1) % num_rows); + } + } + return knn_graph; +} + +TEST(CagraOptimize, HostToHostOptimizesGraph) +{ + raft::resources res; + + constexpr int64_t num_rows = 8; + constexpr int64_t kin = 8; + constexpr int64_t kout = 4; + + auto knn_graph = make_ring_knn_host(num_rows, kin); + auto optimized_graph = raft::make_host_matrix(num_rows, kout); + + // Test the optimize API + cuvs::neighbors::cagra::helpers::optimize(res, knn_graph.view(), optimized_graph.view()); + + // Check basic invariants + ASSERT_EQ(optimized_graph.extent(0), num_rows); + ASSERT_EQ(optimized_graph.extent(1), kout); + + // Check that all neighbors are valid indices + for (int64_t i = 0; i < num_rows; ++i) { + for (int64_t j = 0; j < kout; ++j) { + EXPECT_LT(optimized_graph(i, j), static_cast(num_rows)); + } + } +} + +} // namespace From d0a83d43f7be934f0c4dcef585302a7fad287df0 Mon Sep 17 00:00:00 2001 From: Paul Taylor <178183+trxcllnt@users.noreply.github.com> Date: Tue, 19 Aug 2025 11:14:50 -0700 Subject: [PATCH 38/93] Use build cluster in devcontainers (#1240) RAPIDS has deployed an autoscaling cloud build cluster that can be used to accelerate building large RAPIDS projects. This contributes to https://github.com/rapidsai/build-planning/issues/209. Authors: - Paul Taylor (https://github.com/trxcllnt) Approvers: - Bradley Dice (https://github.com/bdice) URL: https://github.com/rapidsai/cuvs/pull/1240 --- .devcontainer/Dockerfile | 34 +++++++++++++++++-- .../cuda12.9-conda/devcontainer.json | 4 ++- .devcontainer/cuda12.9-pip/devcontainer.json | 6 ++-- .github/workflows/pr.yaml | 15 +++++--- 4 files changed, 50 insertions(+), 9 deletions(-) diff --git a/.devcontainer/Dockerfile b/.devcontainer/Dockerfile index 77b90fa205..31e3837840 100644 --- a/.devcontainer/Dockerfile +++ b/.devcontainer/Dockerfile @@ -21,6 +21,8 @@ ENV DEFAULT_CONDA_ENV=rapids FROM ${PYTHON_PACKAGE_MANAGER}-base +ARG TARGETARCH + ARG CUDA ENV CUDAARCHS="RAPIDS" ENV CUDA_VERSION="${CUDA_VERSION:-${CUDA}}" @@ -32,7 +34,35 @@ ENV PYTHONSAFEPATH="1" ENV PYTHONUNBUFFERED="1" ENV PYTHONDONTWRITEBYTECODE="1" +ENV HISTFILE="/home/coder/.cache/._bash_history" + +### +# sccache configuration +### +ENV AWS_ROLE_ARN="arn:aws:iam::279114543810:role/nv-gha-token-sccache-devs" ENV SCCACHE_REGION="us-east-2" ENV SCCACHE_BUCKET="rapids-sccache-devs" -ENV VAULT_HOST="https://vault.ops.k8s.rapids.ai" -ENV HISTFILE="/home/coder/.cache/._bash_history" +# 2hr (1 minute longer than sccache-dist request timeout) +ENV SCCACHE_IDLE_TIMEOUT=7200 + +### +# sccache-dist configuration +### +# Enable sccache-dist by default +ENV DEVCONTAINER_UTILS_ENABLE_SCCACHE_DIST=1 +# Compile locally if max retries exceeded +ENV SCCACHE_DIST_FALLBACK_TO_LOCAL_COMPILE=true +# Retry transient errors 4 times (for a total of 5 attempts) +ENV SCCACHE_DIST_MAX_RETRIES=4 +ENV SCCACHE_DIST_CONNECT_TIMEOUT=30 +ENV SCCACHE_DIST_CONNECTION_POOL=false +# 1hr 59min (to accommodate debug builds) +ENV SCCACHE_DIST_REQUEST_TIMEOUT=7140 +ENV SCCACHE_DIST_KEEPALIVE_ENABLED=true +ENV SCCACHE_DIST_KEEPALIVE_INTERVAL=20 +ENV SCCACHE_DIST_KEEPALIVE_TIMEOUT=600 +ENV SCCACHE_DIST_URL="https://${TARGETARCH}.linux.sccache.rapids.nvidia.com" + +# Build as much in parallel as possible +ENV INFER_NUM_DEVICE_ARCHITECTURES=1 +ENV MAX_DEVICE_OBJ_TO_COMPILE_IN_PARALLEL=20 diff --git a/.devcontainer/cuda12.9-conda/devcontainer.json b/.devcontainer/cuda12.9-conda/devcontainer.json index 3c1da8fe73..be50d56c15 100644 --- a/.devcontainer/cuda12.9-conda/devcontainer.json +++ b/.devcontainer/cuda12.9-conda/devcontainer.json @@ -11,7 +11,9 @@ "runArgs": [ "--rm", "--name", - "${localEnv:USER:anon}-rapids-${localWorkspaceFolderBasename}-25.10-cuda12.9-conda" + "${localEnv:USER:anon}-rapids-${localWorkspaceFolderBasename}-25.10-cuda12.9-conda", + "--ulimit", + "nofile=500000" ], "hostRequirements": {"gpu": "optional"}, "features": { diff --git a/.devcontainer/cuda12.9-pip/devcontainer.json b/.devcontainer/cuda12.9-pip/devcontainer.json index 0305705fb9..abb88bb356 100644 --- a/.devcontainer/cuda12.9-pip/devcontainer.json +++ b/.devcontainer/cuda12.9-pip/devcontainer.json @@ -5,13 +5,15 @@ "args": { "CUDA": "12.9", "PYTHON_PACKAGE_MANAGER": "pip", - "BASE": "rapidsai/devcontainers:25.10-cpp-cuda12.9-ucx1.18.0-openmpi5.0.7" + "BASE": "rapidsai/devcontainers:25.10-cpp-cuda12.9-ucx1.19.0-openmpi5.0.7" } }, "runArgs": [ "--rm", "--name", - "${localEnv:USER:anon}-rapids-${localWorkspaceFolderBasename}-25.10-cuda12.9-pip" + "${localEnv:USER:anon}-rapids-${localWorkspaceFolderBasename}-25.10-cuda12.9-pip", + "--ulimit", + "nofile=500000" ], "hostRequirements": {"gpu": "optional"}, "features": { diff --git a/.github/workflows/pr.yaml b/.github/workflows/pr.yaml index 89795ed967..41fbc93a84 100644 --- a/.github/workflows/pr.yaml +++ b/.github/workflows/pr.yaml @@ -223,12 +223,19 @@ jobs: needs: telemetry-setup uses: rapidsai/shared-workflows/.github/workflows/build-in-devcontainer.yaml@branch-25.10 with: - arch: '["amd64"]' + arch: '["amd64", "arm64"]' cuda: '["12.9"]' + node_type: "cpu8" + rapids-aux-secret-1: GIST_REPO_READ_ORG_GITHUB_TOKEN + env: | + SCCACHE_DIST_MAX_RETRIES=inf + SCCACHE_SERVER_LOG=sccache=debug + SCCACHE_DIST_FALLBACK_TO_LOCAL_COMPILE=false + SCCACHE_DIST_AUTH_TOKEN_VAR=RAPIDS_AUX_SECRET_1 build_command: | - sccache -z; - build-all --verbose; - sccache -s; + sccache --zero-stats; + build-all -j0 --verbose 2>&1 | tee telemetry-artifacts/build.log; + sccache --show-adv-stats | tee telemetry-artifacts/sccache-stats.txt; telemetry-summarize: # This job must use a self-hosted runner to record telemetry traces. runs-on: linux-amd64-cpu4 From 989c994e3d15d3b81c86aa1d5d10cb95a9cb9e88 Mon Sep 17 00:00:00 2001 From: Puneet Ahuja <167976912+punAhuja@users.noreply.github.com> Date: Fri, 22 Aug 2025 09:33:41 +0530 Subject: [PATCH 39/93] [Java]Binary and scalar quantization (#1104) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Java support for Binary and Scalar Quantization -Incorporated changes for indexing and searching on quantized dataset -Added tests to test quantization, indexing and searching on quantized dataset, with quantized query Authors: - Puneet Ahuja (https://github.com/punAhuja) - Lorenzo Dematté (https://github.com/ldematte) Approvers: - MithunR (https://github.com/mythrocks) - Corey J. Nolet (https://github.com/cjnolet) URL: https://github.com/rapidsai/cuvs/pull/1104 --- .../cuvs/preprocessing/quantize/binary.h | 6 +- .../java/com/nvidia/cuvs/BinaryQuantizer.java | 125 ++++ .../main/java/com/nvidia/cuvs/CagraIndex.java | 8 + .../main/java/com/nvidia/cuvs/CagraQuery.java | 108 +++- .../main/java/com/nvidia/cuvs/CuVSMatrix.java | 18 + .../java/com/nvidia/cuvs/CuVSQuantizer.java | 71 +++ .../com/nvidia/cuvs/Scalar8BitQuantizer.java | 123 ++++ .../com/nvidia/cuvs/spi/CuVSProvider.java | 27 + .../nvidia/cuvs/spi/UnsupportedProvider.java | 45 +- .../cuvs/internal/BinaryQuantizerImpl.java | 324 ++++++++++ .../nvidia/cuvs/internal/CagraIndexImpl.java | 130 +++- .../cuvs/internal/CuVSHostMatrixImpl.java | 7 + .../cuvs/internal/CuVSMatrixBaseImpl.java | 6 + .../internal/Scalar8BitQuantizerImpl.java | 461 ++++++++++++++ .../com/nvidia/cuvs/spi/JDKProvider.java | 41 ++ .../java/com/nvidia/cuvs/QuantizationIT.java | 588 ++++++++++++++++++ java/panama-bindings/headers.h | 2 + 17 files changed, 2053 insertions(+), 37 deletions(-) create mode 100644 java/cuvs-java/src/main/java/com/nvidia/cuvs/BinaryQuantizer.java create mode 100644 java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSQuantizer.java create mode 100644 java/cuvs-java/src/main/java/com/nvidia/cuvs/Scalar8BitQuantizer.java create mode 100644 java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/BinaryQuantizerImpl.java create mode 100644 java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/Scalar8BitQuantizerImpl.java create mode 100644 java/cuvs-java/src/test/java/com/nvidia/cuvs/QuantizationIT.java diff --git a/cpp/include/cuvs/preprocessing/quantize/binary.h b/cpp/include/cuvs/preprocessing/quantize/binary.h index 9a5834ef3e..ec58bae3bd 100644 --- a/cpp/include/cuvs/preprocessing/quantize/binary.h +++ b/cpp/include/cuvs/preprocessing/quantize/binary.h @@ -30,7 +30,7 @@ extern "C" { * and sampling_median thresholds are calculated separately for each dimension. * */ -enum cuvsBinaryQuantizerThreshold { ZERO, MEAN, SAMPLING_MEDIAN }; +typedef enum { ZERO, MEAN, SAMPLING_MEDIAN } cuvsBinaryQuantizerThreshold; /** * @brief Binary quantizer parameters. @@ -39,12 +39,12 @@ struct cuvsBinaryQuantizerParams { /* * specifies the threshold to set a bit in cuvsBinaryQuantizerTransform */ - cuvsBinaryQuantizerThreshold threshold = MEAN; + cuvsBinaryQuantizerThreshold threshold; /* * specifies the sampling ratio */ - float sampling_ratio = 0.1; + float sampling_ratio; }; typedef struct cuvsBinaryQuantizerParams* cuvsBinaryQuantizerParams_t; diff --git a/java/cuvs-java/src/main/java/com/nvidia/cuvs/BinaryQuantizer.java b/java/cuvs-java/src/main/java/com/nvidia/cuvs/BinaryQuantizer.java new file mode 100644 index 0000000000..82584d6cac --- /dev/null +++ b/java/cuvs-java/src/main/java/com/nvidia/cuvs/BinaryQuantizer.java @@ -0,0 +1,125 @@ +/* + * Copyright (c) 2025, NVIDIA CORPORATION. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.nvidia.cuvs; + +import com.nvidia.cuvs.CuVSMatrix.DataType; +import com.nvidia.cuvs.spi.CuVSProvider; + +/** Binary quantizer that transforms float datasets into packed binary datasets. */ +public class BinaryQuantizer implements CuVSQuantizer { + private final CuVSResources resources; + private final ThresholdType thresholdType; + + /** Output data type (BYTE for packed binary data). */ + private final DataType outputDataType = DataType.BYTE; + + private final Object impl; + + /** Creates a binary quantizer with training data. */ + public BinaryQuantizer(CuVSResources resources, CuVSMatrix trainingDataset) throws Throwable { + this(resources, trainingDataset, ThresholdType.MEAN); + } + + /** Creates a binary quantizer with specified threshold type and training data. */ + public BinaryQuantizer( + CuVSResources resources, CuVSMatrix trainingDataset, ThresholdType thresholdType) + throws Throwable { + if (trainingDataset == null) { + throw new IllegalArgumentException("Training dataset cannot be null"); + } + if (trainingDataset.dataType() != DataType.FLOAT) { + throw new IllegalArgumentException( + "Training dataset must have FLOAT data type, got " + trainingDataset.dataType()); + } + this.resources = resources; + this.thresholdType = thresholdType; + this.impl = + CuVSProvider.provider() + .createBinaryQuantizerImpl(resources, trainingDataset, thresholdType); + } + + /** Threshold types for binary quantization. */ + public enum ThresholdType { + /** Zero threshold */ + ZERO(0), + /** Mean threshold */ + MEAN(1), + /** Sampling median threshold */ + SAMPLING_MEDIAN(2); + + private final int value; + + ThresholdType(int value) { + this.value = value; + } + + public int getValue() { + return value; + } + } + + /** Returns the output data type (BYTE). */ + @Override + public DataType outputDataType() { + return outputDataType; + } + + /** Gets the threshold type used for binary quantization. */ + public ThresholdType getThresholdType() { + return thresholdType; + } + + @Override + public CuVSMatrix transform(CuVSMatrix input) throws Throwable { + if (input.dataType() != DataType.FLOAT) { + throw new IllegalArgumentException( + "BinaryQuantizer requires FLOAT input, got " + input.dataType()); + } + + CuVSMatrix result = CuVSProvider.provider().transformBinaryWithImpl(impl, input); + + if (result.dataType() != DataType.BYTE) { + throw new IllegalStateException( + "Expected BYTE output from binary quantization, got " + result.dataType()); + } + + return result; + } + + @Override + public void train(CuVSMatrix trainingData) throws Throwable { + throw new UnsupportedOperationException("Training handled during construction"); + } + + @Override + public CuVSMatrix inverseTransform(CuVSMatrix quantizedData) throws Throwable { + throw new UnsupportedOperationException( + "Binary quantization does not support inverse transformation"); + } + + @Override + public void close() throws Exception { + try { + CuVSProvider.provider().closeBinaryQuantizer(impl); + } catch (Throwable t) { + if (t instanceof Exception) { + throw (Exception) t; + } else { + throw new RuntimeException("Error closing BinaryQuantizer", t); + } + } + } +} diff --git a/java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraIndex.java b/java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraIndex.java index e3b885a333..c0032f7ea4 100644 --- a/java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraIndex.java +++ b/java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraIndex.java @@ -247,6 +247,14 @@ interface Builder { */ Builder withIndexParams(CagraIndexParams cagraIndexParameters); + /** + * Sets the quantizer to use for building the index with quantized data. + * + * @param quantizer The quantizer to apply to the dataset during index construction + * @return An instance of this Builder + */ + Builder withQuantizer(CuVSQuantizer quantizer); + /** * Builds and returns an instance of CagraIndex. * diff --git a/java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraQuery.java b/java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraQuery.java index f81ec3191c..c976949bf7 100644 --- a/java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraQuery.java +++ b/java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraQuery.java @@ -15,14 +15,15 @@ */ package com.nvidia.cuvs; +import com.nvidia.cuvs.CuVSMatrix.DataType; import java.util.Arrays; import java.util.BitSet; import java.util.Objects; import java.util.function.LongToIntFunction; /** - * CagraQuery holds the CagraSearchParams and the query vectors to be used while - * invoking search. + * CagraQuery holds the search parameters plus either raw float[][] vectors + * or a quantized {@link CuVSMatrix} for querying a CAGRA index. * *

Thread Safety: Each CagraQuery instance should use its own * CuVSResources object that is not shared with other threads. Sharing CuVSResources @@ -33,8 +34,9 @@ public class CagraQuery { private final CagraSearchParams cagraSearchParameters; - private final LongToIntFunction mapping; private final float[][] queryVectors; + private final CuVSMatrix quantizedQueries; + private final LongToIntFunction mapping; private final int topK; private final BitSet prefilter; private final int numDocs; @@ -47,23 +49,25 @@ public class CagraQuery { * @param cagraSearchParameters an instance of {@link CagraSearchParams} holding * the search parameters * @param queryVectors 2D float query vector array + * @param quantizedQueries 2D quantized query vector array * @param mapping a function mapping ordinals (neighbor IDs) to custom user IDs * @param topK the top k results to return * @param prefilter A single BitSet to use as filter while searching the CAGRA index * @param numDocs Total number of dataset vectors; used to align the prefilter correctly * @param resources CuVSResources instance to use for this query */ - public CagraQuery( + private CagraQuery( CagraSearchParams cagraSearchParameters, float[][] queryVectors, + CuVSMatrix quantizedQueries, LongToIntFunction mapping, int topK, BitSet prefilter, int numDocs, CuVSResources resources) { - super(); this.cagraSearchParameters = cagraSearchParameters; this.queryVectors = queryVectors; + this.quantizedQueries = quantizedQueries; this.mapping = mapping; this.topK = topK; this.prefilter = prefilter; @@ -71,6 +75,11 @@ public CagraQuery( this.resources = resources; } + /** Start building a new CagraQuery. */ + public static Builder newBuilder(CuVSResources resources) { + return new Builder(resources); + } + /** * Gets the instance of CagraSearchParams initially set. * @@ -81,14 +90,35 @@ public CagraSearchParams getCagraSearchParameters() { } /** - * Gets the query vector 2D float array. - * - * @return 2D float array + * If this query was built without a quantizer, returns the original float vectors. + * Otherwise returns null. */ public float[][] getQueryVectors() { return queryVectors; } + /** + * If this query was built with a quantizer, returns the quantized Dataset. + * Otherwise returns null. + */ + public CuVSMatrix getQuantizedQueries() { + return quantizedQueries; + } + + /** True if this query carries a quantized Dataset instead of float[][] */ + public boolean hasQuantizedQueries() { + return quantizedQueries != null; + } + + /** + * Returns the data type of the query payload: + * - 32 for float32 queries + * - 8 for quantized queries + */ + public DataType getQueryDataType() { + return quantizedQueries != null ? quantizedQueries.dataType() : DataType.FLOAT; + } + /** * Gets the function mapping ordinals (neighbor IDs) to custom user IDs */ @@ -134,10 +164,13 @@ public CuVSResources getResources() { @Override public String toString() { - return "CuVSQuery [cagraSearchParameters=" + return "CagraQuery[" + + "params=" + cagraSearchParameters - + ", queryVectors=" - + Arrays.toString(queryVectors) + + ", floatVectors=" + + (queryVectors != null ? Arrays.toString(queryVectors) : "null") + + ", quantized=" + + (quantizedQueries != null ? ("Dataset@" + quantizedQueries.dataType() + "-bit") : "false") + ", mapping=" + mapping + ", topK=" @@ -156,6 +189,7 @@ public static class Builder { private int topK = 2; private BitSet prefilter; private int numDocs; + private CuVSQuantizer quantizer; private final CuVSResources resources; /** @@ -174,12 +208,12 @@ public Builder(CuVSResources resources) { /** * Sets the instance of configured CagraSearchParams to be passed for search. * - * @param cagraSearchParams an instance of the configured CagraSearchParams to - * be used for this query + * @param params an instance of the configured CagraSearchParams to + * be used for this query * @return an instance of this Builder */ - public Builder withSearchParams(CagraSearchParams cagraSearchParams) { - this.cagraSearchParams = cagraSearchParams; + public Builder withSearchParams(CagraSearchParams params) { + this.cagraSearchParams = params; return this; } @@ -233,13 +267,47 @@ public Builder withPrefilter(BitSet prefilter, int numDocs) { } /** - * Builds an instance of CuVSQuery. - * - * @return an instance of CuVSQuery + * Specify a quantizer to automatically transform the float[][] queryVectors + * into a quantized {@link CuVSMatrix} using the same quantizer used for training. + */ + public Builder withQuantizer(CuVSQuantizer quantizer) { + this.quantizer = quantizer; + return this; + } + + /** + * Builds the CagraQuery. If a quantizer was provided, queryVectors is ignored + * and a quantized Dataset is produced instead. */ - public CagraQuery build() { + public CagraQuery build() throws Throwable { + if (queryVectors == null) { + throw new IllegalArgumentException("Query vectors must be provided"); + } + + CuVSMatrix quantized = null; + float[][] floatsForQuery = queryVectors; + + if (quantizer != null) { + // wrap float[][] in a CuVSMatrix and quantize + try (CuVSMatrix tmp = CuVSMatrix.ofArray(queryVectors)) { + if (tmp.dataType() != DataType.FLOAT) { + throw new IllegalArgumentException( + "Query quantization requires FLOAT input, got " + tmp.dataType()); + } + quantized = quantizer.transform(tmp); + } + floatsForQuery = null; + } + return new CagraQuery( - cagraSearchParams, queryVectors, mapping, topK, prefilter, numDocs, resources); + cagraSearchParams, + floatsForQuery, + quantized, + mapping, + topK, + prefilter, + numDocs, + resources); } } } diff --git a/java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSMatrix.java b/java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSMatrix.java index f8e60c5555..5143c8bbc7 100644 --- a/java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSMatrix.java +++ b/java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSMatrix.java @@ -69,6 +69,17 @@ static CuVSMatrix ofArray(byte[][] vectors) { return CuVSProvider.provider().newMatrixFromArray(vectors); } + /** + * Returns the data type of this matrix. + * + *

DataType provides the fundamental type information needed for all operations. + * FLOAT represents 32-bit floating point data, INT represents 32-bit integer data, + * and BYTE represents 8-bit data (used for quantized values). + * + * @return The DataType of this matrix + */ + DataType dataType(); + interface Builder { /** * Add a single vector to the dataset. @@ -152,6 +163,13 @@ static CuVSMatrix.Builder builder(int size, int columns, DataType dataType) { */ void toArray(byte[][] array); + /** + * Gets where the content is stored + * + * @return MemoryKind whether HOST or DEVICE + */ + CuVSMatrix.MemoryKind memoryKind(); + @Override void close(); } diff --git a/java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSQuantizer.java b/java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSQuantizer.java new file mode 100644 index 0000000000..2c86eff0f1 --- /dev/null +++ b/java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSQuantizer.java @@ -0,0 +1,71 @@ +/* + * Copyright (c) 2025, NVIDIA CORPORATION. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.nvidia.cuvs; + +import com.nvidia.cuvs.CuVSMatrix.DataType; + +/** + * Base interface for all cuVS quantizers providing unified quantization operations. + * + */ +public interface CuVSQuantizer extends AutoCloseable { + + /** + * Returns the data type of quantized data produced by this quantizer. + * + * @return The DataType of the quantized output (BYTE for quantized data) + */ + DataType outputDataType(); + + /** + * Transforms a float32 dataset into a quantized dataset. + * + * @param input The input dataset of float32 vectors (must have precision() == 32) + * @return A new Dataset containing the quantized vectors with this quantizer's precision + */ + CuVSMatrix transform(CuVSMatrix input) throws Throwable; + + /** + * Optional training method - only applies to quantizers that require training. + * + * @param trainingData The dataset to train on (must have precision() == 32) + * @throws Throwable if an error occurs during training + */ + default void train(CuVSMatrix trainingData) throws Throwable { + throw new UnsupportedOperationException( + "Training not supported for " + getClass().getSimpleName()); + } + + /** + * Optional inverse transformation - only applies to quantizers that support reconstruction + * + * @param quantizedData The quantized dataset to inverse transform + * (must have precision() matching this quantizer's precision()) + * @return A float[][] array containing the de-quantized vectors + */ + default CuVSMatrix inverseTransform(CuVSMatrix quantizedData) throws Throwable { + throw new UnsupportedOperationException( + "Inverse transform not supported for " + getClass().getSimpleName()); + } + + /** + * Closes this quantizer and releases any associated resources. + * + * @throws Exception if an error occurs while closing + */ + @Override + void close() throws Exception; +} diff --git a/java/cuvs-java/src/main/java/com/nvidia/cuvs/Scalar8BitQuantizer.java b/java/cuvs-java/src/main/java/com/nvidia/cuvs/Scalar8BitQuantizer.java new file mode 100644 index 0000000000..056a7ad468 --- /dev/null +++ b/java/cuvs-java/src/main/java/com/nvidia/cuvs/Scalar8BitQuantizer.java @@ -0,0 +1,123 @@ +/* + * Copyright (c) 2025, NVIDIA CORPORATION. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.nvidia.cuvs; + +import com.nvidia.cuvs.CuVSMatrix.DataType; +import com.nvidia.cuvs.spi.CuVSProvider; + +/** + * Scalar 8-bit quantizer implementation that transforms 32-bit float datasets into 8-bit integer datasets. + * + *

This quantizer reduces each float32 value to an 8-bit signed integer, providing significant + * compression while maintaining reasonable precision. The quantizer requires training on a representative + * dataset to compute optimal quantization parameters (min/max values per dimension). + * + *

The quantizer supports both transformation and inverse transformation, allowing approximate + * recovery of the original float32 values from the quantized representation. + * + * @since 25.08 + */ +public class Scalar8BitQuantizer implements CuVSQuantizer { + private final Object impl; + + /** + * The data type used by this quantizer (BYTE for 8-bit signed integers). + */ + private final DataType outputDataType = DataType.BYTE; + + /** + * Creates a new scalar 8-bit quantizer with CuVSMatrix training data. + * + *

The training dataset is used to compute quantization parameters (min/max values) + * for each dimension. The quantizer will be immediately ready for use after construction. + * + * @param resources The CuVS resources to use for quantization operations + * @param trainingDataset A CuVSMatrix containing training vectors (must have FLOAT data type) + * @throws Throwable if an error occurs during quantizer training + * @throws IllegalArgumentException if trainingDataset is null or doesn't have FLOAT data type + */ + public Scalar8BitQuantizer(CuVSResources resources, CuVSMatrix trainingDataset) throws Throwable { + if (trainingDataset == null) { + throw new IllegalArgumentException("Training dataset cannot be null"); + } + if (trainingDataset.dataType() != DataType.FLOAT) { + throw new IllegalArgumentException( + "Training dataset must have FLOAT data type, got " + trainingDataset.dataType()); + } + this.impl = CuVSProvider.provider().createScalar8BitQuantizerImpl(resources, trainingDataset); + } + + /** + * Returns the data type of quantized data produced by this quantizer. + * + * @return The DataType (BYTE for this scalar quantizer) + */ + @Override + public DataType outputDataType() { + return outputDataType; + } + + @Override + public CuVSMatrix transform(CuVSMatrix input) throws Throwable { + // Validate input data type + if (input.dataType() != DataType.FLOAT) { + throw new IllegalArgumentException( + "Scalar8BitQuantizer requires FLOAT input, got " + input.dataType()); + } + + CuVSMatrix result = CuVSProvider.provider().transformScalar8Bit(impl, input); + + // Validate output data type + if (result.dataType() != DataType.BYTE) { + throw new IllegalStateException( + "Expected BYTE output from scalar quantization, got " + result.dataType()); + } + + return result; + } + + @Override + public void train(CuVSMatrix trainingData) throws Throwable { + // Training is handled during construction for ScalarQuantizer + // This method could be implemented if you want to support re-training + throw new UnsupportedOperationException("Training handled during construction"); + } + + @Override + public CuVSMatrix inverseTransform(CuVSMatrix quantizedData) throws Throwable { + // Validate input data type for inverse transform + if (quantizedData.dataType() != DataType.BYTE) { + throw new IllegalArgumentException( + "Inverse transform requires BYTE input, got " + quantizedData.dataType()); + } + + return CuVSProvider.provider().inverseTransformScalar8Bit(impl, quantizedData); + } + + @Override + public void close() throws Exception { + try { + CuVSProvider.provider().closeScalar8BitQuantizer(impl); + } catch (Throwable t) { + if (t instanceof Exception) { + throw (Exception) t; + } else { + // Wrap non-Exception throwables (like Error) in RuntimeException + throw new RuntimeException("Error closing Scalar8BitQuantizer", t); + } + } + } +} 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 e528beff1a..316406b24e 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 @@ -15,6 +15,7 @@ */ package com.nvidia.cuvs.spi; +import com.nvidia.cuvs.BinaryQuantizer; import com.nvidia.cuvs.BruteForceIndex; import com.nvidia.cuvs.CagraIndex; import com.nvidia.cuvs.CagraMergeParams; @@ -120,6 +121,32 @@ default CagraIndex mergeCagraIndexes(CagraIndex[] indexes, CagraMergeParams merg return mergeCagraIndexes(indexes); } + /** Creates a new Scalar8BitQuantizer implementation. */ + Object createScalar8BitQuantizerImpl(CuVSResources resources, CuVSMatrix trainingDataset) + throws Throwable; + + /** Performs inverse transform using a Scalar8BitQuantizer. */ + CuVSMatrix inverseTransformScalar8Bit(Object impl, CuVSMatrix quantizedData) throws Throwable; + + /** Transforms dataset using Scalar8BitQuantizer */ + CuVSMatrix transformScalar8Bit(Object impl, CuVSMatrix input) throws Throwable; + + /** Closes Scalar8BitQuantizer implementation */ + void closeScalar8BitQuantizer(Object impl) throws Throwable; + + /** Creates a new BinaryQuantizer implementation. */ + Object createBinaryQuantizerImpl( + CuVSResources resources, + CuVSMatrix trainingDataset, + BinaryQuantizer.ThresholdType thresholdType) + throws Throwable; + + /** Performs transform using a BinaryQuantizer. */ + CuVSMatrix transformBinaryWithImpl(Object impl, CuVSMatrix input) throws Throwable; + + /** Closes BinaryQuantizer implementation. */ + void closeBinaryQuantizer(Object impl) throws Throwable; + /** Retrieves the system-wide provider. */ static CuVSProvider provider() { return CuVSServiceProvider.Holder.INSTANCE; diff --git a/java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/UnsupportedProvider.java b/java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/UnsupportedProvider.java index 8f65bf7068..cd6604663b 100644 --- a/java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/UnsupportedProvider.java +++ b/java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/UnsupportedProvider.java @@ -15,6 +15,7 @@ */ package com.nvidia.cuvs.spi; +import com.nvidia.cuvs.BinaryQuantizer; import com.nvidia.cuvs.BruteForceIndex; import com.nvidia.cuvs.CagraIndex; import com.nvidia.cuvs.CuVSMatrix; @@ -25,7 +26,8 @@ import java.nio.file.Path; /** - * A provider that unconditionally throws UnsupportedOperationException. + * A provider that throws UnsupportedOperationException for all operations. + * Used as a fallback when no proper implementation is available. */ final class UnsupportedProvider implements CuVSProvider { @@ -84,4 +86,45 @@ public CuVSMatrix newMatrixFromArray(int[][] vectors) { public CuVSMatrix newMatrixFromArray(byte[][] vectors) { throw new UnsupportedOperationException(); } + + @Override + public Object createScalar8BitQuantizerImpl(CuVSResources resources, CuVSMatrix trainingDataset) + throws Throwable { + throw new UnsupportedOperationException("CuVS is not supported on this platform"); + } + + @Override + public CuVSMatrix inverseTransformScalar8Bit(Object impl, CuVSMatrix quantizedData) + throws Throwable { + throw new UnsupportedOperationException("CuVS is not supported on this platform"); + } + + @Override + public Object createBinaryQuantizerImpl( + CuVSResources resources, + CuVSMatrix trainingDataset, + BinaryQuantizer.ThresholdType thresholdType) + throws Throwable { + throw new UnsupportedOperationException("CuVS is not supported on this platform"); + } + + @Override + public CuVSMatrix transformBinaryWithImpl(Object impl, CuVSMatrix input) throws Throwable { + throw new UnsupportedOperationException("CuVS is not supported on this platform"); + } + + @Override + public void closeBinaryQuantizer(Object impl) throws Throwable { + throw new UnsupportedOperationException("CuVS is not supported on this platform"); + } + + @Override + public CuVSMatrix transformScalar8Bit(Object impl, CuVSMatrix input) throws Throwable { + throw new UnsupportedOperationException("CuVS is not supported on this platform"); + } + + @Override + public void closeScalar8BitQuantizer(Object impl) throws Throwable { + throw new UnsupportedOperationException("CuVS is not supported on this platform"); + } } diff --git a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/BinaryQuantizerImpl.java b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/BinaryQuantizerImpl.java new file mode 100644 index 0000000000..704fb45738 --- /dev/null +++ b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/BinaryQuantizerImpl.java @@ -0,0 +1,324 @@ +/* + * Copyright (c) 2025, NVIDIA CORPORATION. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.nvidia.cuvs.internal; + +import static com.nvidia.cuvs.internal.common.CloseableRMMAllocation.allocateRMMSegment; +import static com.nvidia.cuvs.internal.common.LinkerHelper.C_FLOAT_BYTE_SIZE; +import static com.nvidia.cuvs.internal.common.LinkerHelper.C_POINTER; +import static com.nvidia.cuvs.internal.common.Util.CudaMemcpyKind.*; +import static com.nvidia.cuvs.internal.common.Util.buildMemorySegment; +import static com.nvidia.cuvs.internal.common.Util.checkCuVSError; +import static com.nvidia.cuvs.internal.common.Util.cudaMemcpy; +import static com.nvidia.cuvs.internal.common.Util.prepareTensor; +import static com.nvidia.cuvs.internal.panama.headers_h.cuvsBinaryQuantizerCreate; +import static com.nvidia.cuvs.internal.panama.headers_h.cuvsBinaryQuantizerDestroy; +import static com.nvidia.cuvs.internal.panama.headers_h.cuvsBinaryQuantizerParamsCreate; +import static com.nvidia.cuvs.internal.panama.headers_h.cuvsBinaryQuantizerParamsDestroy; +import static com.nvidia.cuvs.internal.panama.headers_h.cuvsBinaryQuantizerParams_t; +import static com.nvidia.cuvs.internal.panama.headers_h.cuvsBinaryQuantizerTrain; +import static com.nvidia.cuvs.internal.panama.headers_h.cuvsBinaryQuantizerTransformWithParams; +import static com.nvidia.cuvs.internal.panama.headers_h.cuvsBinaryQuantizer_t; +import static com.nvidia.cuvs.internal.panama.headers_h.cuvsStreamSync; +import static com.nvidia.cuvs.internal.panama.headers_h.kDLCPU; +import static com.nvidia.cuvs.internal.panama.headers_h.kDLCUDA; +import static com.nvidia.cuvs.internal.panama.headers_h.kDLFloat; +import static com.nvidia.cuvs.internal.panama.headers_h.kDLUInt; + +import com.nvidia.cuvs.BinaryQuantizer; +import com.nvidia.cuvs.CuVSMatrix; +import com.nvidia.cuvs.CuVSMatrix.DataType; +import com.nvidia.cuvs.CuVSMatrix.MemoryKind; +import com.nvidia.cuvs.CuVSResources; +import com.nvidia.cuvs.internal.common.Util; +import java.lang.foreign.Arena; +import java.lang.foreign.MemorySegment; +import java.lang.foreign.ValueLayout; + +/** + * Implementation of binary quantizer functionality using Panama Foreign Function Interface. + */ +public class BinaryQuantizerImpl { + + public static CuVSMatrix transform( + CuVSResources cuvsResources, float[][] dataset, BinaryQuantizer.ThresholdType thresholdType) + throws Throwable { + + try (var localArena = Arena.ofConfined(); + var resourcesAccessor = cuvsResources.access()) { + + long rows = dataset.length; + long cols = rows > 0 ? dataset[0].length : 0; + + MemorySegment datasetMemSegment = buildMemorySegment(localArena, dataset); + long cuvsResourcesPtr = resourcesAccessor.handle(); + + return performTransformGPU( + cuvsResourcesPtr, + localArena, + rows, + cols, + datasetMemSegment, + HOST_TO_DEVICE, + thresholdType); + } + } + + public static CuVSMatrix transform( + CuVSResources cuvsResources, CuVSMatrix dataset, BinaryQuantizer.ThresholdType thresholdType) + throws Throwable { + if (dataset.dataType() != DataType.FLOAT) { + throw new IllegalArgumentException( + "BinaryQuantizer requires FLOAT input, got " + dataset.dataType()); + } + + try (var localArena = Arena.ofConfined(); + var resourcesAccessor = cuvsResources.access()) { + + long rows = dataset.size(); + long cols = dataset.columns(); + + MemorySegment datasetMemSegment = ((CuVSMatrixBaseImpl) dataset).memorySegment(); + long cuvsResourcesPtr = resourcesAccessor.handle(); + + // TODO: Currently all datasets are HOST memory kind, preventing testing of GPU quantization + // flow. + // Future improvement: Support GPU-native dataset creation in Java API to enable direct GPU + // quantization testing. + if (dataset.memoryKind() == MemoryKind.HOST) { + return performTransformHost( + cuvsResourcesPtr, localArena, rows, cols, datasetMemSegment, thresholdType); + } else { + CuVSMatrix result = + performTransformGPU( + cuvsResourcesPtr, + localArena, + rows, + cols, + datasetMemSegment, + INFER_DIRECTION, + thresholdType); + + if (result.dataType() != DataType.BYTE) { + throw new IllegalStateException( + "Expected BYTE output from binary quantization, got " + result.dataType()); + } + return result; + } + } + } + + private static CuVSMatrix performTransformGPU( + long cuvsResourcesPtr, + Arena localArena, + long rows, + long cols, + MemorySegment datasetMemSegment, + Util.CudaMemcpyKind memcpyDirection, + BinaryQuantizer.ThresholdType thresholdType) + throws Throwable { + + final long BUFFER_SIZE_BYTES = 1024 * 1024; + final long FLOAT_SIZE = C_FLOAT_BYTE_SIZE; + final long FLOATS_PER_BUFFER = BUFFER_SIZE_BYTES / FLOAT_SIZE; + final long ROWS_PER_BUFFER = Math.max(1, FLOATS_PER_BUFFER / cols); + + var result = new CuVSHostMatrixArenaImpl(rows, cols, DataType.BYTE); + + long maxDatasetBytes = FLOAT_SIZE * ROWS_PER_BUFFER * cols; + long maxOutputBytes = ROWS_PER_BUFFER * cols; + + // Allocate quantizer objects once + MemorySegment paramsSegment = localArena.allocate(cuvsBinaryQuantizerParams_t); + int returnValue = cuvsBinaryQuantizerParamsCreate(paramsSegment); + checkCuVSError(returnValue, "cuvsBinaryQuantizerParamsCreate"); + MemorySegment paramsPtr = paramsSegment.get(C_POINTER, 0); + setBinaryQuantizerThreshold(paramsPtr, thresholdType); + + MemorySegment quantizerSegment = localArena.allocate(cuvsBinaryQuantizer_t); + returnValue = cuvsBinaryQuantizerCreate(quantizerSegment); + checkCuVSError(returnValue, "cuvsBinaryQuantizerCreate"); + MemorySegment quantizerPtr = quantizerSegment.get(C_POINTER, 0); + + try (var datasetDP = allocateRMMSegment(cuvsResourcesPtr, maxDatasetBytes); + var outputDP = allocateRMMSegment(cuvsResourcesPtr, maxOutputBytes)) { + + for (long startRow = 0; startRow < rows; startRow += ROWS_PER_BUFFER) { + long endRow = Math.min(startRow + ROWS_PER_BUFFER, rows); + long chunkRows = endRow - startRow; + long actualDatasetBytes = FLOAT_SIZE * chunkRows * cols; + long actualOutputBytes = chunkRows * cols; + + MemorySegment datasetPtr = datasetDP.handle().asSlice(0, actualDatasetBytes); + MemorySegment outputPtr = outputDP.handle().asSlice(0, actualOutputBytes); + MemorySegment chunkData = + datasetMemSegment.asSlice(startRow * cols * FLOAT_SIZE, actualDatasetBytes); + + cudaMemcpy(datasetPtr, chunkData, actualDatasetBytes, memcpyDirection); + + returnValue = cuvsStreamSync(cuvsResourcesPtr); + checkCuVSError(returnValue, "cuvsStreamSync before transform"); + + long datasetShape[] = {chunkRows, cols}; + long outputShape[] = {chunkRows, cols}; + + MemorySegment datasetTensor = + prepareTensor(localArena, datasetPtr, datasetShape, kDLFloat(), 32, kDLCUDA(), 1); + MemorySegment outputTensor = + prepareTensor(localArena, outputPtr, outputShape, kDLUInt(), 8, kDLCUDA(), 1); + + returnValue = + cuvsBinaryQuantizerTrain(cuvsResourcesPtr, paramsPtr, datasetTensor, quantizerPtr); + checkCuVSError(returnValue, "cuvsBinaryQuantizerTrain (GPU)"); + + returnValue = + cuvsBinaryQuantizerTransformWithParams( + cuvsResourcesPtr, quantizerPtr, datasetTensor, outputTensor); + checkCuVSError(returnValue, "cuvsBinaryQuantizerTransformWithParams (GPU)"); + + returnValue = cuvsStreamSync(cuvsResourcesPtr); + checkCuVSError(returnValue, "cuvsStreamSync after transform"); + + MemorySegment resultChunk = + ((CuVSMatrixBaseImpl) result) + .memorySegment() + .asSlice(startRow * cols, actualOutputBytes); + cudaMemcpy(resultChunk, outputPtr, actualOutputBytes, DEVICE_TO_HOST); + } + } finally { + cuvsBinaryQuantizerDestroy(quantizerPtr); + cuvsBinaryQuantizerParamsDestroy(paramsPtr); + } + + return result; + } + + private static CuVSMatrix performTransformHost( + long cuvsResourcesPtr, + Arena localArena, + long rows, + long cols, + MemorySegment datasetMemSegment, + BinaryQuantizer.ThresholdType thresholdType) + throws Throwable { + + var result = new CuVSHostMatrixArenaImpl(rows, cols, DataType.BYTE); + + MemorySegment paramsSegment = localArena.allocate(cuvsBinaryQuantizerParams_t); + int returnValue = cuvsBinaryQuantizerParamsCreate(paramsSegment); + checkCuVSError(returnValue, "cuvsBinaryQuantizerParamsCreate"); + + MemorySegment paramsPtr = paramsSegment.get(C_POINTER, 0); + setBinaryQuantizerThreshold(paramsPtr, thresholdType); + + MemorySegment quantizerSegment = localArena.allocate(cuvsBinaryQuantizer_t); + returnValue = cuvsBinaryQuantizerCreate(quantizerSegment); + checkCuVSError(returnValue, "cuvsBinaryQuantizerCreate"); + + MemorySegment quantizerPtr = quantizerSegment.get(C_POINTER, 0); + + try { + long[] datasetShape = {rows, cols}; + MemorySegment datasetTensor = + prepareTensor(localArena, datasetMemSegment, datasetShape, kDLFloat(), 32, kDLCPU(), 1); + + returnValue = + cuvsBinaryQuantizerTrain(cuvsResourcesPtr, paramsPtr, datasetTensor, quantizerPtr); + checkCuVSError(returnValue, "cuvsBinaryQuantizerTrain (host)"); + + MemorySegment outputTensor = + prepareTensor( + localArena, + ((CuVSMatrixBaseImpl) result).memorySegment(), + datasetShape, + kDLUInt(), + 8, + kDLCPU(), + 1); + + returnValue = + cuvsBinaryQuantizerTransformWithParams( + cuvsResourcesPtr, quantizerPtr, datasetTensor, outputTensor); + checkCuVSError(returnValue, "cuvsBinaryQuantizerTransformWithParams (host)"); + + return result; + + } finally { + cuvsBinaryQuantizerDestroy(quantizerPtr); + cuvsBinaryQuantizerParamsDestroy(paramsPtr); + } + } + + private static void setBinaryQuantizerThreshold( + MemorySegment paramsPtr, BinaryQuantizer.ThresholdType thresholdType) { + paramsPtr.set(ValueLayout.JAVA_INT, 0, thresholdType.getValue()); + } + + /** Creates a new BinaryQuantizer implementation. */ + public static Object create( + CuVSResources resources, + CuVSMatrix trainingDataset, + BinaryQuantizer.ThresholdType thresholdType) + throws Throwable { + return new BinaryQuantizerImpl(resources, trainingDataset, thresholdType); + } + + /** Performs transform using a BinaryQuantizer implementation. */ + public static CuVSMatrix transformWithImpl(Object impl, CuVSMatrix input) throws Throwable { + if (impl instanceof BinaryQuantizerImpl) { + return ((BinaryQuantizerImpl) impl).transform(input); + } else { + throw new IllegalArgumentException("Invalid implementation object"); + } + } + + /** Closes BinaryQuantizer implementation. */ + public static void close(Object impl) throws Throwable { + if (impl instanceof BinaryQuantizerImpl) { + ((BinaryQuantizerImpl) impl).close(); + } + } + + private final CuVSResources resources; + private final CuVSMatrix trainingDataset; + private final BinaryQuantizer.ThresholdType thresholdType; + private boolean isClosed = false; + + private BinaryQuantizerImpl( + CuVSResources resources, + CuVSMatrix trainingDataset, + BinaryQuantizer.ThresholdType thresholdType) { + this.resources = resources; + this.trainingDataset = trainingDataset; + this.thresholdType = thresholdType; + } + + private CuVSMatrix transform(CuVSMatrix input) throws Throwable { + if (isClosed) { + throw new IllegalStateException("BinaryQuantizerImpl has been closed"); + } + + if (input.dataType() != DataType.FLOAT) { + throw new IllegalArgumentException( + "BinaryQuantizer requires FLOAT input, got " + input.dataType()); + } + + return transform(resources, input, thresholdType); + } + + private void close() throws Throwable { + isClosed = true; + } +} diff --git a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CagraIndexImpl.java b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CagraIndexImpl.java index 0d5fc8d90f..b97c65cf6d 100644 --- a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CagraIndexImpl.java +++ b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CagraIndexImpl.java @@ -30,6 +30,7 @@ import static com.nvidia.cuvs.internal.common.Util.prepareTensor; import static com.nvidia.cuvs.internal.panama.headers_h.*; +import com.nvidia.cuvs.BinaryQuantizer; import com.nvidia.cuvs.CagraCompressionParams; import com.nvidia.cuvs.CagraIndex; import com.nvidia.cuvs.CagraIndexParams; @@ -40,6 +41,8 @@ import com.nvidia.cuvs.CuVSIvfPqIndexParams; import com.nvidia.cuvs.CuVSIvfPqSearchParams; import com.nvidia.cuvs.CuVSMatrix; +import com.nvidia.cuvs.CuVSMatrix.DataType; +import com.nvidia.cuvs.CuVSQuantizer; import com.nvidia.cuvs.CuVSResources; import com.nvidia.cuvs.SearchResults; import com.nvidia.cuvs.internal.common.CloseableHandle; @@ -81,6 +84,7 @@ public class CagraIndexImpl implements CagraIndex { private final CuVSResources resources; private final IndexReference cagraIndexReference; + private final CuVSQuantizer indexQuantizer; private boolean destroyed; /** @@ -95,6 +99,7 @@ private CagraIndexImpl( CagraIndexParams indexParameters, CuVSMatrix dataset, CuVSResources resources) { Objects.requireNonNull(dataset); this.resources = resources; + this.indexQuantizer = null; assert dataset instanceof CuVSMatrixBaseImpl; this.cagraIndexReference = build(indexParameters, (CuVSMatrixBaseImpl) dataset); } @@ -107,6 +112,7 @@ private CagraIndexImpl( */ private CagraIndexImpl(InputStream inputStream, CuVSResources resources) throws Throwable { this.resources = resources; + this.indexQuantizer = null; this.cagraIndexReference = deserialize(inputStream); } @@ -119,6 +125,7 @@ private CagraIndexImpl(InputStream inputStream, CuVSResources resources) throws */ private CagraIndexImpl(IndexReference indexReference, CuVSResources resources) { this.resources = resources; + this.indexQuantizer = null; this.cagraIndexReference = indexReference; this.destroyed = false; } @@ -135,11 +142,13 @@ private CagraIndexImpl( CagraIndexParams.CuvsDistanceType metric, CuVSMatrix graph, CuVSMatrix dataset, - CuVSResources resources) { + CuVSResources resources, + CuVSQuantizer quantizer) { Objects.requireNonNull(graph); Objects.requireNonNull(dataset); this.resources = resources; + this.indexQuantizer = quantizer; assert graph instanceof CuVSMatrixBaseImpl; assert dataset instanceof CuVSMatrixBaseImpl; @@ -181,6 +190,7 @@ public void close() { private IndexReference build(CagraIndexParams indexParameters, CuVSMatrixBaseImpl dataset) { long rows = dataset.size(); long cols = dataset.columns(); + DataType datasetDataType = dataset.dataType(); try (var indexParams = segmentFromIndexParams(indexParameters); var localArena = Arena.ofConfined()) { @@ -195,7 +205,14 @@ private IndexReference build(CagraIndexParams indexParameters, CuVSMatrixBaseImp long[] datasetShape = {rows, cols}; MemorySegment datasetTensor = - prepareTensor(localArena, dataSeg, datasetShape, kDLFloat(), 32, kDLCPU(), 1); + prepareTensor( + localArena, + dataSeg, + datasetShape, + getDataTypeCode(datasetDataType, this.indexQuantizer), + getPrecisionFromDataType(datasetDataType), + kDLCPU(), + 1); var index = createCagraIndex(); @@ -256,20 +273,39 @@ private static MemorySegment createCagraIndex() { public SearchResults search(CagraQuery query) throws Throwable { try (var localArena = Arena.ofConfined()) { checkNotDestroyed(); + + // Extract query information upfront to avoid "mutually-null" fields + long numQueries; + int vectorDimension; + DataType queryDataType; + MemorySegment floatsSeg; + + if (query.hasQuantizedQueries()) { + CuVSMatrix quantizedQueries = query.getQuantizedQueries(); + numQueries = quantizedQueries.size(); + vectorDimension = (int) quantizedQueries.columns(); + queryDataType = quantizedQueries.dataType(); + floatsSeg = ((CuVSMatrixBaseImpl) quantizedQueries).memorySegment(); + } else { + float[][] queryVectors = query.getQueryVectors(); + numQueries = queryVectors.length; + vectorDimension = numQueries > 0 ? queryVectors[0].length : 0; + queryDataType = DataType.FLOAT; // Float queries are always FLOAT + floatsSeg = buildMemorySegment(localArena, queryVectors); + } + int topK = query.getTopK(); - long numQueries = query.getQueryVectors().length; + + // Calculate bytes based on actual data type + long queriesBytes = getQueryBytesSize(numQueries, vectorDimension, queryDataType); long numBlocks = topK * numQueries; - int vectorDimension = numQueries > 0 ? query.getQueryVectors()[0].length : 0; + long neighborsBytes = C_INT_BYTE_SIZE * numQueries * topK; + long distancesBytes = C_FLOAT_BYTE_SIZE * numQueries * topK; SequenceLayout neighborsSequenceLayout = MemoryLayout.sequenceLayout(numBlocks, C_INT); SequenceLayout distancesSequenceLayout = MemoryLayout.sequenceLayout(numBlocks, C_FLOAT); MemorySegment neighborsMemorySegment = localArena.allocate(neighborsSequenceLayout); MemorySegment distancesMemorySegment = localArena.allocate(distancesSequenceLayout); - MemorySegment floatsSeg = buildMemorySegment(localArena, query.getQueryVectors()); - - final long queriesBytes = C_FLOAT_BYTE_SIZE * numQueries * vectorDimension; - final long neighborsBytes = C_INT_BYTE_SIZE * numQueries * topK; - final long distancesBytes = C_FLOAT_BYTE_SIZE * numQueries * topK; final boolean hasPreFilter = query.getPrefilter() != null; final BitSet[] prefilters = hasPreFilter ? new BitSet[] {query.getPrefilter()} : EMPTY_PREFILTER_BITSET; @@ -293,7 +329,13 @@ public SearchResults search(CagraQuery query) throws Throwable { long[] queriesShape = {numQueries, vectorDimension}; MemorySegment queriesTensor = prepareTensor( - localArena, queriesDP.handle(), queriesShape, kDLFloat(), 32, kDLCUDA(), 1); + localArena, + queriesDP.handle(), + queriesShape, + getDataTypeCode(queryDataType, this.indexQuantizer), + getPrecisionFromDataType(queryDataType), + kDLCUDA(), + 1); long[] neighborsShape = {numQueries, topK}; MemorySegment neighborsTensor = prepareTensor( @@ -773,6 +815,51 @@ private static CloseableHandle createMergeParamsSegment(CagraMergeParams mergePa return new CompositeCloseableHandle(seg, handles); } + // Helper method to map DataType to data type codes + private int getDataTypeCode(DataType dataType, CuVSQuantizer quantizer) { + switch (dataType) { + case FLOAT: + return kDLFloat(); + case BYTE: + if (quantizer instanceof BinaryQuantizer) { + return kDLUInt(); + } else { + return kDLInt(); + } + case INT: + return kDLInt(); + default: + throw new IllegalArgumentException("Unsupported data type: " + dataType); + } + } + + // Helper method to convert DataType to precision + private int getPrecisionFromDataType(DataType dataType) { + switch (dataType) { + case FLOAT: + return 32; + case BYTE: + return 8; + case INT: + return 32; + default: + throw new IllegalArgumentException("Unsupported data type: " + dataType); + } + } + + private long getQueryBytesSize(long numQueries, int vectorDimension, DataType dataType) { + switch (dataType) { + case FLOAT: + return C_FLOAT_BYTE_SIZE * numQueries * vectorDimension; + case BYTE: + return numQueries * vectorDimension; + case INT: + return C_INT_BYTE_SIZE * numQueries * vectorDimension; + default: + throw new IllegalArgumentException("Unsupported query data type: " + dataType); + } + } + /** * Builder helps configure and create an instance of {@link CagraIndex}. */ @@ -782,6 +869,7 @@ public static class Builder implements CagraIndex.Builder { private CagraIndexParams cagraIndexParams; private final CuVSResources cuvsResources; private InputStream inputStream; + private CuVSQuantizer quantizer; private CuVSMatrix graph; public Builder(CuVSResources cuvsResources) { @@ -818,11 +906,28 @@ public Builder withIndexParams(CagraIndexParams cagraIndexParameters) { return this; } + @Override + public Builder withQuantizer(CuVSQuantizer quantizer) { + this.quantizer = quantizer; + return this; + } + @Override public CagraIndexImpl build() throws Throwable { if (inputStream != null) { return new CagraIndexImpl(inputStream, cuvsResources); } else { + if (quantizer != null && dataset != null) { + if (dataset.dataType() != DataType.FLOAT) { + throw new IllegalArgumentException( + "Quantizer input requires FLOAT data type, dataset has " + dataset.dataType()); + } + + CuVSMatrix quantizedDataset = quantizer.transform(dataset); + + return new CagraIndexImpl(cagraIndexParams, quantizedDataset, cuvsResources); + } + if (graph != null) { if (cagraIndexParams == null || dataset == null) { throw new IllegalArgumentException( @@ -830,10 +935,9 @@ public CagraIndexImpl build() throws Throwable { + "you must specify the original dataset and the metric used."); } return new CagraIndexImpl( - cagraIndexParams.getCuvsDistanceType(), graph, dataset, cuvsResources); - } else { - return new CagraIndexImpl(cagraIndexParams, dataset, cuvsResources); + cagraIndexParams.getCuvsDistanceType(), graph, dataset, cuvsResources, null); } + return new CagraIndexImpl(cagraIndexParams, dataset, cuvsResources); } } } diff --git a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CuVSHostMatrixImpl.java b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CuVSHostMatrixImpl.java index c2895282ca..ffc4efcd95 100644 --- a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CuVSHostMatrixImpl.java +++ b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CuVSHostMatrixImpl.java @@ -22,6 +22,8 @@ import static com.nvidia.cuvs.internal.panama.headers_h.*; import com.nvidia.cuvs.CuVSHostMatrix; +import com.nvidia.cuvs.CuVSMatrix.DataType; +import com.nvidia.cuvs.CuVSMatrix.MemoryKind; import com.nvidia.cuvs.RowView; import java.lang.foreign.*; import java.lang.invoke.VarHandle; @@ -118,6 +120,11 @@ public void toArray(byte[][] array) { } } + @Override + public MemoryKind memoryKind() { + return MemoryKind.HOST; + } + @Override public void close() {} diff --git a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CuVSMatrixBaseImpl.java b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CuVSMatrixBaseImpl.java index 120960cfc4..fc0461ef22 100644 --- a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CuVSMatrixBaseImpl.java +++ b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CuVSMatrixBaseImpl.java @@ -18,6 +18,7 @@ import static com.nvidia.cuvs.internal.panama.headers_h.*; import com.nvidia.cuvs.CuVSMatrix; +import com.nvidia.cuvs.CuVSMatrix.DataType; import java.lang.foreign.Arena; import java.lang.foreign.MemorySegment; @@ -49,6 +50,11 @@ public MemorySegment memorySegment() { return memorySegment; } + @Override + public DataType dataType() { + return dataType; + } + protected int bits() { return switch (dataType) { case FLOAT, INT, UINT -> 32; diff --git a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/Scalar8BitQuantizerImpl.java b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/Scalar8BitQuantizerImpl.java new file mode 100644 index 0000000000..de49005cb2 --- /dev/null +++ b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/Scalar8BitQuantizerImpl.java @@ -0,0 +1,461 @@ +/* + * Copyright (c) 2025, NVIDIA CORPORATION. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.nvidia.cuvs.internal; + +import static com.nvidia.cuvs.internal.common.CloseableRMMAllocation.allocateRMMSegment; +import static com.nvidia.cuvs.internal.common.LinkerHelper.C_CHAR; +import static com.nvidia.cuvs.internal.common.LinkerHelper.C_FLOAT_BYTE_SIZE; +import static com.nvidia.cuvs.internal.common.LinkerHelper.C_POINTER; +import static com.nvidia.cuvs.internal.common.Util.CudaMemcpyKind.*; +import static com.nvidia.cuvs.internal.common.Util.checkCuVSError; +import static com.nvidia.cuvs.internal.common.Util.cudaMemcpy; +import static com.nvidia.cuvs.internal.common.Util.prepareTensor; +import static com.nvidia.cuvs.internal.panama.headers_h.cuvsScalarQuantizerCreate; +import static com.nvidia.cuvs.internal.panama.headers_h.cuvsScalarQuantizerDestroy; +import static com.nvidia.cuvs.internal.panama.headers_h.cuvsScalarQuantizerInverseTransform; +import static com.nvidia.cuvs.internal.panama.headers_h.cuvsScalarQuantizerParamsCreate; +import static com.nvidia.cuvs.internal.panama.headers_h.cuvsScalarQuantizerParamsDestroy; +import static com.nvidia.cuvs.internal.panama.headers_h.cuvsScalarQuantizerParams_t; +import static com.nvidia.cuvs.internal.panama.headers_h.cuvsScalarQuantizerTrain; +import static com.nvidia.cuvs.internal.panama.headers_h.cuvsScalarQuantizerTransform; +import static com.nvidia.cuvs.internal.panama.headers_h.cuvsScalarQuantizer_t; +import static com.nvidia.cuvs.internal.panama.headers_h.cuvsStreamSync; +import static com.nvidia.cuvs.internal.panama.headers_h.kDLCPU; +import static com.nvidia.cuvs.internal.panama.headers_h.kDLCUDA; +import static com.nvidia.cuvs.internal.panama.headers_h.kDLFloat; +import static com.nvidia.cuvs.internal.panama.headers_h.kDLInt; + +import com.nvidia.cuvs.CuVSMatrix; +import com.nvidia.cuvs.CuVSMatrix.DataType; +import com.nvidia.cuvs.CuVSMatrix.MemoryKind; +import com.nvidia.cuvs.CuVSResources; +import com.nvidia.cuvs.internal.panama.cuvsScalarQuantizerParams; +import java.lang.foreign.Arena; +import java.lang.foreign.MemorySegment; +import java.util.Objects; + +/** + * Implementation of Scalar8BitQuantizer using Panama Foreign Function Interface. + */ +public class Scalar8BitQuantizerImpl { + + private static final long C_BYTE_SIZE = C_CHAR.byteSize(); + + private final CuVSResources resources; + private final MemorySegment quantizerSegment; + private boolean destroyed; + + private Scalar8BitQuantizerImpl(CuVSResources resources, MemorySegment quantizerSegment) { + this.resources = resources; + this.quantizerSegment = quantizerSegment; + } + + private void checkNotDestroyed() { + if (destroyed) { + throw new IllegalStateException("Scalar8BitQuantizer has been destroyed"); + } + } + + public void destroy() throws Throwable { + checkNotDestroyed(); + try { + int returnValue = cuvsScalarQuantizerDestroy(quantizerSegment); + checkCuVSError(returnValue, "cuvsScalarQuantizerDestroy"); + } finally { + destroyed = true; + } + } + + public void close() throws Throwable { + destroy(); + } + + public CuVSMatrix transform(CuVSMatrix dataset) throws Throwable { + checkNotDestroyed(); + + if (dataset.dataType() != DataType.FLOAT) { + throw new IllegalArgumentException( + "Scalar8BitQuantizer requires FLOAT input, got " + dataset.dataType()); + } + + try (Arena resultArena = Arena.ofShared(); + var localArena = Arena.ofConfined(); + var resourcesAccessor = resources.access()) { + + long rows = dataset.size(); + long cols = dataset.columns(); + + MemorySegment datasetMemSegment = ((CuVSMatrixBaseImpl) dataset).memorySegment(); + long cuvsResourcesPtr = resourcesAccessor.handle(); + + // TODO: Currently all datasets are HOST memory kind, preventing testing of GPU quantization + // flow. + // Future improvement: Support GPU-native dataset creation in Java API to enable direct GPU + // quantization testing. + if (dataset.memoryKind() == MemoryKind.HOST) { + return performTransformHost( + cuvsResourcesPtr, quantizerSegment, localArena, rows, cols, datasetMemSegment); + } + + return performTransformGPU( + cuvsResourcesPtr, + quantizerSegment, + localArena, + resultArena, + rows, + cols, + datasetMemSegment); + } catch (Throwable t) { + throw t; + } + } + + /** + * GPU transform with buffered processing for large datasets. + */ + private CuVSMatrix performTransformGPU( + long cuvsResourcesPtr, + MemorySegment quantizerSeg, + Arena localArena, + Arena resultArena, + long rows, + long cols, + MemorySegment datasetMemSegment) + throws Throwable { + + final long BUFFER_SIZE_BYTES = 1024 * 1024; + final long FLOAT_SIZE = C_FLOAT_BYTE_SIZE; + final long FLOATS_PER_BUFFER = BUFFER_SIZE_BYTES / FLOAT_SIZE; + final long ROWS_PER_BUFFER = Math.max(1, FLOATS_PER_BUFFER / cols); + + CuVSMatrix result = new CuVSHostMatrixArenaImpl(rows, cols, DataType.BYTE); + + // Allocate buffers outside the loop (addressing reviewer feedback) + long maxDatasetBytes = FLOAT_SIZE * ROWS_PER_BUFFER * cols; + long maxOutputBytes = C_BYTE_SIZE * ROWS_PER_BUFFER * cols; + + try (var datasetDP = allocateRMMSegment(cuvsResourcesPtr, maxDatasetBytes); + var outputDP = allocateRMMSegment(cuvsResourcesPtr, maxOutputBytes)) { + + for (long startRow = 0; startRow < rows; startRow += ROWS_PER_BUFFER) { + long endRow = Math.min(startRow + ROWS_PER_BUFFER, rows); + long chunkRows = endRow - startRow; + long actualDatasetBytes = FLOAT_SIZE * chunkRows * cols; + long actualOutputBytes = C_BYTE_SIZE * chunkRows * cols; + + // Get slices of pre-allocated buffers + MemorySegment datasetPtr = datasetDP.handle().asSlice(0, actualDatasetBytes); + MemorySegment outputPtr = outputDP.handle().asSlice(0, actualOutputBytes); + + MemorySegment chunkData = + datasetMemSegment.asSlice(startRow * cols * FLOAT_SIZE, actualDatasetBytes); + + cudaMemcpy(datasetPtr, chunkData, actualDatasetBytes, INFER_DIRECTION); + + long datasetShape[] = {chunkRows, cols}; + MemorySegment datasetTensor = + prepareTensor(localArena, datasetPtr, datasetShape, kDLFloat(), 32, kDLCUDA(), 1); + MemorySegment outputTensor = + prepareTensor(localArena, outputPtr, datasetShape, kDLInt(), 8, kDLCUDA(), 1); + + int returnValue = + cuvsScalarQuantizerTransform( + cuvsResourcesPtr, quantizerSeg, datasetTensor, outputTensor); + checkCuVSError(returnValue, "cuvsScalarQuantizerTransform"); + + returnValue = cuvsStreamSync(cuvsResourcesPtr); + checkCuVSError(returnValue, "cuvsStreamSync"); + + MemorySegment resultChunk = + ((CuVSMatrixBaseImpl) result) + .memorySegment() + .asSlice(startRow * cols, actualOutputBytes); + cudaMemcpy(resultChunk, outputPtr, actualOutputBytes, DEVICE_TO_HOST); + } + } // CloseableRMMAllocation automatically closes here + + if (result.dataType() != DataType.BYTE) { + throw new IllegalStateException( + "Expected BYTE output from scalar quantization, got " + result.dataType()); + } + + return result; + } + + private MemorySegment buildMemorySegmentFromBytes(Arena arena, byte[][] data) { + long rows = data.length; + long cols = rows > 0 ? data[0].length : 0; + long totalBytes = rows * cols * C_BYTE_SIZE; + + MemorySegment segment = arena.allocate(totalBytes); + for (int i = 0; i < rows; i++) { + for (int j = 0; j < cols; j++) { + segment.set(C_CHAR, (i * cols + j) * C_BYTE_SIZE, data[i][j]); + } + } + return segment; + } + + public static Scalar8BitQuantizerImpl create(CuVSResources resources, CuVSMatrix trainingDataset) + throws Throwable { + Objects.requireNonNull(resources); + Objects.requireNonNull(trainingDataset); + + try (var localArena = Arena.ofConfined(); + var resourcesAccessor = resources.access()) { + + long rows = trainingDataset.size(); + long cols = trainingDataset.columns(); + MemorySegment datasetMemSegment = ((CuVSMatrixBaseImpl) trainingDataset).memorySegment(); + long cuvsResourcesPtr = resourcesAccessor.handle(); + + MemorySegment paramsSegment = localArena.allocate(cuvsScalarQuantizerParams_t); + int returnValue = cuvsScalarQuantizerParamsCreate(paramsSegment); + checkCuVSError(returnValue, "cuvsScalarQuantizerParamsCreate"); + + MemorySegment paramsPtr = paramsSegment.get(C_POINTER, 0); + cuvsScalarQuantizerParams.quantile(paramsPtr, 0.99f); + + MemorySegment quantizerSegment = localArena.allocate(cuvsScalarQuantizer_t); + returnValue = cuvsScalarQuantizerCreate(quantizerSegment); + checkCuVSError(returnValue, "cuvsScalarQuantizerCreate"); + + MemorySegment quantizerPtr = quantizerSegment.get(C_POINTER, 0); + + if (trainingDataset.memoryKind() == MemoryKind.HOST) { + long datasetShape[] = {rows, cols}; + MemorySegment datasetTensor = + prepareTensor(localArena, datasetMemSegment, datasetShape, kDLFloat(), 32, kDLCPU(), 1); + + returnValue = + cuvsScalarQuantizerTrain(cuvsResourcesPtr, paramsPtr, datasetTensor, quantizerPtr); + checkCuVSError(returnValue, "cuvsScalarQuantizerTrain"); + } else { + performGPUTraining( + cuvsResourcesPtr, paramsPtr, quantizerPtr, localArena, rows, cols, datasetMemSegment); + } + + returnValue = cuvsScalarQuantizerParamsDestroy(paramsPtr); + checkCuVSError(returnValue, "cuvsScalarQuantizerParamsDestroy"); + + return new Scalar8BitQuantizerImpl(resources, quantizerPtr); + } + } + + public CuVSMatrix inverseTransform(CuVSMatrix quantizedData) throws Throwable { + checkNotDestroyed(); + + if (quantizedData.dataType() != DataType.BYTE) { + throw new IllegalArgumentException( + "Inverse transform requires BYTE input, got " + quantizedData.dataType()); + } + + try (Arena resultArena = Arena.ofShared(); + var localArena = Arena.ofConfined(); + var resourcesAccessor = resources.access()) { + + long rows = quantizedData.size(); + long cols = quantizedData.columns(); + MemorySegment quantizedMemSegment = ((CuVSMatrixBaseImpl) quantizedData).memorySegment(); + long cuvsResourcesPtr = resourcesAccessor.handle(); + + // TODO: Currently all datasets are HOST memory kind, preventing testing of GPU quantization + // flow. + // Future improvement: Support GPU-native dataset creation in Java API to enable direct GPU + // quantization testing. + if (quantizedData.memoryKind() == MemoryKind.HOST) { + return performInverseTransformHost( + cuvsResourcesPtr, quantizerSegment, localArena, rows, cols, quantizedMemSegment); + } + + return performInverseTransformGPU( + cuvsResourcesPtr, + quantizerSegment, + localArena, + resultArena, + rows, + cols, + quantizedMemSegment); + } + } + + private CuVSMatrix performInverseTransformGPU( + long cuvsResourcesPtr, + MemorySegment quantizerSeg, + Arena localArena, + Arena resultArena, + long rows, + long cols, + MemorySegment quantizedMemSegment) + throws Throwable { + + final long BUFFER_SIZE_BYTES = 1024 * 1024; + final long BYTE_SIZE = C_BYTE_SIZE; + final long BYTES_PER_BUFFER = BUFFER_SIZE_BYTES; + final long ROWS_PER_BUFFER = Math.max(1, BYTES_PER_BUFFER / (cols * BYTE_SIZE)); + + CuVSMatrix result = new CuVSHostMatrixArenaImpl(rows, cols, DataType.FLOAT); + + // Allocate buffers outside the loop (addressing reviewer feedback) + long maxQuantizedBytes = BYTE_SIZE * ROWS_PER_BUFFER * cols; + long maxOutputBytes = C_FLOAT_BYTE_SIZE * ROWS_PER_BUFFER * cols; + + try (var quantizedDP = allocateRMMSegment(cuvsResourcesPtr, maxQuantizedBytes); + var outputDP = allocateRMMSegment(cuvsResourcesPtr, maxOutputBytes)) { + + for (long startRow = 0; startRow < rows; startRow += ROWS_PER_BUFFER) { + long endRow = Math.min(startRow + ROWS_PER_BUFFER, rows); + long chunkRows = endRow - startRow; + long actualQuantizedBytes = BYTE_SIZE * chunkRows * cols; + long actualOutputBytes = C_FLOAT_BYTE_SIZE * chunkRows * cols; + + // Get slices of pre-allocated buffers + MemorySegment quantizedPtr = quantizedDP.handle().asSlice(0, actualQuantizedBytes); + MemorySegment outputPtr = outputDP.handle().asSlice(0, actualOutputBytes); + + MemorySegment chunkData = + quantizedMemSegment.asSlice(startRow * cols * BYTE_SIZE, actualQuantizedBytes); + cudaMemcpy(quantizedPtr, chunkData, actualQuantizedBytes, INFER_DIRECTION); + + long datasetShape[] = {chunkRows, cols}; + MemorySegment quantizedTensor = + prepareTensor(localArena, quantizedPtr, datasetShape, kDLInt(), 8, kDLCUDA(), 1); + MemorySegment outputTensor = + prepareTensor(localArena, outputPtr, datasetShape, kDLFloat(), 32, kDLCUDA(), 1); + + int returnValue = + cuvsScalarQuantizerInverseTransform( + cuvsResourcesPtr, quantizerSeg, quantizedTensor, outputTensor); + checkCuVSError(returnValue, "cuvsScalarQuantizerInverseTransform"); + + returnValue = cuvsStreamSync(cuvsResourcesPtr); + checkCuVSError(returnValue, "cuvsStreamSync"); + + MemorySegment resultChunk = + ((CuVSMatrixBaseImpl) result) + .memorySegment() + .asSlice(startRow * cols * C_FLOAT_BYTE_SIZE, actualOutputBytes); + cudaMemcpy(resultChunk, outputPtr, actualOutputBytes, DEVICE_TO_HOST); + } + } // CloseableRMMAllocation automatically closes here + + return result; + } + + private static void performGPUTraining( + long cuvsResourcesPtr, + MemorySegment paramsPtr, + MemorySegment quantizerPtr, + Arena localArena, + long rows, + long cols, + MemorySegment datasetMemSegment) + throws Throwable { + + final long BUFFER_SIZE_BYTES = 1024 * 1024; + final long FLOAT_SIZE = C_FLOAT_BYTE_SIZE; + final long FLOATS_PER_BUFFER = BUFFER_SIZE_BYTES / FLOAT_SIZE; + final long ROWS_PER_BUFFER = Math.max(1, FLOATS_PER_BUFFER / cols); + + // Allocate buffer outside the loop (addressing reviewer feedback) + long maxDatasetBytes = FLOAT_SIZE * ROWS_PER_BUFFER * cols; + + try (var datasetDP = allocateRMMSegment(cuvsResourcesPtr, maxDatasetBytes)) { + for (long startRow = 0; startRow < rows; startRow += ROWS_PER_BUFFER) { + long endRow = Math.min(startRow + ROWS_PER_BUFFER, rows); + long chunkRows = endRow - startRow; + long actualDatasetBytes = FLOAT_SIZE * chunkRows * cols; + + // Get slice of pre-allocated buffer + MemorySegment datasetPtr = datasetDP.handle().asSlice(0, actualDatasetBytes); + + MemorySegment chunkData = + datasetMemSegment.asSlice(startRow * cols * FLOAT_SIZE, actualDatasetBytes); + cudaMemcpy(datasetPtr, chunkData, actualDatasetBytes, INFER_DIRECTION); + + long datasetShape[] = {chunkRows, cols}; + MemorySegment datasetTensor = + prepareTensor(localArena, datasetPtr, datasetShape, kDLFloat(), 32, kDLCUDA(), 1); + + int returnValue = + cuvsScalarQuantizerTrain(cuvsResourcesPtr, paramsPtr, datasetTensor, quantizerPtr); + checkCuVSError(returnValue, "cuvsScalarQuantizerTrain"); + + returnValue = cuvsStreamSync(cuvsResourcesPtr); + checkCuVSError(returnValue, "cuvsStreamSync"); + } + } // CloseableRMMAllocation automatically closes here + } + + private static CuVSMatrix performTransformHost( + long resPtr, + MemorySegment quantizerSeg, + Arena localArena, + long rows, + long cols, + MemorySegment dataSeg) + throws Throwable { + + var result = new CuVSHostMatrixArenaImpl(rows, cols, DataType.BYTE); + + long[] shape = {rows, cols}; + MemorySegment inTensor = prepareTensor(localArena, dataSeg, shape, kDLFloat(), 32, kDLCPU(), 1); + MemorySegment outTensor = + prepareTensor( + localArena, + ((CuVSMatrixBaseImpl) result).memorySegment(), + shape, + kDLInt(), + 8, + kDLCPU(), + 1); + + int rv = cuvsScalarQuantizerTransform(resPtr, quantizerSeg, inTensor, outTensor); + checkCuVSError(rv, "cuvsScalarQuantizerTransform (host)"); + + return result; + } + + private static CuVSMatrix performInverseTransformHost( + long resPtr, + MemorySegment quantizerSeg, + Arena localArena, + long rows, + long cols, + MemorySegment quantizedSeg) + throws Throwable { + + var result = new CuVSHostMatrixArenaImpl(rows, cols, DataType.FLOAT); + + long[] shape = {rows, cols}; + MemorySegment inTensor = + prepareTensor(localArena, quantizedSeg, shape, kDLInt(), 8, kDLCPU(), 1); + MemorySegment outTensor = + prepareTensor( + localArena, + ((CuVSMatrixBaseImpl) result).memorySegment(), + shape, + kDLFloat(), + 32, + kDLCPU(), + 1); + + int rv = cuvsScalarQuantizerInverseTransform(resPtr, quantizerSeg, inTensor, outTensor); + checkCuVSError(rv, "cuvsScalarQuantizerInverseTransform (host)"); + + return result; + } +} diff --git a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/spi/JDKProvider.java b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/spi/JDKProvider.java index 668befe236..0c29188456 100644 --- a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/spi/JDKProvider.java +++ b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/spi/JDKProvider.java @@ -188,4 +188,45 @@ public CuVSMatrix newMatrixFromArray(byte[][] vectors) { Util.copy(dataset.memorySegment(), vectors); return dataset; } + + @Override + public Object createScalar8BitQuantizerImpl(CuVSResources resources, CuVSMatrix trainingDataset) + throws Throwable { + return Scalar8BitQuantizerImpl.create(resources, trainingDataset); + } + + @Override + public CuVSMatrix transformScalar8Bit(Object impl, CuVSMatrix input) throws Throwable { + return ((Scalar8BitQuantizerImpl) impl).transform(input); + } + + @Override + public CuVSMatrix inverseTransformScalar8Bit(Object impl, CuVSMatrix quantizedData) + throws Throwable { + return ((Scalar8BitQuantizerImpl) impl).inverseTransform(quantizedData); + } + + @Override + public Object createBinaryQuantizerImpl( + CuVSResources resources, + CuVSMatrix trainingDataset, + BinaryQuantizer.ThresholdType thresholdType) + throws Throwable { + return BinaryQuantizerImpl.create(resources, trainingDataset, thresholdType); + } + + @Override + public CuVSMatrix transformBinaryWithImpl(Object impl, CuVSMatrix input) throws Throwable { + return BinaryQuantizerImpl.transformWithImpl(impl, input); + } + + @Override + public void closeBinaryQuantizer(Object impl) throws Throwable { + BinaryQuantizerImpl.close(impl); + } + + @Override + public void closeScalar8BitQuantizer(Object impl) throws Throwable { + ((Scalar8BitQuantizerImpl) impl).close(); + } } diff --git a/java/cuvs-java/src/test/java/com/nvidia/cuvs/QuantizationIT.java b/java/cuvs-java/src/test/java/com/nvidia/cuvs/QuantizationIT.java new file mode 100644 index 0000000000..8189fd9139 --- /dev/null +++ b/java/cuvs-java/src/test/java/com/nvidia/cuvs/QuantizationIT.java @@ -0,0 +1,588 @@ +/* + * Copyright (c) 2025, NVIDIA CORPORATION. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.nvidia.cuvs; + +import static com.carrotsearch.randomizedtesting.RandomizedTest.assumeTrue; +import static org.junit.Assert.*; + +import com.carrotsearch.randomizedtesting.RandomizedRunner; +import com.nvidia.cuvs.CagraIndexParams.CagraGraphBuildAlgo; +import com.nvidia.cuvs.CagraIndexParams.CuvsDistanceType; +import com.nvidia.cuvs.CuVSMatrix.DataType; +import java.lang.invoke.MethodHandles; +import java.util.Map; +import org.junit.Before; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +@RunWith(RandomizedRunner.class) +public class QuantizationIT extends CuVSTestCase { + + private static final Logger log = LoggerFactory.getLogger(MethodHandles.lookup().lookupClass()); + + @Before + public void setup() { + assumeTrue("not supported on " + System.getProperty("os.name"), isLinuxAmd64()); + initializeRandom(); + log.info("Random context initialized for quantization test."); + } + + private static float[][] createSimpleDataset() { + return new float[][] { + {1.0f, 2.0f, 3.0f}, + {4.0f, 5.0f, 6.0f}, + {7.0f, 8.0f, 9.0f}, + {10.0f, 11.0f, 12.0f}, + {13.0f, 14.0f, 15.0f} + }; + } + + private static float[][] createSimpleQueries() { + return new float[][] { + {1.1f, 2.1f, 3.1f}, + {7.1f, 8.1f, 9.1f} + }; + } + + @Test + public void testScalarQuantizerBasicFlow() throws Throwable { + float[][] dataset = createSimpleDataset(); + float[][] queries = createSimpleQueries(); + + try (CuVSResources resources = CheckedCuVSResources.create()) { + CuVSMatrix trainingDataset = CuVSMatrix.ofArray(dataset); + assertEquals(DataType.FLOAT, trainingDataset.dataType()); + assertEquals(dataset.length, trainingDataset.size()); + assertEquals(dataset[0].length, trainingDataset.columns()); + + Scalar8BitQuantizer quantizer = new Scalar8BitQuantizer(resources, trainingDataset); + assertEquals(DataType.BYTE, quantizer.outputDataType()); + log.info("Created scalar quantizer with BYTE data type"); + + CagraIndexParams indexParams = + new CagraIndexParams.Builder() + .withCagraGraphBuildAlgo(CagraGraphBuildAlgo.NN_DESCENT) + .withGraphDegree(2) + .withIntermediateGraphDegree(3) + .withNumWriterThreads(1) + .withMetric(CuvsDistanceType.L2Expanded) + .build(); + + CagraIndex index = + CagraIndex.newBuilder(resources) + .withDataset(trainingDataset) + .withQuantizer(quantizer) + .withIndexParams(indexParams) + .build(); + + log.info("Built index with quantized dataset"); + + CagraSearchParams searchParams = new CagraSearchParams.Builder().build(); + + CagraQuery query = + CagraQuery.newBuilder(resources) + .withQueryVectors(queries) + .withQuantizer(quantizer) + .withSearchParams(searchParams) + .withTopK(3) + .build(); + + assertTrue("Query should have quantized vectors", query.hasQuantizedQueries()); + assertEquals(DataType.BYTE, query.getQueryDataType()); + log.info("Created quantized query"); + + // Perform search + SearchResults results = index.search(query); + assertNotNull("Search results should not be null", results); + + // Verify results + assertEquals( + "Should have results for all queries", queries.length, results.getResults().size()); + + for (int i = 0; i < results.getResults().size(); i++) { + Map queryResults = results.getResults().get(i); + assertFalse("Query " + i + " should return some neighbors", queryResults.isEmpty()); + + // Verify all returned IDs are within valid range + for (Integer id : queryResults.keySet()) { + assertTrue("Returned ID should be valid", id >= 0 && id < dataset.length); + } + } + + log.info( + "Search completed successfully with {} queries returning neighbors", + results.getResults().size()); + + // Cleanup + index.close(); + quantizer.close(); + trainingDataset.close(); + } + } + + @Test + public void testBinaryQuantizerBasicFlow() throws Throwable { + float[][] dataset = createSimpleDataset(); + float[][] queries = createSimpleQueries(); + + try (CuVSResources resources = CheckedCuVSResources.create()) { + CuVSMatrix trainingDataset = CuVSMatrix.ofArray(dataset); + assertEquals(DataType.FLOAT, trainingDataset.dataType()); + + BinaryQuantizer quantizer = new BinaryQuantizer(resources, trainingDataset); + assertEquals(DataType.BYTE, quantizer.outputDataType()); + log.info("Created binary quantizer"); + + CagraIndexParams indexParams = + new CagraIndexParams.Builder() + .withCagraGraphBuildAlgo(CagraGraphBuildAlgo.NN_DESCENT) + .withGraphDegree(2) + .withIntermediateGraphDegree(3) + .withNumWriterThreads(1) + .withMetric(CuvsDistanceType.L2Expanded) + .build(); + + CagraIndex index = + CagraIndex.newBuilder(resources) + .withDataset(trainingDataset) + .withQuantizer(quantizer) + .withIndexParams(indexParams) + .build(); + + log.info("Built index with binary quantized dataset"); + + CagraSearchParams searchParams = new CagraSearchParams.Builder().build(); + + CagraQuery query = + CagraQuery.newBuilder(resources) + .withQueryVectors(queries) + .withQuantizer(quantizer) + .withSearchParams(searchParams) + .withTopK(3) + .build(); + + assertTrue("Query should have quantized vectors", query.hasQuantizedQueries()); + assertEquals(DataType.BYTE, query.getQueryDataType()); + + // Perform search + SearchResults results = index.search(query); + assertNotNull("Search results should not be null", results); + + // Verify results + assertEquals( + "Should have results for all queries", queries.length, results.getResults().size()); + + for (int i = 0; i < results.getResults().size(); i++) { + Map queryResults = results.getResults().get(i); + assertFalse("Query " + i + " should return some neighbors", queryResults.isEmpty()); + + // Verify all returned IDs are within valid range + for (Integer id : queryResults.keySet()) { + assertTrue("Returned ID should be valid", id >= 0 && id < dataset.length); + } + } + + log.info("Binary quantized search completed successfully"); + + CuVSMatrix quantizedQueries = query.getQuantizedQueries(); + try { + quantizer.inverseTransform(quantizedQueries); + fail("Expected UnsupportedOperationException to be thrown"); + } catch (UnsupportedOperationException e) { + assertTrue( + "Exception message should mention inverse transform", + e.getMessage().contains("inverse")); + } + + // Cleanup + index.close(); + quantizer.close(); + trainingDataset.close(); + quantizedQueries.close(); + } + } + + @Test + public void testScalarQuantizerInverseTransform() throws Throwable { + float[][] dataset = createSimpleDataset(); + + try (CuVSResources resources = CheckedCuVSResources.create()) { + CuVSMatrix trainingDataset = CuVSMatrix.ofArray(dataset); + Scalar8BitQuantizer quantizer = new Scalar8BitQuantizer(resources, trainingDataset); + + // Transform and inverse transform + CuVSMatrix quantized = quantizer.transform(trainingDataset); + assertEquals(DataType.BYTE, quantized.dataType()); + + CuVSMatrix recovered = quantizer.inverseTransform(quantized); + assertEquals(DataType.FLOAT, recovered.dataType()); + assertEquals(trainingDataset.size(), recovered.size()); + assertEquals(trainingDataset.columns(), recovered.columns()); + + log.info("Inverse transform completed successfully"); + + // Verify quantization worked + assertEquals(DataType.BYTE, quantized.dataType()); + assertEquals(DataType.FLOAT, recovered.dataType()); + + // Cleanup + quantizer.close(); + trainingDataset.close(); + quantized.close(); + recovered.close(); + } + } + + @Test + public void testCPUBinaryQuantization() throws Throwable { + float[][] dataset = createSimpleDataset(); + + try (CuVSResources resources = CheckedCuVSResources.create()) { + CuVSMatrix inputDataset = CuVSMatrix.ofArray(dataset); + assertEquals(CuVSMatrix.MemoryKind.HOST, inputDataset.memoryKind()); + + BinaryQuantizer quantizer = new BinaryQuantizer(resources, inputDataset); + + CuVSMatrix quantized = quantizer.transform(inputDataset); + assertEquals(DataType.BYTE, quantized.dataType()); + assertEquals(CuVSMatrix.MemoryKind.HOST, quantized.memoryKind()); + + try { + quantizer.inverseTransform(quantized); + fail("Expected UnsupportedOperationException to be thrown"); + } catch (UnsupportedOperationException e) { + assertTrue( + "Exception message should mention inverse transform", + e.getMessage().contains("inverse")); + } + + quantizer.close(); + inputDataset.close(); + quantized.close(); + + log.info("✓ Binary CPU quantization test passed"); + } + } + + @Test + public void testRandomizedScalarQuantizedIndexingAndSearch() throws Throwable { + int maxDatasetSize = 60; + int maxDimensions = 20; + int maxQueries = 5; + int maxTopK = 4; + + for (int iteration = 0; iteration < 5; iteration++) { + int datasetSize = random.nextInt(maxDatasetSize) + maxTopK; + int dims = random.nextInt(maxDimensions) + 2; + int numQueries = random.nextInt(maxQueries) + 1; + int topK = Math.min(random.nextInt(maxTopK) + 1, datasetSize); + + // Generate dataset and query + float[][] dataset = new float[datasetSize][dims]; + float[][] queries = new float[numQueries][dims]; + for (int i = 0; i < datasetSize; i++) + for (int d = 0; d < dims; d++) dataset[i][d] = (float) random.nextGaussian(); + for (int i = 0; i < numQueries; i++) + for (int d = 0; d < dims; d++) queries[i][d] = (float) random.nextGaussian(); + + try (CuVSResources resources = CheckedCuVSResources.create()) { + CuVSMatrix trainingDataset = CuVSMatrix.ofArray(dataset); + Scalar8BitQuantizer quantizer = new Scalar8BitQuantizer(resources, trainingDataset); + + CagraIndexParams indexParams = + new CagraIndexParams.Builder() + .withCagraGraphBuildAlgo(CagraGraphBuildAlgo.NN_DESCENT) + .withGraphDegree(Math.min(4, datasetSize - 1)) + .withIntermediateGraphDegree(Math.min(5, datasetSize)) + .withNumWriterThreads(1) + .withMetric(CuvsDistanceType.L2Expanded) + .build(); + + CagraIndex index = + CagraIndex.newBuilder(resources) + .withDataset(trainingDataset) + .withQuantizer(quantizer) + .withIndexParams(indexParams) + .build(); + + CagraSearchParams searchParams = new CagraSearchParams.Builder().build(); + CagraQuery query = + CagraQuery.newBuilder(resources) + .withQueryVectors(queries) + .withQuantizer(quantizer) + .withSearchParams(searchParams) + .withTopK(topK) + .build(); + + SearchResults results = index.search(query); + assertNotNull("Results must not be null", results); + assertEquals("One result per query", numQueries, results.getResults().size()); + + for (int qi = 0; qi < numQueries; qi++) { + Map result = results.getResults().get(qi); + // It is possible for result to have < topK entries if datasetSize == topK + assertFalse("Each query should return non-empty result", result.isEmpty()); + for (Integer id : result.keySet()) { + assertTrue("ID in result should be in range", id >= 0 && id < datasetSize); + } + } + + quantizer.close(); + trainingDataset.close(); + index.close(); + } + } + } + + @Test + public void testBinaryQuantizationWithZeroThreshold() throws Throwable { + float[][] dataset = + new float[][] { + {-2.0f, -1.0f, 0.0f, 1.0f, 2.0f}, // crosses zero + {0.1f, 0.5f, 1.0f, 1.5f, 2.0f}, // all positive + {-2.0f, -1.5f, -1.0f, -0.5f, 0.0f}, // mostly negative + {-1.0f, 0.0f, 1.0f, 2.0f, 3.0f} // mixed + }; + + float[][] queries = + new float[][] { + {0.5f, 0.5f, 0.5f, 0.5f, 0.5f}, + {-0.5f, -0.5f, -0.5f, -0.5f, -0.5f} + }; + + try (CuVSResources resources = CheckedCuVSResources.create()) { + CuVSMatrix trainingDataset = CuVSMatrix.ofArray(dataset); + + BinaryQuantizer quantizer = + new BinaryQuantizer(resources, trainingDataset, BinaryQuantizer.ThresholdType.ZERO); + assertEquals(DataType.BYTE, quantizer.outputDataType()); + assertEquals(BinaryQuantizer.ThresholdType.ZERO, quantizer.getThresholdType()); + + CagraIndexParams indexParams = + new CagraIndexParams.Builder() + .withCagraGraphBuildAlgo(CagraGraphBuildAlgo.NN_DESCENT) + .withGraphDegree(2) + .withIntermediateGraphDegree(3) + .withNumWriterThreads(1) + .withMetric(CuvsDistanceType.L2Expanded) + .build(); + + CagraIndex index = + CagraIndex.newBuilder(resources) + .withDataset(trainingDataset) + .withQuantizer(quantizer) + .withIndexParams(indexParams) + .build(); + + log.info("Built index with ZERO threshold binary quantization"); + + CagraSearchParams searchParams = new CagraSearchParams.Builder().build(); + CagraQuery query = + CagraQuery.newBuilder(resources) + .withQueryVectors(queries) + .withQuantizer(quantizer) + .withSearchParams(searchParams) + .withTopK(2) + .build(); + + assertTrue("Query should have quantized vectors", query.hasQuantizedQueries()); + assertEquals(DataType.BYTE, query.getQueryDataType()); + + SearchResults results = index.search(query); + assertNotNull("Search results should not be null", results); + assertEquals( + "Should have results for all queries", queries.length, results.getResults().size()); + + CuVSMatrix quantizedDataset = quantizer.transform(trainingDataset); + byte[][] result = new byte[(int) quantizedDataset.size()][(int) quantizedDataset.columns()]; + quantizedDataset.toArray(result); + + byte firstByte = result[0][0]; + int[] unpackedBits = new int[5]; + for (int i = 0; i < 5; i++) { + unpackedBits[i] = (firstByte >> i) & 1; + } + + int[] expectedRow0 = {0, 0, 0, 1, 1}; + assertArrayEquals("ZERO threshold failed for row 0", expectedRow0, unpackedBits); + + index.close(); + quantizer.close(); + trainingDataset.close(); + quantizedDataset.close(); + + log.info("✓ ZERO threshold binary quantization test passed"); + } + } + + @Test + public void testBinaryQuantizationWithMeanThreshold() throws Throwable { + float[][] dataset = + new float[][] { + {1.0f, 2.0f, 3.0f, 4.0f, 5.0f}, // mean = 3.0 + {0.0f, 1.0f, 2.0f, 3.0f, 4.0f}, // mean = 2.0 + {2.0f, 4.0f, 6.0f, 8.0f, 10.0f}, // mean = 6.0 + {-1.0f, 0.0f, 1.0f, 2.0f, 3.0f} // mean = 1.0 + }; + + float[][] queries = + new float[][] { + {2.5f, 2.5f, 2.5f, 2.5f, 2.5f}, + {1.0f, 3.0f, 5.0f, 7.0f, 9.0f} + }; + + try (CuVSResources resources = CheckedCuVSResources.create()) { + CuVSMatrix trainingDataset = CuVSMatrix.ofArray(dataset); + + BinaryQuantizer quantizer = + new BinaryQuantizer(resources, trainingDataset, BinaryQuantizer.ThresholdType.MEAN); + assertEquals(BinaryQuantizer.ThresholdType.MEAN, quantizer.getThresholdType()); + + CagraIndexParams indexParams = + new CagraIndexParams.Builder() + .withCagraGraphBuildAlgo(CagraGraphBuildAlgo.NN_DESCENT) + .withGraphDegree(2) + .withIntermediateGraphDegree(3) + .withNumWriterThreads(1) + .withMetric(CuvsDistanceType.L2Expanded) + .build(); + + CagraIndex index = + CagraIndex.newBuilder(resources) + .withDataset(trainingDataset) + .withQuantizer(quantizer) + .withIndexParams(indexParams) + .build(); + + log.info("Built index with MEAN threshold binary quantization"); + + // Perform search + CagraSearchParams searchParams = new CagraSearchParams.Builder().build(); + CagraQuery query = + CagraQuery.newBuilder(resources) + .withQueryVectors(queries) + .withQuantizer(quantizer) + .withSearchParams(searchParams) + .withTopK(2) + .build(); + + SearchResults results = index.search(query); + assertNotNull("Search results should not be null", results); + assertEquals( + "Should have results for all queries", queries.length, results.getResults().size()); + + // Calculate overall mean: (3.0 + 2.0 + 6.0 + 1.0) / 4 = 3.0 + // Verify quantization behavior - values > mean should become 1, values <= mean should become + // 0 + CuVSMatrix quantizedDataset = quantizer.transform(trainingDataset); + byte[][] result1 = new byte[(int) quantizedDataset.size()][(int) quantizedDataset.columns()]; + + // Copy data from CuVSMatrix to array + quantizedDataset.toArray(result1); + + // Verify that MEAN threshold produces different results than ZERO threshold + BinaryQuantizer zeroQuantizer = + new BinaryQuantizer(resources, trainingDataset, BinaryQuantizer.ThresholdType.ZERO); + CuVSMatrix zeroQuantized = zeroQuantizer.transform(trainingDataset); + byte[][] zeroData = new byte[(int) zeroQuantized.size()][(int) zeroQuantized.columns()]; + zeroQuantized.toArray(zeroData); + + boolean isDifferent = false; + for (int i = 0; i < result1.length && !isDifferent; i++) { + for (int j = 0; j < result1[i].length && !isDifferent; j++) { + if (result1[i][j] != zeroData[i][j]) { + isDifferent = true; + } + } + } + assertTrue("MEAN and ZERO thresholds should produce different results", isDifferent); + + // Cleanup + index.close(); + quantizer.close(); + zeroQuantizer.close(); + trainingDataset.close(); + quantizedDataset.close(); + zeroQuantized.close(); + + log.info("✓ MEAN threshold binary quantization test passed"); + } + } + + @Test + public void testBinaryQuantizationWithSamplingMedianThreshold() throws Throwable { + float[][] dataset = createSimpleDataset(); + float[][] queries = createSimpleQueries(); + + try (CuVSResources resources = CheckedCuVSResources.create()) { + CuVSMatrix trainingDataset = CuVSMatrix.ofArray(dataset); + assertEquals(DataType.FLOAT, trainingDataset.dataType()); + + BinaryQuantizer quantizer = + new BinaryQuantizer( + resources, trainingDataset, BinaryQuantizer.ThresholdType.SAMPLING_MEDIAN); + assertEquals(DataType.BYTE, quantizer.outputDataType()); + log.info("Created binary quantizer with sampling median threshold"); + + CagraIndexParams indexParams = + new CagraIndexParams.Builder() + .withCagraGraphBuildAlgo(CagraGraphBuildAlgo.NN_DESCENT) + .withGraphDegree(2) + .withIntermediateGraphDegree(3) + .withNumWriterThreads(1) + .withMetric(CuvsDistanceType.L2Expanded) + .build(); + + CagraIndex index = + CagraIndex.newBuilder(resources) + .withDataset(trainingDataset) + .withQuantizer(quantizer) + .withIndexParams(indexParams) + .build(); + + log.info("Built index with binary quantized dataset"); + + CagraSearchParams searchParams = new CagraSearchParams.Builder().build(); + + CagraQuery query = + CagraQuery.newBuilder(resources) + .withQueryVectors(queries) + .withQuantizer(quantizer) + .withSearchParams(searchParams) + .withTopK(3) + .build(); + + assertTrue("Query should have quantized vectors", query.hasQuantizedQueries()); + assertEquals(DataType.BYTE, query.getQueryDataType()); + + SearchResults results = index.search(query); + assertNotNull("Search results should not be null", results); + + assertEquals( + "Should return results for each query", queries.length, results.getResults().size()); + + for (int i = 0; i < results.getResults().size(); i++) { + Map result = results.getResults().get(i); + assertFalse("Each query should return non-empty result", result.isEmpty()); + assertTrue("Should return at most topK results", result.size() <= 3); + } + + log.info("✓ Binary quantization with sampling median threshold test passed"); + } + } +} diff --git a/java/panama-bindings/headers.h b/java/panama-bindings/headers.h index 57a13cf1b1..ba4cf8c203 100644 --- a/java/panama-bindings/headers.h +++ b/java/panama-bindings/headers.h @@ -22,6 +22,8 @@ #include #include #include +#include +#include #include #include From 3f953364fb7aa6162adafbd70070c1fd39e6f80f Mon Sep 17 00:00:00 2001 From: Ben Frederickson Date: Thu, 21 Aug 2025 21:40:07 -0700 Subject: [PATCH 40/93] Add rust bindings for kmeans (#1266) Authors: - Ben Frederickson (https://github.com/benfred) Approvers: - Corey J. Nolet (https://github.com/cjnolet) URL: https://github.com/rapidsai/cuvs/pull/1266 --- rust/cuvs-sys/cuvs_c_wrapper.h | 3 +- rust/cuvs/src/cluster/kmeans/mod.rs | 212 +++++++++++++++++++++++++ rust/cuvs/src/cluster/kmeans/params.rs | 152 ++++++++++++++++++ rust/cuvs/src/cluster/mod.rs | 17 ++ rust/cuvs/src/dlpack.rs | 8 +- rust/cuvs/src/lib.rs | 1 + 6 files changed, 390 insertions(+), 3 deletions(-) create mode 100644 rust/cuvs/src/cluster/kmeans/mod.rs create mode 100644 rust/cuvs/src/cluster/kmeans/params.rs create mode 100644 rust/cuvs/src/cluster/mod.rs diff --git a/rust/cuvs-sys/cuvs_c_wrapper.h b/rust/cuvs-sys/cuvs_c_wrapper.h index 0c66928bcf..aeac04f54e 100644 --- a/rust/cuvs-sys/cuvs_c_wrapper.h +++ b/rust/cuvs-sys/cuvs_c_wrapper.h @@ -1,5 +1,5 @@ /* - * Copyright (c) 2024, NVIDIA CORPORATION. + * Copyright (c) 2024-2025, NVIDIA CORPORATION. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -17,6 +17,7 @@ // wrapper file containing all the C-API's we should automatically be creating rust // bindings for #include +#include #include #include #include diff --git a/rust/cuvs/src/cluster/kmeans/mod.rs b/rust/cuvs/src/cluster/kmeans/mod.rs new file mode 100644 index 0000000000..2ebac1cdf6 --- /dev/null +++ b/rust/cuvs/src/cluster/kmeans/mod.rs @@ -0,0 +1,212 @@ +/* + * Copyright (c) 2025, NVIDIA CORPORATION. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +//! Kmeans clustering API's +//! +//! Example: +//! ``` +//! +//! use cuvs::cluster::kmeans; +//! use cuvs::{ManagedTensor, Resources, Result}; +//! +//! use ndarray_rand::rand_distr::Uniform; +//! use ndarray_rand::RandomExt; +//! +//! fn kmeans_example() -> Result<()> { +//! let res = Resources::new()?; +//! +//! // Create a new random dataset to index +//! let n_datapoints = 65536; +//! let n_features = 512; +//! let n_clusters = 8; +//! let dataset = +//! ndarray::Array::::random((n_datapoints, n_features), Uniform::new(0., 1.0)); +//! let dataset = ManagedTensor::from(&dataset).to_device(&res)?; +//! +//! let centroids_host = ndarray::Array::::zeros((n_clusters, n_features)); +//! let mut centroids = ManagedTensor::from(¢roids_host).to_device(&res)?; +//! +//! // find the centroids with the kmeans index +//! let kmeans_params = kmeans::Params::new()?.set_n_clusters(n_clusters as i32); +//! let (inertia, n_iter) = kmeans::fit(&res, &kmeans_params, &dataset, &None, &mut centroids)?; +//! Ok(()) +//! } +//! ``` + +mod params; + +pub use params::Params; + +use crate::dlpack::ManagedTensor; +use crate::error::{check_cuvs, Result}; +use crate::resources::Resources; + +/// Find clusters with the k-means algorithm +/// +/// # Arguments +/// +/// * `res` - Resources to use +/// * `params` - Parameters to use to fit KMeans model +/// * `x` - A matrix in device memory - shape (m, k) +/// * `sample_weight` - Optional device matrix shape (n_clusters, 1) +/// * `centroids` - Output device matrix, that has the centroids for each cluster +/// shape (n_clusters, k) +pub fn fit( + res: &Resources, + params: &Params, + x: &ManagedTensor, + sample_weight: &Option, + centroids: &mut ManagedTensor, +) -> Result<(f64, i32)> { + let mut inertia: f64 = 0.0; + let mut niter: i32 = 0; + + unsafe { + let sample_weight_dlpack = match sample_weight { + Some(tensor) => tensor.as_ptr(), + None => std::ptr::null_mut(), + }; + check_cuvs(ffi::cuvsKMeansFit( + res.0, + params.0, + x.as_ptr(), + sample_weight_dlpack, + centroids.as_ptr(), + &mut inertia as *mut f64, + &mut niter as *mut i32, + ))?; + } + Ok((inertia, niter)) +} + +/// Predict clusters with the k-means algorithm +/// +/// # Arguments +/// +/// * `res` - Resources to use +/// * `params` - Parameters to use to fit KMeans model +/// * `x` - Input matrix in device memory - shape (m, k) +/// * `sample_weight` - Optional device matrix shape (n_clusters, 1) +/// * `centroids` - Centroids calculated by fit in device memory, shape (n_clusters, k) +/// * `labels` - preallocated CUDA array interface matrix shape (m, 1) to hold the output labels +/// * `normalize_weight` - whether or not to normalize the weights +pub fn predict( + res: &Resources, + params: &Params, + x: &ManagedTensor, + sample_weight: &Option, + centroids: &ManagedTensor, + labels: &mut ManagedTensor, + normalize_weight: bool, +) -> Result { + let mut inertia: f64 = 0.0; + + unsafe { + let sample_weight_dlpack = match sample_weight { + Some(tensor) => tensor.as_ptr(), + None => std::ptr::null_mut(), + }; + check_cuvs(ffi::cuvsKMeansPredict( + res.0, + params.0, + x.as_ptr(), + sample_weight_dlpack, + centroids.as_ptr(), + labels.as_ptr(), + normalize_weight, + &mut inertia as *mut f64, + ))?; + } + Ok(inertia) +} + +/// Compute cluster cost given an input matrix and existing centroids +/// # Arguments +/// +/// * `res` - Resources to use +/// * `x` - Input matrix in device memory - shape (m, k) +/// * `centroids` - Centroids calculated by fit in device memory, shape (n_clusters, k) +pub fn cluster_cost(res: &Resources, x: &ManagedTensor, centroids: &ManagedTensor) -> Result { + let mut inertia: f64 = 0.0; + + unsafe { + check_cuvs(ffi::cuvsKMeansClusterCost( + res.0, + x.as_ptr(), + centroids.as_ptr(), + &mut inertia as *mut f64, + ))?; + } + Ok(inertia) +} + +#[cfg(test)] +mod tests { + use super::*; + use ndarray_rand::rand_distr::Uniform; + use ndarray_rand::RandomExt; + + #[test] + fn test_kmeans() { + let res = Resources::new().unwrap(); + + let n_clusters = 4; + + // Create a new random dataset to index + let n_datapoints = 256; + let n_features = 16; + let dataset = + ndarray::Array::::random((n_datapoints, n_features), Uniform::new(0., 1.0)); + let dataset = ManagedTensor::from(&dataset).to_device(&res).unwrap(); + + let centroids_host = ndarray::Array::::zeros((n_clusters, n_features)); + let mut centroids = ManagedTensor::from(¢roids_host) + .to_device(&res) + .unwrap(); + + let params = Params::new().unwrap().set_n_clusters(n_clusters as i32); + + // compute the inertia, before fitting centroids + let original_inertia = cluster_cost(&res, &dataset, ¢roids).unwrap(); + + // fit the centroids, make sure that inertia has gone down + let (inertia, n_iter) = fit(&res, ¶ms, &dataset, &None, &mut centroids).unwrap(); + + assert!(inertia < original_inertia); + assert!(n_iter >= 1); + + let mut labels_host = ndarray::Array::::zeros((n_clusters,)); + let mut labels = ManagedTensor::from(&labels_host).to_device(&res).unwrap(); + + // make sure the prediction for each centroid is the centroid itself + predict( + &res, + ¶ms, + ¢roids, + &None, + ¢roids, + &mut labels, + false, + ) + .unwrap(); + + labels.to_host(&res, &mut labels_host).unwrap(); + assert_eq!(labels_host[[0,]], 0); + assert_eq!(labels_host[[1,]], 1); + assert_eq!(labels_host[[2,]], 2); + assert_eq!(labels_host[[3,]], 3); + } +} diff --git a/rust/cuvs/src/cluster/kmeans/params.rs b/rust/cuvs/src/cluster/kmeans/params.rs new file mode 100644 index 0000000000..bbefe19351 --- /dev/null +++ b/rust/cuvs/src/cluster/kmeans/params.rs @@ -0,0 +1,152 @@ +/* + * Copyright (c) 2025, NVIDIA CORPORATION. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +use crate::distance_type::DistanceType; +use crate::error::{check_cuvs, Result}; +use std::fmt; +use std::io::{stderr, Write}; + +pub struct Params(pub ffi::cuvsKMeansParams_t); + +impl Params { + /// Returns a new Params + pub fn new() -> Result { + unsafe { + let mut params = std::mem::MaybeUninit::::uninit(); + check_cuvs(ffi::cuvsKMeansParamsCreate(params.as_mut_ptr()))?; + Ok(Params(params.assume_init())) + } + } + + /// DistanceType to use for fitting kmeans + pub fn set_metric(self, metric: DistanceType) -> Params { + unsafe { + (*self.0).metric = metric; + } + self + } + + /// The number of clusters to form as well as the number of centroids to generate (default:8). + pub fn set_n_clusters(self, n_clusters: i32) -> Params { + unsafe { + (*self.0).n_clusters = n_clusters; + } + self + } + + /// Maximum number of iterations of the k-means algorithm for a single run. + pub fn set_max_iter(self, max_iter: i32) -> Params { + unsafe { + (*self.0).max_iter = max_iter; + } + self + } + + /// Relative tolerance with regards to inertia to declare convergence. + pub fn set_tol(self, tol: f64) -> Params { + unsafe { + (*self.0).tol = tol; + } + self + } + + /// Number of instance k-means algorithm will be run with different seeds. + pub fn set_n_init(self, n_init: i32) -> Params { + unsafe { + (*self.0).n_init = n_init; + } + self + } + + /// Oversampling factor for use in the k-means|| algorithm + pub fn set_oversampling_factor(self, oversampling_factor: f64) -> Params { + unsafe { + (*self.0).oversampling_factor = oversampling_factor; + } + self + } + + /** + * batch_samples and batch_centroids are used to tile 1NN computation which is + * useful to optimize/control the memory footprint + * Default tile is [batch_samples x n_clusters] i.e. when batch_centroids is 0 + * then don't tile the centroids. + */ + pub fn set_batch_samples(self, batch_samples: i32) -> Params { + unsafe { + (*self.0).batch_samples = batch_samples; + } + self + } + /// if 0 then batch_centroids = n_clusters + pub fn set_batch_centroids(self, batch_centroids: i32) -> Params { + unsafe { + (*self.0).batch_centroids = batch_centroids; + } + self + } + + /// Whether to use hierarchical (balanced) kmeans or not + pub fn set_hierarchical(self, hierarchical: bool) -> Params { + unsafe { + (*self.0).hierarchical = hierarchical; + } + self + } + + /// For hierarchical k-means , defines the number of training iterations + pub fn set_hierarchical_n_iters(self, hierarchical_n_iters: i32) -> Params { + unsafe { + (*self.0).hierarchical_n_iters = hierarchical_n_iters; + } + self + } +} + +impl fmt::Debug for Params { + fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { + // custom debug trait here, default value will show the pointer address + // for the inner params object which isn't that useful. + write!(f, "Params({:?})", unsafe { *self.0 }) + } +} + +impl Drop for Params { + fn drop(&mut self) { + if let Err(e) = check_cuvs(unsafe { ffi::cuvsKMeansParamsDestroy(self.0) }) { + write!(stderr(), "failed to call cuvsKMeansParamsDestroy {:?}", e) + .expect("failed to write to stderr"); + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_params() { + let params = Params::new() + .unwrap() + .set_n_clusters(128) + .set_hierarchical(true); + + unsafe { + assert_eq!((*params.0).n_clusters, 128); + assert_eq!((*params.0).hierarchical, true); + } + } +} diff --git a/rust/cuvs/src/cluster/mod.rs b/rust/cuvs/src/cluster/mod.rs new file mode 100644 index 0000000000..2972b3431d --- /dev/null +++ b/rust/cuvs/src/cluster/mod.rs @@ -0,0 +1,17 @@ +/* + * Copyright (c) 2025, NVIDIA CORPORATION. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +pub mod kmeans; diff --git a/rust/cuvs/src/dlpack.rs b/rust/cuvs/src/dlpack.rs index b74484a5c9..dfc730c24d 100644 --- a/rust/cuvs/src/dlpack.rs +++ b/rust/cuvs/src/dlpack.rs @@ -1,5 +1,5 @@ /* - * Copyright (c) 2024, NVIDIA CORPORATION. + * Copyright (c) 2024-2025, NVIDIA CORPORATION. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -51,7 +51,7 @@ impl ManagedTensor { res.get_cuda_stream()?, ))?; - let mut ret = self.0.clone(); + let mut ret = self.0; ret.dl_tensor.data = device_data; ret.deleter = Some(rmm_free_tensor); ret.dl_tensor.device.device_type = ffi::DLDeviceType::kDLCUDA; @@ -205,5 +205,9 @@ mod tests { // make sure we can get the shape ok assert_eq!(unsafe { *tensor.shape }, 8); assert_eq!(unsafe { *tensor.shape.add(1) }, 4); + + let arr = ndarray::Array::::zeros((8,)); + let tensor = unsafe { (*(ManagedTensor::from(&arr).as_ptr())).dl_tensor }; + assert_eq!(tensor.ndim, 1); } } diff --git a/rust/cuvs/src/lib.rs b/rust/cuvs/src/lib.rs index c2059b6097..085bbbb801 100644 --- a/rust/cuvs/src/lib.rs +++ b/rust/cuvs/src/lib.rs @@ -20,6 +20,7 @@ //! approximate nearest neighbors search on the GPU. pub mod brute_force; pub mod cagra; +pub mod cluster; pub mod distance; pub mod distance_type; mod dlpack; From 8a1f2f6a137bbcfb1b3ce4120783df327dfcf285 Mon Sep 17 00:00:00 2001 From: MithunR Date: Fri, 22 Aug 2025 12:02:58 -0700 Subject: [PATCH 41/93] Revert "[Java]Binary and scalar quantization (#1104)" (#1274) This reverts commit 989c994e3d15d3b81c86aa1d5d10cb95a9cb9e88. The changes in #1104 were merged by accident, ahead of higher priority features. This change is to revert #1104, to reduce rework on the other PRs. Java support for Quantization will be revisited when prioritized. The changes in #1104 can be rebased then. Authors: - MithunR (https://github.com/mythrocks) Approvers: - Corey J. Nolet (https://github.com/cjnolet) URL: https://github.com/rapidsai/cuvs/pull/1274 --- .../cuvs/preprocessing/quantize/binary.h | 6 +- .../java/com/nvidia/cuvs/BinaryQuantizer.java | 125 ---- .../main/java/com/nvidia/cuvs/CagraIndex.java | 8 - .../main/java/com/nvidia/cuvs/CagraQuery.java | 108 +--- .../main/java/com/nvidia/cuvs/CuVSMatrix.java | 18 - .../java/com/nvidia/cuvs/CuVSQuantizer.java | 71 --- .../com/nvidia/cuvs/Scalar8BitQuantizer.java | 123 ---- .../com/nvidia/cuvs/spi/CuVSProvider.java | 27 - .../nvidia/cuvs/spi/UnsupportedProvider.java | 45 +- .../cuvs/internal/BinaryQuantizerImpl.java | 324 ---------- .../nvidia/cuvs/internal/CagraIndexImpl.java | 130 +--- .../cuvs/internal/CuVSHostMatrixImpl.java | 7 - .../cuvs/internal/CuVSMatrixBaseImpl.java | 6 - .../internal/Scalar8BitQuantizerImpl.java | 461 -------------- .../com/nvidia/cuvs/spi/JDKProvider.java | 41 -- .../java/com/nvidia/cuvs/QuantizationIT.java | 588 ------------------ java/panama-bindings/headers.h | 2 - 17 files changed, 37 insertions(+), 2053 deletions(-) delete mode 100644 java/cuvs-java/src/main/java/com/nvidia/cuvs/BinaryQuantizer.java delete mode 100644 java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSQuantizer.java delete mode 100644 java/cuvs-java/src/main/java/com/nvidia/cuvs/Scalar8BitQuantizer.java delete mode 100644 java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/BinaryQuantizerImpl.java delete mode 100644 java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/Scalar8BitQuantizerImpl.java delete mode 100644 java/cuvs-java/src/test/java/com/nvidia/cuvs/QuantizationIT.java diff --git a/cpp/include/cuvs/preprocessing/quantize/binary.h b/cpp/include/cuvs/preprocessing/quantize/binary.h index ec58bae3bd..9a5834ef3e 100644 --- a/cpp/include/cuvs/preprocessing/quantize/binary.h +++ b/cpp/include/cuvs/preprocessing/quantize/binary.h @@ -30,7 +30,7 @@ extern "C" { * and sampling_median thresholds are calculated separately for each dimension. * */ -typedef enum { ZERO, MEAN, SAMPLING_MEDIAN } cuvsBinaryQuantizerThreshold; +enum cuvsBinaryQuantizerThreshold { ZERO, MEAN, SAMPLING_MEDIAN }; /** * @brief Binary quantizer parameters. @@ -39,12 +39,12 @@ struct cuvsBinaryQuantizerParams { /* * specifies the threshold to set a bit in cuvsBinaryQuantizerTransform */ - cuvsBinaryQuantizerThreshold threshold; + cuvsBinaryQuantizerThreshold threshold = MEAN; /* * specifies the sampling ratio */ - float sampling_ratio; + float sampling_ratio = 0.1; }; typedef struct cuvsBinaryQuantizerParams* cuvsBinaryQuantizerParams_t; diff --git a/java/cuvs-java/src/main/java/com/nvidia/cuvs/BinaryQuantizer.java b/java/cuvs-java/src/main/java/com/nvidia/cuvs/BinaryQuantizer.java deleted file mode 100644 index 82584d6cac..0000000000 --- a/java/cuvs-java/src/main/java/com/nvidia/cuvs/BinaryQuantizer.java +++ /dev/null @@ -1,125 +0,0 @@ -/* - * Copyright (c) 2025, NVIDIA CORPORATION. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ -package com.nvidia.cuvs; - -import com.nvidia.cuvs.CuVSMatrix.DataType; -import com.nvidia.cuvs.spi.CuVSProvider; - -/** Binary quantizer that transforms float datasets into packed binary datasets. */ -public class BinaryQuantizer implements CuVSQuantizer { - private final CuVSResources resources; - private final ThresholdType thresholdType; - - /** Output data type (BYTE for packed binary data). */ - private final DataType outputDataType = DataType.BYTE; - - private final Object impl; - - /** Creates a binary quantizer with training data. */ - public BinaryQuantizer(CuVSResources resources, CuVSMatrix trainingDataset) throws Throwable { - this(resources, trainingDataset, ThresholdType.MEAN); - } - - /** Creates a binary quantizer with specified threshold type and training data. */ - public BinaryQuantizer( - CuVSResources resources, CuVSMatrix trainingDataset, ThresholdType thresholdType) - throws Throwable { - if (trainingDataset == null) { - throw new IllegalArgumentException("Training dataset cannot be null"); - } - if (trainingDataset.dataType() != DataType.FLOAT) { - throw new IllegalArgumentException( - "Training dataset must have FLOAT data type, got " + trainingDataset.dataType()); - } - this.resources = resources; - this.thresholdType = thresholdType; - this.impl = - CuVSProvider.provider() - .createBinaryQuantizerImpl(resources, trainingDataset, thresholdType); - } - - /** Threshold types for binary quantization. */ - public enum ThresholdType { - /** Zero threshold */ - ZERO(0), - /** Mean threshold */ - MEAN(1), - /** Sampling median threshold */ - SAMPLING_MEDIAN(2); - - private final int value; - - ThresholdType(int value) { - this.value = value; - } - - public int getValue() { - return value; - } - } - - /** Returns the output data type (BYTE). */ - @Override - public DataType outputDataType() { - return outputDataType; - } - - /** Gets the threshold type used for binary quantization. */ - public ThresholdType getThresholdType() { - return thresholdType; - } - - @Override - public CuVSMatrix transform(CuVSMatrix input) throws Throwable { - if (input.dataType() != DataType.FLOAT) { - throw new IllegalArgumentException( - "BinaryQuantizer requires FLOAT input, got " + input.dataType()); - } - - CuVSMatrix result = CuVSProvider.provider().transformBinaryWithImpl(impl, input); - - if (result.dataType() != DataType.BYTE) { - throw new IllegalStateException( - "Expected BYTE output from binary quantization, got " + result.dataType()); - } - - return result; - } - - @Override - public void train(CuVSMatrix trainingData) throws Throwable { - throw new UnsupportedOperationException("Training handled during construction"); - } - - @Override - public CuVSMatrix inverseTransform(CuVSMatrix quantizedData) throws Throwable { - throw new UnsupportedOperationException( - "Binary quantization does not support inverse transformation"); - } - - @Override - public void close() throws Exception { - try { - CuVSProvider.provider().closeBinaryQuantizer(impl); - } catch (Throwable t) { - if (t instanceof Exception) { - throw (Exception) t; - } else { - throw new RuntimeException("Error closing BinaryQuantizer", t); - } - } - } -} diff --git a/java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraIndex.java b/java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraIndex.java index c0032f7ea4..e3b885a333 100644 --- a/java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraIndex.java +++ b/java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraIndex.java @@ -247,14 +247,6 @@ interface Builder { */ Builder withIndexParams(CagraIndexParams cagraIndexParameters); - /** - * Sets the quantizer to use for building the index with quantized data. - * - * @param quantizer The quantizer to apply to the dataset during index construction - * @return An instance of this Builder - */ - Builder withQuantizer(CuVSQuantizer quantizer); - /** * Builds and returns an instance of CagraIndex. * diff --git a/java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraQuery.java b/java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraQuery.java index c976949bf7..f81ec3191c 100644 --- a/java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraQuery.java +++ b/java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraQuery.java @@ -15,15 +15,14 @@ */ package com.nvidia.cuvs; -import com.nvidia.cuvs.CuVSMatrix.DataType; import java.util.Arrays; import java.util.BitSet; import java.util.Objects; import java.util.function.LongToIntFunction; /** - * CagraQuery holds the search parameters plus either raw float[][] vectors - * or a quantized {@link CuVSMatrix} for querying a CAGRA index. + * CagraQuery holds the CagraSearchParams and the query vectors to be used while + * invoking search. * *

Thread Safety: Each CagraQuery instance should use its own * CuVSResources object that is not shared with other threads. Sharing CuVSResources @@ -34,9 +33,8 @@ public class CagraQuery { private final CagraSearchParams cagraSearchParameters; - private final float[][] queryVectors; - private final CuVSMatrix quantizedQueries; private final LongToIntFunction mapping; + private final float[][] queryVectors; private final int topK; private final BitSet prefilter; private final int numDocs; @@ -49,25 +47,23 @@ public class CagraQuery { * @param cagraSearchParameters an instance of {@link CagraSearchParams} holding * the search parameters * @param queryVectors 2D float query vector array - * @param quantizedQueries 2D quantized query vector array * @param mapping a function mapping ordinals (neighbor IDs) to custom user IDs * @param topK the top k results to return * @param prefilter A single BitSet to use as filter while searching the CAGRA index * @param numDocs Total number of dataset vectors; used to align the prefilter correctly * @param resources CuVSResources instance to use for this query */ - private CagraQuery( + public CagraQuery( CagraSearchParams cagraSearchParameters, float[][] queryVectors, - CuVSMatrix quantizedQueries, LongToIntFunction mapping, int topK, BitSet prefilter, int numDocs, CuVSResources resources) { + super(); this.cagraSearchParameters = cagraSearchParameters; this.queryVectors = queryVectors; - this.quantizedQueries = quantizedQueries; this.mapping = mapping; this.topK = topK; this.prefilter = prefilter; @@ -75,11 +71,6 @@ private CagraQuery( this.resources = resources; } - /** Start building a new CagraQuery. */ - public static Builder newBuilder(CuVSResources resources) { - return new Builder(resources); - } - /** * Gets the instance of CagraSearchParams initially set. * @@ -90,35 +81,14 @@ public CagraSearchParams getCagraSearchParameters() { } /** - * If this query was built without a quantizer, returns the original float vectors. - * Otherwise returns null. + * Gets the query vector 2D float array. + * + * @return 2D float array */ public float[][] getQueryVectors() { return queryVectors; } - /** - * If this query was built with a quantizer, returns the quantized Dataset. - * Otherwise returns null. - */ - public CuVSMatrix getQuantizedQueries() { - return quantizedQueries; - } - - /** True if this query carries a quantized Dataset instead of float[][] */ - public boolean hasQuantizedQueries() { - return quantizedQueries != null; - } - - /** - * Returns the data type of the query payload: - * - 32 for float32 queries - * - 8 for quantized queries - */ - public DataType getQueryDataType() { - return quantizedQueries != null ? quantizedQueries.dataType() : DataType.FLOAT; - } - /** * Gets the function mapping ordinals (neighbor IDs) to custom user IDs */ @@ -164,13 +134,10 @@ public CuVSResources getResources() { @Override public String toString() { - return "CagraQuery[" - + "params=" + return "CuVSQuery [cagraSearchParameters=" + cagraSearchParameters - + ", floatVectors=" - + (queryVectors != null ? Arrays.toString(queryVectors) : "null") - + ", quantized=" - + (quantizedQueries != null ? ("Dataset@" + quantizedQueries.dataType() + "-bit") : "false") + + ", queryVectors=" + + Arrays.toString(queryVectors) + ", mapping=" + mapping + ", topK=" @@ -189,7 +156,6 @@ public static class Builder { private int topK = 2; private BitSet prefilter; private int numDocs; - private CuVSQuantizer quantizer; private final CuVSResources resources; /** @@ -208,12 +174,12 @@ public Builder(CuVSResources resources) { /** * Sets the instance of configured CagraSearchParams to be passed for search. * - * @param params an instance of the configured CagraSearchParams to - * be used for this query + * @param cagraSearchParams an instance of the configured CagraSearchParams to + * be used for this query * @return an instance of this Builder */ - public Builder withSearchParams(CagraSearchParams params) { - this.cagraSearchParams = params; + public Builder withSearchParams(CagraSearchParams cagraSearchParams) { + this.cagraSearchParams = cagraSearchParams; return this; } @@ -267,47 +233,13 @@ public Builder withPrefilter(BitSet prefilter, int numDocs) { } /** - * Specify a quantizer to automatically transform the float[][] queryVectors - * into a quantized {@link CuVSMatrix} using the same quantizer used for training. - */ - public Builder withQuantizer(CuVSQuantizer quantizer) { - this.quantizer = quantizer; - return this; - } - - /** - * Builds the CagraQuery. If a quantizer was provided, queryVectors is ignored - * and a quantized Dataset is produced instead. + * Builds an instance of CuVSQuery. + * + * @return an instance of CuVSQuery */ - public CagraQuery build() throws Throwable { - if (queryVectors == null) { - throw new IllegalArgumentException("Query vectors must be provided"); - } - - CuVSMatrix quantized = null; - float[][] floatsForQuery = queryVectors; - - if (quantizer != null) { - // wrap float[][] in a CuVSMatrix and quantize - try (CuVSMatrix tmp = CuVSMatrix.ofArray(queryVectors)) { - if (tmp.dataType() != DataType.FLOAT) { - throw new IllegalArgumentException( - "Query quantization requires FLOAT input, got " + tmp.dataType()); - } - quantized = quantizer.transform(tmp); - } - floatsForQuery = null; - } - + public CagraQuery build() { return new CagraQuery( - cagraSearchParams, - floatsForQuery, - quantized, - mapping, - topK, - prefilter, - numDocs, - resources); + cagraSearchParams, queryVectors, mapping, topK, prefilter, numDocs, resources); } } } diff --git a/java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSMatrix.java b/java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSMatrix.java index 5143c8bbc7..f8e60c5555 100644 --- a/java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSMatrix.java +++ b/java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSMatrix.java @@ -69,17 +69,6 @@ static CuVSMatrix ofArray(byte[][] vectors) { return CuVSProvider.provider().newMatrixFromArray(vectors); } - /** - * Returns the data type of this matrix. - * - *

DataType provides the fundamental type information needed for all operations. - * FLOAT represents 32-bit floating point data, INT represents 32-bit integer data, - * and BYTE represents 8-bit data (used for quantized values). - * - * @return The DataType of this matrix - */ - DataType dataType(); - interface Builder { /** * Add a single vector to the dataset. @@ -163,13 +152,6 @@ static CuVSMatrix.Builder builder(int size, int columns, DataType dataType) { */ void toArray(byte[][] array); - /** - * Gets where the content is stored - * - * @return MemoryKind whether HOST or DEVICE - */ - CuVSMatrix.MemoryKind memoryKind(); - @Override void close(); } diff --git a/java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSQuantizer.java b/java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSQuantizer.java deleted file mode 100644 index 2c86eff0f1..0000000000 --- a/java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSQuantizer.java +++ /dev/null @@ -1,71 +0,0 @@ -/* - * Copyright (c) 2025, NVIDIA CORPORATION. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ -package com.nvidia.cuvs; - -import com.nvidia.cuvs.CuVSMatrix.DataType; - -/** - * Base interface for all cuVS quantizers providing unified quantization operations. - * - */ -public interface CuVSQuantizer extends AutoCloseable { - - /** - * Returns the data type of quantized data produced by this quantizer. - * - * @return The DataType of the quantized output (BYTE for quantized data) - */ - DataType outputDataType(); - - /** - * Transforms a float32 dataset into a quantized dataset. - * - * @param input The input dataset of float32 vectors (must have precision() == 32) - * @return A new Dataset containing the quantized vectors with this quantizer's precision - */ - CuVSMatrix transform(CuVSMatrix input) throws Throwable; - - /** - * Optional training method - only applies to quantizers that require training. - * - * @param trainingData The dataset to train on (must have precision() == 32) - * @throws Throwable if an error occurs during training - */ - default void train(CuVSMatrix trainingData) throws Throwable { - throw new UnsupportedOperationException( - "Training not supported for " + getClass().getSimpleName()); - } - - /** - * Optional inverse transformation - only applies to quantizers that support reconstruction - * - * @param quantizedData The quantized dataset to inverse transform - * (must have precision() matching this quantizer's precision()) - * @return A float[][] array containing the de-quantized vectors - */ - default CuVSMatrix inverseTransform(CuVSMatrix quantizedData) throws Throwable { - throw new UnsupportedOperationException( - "Inverse transform not supported for " + getClass().getSimpleName()); - } - - /** - * Closes this quantizer and releases any associated resources. - * - * @throws Exception if an error occurs while closing - */ - @Override - void close() throws Exception; -} diff --git a/java/cuvs-java/src/main/java/com/nvidia/cuvs/Scalar8BitQuantizer.java b/java/cuvs-java/src/main/java/com/nvidia/cuvs/Scalar8BitQuantizer.java deleted file mode 100644 index 056a7ad468..0000000000 --- a/java/cuvs-java/src/main/java/com/nvidia/cuvs/Scalar8BitQuantizer.java +++ /dev/null @@ -1,123 +0,0 @@ -/* - * Copyright (c) 2025, NVIDIA CORPORATION. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ -package com.nvidia.cuvs; - -import com.nvidia.cuvs.CuVSMatrix.DataType; -import com.nvidia.cuvs.spi.CuVSProvider; - -/** - * Scalar 8-bit quantizer implementation that transforms 32-bit float datasets into 8-bit integer datasets. - * - *

This quantizer reduces each float32 value to an 8-bit signed integer, providing significant - * compression while maintaining reasonable precision. The quantizer requires training on a representative - * dataset to compute optimal quantization parameters (min/max values per dimension). - * - *

The quantizer supports both transformation and inverse transformation, allowing approximate - * recovery of the original float32 values from the quantized representation. - * - * @since 25.08 - */ -public class Scalar8BitQuantizer implements CuVSQuantizer { - private final Object impl; - - /** - * The data type used by this quantizer (BYTE for 8-bit signed integers). - */ - private final DataType outputDataType = DataType.BYTE; - - /** - * Creates a new scalar 8-bit quantizer with CuVSMatrix training data. - * - *

The training dataset is used to compute quantization parameters (min/max values) - * for each dimension. The quantizer will be immediately ready for use after construction. - * - * @param resources The CuVS resources to use for quantization operations - * @param trainingDataset A CuVSMatrix containing training vectors (must have FLOAT data type) - * @throws Throwable if an error occurs during quantizer training - * @throws IllegalArgumentException if trainingDataset is null or doesn't have FLOAT data type - */ - public Scalar8BitQuantizer(CuVSResources resources, CuVSMatrix trainingDataset) throws Throwable { - if (trainingDataset == null) { - throw new IllegalArgumentException("Training dataset cannot be null"); - } - if (trainingDataset.dataType() != DataType.FLOAT) { - throw new IllegalArgumentException( - "Training dataset must have FLOAT data type, got " + trainingDataset.dataType()); - } - this.impl = CuVSProvider.provider().createScalar8BitQuantizerImpl(resources, trainingDataset); - } - - /** - * Returns the data type of quantized data produced by this quantizer. - * - * @return The DataType (BYTE for this scalar quantizer) - */ - @Override - public DataType outputDataType() { - return outputDataType; - } - - @Override - public CuVSMatrix transform(CuVSMatrix input) throws Throwable { - // Validate input data type - if (input.dataType() != DataType.FLOAT) { - throw new IllegalArgumentException( - "Scalar8BitQuantizer requires FLOAT input, got " + input.dataType()); - } - - CuVSMatrix result = CuVSProvider.provider().transformScalar8Bit(impl, input); - - // Validate output data type - if (result.dataType() != DataType.BYTE) { - throw new IllegalStateException( - "Expected BYTE output from scalar quantization, got " + result.dataType()); - } - - return result; - } - - @Override - public void train(CuVSMatrix trainingData) throws Throwable { - // Training is handled during construction for ScalarQuantizer - // This method could be implemented if you want to support re-training - throw new UnsupportedOperationException("Training handled during construction"); - } - - @Override - public CuVSMatrix inverseTransform(CuVSMatrix quantizedData) throws Throwable { - // Validate input data type for inverse transform - if (quantizedData.dataType() != DataType.BYTE) { - throw new IllegalArgumentException( - "Inverse transform requires BYTE input, got " + quantizedData.dataType()); - } - - return CuVSProvider.provider().inverseTransformScalar8Bit(impl, quantizedData); - } - - @Override - public void close() throws Exception { - try { - CuVSProvider.provider().closeScalar8BitQuantizer(impl); - } catch (Throwable t) { - if (t instanceof Exception) { - throw (Exception) t; - } else { - // Wrap non-Exception throwables (like Error) in RuntimeException - throw new RuntimeException("Error closing Scalar8BitQuantizer", t); - } - } - } -} 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 316406b24e..e528beff1a 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 @@ -15,7 +15,6 @@ */ package com.nvidia.cuvs.spi; -import com.nvidia.cuvs.BinaryQuantizer; import com.nvidia.cuvs.BruteForceIndex; import com.nvidia.cuvs.CagraIndex; import com.nvidia.cuvs.CagraMergeParams; @@ -121,32 +120,6 @@ default CagraIndex mergeCagraIndexes(CagraIndex[] indexes, CagraMergeParams merg return mergeCagraIndexes(indexes); } - /** Creates a new Scalar8BitQuantizer implementation. */ - Object createScalar8BitQuantizerImpl(CuVSResources resources, CuVSMatrix trainingDataset) - throws Throwable; - - /** Performs inverse transform using a Scalar8BitQuantizer. */ - CuVSMatrix inverseTransformScalar8Bit(Object impl, CuVSMatrix quantizedData) throws Throwable; - - /** Transforms dataset using Scalar8BitQuantizer */ - CuVSMatrix transformScalar8Bit(Object impl, CuVSMatrix input) throws Throwable; - - /** Closes Scalar8BitQuantizer implementation */ - void closeScalar8BitQuantizer(Object impl) throws Throwable; - - /** Creates a new BinaryQuantizer implementation. */ - Object createBinaryQuantizerImpl( - CuVSResources resources, - CuVSMatrix trainingDataset, - BinaryQuantizer.ThresholdType thresholdType) - throws Throwable; - - /** Performs transform using a BinaryQuantizer. */ - CuVSMatrix transformBinaryWithImpl(Object impl, CuVSMatrix input) throws Throwable; - - /** Closes BinaryQuantizer implementation. */ - void closeBinaryQuantizer(Object impl) throws Throwable; - /** Retrieves the system-wide provider. */ static CuVSProvider provider() { return CuVSServiceProvider.Holder.INSTANCE; diff --git a/java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/UnsupportedProvider.java b/java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/UnsupportedProvider.java index cd6604663b..8f65bf7068 100644 --- a/java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/UnsupportedProvider.java +++ b/java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/UnsupportedProvider.java @@ -15,7 +15,6 @@ */ package com.nvidia.cuvs.spi; -import com.nvidia.cuvs.BinaryQuantizer; import com.nvidia.cuvs.BruteForceIndex; import com.nvidia.cuvs.CagraIndex; import com.nvidia.cuvs.CuVSMatrix; @@ -26,8 +25,7 @@ import java.nio.file.Path; /** - * A provider that throws UnsupportedOperationException for all operations. - * Used as a fallback when no proper implementation is available. + * A provider that unconditionally throws UnsupportedOperationException. */ final class UnsupportedProvider implements CuVSProvider { @@ -86,45 +84,4 @@ public CuVSMatrix newMatrixFromArray(int[][] vectors) { public CuVSMatrix newMatrixFromArray(byte[][] vectors) { throw new UnsupportedOperationException(); } - - @Override - public Object createScalar8BitQuantizerImpl(CuVSResources resources, CuVSMatrix trainingDataset) - throws Throwable { - throw new UnsupportedOperationException("CuVS is not supported on this platform"); - } - - @Override - public CuVSMatrix inverseTransformScalar8Bit(Object impl, CuVSMatrix quantizedData) - throws Throwable { - throw new UnsupportedOperationException("CuVS is not supported on this platform"); - } - - @Override - public Object createBinaryQuantizerImpl( - CuVSResources resources, - CuVSMatrix trainingDataset, - BinaryQuantizer.ThresholdType thresholdType) - throws Throwable { - throw new UnsupportedOperationException("CuVS is not supported on this platform"); - } - - @Override - public CuVSMatrix transformBinaryWithImpl(Object impl, CuVSMatrix input) throws Throwable { - throw new UnsupportedOperationException("CuVS is not supported on this platform"); - } - - @Override - public void closeBinaryQuantizer(Object impl) throws Throwable { - throw new UnsupportedOperationException("CuVS is not supported on this platform"); - } - - @Override - public CuVSMatrix transformScalar8Bit(Object impl, CuVSMatrix input) throws Throwable { - throw new UnsupportedOperationException("CuVS is not supported on this platform"); - } - - @Override - public void closeScalar8BitQuantizer(Object impl) throws Throwable { - throw new UnsupportedOperationException("CuVS is not supported on this platform"); - } } diff --git a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/BinaryQuantizerImpl.java b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/BinaryQuantizerImpl.java deleted file mode 100644 index 704fb45738..0000000000 --- a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/BinaryQuantizerImpl.java +++ /dev/null @@ -1,324 +0,0 @@ -/* - * Copyright (c) 2025, NVIDIA CORPORATION. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ -package com.nvidia.cuvs.internal; - -import static com.nvidia.cuvs.internal.common.CloseableRMMAllocation.allocateRMMSegment; -import static com.nvidia.cuvs.internal.common.LinkerHelper.C_FLOAT_BYTE_SIZE; -import static com.nvidia.cuvs.internal.common.LinkerHelper.C_POINTER; -import static com.nvidia.cuvs.internal.common.Util.CudaMemcpyKind.*; -import static com.nvidia.cuvs.internal.common.Util.buildMemorySegment; -import static com.nvidia.cuvs.internal.common.Util.checkCuVSError; -import static com.nvidia.cuvs.internal.common.Util.cudaMemcpy; -import static com.nvidia.cuvs.internal.common.Util.prepareTensor; -import static com.nvidia.cuvs.internal.panama.headers_h.cuvsBinaryQuantizerCreate; -import static com.nvidia.cuvs.internal.panama.headers_h.cuvsBinaryQuantizerDestroy; -import static com.nvidia.cuvs.internal.panama.headers_h.cuvsBinaryQuantizerParamsCreate; -import static com.nvidia.cuvs.internal.panama.headers_h.cuvsBinaryQuantizerParamsDestroy; -import static com.nvidia.cuvs.internal.panama.headers_h.cuvsBinaryQuantizerParams_t; -import static com.nvidia.cuvs.internal.panama.headers_h.cuvsBinaryQuantizerTrain; -import static com.nvidia.cuvs.internal.panama.headers_h.cuvsBinaryQuantizerTransformWithParams; -import static com.nvidia.cuvs.internal.panama.headers_h.cuvsBinaryQuantizer_t; -import static com.nvidia.cuvs.internal.panama.headers_h.cuvsStreamSync; -import static com.nvidia.cuvs.internal.panama.headers_h.kDLCPU; -import static com.nvidia.cuvs.internal.panama.headers_h.kDLCUDA; -import static com.nvidia.cuvs.internal.panama.headers_h.kDLFloat; -import static com.nvidia.cuvs.internal.panama.headers_h.kDLUInt; - -import com.nvidia.cuvs.BinaryQuantizer; -import com.nvidia.cuvs.CuVSMatrix; -import com.nvidia.cuvs.CuVSMatrix.DataType; -import com.nvidia.cuvs.CuVSMatrix.MemoryKind; -import com.nvidia.cuvs.CuVSResources; -import com.nvidia.cuvs.internal.common.Util; -import java.lang.foreign.Arena; -import java.lang.foreign.MemorySegment; -import java.lang.foreign.ValueLayout; - -/** - * Implementation of binary quantizer functionality using Panama Foreign Function Interface. - */ -public class BinaryQuantizerImpl { - - public static CuVSMatrix transform( - CuVSResources cuvsResources, float[][] dataset, BinaryQuantizer.ThresholdType thresholdType) - throws Throwable { - - try (var localArena = Arena.ofConfined(); - var resourcesAccessor = cuvsResources.access()) { - - long rows = dataset.length; - long cols = rows > 0 ? dataset[0].length : 0; - - MemorySegment datasetMemSegment = buildMemorySegment(localArena, dataset); - long cuvsResourcesPtr = resourcesAccessor.handle(); - - return performTransformGPU( - cuvsResourcesPtr, - localArena, - rows, - cols, - datasetMemSegment, - HOST_TO_DEVICE, - thresholdType); - } - } - - public static CuVSMatrix transform( - CuVSResources cuvsResources, CuVSMatrix dataset, BinaryQuantizer.ThresholdType thresholdType) - throws Throwable { - if (dataset.dataType() != DataType.FLOAT) { - throw new IllegalArgumentException( - "BinaryQuantizer requires FLOAT input, got " + dataset.dataType()); - } - - try (var localArena = Arena.ofConfined(); - var resourcesAccessor = cuvsResources.access()) { - - long rows = dataset.size(); - long cols = dataset.columns(); - - MemorySegment datasetMemSegment = ((CuVSMatrixBaseImpl) dataset).memorySegment(); - long cuvsResourcesPtr = resourcesAccessor.handle(); - - // TODO: Currently all datasets are HOST memory kind, preventing testing of GPU quantization - // flow. - // Future improvement: Support GPU-native dataset creation in Java API to enable direct GPU - // quantization testing. - if (dataset.memoryKind() == MemoryKind.HOST) { - return performTransformHost( - cuvsResourcesPtr, localArena, rows, cols, datasetMemSegment, thresholdType); - } else { - CuVSMatrix result = - performTransformGPU( - cuvsResourcesPtr, - localArena, - rows, - cols, - datasetMemSegment, - INFER_DIRECTION, - thresholdType); - - if (result.dataType() != DataType.BYTE) { - throw new IllegalStateException( - "Expected BYTE output from binary quantization, got " + result.dataType()); - } - return result; - } - } - } - - private static CuVSMatrix performTransformGPU( - long cuvsResourcesPtr, - Arena localArena, - long rows, - long cols, - MemorySegment datasetMemSegment, - Util.CudaMemcpyKind memcpyDirection, - BinaryQuantizer.ThresholdType thresholdType) - throws Throwable { - - final long BUFFER_SIZE_BYTES = 1024 * 1024; - final long FLOAT_SIZE = C_FLOAT_BYTE_SIZE; - final long FLOATS_PER_BUFFER = BUFFER_SIZE_BYTES / FLOAT_SIZE; - final long ROWS_PER_BUFFER = Math.max(1, FLOATS_PER_BUFFER / cols); - - var result = new CuVSHostMatrixArenaImpl(rows, cols, DataType.BYTE); - - long maxDatasetBytes = FLOAT_SIZE * ROWS_PER_BUFFER * cols; - long maxOutputBytes = ROWS_PER_BUFFER * cols; - - // Allocate quantizer objects once - MemorySegment paramsSegment = localArena.allocate(cuvsBinaryQuantizerParams_t); - int returnValue = cuvsBinaryQuantizerParamsCreate(paramsSegment); - checkCuVSError(returnValue, "cuvsBinaryQuantizerParamsCreate"); - MemorySegment paramsPtr = paramsSegment.get(C_POINTER, 0); - setBinaryQuantizerThreshold(paramsPtr, thresholdType); - - MemorySegment quantizerSegment = localArena.allocate(cuvsBinaryQuantizer_t); - returnValue = cuvsBinaryQuantizerCreate(quantizerSegment); - checkCuVSError(returnValue, "cuvsBinaryQuantizerCreate"); - MemorySegment quantizerPtr = quantizerSegment.get(C_POINTER, 0); - - try (var datasetDP = allocateRMMSegment(cuvsResourcesPtr, maxDatasetBytes); - var outputDP = allocateRMMSegment(cuvsResourcesPtr, maxOutputBytes)) { - - for (long startRow = 0; startRow < rows; startRow += ROWS_PER_BUFFER) { - long endRow = Math.min(startRow + ROWS_PER_BUFFER, rows); - long chunkRows = endRow - startRow; - long actualDatasetBytes = FLOAT_SIZE * chunkRows * cols; - long actualOutputBytes = chunkRows * cols; - - MemorySegment datasetPtr = datasetDP.handle().asSlice(0, actualDatasetBytes); - MemorySegment outputPtr = outputDP.handle().asSlice(0, actualOutputBytes); - MemorySegment chunkData = - datasetMemSegment.asSlice(startRow * cols * FLOAT_SIZE, actualDatasetBytes); - - cudaMemcpy(datasetPtr, chunkData, actualDatasetBytes, memcpyDirection); - - returnValue = cuvsStreamSync(cuvsResourcesPtr); - checkCuVSError(returnValue, "cuvsStreamSync before transform"); - - long datasetShape[] = {chunkRows, cols}; - long outputShape[] = {chunkRows, cols}; - - MemorySegment datasetTensor = - prepareTensor(localArena, datasetPtr, datasetShape, kDLFloat(), 32, kDLCUDA(), 1); - MemorySegment outputTensor = - prepareTensor(localArena, outputPtr, outputShape, kDLUInt(), 8, kDLCUDA(), 1); - - returnValue = - cuvsBinaryQuantizerTrain(cuvsResourcesPtr, paramsPtr, datasetTensor, quantizerPtr); - checkCuVSError(returnValue, "cuvsBinaryQuantizerTrain (GPU)"); - - returnValue = - cuvsBinaryQuantizerTransformWithParams( - cuvsResourcesPtr, quantizerPtr, datasetTensor, outputTensor); - checkCuVSError(returnValue, "cuvsBinaryQuantizerTransformWithParams (GPU)"); - - returnValue = cuvsStreamSync(cuvsResourcesPtr); - checkCuVSError(returnValue, "cuvsStreamSync after transform"); - - MemorySegment resultChunk = - ((CuVSMatrixBaseImpl) result) - .memorySegment() - .asSlice(startRow * cols, actualOutputBytes); - cudaMemcpy(resultChunk, outputPtr, actualOutputBytes, DEVICE_TO_HOST); - } - } finally { - cuvsBinaryQuantizerDestroy(quantizerPtr); - cuvsBinaryQuantizerParamsDestroy(paramsPtr); - } - - return result; - } - - private static CuVSMatrix performTransformHost( - long cuvsResourcesPtr, - Arena localArena, - long rows, - long cols, - MemorySegment datasetMemSegment, - BinaryQuantizer.ThresholdType thresholdType) - throws Throwable { - - var result = new CuVSHostMatrixArenaImpl(rows, cols, DataType.BYTE); - - MemorySegment paramsSegment = localArena.allocate(cuvsBinaryQuantizerParams_t); - int returnValue = cuvsBinaryQuantizerParamsCreate(paramsSegment); - checkCuVSError(returnValue, "cuvsBinaryQuantizerParamsCreate"); - - MemorySegment paramsPtr = paramsSegment.get(C_POINTER, 0); - setBinaryQuantizerThreshold(paramsPtr, thresholdType); - - MemorySegment quantizerSegment = localArena.allocate(cuvsBinaryQuantizer_t); - returnValue = cuvsBinaryQuantizerCreate(quantizerSegment); - checkCuVSError(returnValue, "cuvsBinaryQuantizerCreate"); - - MemorySegment quantizerPtr = quantizerSegment.get(C_POINTER, 0); - - try { - long[] datasetShape = {rows, cols}; - MemorySegment datasetTensor = - prepareTensor(localArena, datasetMemSegment, datasetShape, kDLFloat(), 32, kDLCPU(), 1); - - returnValue = - cuvsBinaryQuantizerTrain(cuvsResourcesPtr, paramsPtr, datasetTensor, quantizerPtr); - checkCuVSError(returnValue, "cuvsBinaryQuantizerTrain (host)"); - - MemorySegment outputTensor = - prepareTensor( - localArena, - ((CuVSMatrixBaseImpl) result).memorySegment(), - datasetShape, - kDLUInt(), - 8, - kDLCPU(), - 1); - - returnValue = - cuvsBinaryQuantizerTransformWithParams( - cuvsResourcesPtr, quantizerPtr, datasetTensor, outputTensor); - checkCuVSError(returnValue, "cuvsBinaryQuantizerTransformWithParams (host)"); - - return result; - - } finally { - cuvsBinaryQuantizerDestroy(quantizerPtr); - cuvsBinaryQuantizerParamsDestroy(paramsPtr); - } - } - - private static void setBinaryQuantizerThreshold( - MemorySegment paramsPtr, BinaryQuantizer.ThresholdType thresholdType) { - paramsPtr.set(ValueLayout.JAVA_INT, 0, thresholdType.getValue()); - } - - /** Creates a new BinaryQuantizer implementation. */ - public static Object create( - CuVSResources resources, - CuVSMatrix trainingDataset, - BinaryQuantizer.ThresholdType thresholdType) - throws Throwable { - return new BinaryQuantizerImpl(resources, trainingDataset, thresholdType); - } - - /** Performs transform using a BinaryQuantizer implementation. */ - public static CuVSMatrix transformWithImpl(Object impl, CuVSMatrix input) throws Throwable { - if (impl instanceof BinaryQuantizerImpl) { - return ((BinaryQuantizerImpl) impl).transform(input); - } else { - throw new IllegalArgumentException("Invalid implementation object"); - } - } - - /** Closes BinaryQuantizer implementation. */ - public static void close(Object impl) throws Throwable { - if (impl instanceof BinaryQuantizerImpl) { - ((BinaryQuantizerImpl) impl).close(); - } - } - - private final CuVSResources resources; - private final CuVSMatrix trainingDataset; - private final BinaryQuantizer.ThresholdType thresholdType; - private boolean isClosed = false; - - private BinaryQuantizerImpl( - CuVSResources resources, - CuVSMatrix trainingDataset, - BinaryQuantizer.ThresholdType thresholdType) { - this.resources = resources; - this.trainingDataset = trainingDataset; - this.thresholdType = thresholdType; - } - - private CuVSMatrix transform(CuVSMatrix input) throws Throwable { - if (isClosed) { - throw new IllegalStateException("BinaryQuantizerImpl has been closed"); - } - - if (input.dataType() != DataType.FLOAT) { - throw new IllegalArgumentException( - "BinaryQuantizer requires FLOAT input, got " + input.dataType()); - } - - return transform(resources, input, thresholdType); - } - - private void close() throws Throwable { - isClosed = true; - } -} diff --git a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CagraIndexImpl.java b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CagraIndexImpl.java index b97c65cf6d..0d5fc8d90f 100644 --- a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CagraIndexImpl.java +++ b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CagraIndexImpl.java @@ -30,7 +30,6 @@ import static com.nvidia.cuvs.internal.common.Util.prepareTensor; import static com.nvidia.cuvs.internal.panama.headers_h.*; -import com.nvidia.cuvs.BinaryQuantizer; import com.nvidia.cuvs.CagraCompressionParams; import com.nvidia.cuvs.CagraIndex; import com.nvidia.cuvs.CagraIndexParams; @@ -41,8 +40,6 @@ import com.nvidia.cuvs.CuVSIvfPqIndexParams; import com.nvidia.cuvs.CuVSIvfPqSearchParams; import com.nvidia.cuvs.CuVSMatrix; -import com.nvidia.cuvs.CuVSMatrix.DataType; -import com.nvidia.cuvs.CuVSQuantizer; import com.nvidia.cuvs.CuVSResources; import com.nvidia.cuvs.SearchResults; import com.nvidia.cuvs.internal.common.CloseableHandle; @@ -84,7 +81,6 @@ public class CagraIndexImpl implements CagraIndex { private final CuVSResources resources; private final IndexReference cagraIndexReference; - private final CuVSQuantizer indexQuantizer; private boolean destroyed; /** @@ -99,7 +95,6 @@ private CagraIndexImpl( CagraIndexParams indexParameters, CuVSMatrix dataset, CuVSResources resources) { Objects.requireNonNull(dataset); this.resources = resources; - this.indexQuantizer = null; assert dataset instanceof CuVSMatrixBaseImpl; this.cagraIndexReference = build(indexParameters, (CuVSMatrixBaseImpl) dataset); } @@ -112,7 +107,6 @@ private CagraIndexImpl( */ private CagraIndexImpl(InputStream inputStream, CuVSResources resources) throws Throwable { this.resources = resources; - this.indexQuantizer = null; this.cagraIndexReference = deserialize(inputStream); } @@ -125,7 +119,6 @@ private CagraIndexImpl(InputStream inputStream, CuVSResources resources) throws */ private CagraIndexImpl(IndexReference indexReference, CuVSResources resources) { this.resources = resources; - this.indexQuantizer = null; this.cagraIndexReference = indexReference; this.destroyed = false; } @@ -142,13 +135,11 @@ private CagraIndexImpl( CagraIndexParams.CuvsDistanceType metric, CuVSMatrix graph, CuVSMatrix dataset, - CuVSResources resources, - CuVSQuantizer quantizer) { + CuVSResources resources) { Objects.requireNonNull(graph); Objects.requireNonNull(dataset); this.resources = resources; - this.indexQuantizer = quantizer; assert graph instanceof CuVSMatrixBaseImpl; assert dataset instanceof CuVSMatrixBaseImpl; @@ -190,7 +181,6 @@ public void close() { private IndexReference build(CagraIndexParams indexParameters, CuVSMatrixBaseImpl dataset) { long rows = dataset.size(); long cols = dataset.columns(); - DataType datasetDataType = dataset.dataType(); try (var indexParams = segmentFromIndexParams(indexParameters); var localArena = Arena.ofConfined()) { @@ -205,14 +195,7 @@ private IndexReference build(CagraIndexParams indexParameters, CuVSMatrixBaseImp long[] datasetShape = {rows, cols}; MemorySegment datasetTensor = - prepareTensor( - localArena, - dataSeg, - datasetShape, - getDataTypeCode(datasetDataType, this.indexQuantizer), - getPrecisionFromDataType(datasetDataType), - kDLCPU(), - 1); + prepareTensor(localArena, dataSeg, datasetShape, kDLFloat(), 32, kDLCPU(), 1); var index = createCagraIndex(); @@ -273,39 +256,20 @@ private static MemorySegment createCagraIndex() { public SearchResults search(CagraQuery query) throws Throwable { try (var localArena = Arena.ofConfined()) { checkNotDestroyed(); - - // Extract query information upfront to avoid "mutually-null" fields - long numQueries; - int vectorDimension; - DataType queryDataType; - MemorySegment floatsSeg; - - if (query.hasQuantizedQueries()) { - CuVSMatrix quantizedQueries = query.getQuantizedQueries(); - numQueries = quantizedQueries.size(); - vectorDimension = (int) quantizedQueries.columns(); - queryDataType = quantizedQueries.dataType(); - floatsSeg = ((CuVSMatrixBaseImpl) quantizedQueries).memorySegment(); - } else { - float[][] queryVectors = query.getQueryVectors(); - numQueries = queryVectors.length; - vectorDimension = numQueries > 0 ? queryVectors[0].length : 0; - queryDataType = DataType.FLOAT; // Float queries are always FLOAT - floatsSeg = buildMemorySegment(localArena, queryVectors); - } - int topK = query.getTopK(); - - // Calculate bytes based on actual data type - long queriesBytes = getQueryBytesSize(numQueries, vectorDimension, queryDataType); + long numQueries = query.getQueryVectors().length; long numBlocks = topK * numQueries; - long neighborsBytes = C_INT_BYTE_SIZE * numQueries * topK; - long distancesBytes = C_FLOAT_BYTE_SIZE * numQueries * topK; + int vectorDimension = numQueries > 0 ? query.getQueryVectors()[0].length : 0; SequenceLayout neighborsSequenceLayout = MemoryLayout.sequenceLayout(numBlocks, C_INT); SequenceLayout distancesSequenceLayout = MemoryLayout.sequenceLayout(numBlocks, C_FLOAT); MemorySegment neighborsMemorySegment = localArena.allocate(neighborsSequenceLayout); MemorySegment distancesMemorySegment = localArena.allocate(distancesSequenceLayout); + MemorySegment floatsSeg = buildMemorySegment(localArena, query.getQueryVectors()); + + final long queriesBytes = C_FLOAT_BYTE_SIZE * numQueries * vectorDimension; + final long neighborsBytes = C_INT_BYTE_SIZE * numQueries * topK; + final long distancesBytes = C_FLOAT_BYTE_SIZE * numQueries * topK; final boolean hasPreFilter = query.getPrefilter() != null; final BitSet[] prefilters = hasPreFilter ? new BitSet[] {query.getPrefilter()} : EMPTY_PREFILTER_BITSET; @@ -329,13 +293,7 @@ public SearchResults search(CagraQuery query) throws Throwable { long[] queriesShape = {numQueries, vectorDimension}; MemorySegment queriesTensor = prepareTensor( - localArena, - queriesDP.handle(), - queriesShape, - getDataTypeCode(queryDataType, this.indexQuantizer), - getPrecisionFromDataType(queryDataType), - kDLCUDA(), - 1); + localArena, queriesDP.handle(), queriesShape, kDLFloat(), 32, kDLCUDA(), 1); long[] neighborsShape = {numQueries, topK}; MemorySegment neighborsTensor = prepareTensor( @@ -815,51 +773,6 @@ private static CloseableHandle createMergeParamsSegment(CagraMergeParams mergePa return new CompositeCloseableHandle(seg, handles); } - // Helper method to map DataType to data type codes - private int getDataTypeCode(DataType dataType, CuVSQuantizer quantizer) { - switch (dataType) { - case FLOAT: - return kDLFloat(); - case BYTE: - if (quantizer instanceof BinaryQuantizer) { - return kDLUInt(); - } else { - return kDLInt(); - } - case INT: - return kDLInt(); - default: - throw new IllegalArgumentException("Unsupported data type: " + dataType); - } - } - - // Helper method to convert DataType to precision - private int getPrecisionFromDataType(DataType dataType) { - switch (dataType) { - case FLOAT: - return 32; - case BYTE: - return 8; - case INT: - return 32; - default: - throw new IllegalArgumentException("Unsupported data type: " + dataType); - } - } - - private long getQueryBytesSize(long numQueries, int vectorDimension, DataType dataType) { - switch (dataType) { - case FLOAT: - return C_FLOAT_BYTE_SIZE * numQueries * vectorDimension; - case BYTE: - return numQueries * vectorDimension; - case INT: - return C_INT_BYTE_SIZE * numQueries * vectorDimension; - default: - throw new IllegalArgumentException("Unsupported query data type: " + dataType); - } - } - /** * Builder helps configure and create an instance of {@link CagraIndex}. */ @@ -869,7 +782,6 @@ public static class Builder implements CagraIndex.Builder { private CagraIndexParams cagraIndexParams; private final CuVSResources cuvsResources; private InputStream inputStream; - private CuVSQuantizer quantizer; private CuVSMatrix graph; public Builder(CuVSResources cuvsResources) { @@ -906,28 +818,11 @@ public Builder withIndexParams(CagraIndexParams cagraIndexParameters) { return this; } - @Override - public Builder withQuantizer(CuVSQuantizer quantizer) { - this.quantizer = quantizer; - return this; - } - @Override public CagraIndexImpl build() throws Throwable { if (inputStream != null) { return new CagraIndexImpl(inputStream, cuvsResources); } else { - if (quantizer != null && dataset != null) { - if (dataset.dataType() != DataType.FLOAT) { - throw new IllegalArgumentException( - "Quantizer input requires FLOAT data type, dataset has " + dataset.dataType()); - } - - CuVSMatrix quantizedDataset = quantizer.transform(dataset); - - return new CagraIndexImpl(cagraIndexParams, quantizedDataset, cuvsResources); - } - if (graph != null) { if (cagraIndexParams == null || dataset == null) { throw new IllegalArgumentException( @@ -935,9 +830,10 @@ public CagraIndexImpl build() throws Throwable { + "you must specify the original dataset and the metric used."); } return new CagraIndexImpl( - cagraIndexParams.getCuvsDistanceType(), graph, dataset, cuvsResources, null); + cagraIndexParams.getCuvsDistanceType(), graph, dataset, cuvsResources); + } else { + return new CagraIndexImpl(cagraIndexParams, dataset, cuvsResources); } - return new CagraIndexImpl(cagraIndexParams, dataset, cuvsResources); } } } diff --git a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CuVSHostMatrixImpl.java b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CuVSHostMatrixImpl.java index ffc4efcd95..c2895282ca 100644 --- a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CuVSHostMatrixImpl.java +++ b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CuVSHostMatrixImpl.java @@ -22,8 +22,6 @@ import static com.nvidia.cuvs.internal.panama.headers_h.*; import com.nvidia.cuvs.CuVSHostMatrix; -import com.nvidia.cuvs.CuVSMatrix.DataType; -import com.nvidia.cuvs.CuVSMatrix.MemoryKind; import com.nvidia.cuvs.RowView; import java.lang.foreign.*; import java.lang.invoke.VarHandle; @@ -120,11 +118,6 @@ public void toArray(byte[][] array) { } } - @Override - public MemoryKind memoryKind() { - return MemoryKind.HOST; - } - @Override public void close() {} diff --git a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CuVSMatrixBaseImpl.java b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CuVSMatrixBaseImpl.java index fc0461ef22..120960cfc4 100644 --- a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CuVSMatrixBaseImpl.java +++ b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CuVSMatrixBaseImpl.java @@ -18,7 +18,6 @@ import static com.nvidia.cuvs.internal.panama.headers_h.*; import com.nvidia.cuvs.CuVSMatrix; -import com.nvidia.cuvs.CuVSMatrix.DataType; import java.lang.foreign.Arena; import java.lang.foreign.MemorySegment; @@ -50,11 +49,6 @@ public MemorySegment memorySegment() { return memorySegment; } - @Override - public DataType dataType() { - return dataType; - } - protected int bits() { return switch (dataType) { case FLOAT, INT, UINT -> 32; diff --git a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/Scalar8BitQuantizerImpl.java b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/Scalar8BitQuantizerImpl.java deleted file mode 100644 index de49005cb2..0000000000 --- a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/Scalar8BitQuantizerImpl.java +++ /dev/null @@ -1,461 +0,0 @@ -/* - * Copyright (c) 2025, NVIDIA CORPORATION. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ -package com.nvidia.cuvs.internal; - -import static com.nvidia.cuvs.internal.common.CloseableRMMAllocation.allocateRMMSegment; -import static com.nvidia.cuvs.internal.common.LinkerHelper.C_CHAR; -import static com.nvidia.cuvs.internal.common.LinkerHelper.C_FLOAT_BYTE_SIZE; -import static com.nvidia.cuvs.internal.common.LinkerHelper.C_POINTER; -import static com.nvidia.cuvs.internal.common.Util.CudaMemcpyKind.*; -import static com.nvidia.cuvs.internal.common.Util.checkCuVSError; -import static com.nvidia.cuvs.internal.common.Util.cudaMemcpy; -import static com.nvidia.cuvs.internal.common.Util.prepareTensor; -import static com.nvidia.cuvs.internal.panama.headers_h.cuvsScalarQuantizerCreate; -import static com.nvidia.cuvs.internal.panama.headers_h.cuvsScalarQuantizerDestroy; -import static com.nvidia.cuvs.internal.panama.headers_h.cuvsScalarQuantizerInverseTransform; -import static com.nvidia.cuvs.internal.panama.headers_h.cuvsScalarQuantizerParamsCreate; -import static com.nvidia.cuvs.internal.panama.headers_h.cuvsScalarQuantizerParamsDestroy; -import static com.nvidia.cuvs.internal.panama.headers_h.cuvsScalarQuantizerParams_t; -import static com.nvidia.cuvs.internal.panama.headers_h.cuvsScalarQuantizerTrain; -import static com.nvidia.cuvs.internal.panama.headers_h.cuvsScalarQuantizerTransform; -import static com.nvidia.cuvs.internal.panama.headers_h.cuvsScalarQuantizer_t; -import static com.nvidia.cuvs.internal.panama.headers_h.cuvsStreamSync; -import static com.nvidia.cuvs.internal.panama.headers_h.kDLCPU; -import static com.nvidia.cuvs.internal.panama.headers_h.kDLCUDA; -import static com.nvidia.cuvs.internal.panama.headers_h.kDLFloat; -import static com.nvidia.cuvs.internal.panama.headers_h.kDLInt; - -import com.nvidia.cuvs.CuVSMatrix; -import com.nvidia.cuvs.CuVSMatrix.DataType; -import com.nvidia.cuvs.CuVSMatrix.MemoryKind; -import com.nvidia.cuvs.CuVSResources; -import com.nvidia.cuvs.internal.panama.cuvsScalarQuantizerParams; -import java.lang.foreign.Arena; -import java.lang.foreign.MemorySegment; -import java.util.Objects; - -/** - * Implementation of Scalar8BitQuantizer using Panama Foreign Function Interface. - */ -public class Scalar8BitQuantizerImpl { - - private static final long C_BYTE_SIZE = C_CHAR.byteSize(); - - private final CuVSResources resources; - private final MemorySegment quantizerSegment; - private boolean destroyed; - - private Scalar8BitQuantizerImpl(CuVSResources resources, MemorySegment quantizerSegment) { - this.resources = resources; - this.quantizerSegment = quantizerSegment; - } - - private void checkNotDestroyed() { - if (destroyed) { - throw new IllegalStateException("Scalar8BitQuantizer has been destroyed"); - } - } - - public void destroy() throws Throwable { - checkNotDestroyed(); - try { - int returnValue = cuvsScalarQuantizerDestroy(quantizerSegment); - checkCuVSError(returnValue, "cuvsScalarQuantizerDestroy"); - } finally { - destroyed = true; - } - } - - public void close() throws Throwable { - destroy(); - } - - public CuVSMatrix transform(CuVSMatrix dataset) throws Throwable { - checkNotDestroyed(); - - if (dataset.dataType() != DataType.FLOAT) { - throw new IllegalArgumentException( - "Scalar8BitQuantizer requires FLOAT input, got " + dataset.dataType()); - } - - try (Arena resultArena = Arena.ofShared(); - var localArena = Arena.ofConfined(); - var resourcesAccessor = resources.access()) { - - long rows = dataset.size(); - long cols = dataset.columns(); - - MemorySegment datasetMemSegment = ((CuVSMatrixBaseImpl) dataset).memorySegment(); - long cuvsResourcesPtr = resourcesAccessor.handle(); - - // TODO: Currently all datasets are HOST memory kind, preventing testing of GPU quantization - // flow. - // Future improvement: Support GPU-native dataset creation in Java API to enable direct GPU - // quantization testing. - if (dataset.memoryKind() == MemoryKind.HOST) { - return performTransformHost( - cuvsResourcesPtr, quantizerSegment, localArena, rows, cols, datasetMemSegment); - } - - return performTransformGPU( - cuvsResourcesPtr, - quantizerSegment, - localArena, - resultArena, - rows, - cols, - datasetMemSegment); - } catch (Throwable t) { - throw t; - } - } - - /** - * GPU transform with buffered processing for large datasets. - */ - private CuVSMatrix performTransformGPU( - long cuvsResourcesPtr, - MemorySegment quantizerSeg, - Arena localArena, - Arena resultArena, - long rows, - long cols, - MemorySegment datasetMemSegment) - throws Throwable { - - final long BUFFER_SIZE_BYTES = 1024 * 1024; - final long FLOAT_SIZE = C_FLOAT_BYTE_SIZE; - final long FLOATS_PER_BUFFER = BUFFER_SIZE_BYTES / FLOAT_SIZE; - final long ROWS_PER_BUFFER = Math.max(1, FLOATS_PER_BUFFER / cols); - - CuVSMatrix result = new CuVSHostMatrixArenaImpl(rows, cols, DataType.BYTE); - - // Allocate buffers outside the loop (addressing reviewer feedback) - long maxDatasetBytes = FLOAT_SIZE * ROWS_PER_BUFFER * cols; - long maxOutputBytes = C_BYTE_SIZE * ROWS_PER_BUFFER * cols; - - try (var datasetDP = allocateRMMSegment(cuvsResourcesPtr, maxDatasetBytes); - var outputDP = allocateRMMSegment(cuvsResourcesPtr, maxOutputBytes)) { - - for (long startRow = 0; startRow < rows; startRow += ROWS_PER_BUFFER) { - long endRow = Math.min(startRow + ROWS_PER_BUFFER, rows); - long chunkRows = endRow - startRow; - long actualDatasetBytes = FLOAT_SIZE * chunkRows * cols; - long actualOutputBytes = C_BYTE_SIZE * chunkRows * cols; - - // Get slices of pre-allocated buffers - MemorySegment datasetPtr = datasetDP.handle().asSlice(0, actualDatasetBytes); - MemorySegment outputPtr = outputDP.handle().asSlice(0, actualOutputBytes); - - MemorySegment chunkData = - datasetMemSegment.asSlice(startRow * cols * FLOAT_SIZE, actualDatasetBytes); - - cudaMemcpy(datasetPtr, chunkData, actualDatasetBytes, INFER_DIRECTION); - - long datasetShape[] = {chunkRows, cols}; - MemorySegment datasetTensor = - prepareTensor(localArena, datasetPtr, datasetShape, kDLFloat(), 32, kDLCUDA(), 1); - MemorySegment outputTensor = - prepareTensor(localArena, outputPtr, datasetShape, kDLInt(), 8, kDLCUDA(), 1); - - int returnValue = - cuvsScalarQuantizerTransform( - cuvsResourcesPtr, quantizerSeg, datasetTensor, outputTensor); - checkCuVSError(returnValue, "cuvsScalarQuantizerTransform"); - - returnValue = cuvsStreamSync(cuvsResourcesPtr); - checkCuVSError(returnValue, "cuvsStreamSync"); - - MemorySegment resultChunk = - ((CuVSMatrixBaseImpl) result) - .memorySegment() - .asSlice(startRow * cols, actualOutputBytes); - cudaMemcpy(resultChunk, outputPtr, actualOutputBytes, DEVICE_TO_HOST); - } - } // CloseableRMMAllocation automatically closes here - - if (result.dataType() != DataType.BYTE) { - throw new IllegalStateException( - "Expected BYTE output from scalar quantization, got " + result.dataType()); - } - - return result; - } - - private MemorySegment buildMemorySegmentFromBytes(Arena arena, byte[][] data) { - long rows = data.length; - long cols = rows > 0 ? data[0].length : 0; - long totalBytes = rows * cols * C_BYTE_SIZE; - - MemorySegment segment = arena.allocate(totalBytes); - for (int i = 0; i < rows; i++) { - for (int j = 0; j < cols; j++) { - segment.set(C_CHAR, (i * cols + j) * C_BYTE_SIZE, data[i][j]); - } - } - return segment; - } - - public static Scalar8BitQuantizerImpl create(CuVSResources resources, CuVSMatrix trainingDataset) - throws Throwable { - Objects.requireNonNull(resources); - Objects.requireNonNull(trainingDataset); - - try (var localArena = Arena.ofConfined(); - var resourcesAccessor = resources.access()) { - - long rows = trainingDataset.size(); - long cols = trainingDataset.columns(); - MemorySegment datasetMemSegment = ((CuVSMatrixBaseImpl) trainingDataset).memorySegment(); - long cuvsResourcesPtr = resourcesAccessor.handle(); - - MemorySegment paramsSegment = localArena.allocate(cuvsScalarQuantizerParams_t); - int returnValue = cuvsScalarQuantizerParamsCreate(paramsSegment); - checkCuVSError(returnValue, "cuvsScalarQuantizerParamsCreate"); - - MemorySegment paramsPtr = paramsSegment.get(C_POINTER, 0); - cuvsScalarQuantizerParams.quantile(paramsPtr, 0.99f); - - MemorySegment quantizerSegment = localArena.allocate(cuvsScalarQuantizer_t); - returnValue = cuvsScalarQuantizerCreate(quantizerSegment); - checkCuVSError(returnValue, "cuvsScalarQuantizerCreate"); - - MemorySegment quantizerPtr = quantizerSegment.get(C_POINTER, 0); - - if (trainingDataset.memoryKind() == MemoryKind.HOST) { - long datasetShape[] = {rows, cols}; - MemorySegment datasetTensor = - prepareTensor(localArena, datasetMemSegment, datasetShape, kDLFloat(), 32, kDLCPU(), 1); - - returnValue = - cuvsScalarQuantizerTrain(cuvsResourcesPtr, paramsPtr, datasetTensor, quantizerPtr); - checkCuVSError(returnValue, "cuvsScalarQuantizerTrain"); - } else { - performGPUTraining( - cuvsResourcesPtr, paramsPtr, quantizerPtr, localArena, rows, cols, datasetMemSegment); - } - - returnValue = cuvsScalarQuantizerParamsDestroy(paramsPtr); - checkCuVSError(returnValue, "cuvsScalarQuantizerParamsDestroy"); - - return new Scalar8BitQuantizerImpl(resources, quantizerPtr); - } - } - - public CuVSMatrix inverseTransform(CuVSMatrix quantizedData) throws Throwable { - checkNotDestroyed(); - - if (quantizedData.dataType() != DataType.BYTE) { - throw new IllegalArgumentException( - "Inverse transform requires BYTE input, got " + quantizedData.dataType()); - } - - try (Arena resultArena = Arena.ofShared(); - var localArena = Arena.ofConfined(); - var resourcesAccessor = resources.access()) { - - long rows = quantizedData.size(); - long cols = quantizedData.columns(); - MemorySegment quantizedMemSegment = ((CuVSMatrixBaseImpl) quantizedData).memorySegment(); - long cuvsResourcesPtr = resourcesAccessor.handle(); - - // TODO: Currently all datasets are HOST memory kind, preventing testing of GPU quantization - // flow. - // Future improvement: Support GPU-native dataset creation in Java API to enable direct GPU - // quantization testing. - if (quantizedData.memoryKind() == MemoryKind.HOST) { - return performInverseTransformHost( - cuvsResourcesPtr, quantizerSegment, localArena, rows, cols, quantizedMemSegment); - } - - return performInverseTransformGPU( - cuvsResourcesPtr, - quantizerSegment, - localArena, - resultArena, - rows, - cols, - quantizedMemSegment); - } - } - - private CuVSMatrix performInverseTransformGPU( - long cuvsResourcesPtr, - MemorySegment quantizerSeg, - Arena localArena, - Arena resultArena, - long rows, - long cols, - MemorySegment quantizedMemSegment) - throws Throwable { - - final long BUFFER_SIZE_BYTES = 1024 * 1024; - final long BYTE_SIZE = C_BYTE_SIZE; - final long BYTES_PER_BUFFER = BUFFER_SIZE_BYTES; - final long ROWS_PER_BUFFER = Math.max(1, BYTES_PER_BUFFER / (cols * BYTE_SIZE)); - - CuVSMatrix result = new CuVSHostMatrixArenaImpl(rows, cols, DataType.FLOAT); - - // Allocate buffers outside the loop (addressing reviewer feedback) - long maxQuantizedBytes = BYTE_SIZE * ROWS_PER_BUFFER * cols; - long maxOutputBytes = C_FLOAT_BYTE_SIZE * ROWS_PER_BUFFER * cols; - - try (var quantizedDP = allocateRMMSegment(cuvsResourcesPtr, maxQuantizedBytes); - var outputDP = allocateRMMSegment(cuvsResourcesPtr, maxOutputBytes)) { - - for (long startRow = 0; startRow < rows; startRow += ROWS_PER_BUFFER) { - long endRow = Math.min(startRow + ROWS_PER_BUFFER, rows); - long chunkRows = endRow - startRow; - long actualQuantizedBytes = BYTE_SIZE * chunkRows * cols; - long actualOutputBytes = C_FLOAT_BYTE_SIZE * chunkRows * cols; - - // Get slices of pre-allocated buffers - MemorySegment quantizedPtr = quantizedDP.handle().asSlice(0, actualQuantizedBytes); - MemorySegment outputPtr = outputDP.handle().asSlice(0, actualOutputBytes); - - MemorySegment chunkData = - quantizedMemSegment.asSlice(startRow * cols * BYTE_SIZE, actualQuantizedBytes); - cudaMemcpy(quantizedPtr, chunkData, actualQuantizedBytes, INFER_DIRECTION); - - long datasetShape[] = {chunkRows, cols}; - MemorySegment quantizedTensor = - prepareTensor(localArena, quantizedPtr, datasetShape, kDLInt(), 8, kDLCUDA(), 1); - MemorySegment outputTensor = - prepareTensor(localArena, outputPtr, datasetShape, kDLFloat(), 32, kDLCUDA(), 1); - - int returnValue = - cuvsScalarQuantizerInverseTransform( - cuvsResourcesPtr, quantizerSeg, quantizedTensor, outputTensor); - checkCuVSError(returnValue, "cuvsScalarQuantizerInverseTransform"); - - returnValue = cuvsStreamSync(cuvsResourcesPtr); - checkCuVSError(returnValue, "cuvsStreamSync"); - - MemorySegment resultChunk = - ((CuVSMatrixBaseImpl) result) - .memorySegment() - .asSlice(startRow * cols * C_FLOAT_BYTE_SIZE, actualOutputBytes); - cudaMemcpy(resultChunk, outputPtr, actualOutputBytes, DEVICE_TO_HOST); - } - } // CloseableRMMAllocation automatically closes here - - return result; - } - - private static void performGPUTraining( - long cuvsResourcesPtr, - MemorySegment paramsPtr, - MemorySegment quantizerPtr, - Arena localArena, - long rows, - long cols, - MemorySegment datasetMemSegment) - throws Throwable { - - final long BUFFER_SIZE_BYTES = 1024 * 1024; - final long FLOAT_SIZE = C_FLOAT_BYTE_SIZE; - final long FLOATS_PER_BUFFER = BUFFER_SIZE_BYTES / FLOAT_SIZE; - final long ROWS_PER_BUFFER = Math.max(1, FLOATS_PER_BUFFER / cols); - - // Allocate buffer outside the loop (addressing reviewer feedback) - long maxDatasetBytes = FLOAT_SIZE * ROWS_PER_BUFFER * cols; - - try (var datasetDP = allocateRMMSegment(cuvsResourcesPtr, maxDatasetBytes)) { - for (long startRow = 0; startRow < rows; startRow += ROWS_PER_BUFFER) { - long endRow = Math.min(startRow + ROWS_PER_BUFFER, rows); - long chunkRows = endRow - startRow; - long actualDatasetBytes = FLOAT_SIZE * chunkRows * cols; - - // Get slice of pre-allocated buffer - MemorySegment datasetPtr = datasetDP.handle().asSlice(0, actualDatasetBytes); - - MemorySegment chunkData = - datasetMemSegment.asSlice(startRow * cols * FLOAT_SIZE, actualDatasetBytes); - cudaMemcpy(datasetPtr, chunkData, actualDatasetBytes, INFER_DIRECTION); - - long datasetShape[] = {chunkRows, cols}; - MemorySegment datasetTensor = - prepareTensor(localArena, datasetPtr, datasetShape, kDLFloat(), 32, kDLCUDA(), 1); - - int returnValue = - cuvsScalarQuantizerTrain(cuvsResourcesPtr, paramsPtr, datasetTensor, quantizerPtr); - checkCuVSError(returnValue, "cuvsScalarQuantizerTrain"); - - returnValue = cuvsStreamSync(cuvsResourcesPtr); - checkCuVSError(returnValue, "cuvsStreamSync"); - } - } // CloseableRMMAllocation automatically closes here - } - - private static CuVSMatrix performTransformHost( - long resPtr, - MemorySegment quantizerSeg, - Arena localArena, - long rows, - long cols, - MemorySegment dataSeg) - throws Throwable { - - var result = new CuVSHostMatrixArenaImpl(rows, cols, DataType.BYTE); - - long[] shape = {rows, cols}; - MemorySegment inTensor = prepareTensor(localArena, dataSeg, shape, kDLFloat(), 32, kDLCPU(), 1); - MemorySegment outTensor = - prepareTensor( - localArena, - ((CuVSMatrixBaseImpl) result).memorySegment(), - shape, - kDLInt(), - 8, - kDLCPU(), - 1); - - int rv = cuvsScalarQuantizerTransform(resPtr, quantizerSeg, inTensor, outTensor); - checkCuVSError(rv, "cuvsScalarQuantizerTransform (host)"); - - return result; - } - - private static CuVSMatrix performInverseTransformHost( - long resPtr, - MemorySegment quantizerSeg, - Arena localArena, - long rows, - long cols, - MemorySegment quantizedSeg) - throws Throwable { - - var result = new CuVSHostMatrixArenaImpl(rows, cols, DataType.FLOAT); - - long[] shape = {rows, cols}; - MemorySegment inTensor = - prepareTensor(localArena, quantizedSeg, shape, kDLInt(), 8, kDLCPU(), 1); - MemorySegment outTensor = - prepareTensor( - localArena, - ((CuVSMatrixBaseImpl) result).memorySegment(), - shape, - kDLFloat(), - 32, - kDLCPU(), - 1); - - int rv = cuvsScalarQuantizerInverseTransform(resPtr, quantizerSeg, inTensor, outTensor); - checkCuVSError(rv, "cuvsScalarQuantizerInverseTransform (host)"); - - return result; - } -} diff --git a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/spi/JDKProvider.java b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/spi/JDKProvider.java index 0c29188456..668befe236 100644 --- a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/spi/JDKProvider.java +++ b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/spi/JDKProvider.java @@ -188,45 +188,4 @@ public CuVSMatrix newMatrixFromArray(byte[][] vectors) { Util.copy(dataset.memorySegment(), vectors); return dataset; } - - @Override - public Object createScalar8BitQuantizerImpl(CuVSResources resources, CuVSMatrix trainingDataset) - throws Throwable { - return Scalar8BitQuantizerImpl.create(resources, trainingDataset); - } - - @Override - public CuVSMatrix transformScalar8Bit(Object impl, CuVSMatrix input) throws Throwable { - return ((Scalar8BitQuantizerImpl) impl).transform(input); - } - - @Override - public CuVSMatrix inverseTransformScalar8Bit(Object impl, CuVSMatrix quantizedData) - throws Throwable { - return ((Scalar8BitQuantizerImpl) impl).inverseTransform(quantizedData); - } - - @Override - public Object createBinaryQuantizerImpl( - CuVSResources resources, - CuVSMatrix trainingDataset, - BinaryQuantizer.ThresholdType thresholdType) - throws Throwable { - return BinaryQuantizerImpl.create(resources, trainingDataset, thresholdType); - } - - @Override - public CuVSMatrix transformBinaryWithImpl(Object impl, CuVSMatrix input) throws Throwable { - return BinaryQuantizerImpl.transformWithImpl(impl, input); - } - - @Override - public void closeBinaryQuantizer(Object impl) throws Throwable { - BinaryQuantizerImpl.close(impl); - } - - @Override - public void closeScalar8BitQuantizer(Object impl) throws Throwable { - ((Scalar8BitQuantizerImpl) impl).close(); - } } diff --git a/java/cuvs-java/src/test/java/com/nvidia/cuvs/QuantizationIT.java b/java/cuvs-java/src/test/java/com/nvidia/cuvs/QuantizationIT.java deleted file mode 100644 index 8189fd9139..0000000000 --- a/java/cuvs-java/src/test/java/com/nvidia/cuvs/QuantizationIT.java +++ /dev/null @@ -1,588 +0,0 @@ -/* - * Copyright (c) 2025, NVIDIA CORPORATION. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ -package com.nvidia.cuvs; - -import static com.carrotsearch.randomizedtesting.RandomizedTest.assumeTrue; -import static org.junit.Assert.*; - -import com.carrotsearch.randomizedtesting.RandomizedRunner; -import com.nvidia.cuvs.CagraIndexParams.CagraGraphBuildAlgo; -import com.nvidia.cuvs.CagraIndexParams.CuvsDistanceType; -import com.nvidia.cuvs.CuVSMatrix.DataType; -import java.lang.invoke.MethodHandles; -import java.util.Map; -import org.junit.Before; -import org.junit.Test; -import org.junit.runner.RunWith; -import org.slf4j.Logger; -import org.slf4j.LoggerFactory; - -@RunWith(RandomizedRunner.class) -public class QuantizationIT extends CuVSTestCase { - - private static final Logger log = LoggerFactory.getLogger(MethodHandles.lookup().lookupClass()); - - @Before - public void setup() { - assumeTrue("not supported on " + System.getProperty("os.name"), isLinuxAmd64()); - initializeRandom(); - log.info("Random context initialized for quantization test."); - } - - private static float[][] createSimpleDataset() { - return new float[][] { - {1.0f, 2.0f, 3.0f}, - {4.0f, 5.0f, 6.0f}, - {7.0f, 8.0f, 9.0f}, - {10.0f, 11.0f, 12.0f}, - {13.0f, 14.0f, 15.0f} - }; - } - - private static float[][] createSimpleQueries() { - return new float[][] { - {1.1f, 2.1f, 3.1f}, - {7.1f, 8.1f, 9.1f} - }; - } - - @Test - public void testScalarQuantizerBasicFlow() throws Throwable { - float[][] dataset = createSimpleDataset(); - float[][] queries = createSimpleQueries(); - - try (CuVSResources resources = CheckedCuVSResources.create()) { - CuVSMatrix trainingDataset = CuVSMatrix.ofArray(dataset); - assertEquals(DataType.FLOAT, trainingDataset.dataType()); - assertEquals(dataset.length, trainingDataset.size()); - assertEquals(dataset[0].length, trainingDataset.columns()); - - Scalar8BitQuantizer quantizer = new Scalar8BitQuantizer(resources, trainingDataset); - assertEquals(DataType.BYTE, quantizer.outputDataType()); - log.info("Created scalar quantizer with BYTE data type"); - - CagraIndexParams indexParams = - new CagraIndexParams.Builder() - .withCagraGraphBuildAlgo(CagraGraphBuildAlgo.NN_DESCENT) - .withGraphDegree(2) - .withIntermediateGraphDegree(3) - .withNumWriterThreads(1) - .withMetric(CuvsDistanceType.L2Expanded) - .build(); - - CagraIndex index = - CagraIndex.newBuilder(resources) - .withDataset(trainingDataset) - .withQuantizer(quantizer) - .withIndexParams(indexParams) - .build(); - - log.info("Built index with quantized dataset"); - - CagraSearchParams searchParams = new CagraSearchParams.Builder().build(); - - CagraQuery query = - CagraQuery.newBuilder(resources) - .withQueryVectors(queries) - .withQuantizer(quantizer) - .withSearchParams(searchParams) - .withTopK(3) - .build(); - - assertTrue("Query should have quantized vectors", query.hasQuantizedQueries()); - assertEquals(DataType.BYTE, query.getQueryDataType()); - log.info("Created quantized query"); - - // Perform search - SearchResults results = index.search(query); - assertNotNull("Search results should not be null", results); - - // Verify results - assertEquals( - "Should have results for all queries", queries.length, results.getResults().size()); - - for (int i = 0; i < results.getResults().size(); i++) { - Map queryResults = results.getResults().get(i); - assertFalse("Query " + i + " should return some neighbors", queryResults.isEmpty()); - - // Verify all returned IDs are within valid range - for (Integer id : queryResults.keySet()) { - assertTrue("Returned ID should be valid", id >= 0 && id < dataset.length); - } - } - - log.info( - "Search completed successfully with {} queries returning neighbors", - results.getResults().size()); - - // Cleanup - index.close(); - quantizer.close(); - trainingDataset.close(); - } - } - - @Test - public void testBinaryQuantizerBasicFlow() throws Throwable { - float[][] dataset = createSimpleDataset(); - float[][] queries = createSimpleQueries(); - - try (CuVSResources resources = CheckedCuVSResources.create()) { - CuVSMatrix trainingDataset = CuVSMatrix.ofArray(dataset); - assertEquals(DataType.FLOAT, trainingDataset.dataType()); - - BinaryQuantizer quantizer = new BinaryQuantizer(resources, trainingDataset); - assertEquals(DataType.BYTE, quantizer.outputDataType()); - log.info("Created binary quantizer"); - - CagraIndexParams indexParams = - new CagraIndexParams.Builder() - .withCagraGraphBuildAlgo(CagraGraphBuildAlgo.NN_DESCENT) - .withGraphDegree(2) - .withIntermediateGraphDegree(3) - .withNumWriterThreads(1) - .withMetric(CuvsDistanceType.L2Expanded) - .build(); - - CagraIndex index = - CagraIndex.newBuilder(resources) - .withDataset(trainingDataset) - .withQuantizer(quantizer) - .withIndexParams(indexParams) - .build(); - - log.info("Built index with binary quantized dataset"); - - CagraSearchParams searchParams = new CagraSearchParams.Builder().build(); - - CagraQuery query = - CagraQuery.newBuilder(resources) - .withQueryVectors(queries) - .withQuantizer(quantizer) - .withSearchParams(searchParams) - .withTopK(3) - .build(); - - assertTrue("Query should have quantized vectors", query.hasQuantizedQueries()); - assertEquals(DataType.BYTE, query.getQueryDataType()); - - // Perform search - SearchResults results = index.search(query); - assertNotNull("Search results should not be null", results); - - // Verify results - assertEquals( - "Should have results for all queries", queries.length, results.getResults().size()); - - for (int i = 0; i < results.getResults().size(); i++) { - Map queryResults = results.getResults().get(i); - assertFalse("Query " + i + " should return some neighbors", queryResults.isEmpty()); - - // Verify all returned IDs are within valid range - for (Integer id : queryResults.keySet()) { - assertTrue("Returned ID should be valid", id >= 0 && id < dataset.length); - } - } - - log.info("Binary quantized search completed successfully"); - - CuVSMatrix quantizedQueries = query.getQuantizedQueries(); - try { - quantizer.inverseTransform(quantizedQueries); - fail("Expected UnsupportedOperationException to be thrown"); - } catch (UnsupportedOperationException e) { - assertTrue( - "Exception message should mention inverse transform", - e.getMessage().contains("inverse")); - } - - // Cleanup - index.close(); - quantizer.close(); - trainingDataset.close(); - quantizedQueries.close(); - } - } - - @Test - public void testScalarQuantizerInverseTransform() throws Throwable { - float[][] dataset = createSimpleDataset(); - - try (CuVSResources resources = CheckedCuVSResources.create()) { - CuVSMatrix trainingDataset = CuVSMatrix.ofArray(dataset); - Scalar8BitQuantizer quantizer = new Scalar8BitQuantizer(resources, trainingDataset); - - // Transform and inverse transform - CuVSMatrix quantized = quantizer.transform(trainingDataset); - assertEquals(DataType.BYTE, quantized.dataType()); - - CuVSMatrix recovered = quantizer.inverseTransform(quantized); - assertEquals(DataType.FLOAT, recovered.dataType()); - assertEquals(trainingDataset.size(), recovered.size()); - assertEquals(trainingDataset.columns(), recovered.columns()); - - log.info("Inverse transform completed successfully"); - - // Verify quantization worked - assertEquals(DataType.BYTE, quantized.dataType()); - assertEquals(DataType.FLOAT, recovered.dataType()); - - // Cleanup - quantizer.close(); - trainingDataset.close(); - quantized.close(); - recovered.close(); - } - } - - @Test - public void testCPUBinaryQuantization() throws Throwable { - float[][] dataset = createSimpleDataset(); - - try (CuVSResources resources = CheckedCuVSResources.create()) { - CuVSMatrix inputDataset = CuVSMatrix.ofArray(dataset); - assertEquals(CuVSMatrix.MemoryKind.HOST, inputDataset.memoryKind()); - - BinaryQuantizer quantizer = new BinaryQuantizer(resources, inputDataset); - - CuVSMatrix quantized = quantizer.transform(inputDataset); - assertEquals(DataType.BYTE, quantized.dataType()); - assertEquals(CuVSMatrix.MemoryKind.HOST, quantized.memoryKind()); - - try { - quantizer.inverseTransform(quantized); - fail("Expected UnsupportedOperationException to be thrown"); - } catch (UnsupportedOperationException e) { - assertTrue( - "Exception message should mention inverse transform", - e.getMessage().contains("inverse")); - } - - quantizer.close(); - inputDataset.close(); - quantized.close(); - - log.info("✓ Binary CPU quantization test passed"); - } - } - - @Test - public void testRandomizedScalarQuantizedIndexingAndSearch() throws Throwable { - int maxDatasetSize = 60; - int maxDimensions = 20; - int maxQueries = 5; - int maxTopK = 4; - - for (int iteration = 0; iteration < 5; iteration++) { - int datasetSize = random.nextInt(maxDatasetSize) + maxTopK; - int dims = random.nextInt(maxDimensions) + 2; - int numQueries = random.nextInt(maxQueries) + 1; - int topK = Math.min(random.nextInt(maxTopK) + 1, datasetSize); - - // Generate dataset and query - float[][] dataset = new float[datasetSize][dims]; - float[][] queries = new float[numQueries][dims]; - for (int i = 0; i < datasetSize; i++) - for (int d = 0; d < dims; d++) dataset[i][d] = (float) random.nextGaussian(); - for (int i = 0; i < numQueries; i++) - for (int d = 0; d < dims; d++) queries[i][d] = (float) random.nextGaussian(); - - try (CuVSResources resources = CheckedCuVSResources.create()) { - CuVSMatrix trainingDataset = CuVSMatrix.ofArray(dataset); - Scalar8BitQuantizer quantizer = new Scalar8BitQuantizer(resources, trainingDataset); - - CagraIndexParams indexParams = - new CagraIndexParams.Builder() - .withCagraGraphBuildAlgo(CagraGraphBuildAlgo.NN_DESCENT) - .withGraphDegree(Math.min(4, datasetSize - 1)) - .withIntermediateGraphDegree(Math.min(5, datasetSize)) - .withNumWriterThreads(1) - .withMetric(CuvsDistanceType.L2Expanded) - .build(); - - CagraIndex index = - CagraIndex.newBuilder(resources) - .withDataset(trainingDataset) - .withQuantizer(quantizer) - .withIndexParams(indexParams) - .build(); - - CagraSearchParams searchParams = new CagraSearchParams.Builder().build(); - CagraQuery query = - CagraQuery.newBuilder(resources) - .withQueryVectors(queries) - .withQuantizer(quantizer) - .withSearchParams(searchParams) - .withTopK(topK) - .build(); - - SearchResults results = index.search(query); - assertNotNull("Results must not be null", results); - assertEquals("One result per query", numQueries, results.getResults().size()); - - for (int qi = 0; qi < numQueries; qi++) { - Map result = results.getResults().get(qi); - // It is possible for result to have < topK entries if datasetSize == topK - assertFalse("Each query should return non-empty result", result.isEmpty()); - for (Integer id : result.keySet()) { - assertTrue("ID in result should be in range", id >= 0 && id < datasetSize); - } - } - - quantizer.close(); - trainingDataset.close(); - index.close(); - } - } - } - - @Test - public void testBinaryQuantizationWithZeroThreshold() throws Throwable { - float[][] dataset = - new float[][] { - {-2.0f, -1.0f, 0.0f, 1.0f, 2.0f}, // crosses zero - {0.1f, 0.5f, 1.0f, 1.5f, 2.0f}, // all positive - {-2.0f, -1.5f, -1.0f, -0.5f, 0.0f}, // mostly negative - {-1.0f, 0.0f, 1.0f, 2.0f, 3.0f} // mixed - }; - - float[][] queries = - new float[][] { - {0.5f, 0.5f, 0.5f, 0.5f, 0.5f}, - {-0.5f, -0.5f, -0.5f, -0.5f, -0.5f} - }; - - try (CuVSResources resources = CheckedCuVSResources.create()) { - CuVSMatrix trainingDataset = CuVSMatrix.ofArray(dataset); - - BinaryQuantizer quantizer = - new BinaryQuantizer(resources, trainingDataset, BinaryQuantizer.ThresholdType.ZERO); - assertEquals(DataType.BYTE, quantizer.outputDataType()); - assertEquals(BinaryQuantizer.ThresholdType.ZERO, quantizer.getThresholdType()); - - CagraIndexParams indexParams = - new CagraIndexParams.Builder() - .withCagraGraphBuildAlgo(CagraGraphBuildAlgo.NN_DESCENT) - .withGraphDegree(2) - .withIntermediateGraphDegree(3) - .withNumWriterThreads(1) - .withMetric(CuvsDistanceType.L2Expanded) - .build(); - - CagraIndex index = - CagraIndex.newBuilder(resources) - .withDataset(trainingDataset) - .withQuantizer(quantizer) - .withIndexParams(indexParams) - .build(); - - log.info("Built index with ZERO threshold binary quantization"); - - CagraSearchParams searchParams = new CagraSearchParams.Builder().build(); - CagraQuery query = - CagraQuery.newBuilder(resources) - .withQueryVectors(queries) - .withQuantizer(quantizer) - .withSearchParams(searchParams) - .withTopK(2) - .build(); - - assertTrue("Query should have quantized vectors", query.hasQuantizedQueries()); - assertEquals(DataType.BYTE, query.getQueryDataType()); - - SearchResults results = index.search(query); - assertNotNull("Search results should not be null", results); - assertEquals( - "Should have results for all queries", queries.length, results.getResults().size()); - - CuVSMatrix quantizedDataset = quantizer.transform(trainingDataset); - byte[][] result = new byte[(int) quantizedDataset.size()][(int) quantizedDataset.columns()]; - quantizedDataset.toArray(result); - - byte firstByte = result[0][0]; - int[] unpackedBits = new int[5]; - for (int i = 0; i < 5; i++) { - unpackedBits[i] = (firstByte >> i) & 1; - } - - int[] expectedRow0 = {0, 0, 0, 1, 1}; - assertArrayEquals("ZERO threshold failed for row 0", expectedRow0, unpackedBits); - - index.close(); - quantizer.close(); - trainingDataset.close(); - quantizedDataset.close(); - - log.info("✓ ZERO threshold binary quantization test passed"); - } - } - - @Test - public void testBinaryQuantizationWithMeanThreshold() throws Throwable { - float[][] dataset = - new float[][] { - {1.0f, 2.0f, 3.0f, 4.0f, 5.0f}, // mean = 3.0 - {0.0f, 1.0f, 2.0f, 3.0f, 4.0f}, // mean = 2.0 - {2.0f, 4.0f, 6.0f, 8.0f, 10.0f}, // mean = 6.0 - {-1.0f, 0.0f, 1.0f, 2.0f, 3.0f} // mean = 1.0 - }; - - float[][] queries = - new float[][] { - {2.5f, 2.5f, 2.5f, 2.5f, 2.5f}, - {1.0f, 3.0f, 5.0f, 7.0f, 9.0f} - }; - - try (CuVSResources resources = CheckedCuVSResources.create()) { - CuVSMatrix trainingDataset = CuVSMatrix.ofArray(dataset); - - BinaryQuantizer quantizer = - new BinaryQuantizer(resources, trainingDataset, BinaryQuantizer.ThresholdType.MEAN); - assertEquals(BinaryQuantizer.ThresholdType.MEAN, quantizer.getThresholdType()); - - CagraIndexParams indexParams = - new CagraIndexParams.Builder() - .withCagraGraphBuildAlgo(CagraGraphBuildAlgo.NN_DESCENT) - .withGraphDegree(2) - .withIntermediateGraphDegree(3) - .withNumWriterThreads(1) - .withMetric(CuvsDistanceType.L2Expanded) - .build(); - - CagraIndex index = - CagraIndex.newBuilder(resources) - .withDataset(trainingDataset) - .withQuantizer(quantizer) - .withIndexParams(indexParams) - .build(); - - log.info("Built index with MEAN threshold binary quantization"); - - // Perform search - CagraSearchParams searchParams = new CagraSearchParams.Builder().build(); - CagraQuery query = - CagraQuery.newBuilder(resources) - .withQueryVectors(queries) - .withQuantizer(quantizer) - .withSearchParams(searchParams) - .withTopK(2) - .build(); - - SearchResults results = index.search(query); - assertNotNull("Search results should not be null", results); - assertEquals( - "Should have results for all queries", queries.length, results.getResults().size()); - - // Calculate overall mean: (3.0 + 2.0 + 6.0 + 1.0) / 4 = 3.0 - // Verify quantization behavior - values > mean should become 1, values <= mean should become - // 0 - CuVSMatrix quantizedDataset = quantizer.transform(trainingDataset); - byte[][] result1 = new byte[(int) quantizedDataset.size()][(int) quantizedDataset.columns()]; - - // Copy data from CuVSMatrix to array - quantizedDataset.toArray(result1); - - // Verify that MEAN threshold produces different results than ZERO threshold - BinaryQuantizer zeroQuantizer = - new BinaryQuantizer(resources, trainingDataset, BinaryQuantizer.ThresholdType.ZERO); - CuVSMatrix zeroQuantized = zeroQuantizer.transform(trainingDataset); - byte[][] zeroData = new byte[(int) zeroQuantized.size()][(int) zeroQuantized.columns()]; - zeroQuantized.toArray(zeroData); - - boolean isDifferent = false; - for (int i = 0; i < result1.length && !isDifferent; i++) { - for (int j = 0; j < result1[i].length && !isDifferent; j++) { - if (result1[i][j] != zeroData[i][j]) { - isDifferent = true; - } - } - } - assertTrue("MEAN and ZERO thresholds should produce different results", isDifferent); - - // Cleanup - index.close(); - quantizer.close(); - zeroQuantizer.close(); - trainingDataset.close(); - quantizedDataset.close(); - zeroQuantized.close(); - - log.info("✓ MEAN threshold binary quantization test passed"); - } - } - - @Test - public void testBinaryQuantizationWithSamplingMedianThreshold() throws Throwable { - float[][] dataset = createSimpleDataset(); - float[][] queries = createSimpleQueries(); - - try (CuVSResources resources = CheckedCuVSResources.create()) { - CuVSMatrix trainingDataset = CuVSMatrix.ofArray(dataset); - assertEquals(DataType.FLOAT, trainingDataset.dataType()); - - BinaryQuantizer quantizer = - new BinaryQuantizer( - resources, trainingDataset, BinaryQuantizer.ThresholdType.SAMPLING_MEDIAN); - assertEquals(DataType.BYTE, quantizer.outputDataType()); - log.info("Created binary quantizer with sampling median threshold"); - - CagraIndexParams indexParams = - new CagraIndexParams.Builder() - .withCagraGraphBuildAlgo(CagraGraphBuildAlgo.NN_DESCENT) - .withGraphDegree(2) - .withIntermediateGraphDegree(3) - .withNumWriterThreads(1) - .withMetric(CuvsDistanceType.L2Expanded) - .build(); - - CagraIndex index = - CagraIndex.newBuilder(resources) - .withDataset(trainingDataset) - .withQuantizer(quantizer) - .withIndexParams(indexParams) - .build(); - - log.info("Built index with binary quantized dataset"); - - CagraSearchParams searchParams = new CagraSearchParams.Builder().build(); - - CagraQuery query = - CagraQuery.newBuilder(resources) - .withQueryVectors(queries) - .withQuantizer(quantizer) - .withSearchParams(searchParams) - .withTopK(3) - .build(); - - assertTrue("Query should have quantized vectors", query.hasQuantizedQueries()); - assertEquals(DataType.BYTE, query.getQueryDataType()); - - SearchResults results = index.search(query); - assertNotNull("Search results should not be null", results); - - assertEquals( - "Should return results for each query", queries.length, results.getResults().size()); - - for (int i = 0; i < results.getResults().size(); i++) { - Map result = results.getResults().get(i); - assertFalse("Each query should return non-empty result", result.isEmpty()); - assertTrue("Should return at most topK results", result.size() <= 3); - } - - log.info("✓ Binary quantization with sampling median threshold test passed"); - } - } -} diff --git a/java/panama-bindings/headers.h b/java/panama-bindings/headers.h index ba4cf8c203..57a13cf1b1 100644 --- a/java/panama-bindings/headers.h +++ b/java/panama-bindings/headers.h @@ -22,8 +22,6 @@ #include #include #include -#include -#include #include #include From 50ed7109abc2bb579e323404a6a7525adc7da275 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Lorenzo=20Dematt=C3=A9?= Date: Fri, 22 Aug 2025 23:50:08 +0200 Subject: [PATCH 42/93] [Java] CuVSMatrix for device memory (#1232) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit This PR introduces implementation classes for `CuVSDeviceMatrix` (a `CuVSMatrix` backed by device memory). It reworks the base implementation classes a bit to increase reuse, and adds benchmarks and tests for the new classes. Benchmarks were used to try out different implementations so the best one could be chosen: - Row access is backed by a buffer of pinned memory - Builders for device memory use `cudaMemcpyAsync` with the `critical` linker option to directly access heap-based memory - `cuvsMatrixCopy` is used across the board, as it has the same performances of the various `cudaMemcpy*` functions. There are some places in the codebase that will benefit from refactoring to use `CuVSDeviceMatrix` (or a generic `CuVSMatrix` plus `toHost`/`toTensor`/`fromTensor` functions); replacing these multiple ad-hoc implementations with `CuVSDeviceMatrix` will be addressed in a follow-up PR. Final numbers: ``` Benchmark (dims) (size) Mode Cnt Score Error Units CuVSDeviceMatrixBenchmarks.matrixCopyDeviceToHost 2048 16384 thrpt 5 70.531 ± 0.322 ops/s CuVSDeviceMatrixBenchmarks.matrixDeviceBuilder 2048 16384 thrpt 5 35.493 ± 0.772 ops/s CuVSDeviceMatrixBenchmarks.matrixReadRowsFromDevice 2048 16384 thrpt 5 83.616 ± 0.745 ops/s ``` With theoretical max for the PCI-E bus of 15.7 GB/s and a data size of 128MB, we get close to 2/3 of the maximum theoretical throughput (see comments for details). Authors: - Lorenzo Dematté (https://github.com/ldematte) - MithunR (https://github.com/mythrocks) Approvers: - MithunR (https://github.com/mythrocks) URL: https://github.com/rapidsai/cuvs/pull/1232 --- java/.gitignore | 1 + .../cuvs/CuVSDeviceMatrixBenchmarks.java | 107 ++++++++ .../main/java/com/nvidia/cuvs/CagraIndex.java | 4 +- .../com/nvidia/cuvs/CuVSDeviceMatrix.java | 23 +- .../main/java/com/nvidia/cuvs/CuVSMatrix.java | 43 ++- .../main/java/com/nvidia/cuvs/RowView.java | 2 +- .../com/nvidia/cuvs/spi/CuVSProvider.java | 36 ++- .../nvidia/cuvs/spi/UnsupportedProvider.java | 30 ++- .../cuvs/internal/BruteForceIndexImpl.java | 26 +- .../nvidia/cuvs/internal/CagraIndexImpl.java | 100 ++----- .../cuvs/internal/CuVSDeviceMatrixImpl.java | 246 ++++++++++++++++++ .../internal/CuVSDeviceMatrixRMMImpl.java | 97 +++++++ .../cuvs/internal/CuVSHostMatrixImpl.java | 157 ++++------- .../cuvs/internal/CuVSMatrixBaseImpl.java | 139 +++++++++- .../nvidia/cuvs/internal/HnswIndexImpl.java | 11 +- .../nvidia/cuvs/internal/SliceRowView.java | 126 +++++++++ .../nvidia/cuvs/internal/TieredIndexImpl.java | 14 +- .../com/nvidia/cuvs/internal/common/Util.java | 131 ++++++++-- .../com/nvidia/cuvs/spi/JDKProvider.java | 149 ++++++++++- .../nvidia/cuvs/BruteForceRandomizedIT.java | 2 +- .../nvidia/cuvs/CagraBuildAndSearchIT.java | 3 +- .../com/nvidia/cuvs/CagraRandomizedIT.java | 2 +- .../java/com/nvidia/cuvs/CuVSMatrixIT.java | 212 ++++++++++++--- .../com/nvidia/cuvs/HnswRandomizedIT.java | 2 +- 24 files changed, 1337 insertions(+), 326 deletions(-) create mode 100644 java/benchmarks/src/main/java/com/nvidia/cuvs/CuVSDeviceMatrixBenchmarks.java create mode 100644 java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CuVSDeviceMatrixImpl.java create mode 100644 java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CuVSDeviceMatrixRMMImpl.java create mode 100644 java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/SliceRowView.java diff --git a/java/.gitignore b/java/.gitignore index a240aca207..2e98aff7ba 100644 --- a/java/.gitignore +++ b/java/.gitignore @@ -1,4 +1,5 @@ *.iml +hs_err*.log target/ jextract-22/ openjdk-22-jextract* diff --git a/java/benchmarks/src/main/java/com/nvidia/cuvs/CuVSDeviceMatrixBenchmarks.java b/java/benchmarks/src/main/java/com/nvidia/cuvs/CuVSDeviceMatrixBenchmarks.java new file mode 100644 index 0000000000..2a0c430fb4 --- /dev/null +++ b/java/benchmarks/src/main/java/com/nvidia/cuvs/CuVSDeviceMatrixBenchmarks.java @@ -0,0 +1,107 @@ +/* + * Copyright (c) 2025, NVIDIA CORPORATION. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.nvidia.cuvs; + +import org.openjdk.jmh.annotations.*; +import org.openjdk.jmh.infra.Blackhole; + +import java.util.Random; + +@Fork(value = 1, warmups = 0) +@State(Scope.Benchmark) +public class CuVSDeviceMatrixBenchmarks { + + @Param({"2048"}) + private int dims; + + @Param({"16384"}) + private int size; + + private static final Random random = new Random(); + + private float[][] data; + + private CuVSResources resources; + private CuVSDeviceMatrix deviceMatrix; + private CuVSHostMatrix hostMatrix; + + private float[][] createRandomData() { + var array = new float[size][dims]; + + for (int i = 0; i < size; ++i) { + for (int j = 0; j < dims; ++j) { + array[i][j] = random.nextFloat(); + } + } + return array; + } + + @Setup + public void initialize() throws Throwable { + data = createRandomData(); + resources = CuVSResources.create(); + + var builder0 = CuVSMatrix.deviceBuilder(resources, size, dims, CuVSMatrix.DataType.FLOAT); + + for (int i = 0; i < size; ++i) { + var array = data[i]; + builder0.addVector(array); + } + + deviceMatrix = builder0.build(); + hostMatrix = CuVSMatrix.hostBuilder(size, dims, CuVSMatrix.DataType.FLOAT).build(); + } + + @TearDown + public void cleanUp() { + if (deviceMatrix != null) { + deviceMatrix.close(); + } + if (hostMatrix != null) { + hostMatrix.close(); + } + if (resources != null) { + resources.close(); + } + } + + @Benchmark + public void matrixReadRowsFromDevice(Blackhole bh) { + for (int i = 0; i < size; ++i) { + bh.consume(deviceMatrix.getRow(i)); + } + } + + @Benchmark + public void matrixCopyDeviceToHost() { + deviceMatrix.toHost(hostMatrix); + } + + @Benchmark + public void matrixDeviceBuilder() throws Throwable { + try (CuVSResources resources = CuVSResources.create()) { + var builder = CuVSMatrix.deviceBuilder(resources, size, dims, CuVSMatrix.DataType.FLOAT); + + for (int i = 0; i < size; ++i) { + var array = data[i]; + builder.addVector(array); + } + CuVSDeviceMatrix matrix = builder.build(); + matrix.close(); + } + } +} diff --git a/java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraIndex.java b/java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraIndex.java index e3b885a333..a3cb68448c 100644 --- a/java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraIndex.java +++ b/java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraIndex.java @@ -53,10 +53,10 @@ public interface CagraIndex extends AutoCloseable { /** Returns the CAGRA graph * - * @return a {@link CuVSMatrix} encapsulating the native int (uint32_t) array used to represent + * @return a {@link CuVSDeviceMatrix} encapsulating the native int (uint32_t) array used to represent * the cagra graph */ - CuVSMatrix getGraph(); + CuVSDeviceMatrix getGraph(); /** * A method to persist a CAGRA index using an instance of {@link OutputStream} diff --git a/java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSDeviceMatrix.java b/java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSDeviceMatrix.java index 23d29fd479..b654163a63 100644 --- a/java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSDeviceMatrix.java +++ b/java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSDeviceMatrix.java @@ -18,4 +18,25 @@ /** * A Dataset implementation backed by device (GPU) memory. */ -public interface CuVSDeviceMatrix extends CuVSMatrix {} +public interface CuVSDeviceMatrix extends CuVSMatrix { + + /** + * Fills the provided, pre-allocated host matrix with data from this device matrix. + * The content of the provided host matrix will be overwritten; the 2 matrices must have the + * same element type and dimension. + * + * @param hostMatrix the host-memory-backed matrix to fill. + */ + void toHost(CuVSHostMatrix hostMatrix); + + /** + * Returns a new, host matrix with data from this device matrix. + * The returned host matrix will need to be managed by the caller, which will be + * responsible to call {@link CuVSMatrix#close()} to free its resources when done. + */ + default CuVSHostMatrix toHost() { + var hostMatrix = CuVSMatrix.hostBuilder(size(), columns(), dataType()).build(); + toHost(hostMatrix); + return hostMatrix; + } +} diff --git a/java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSMatrix.java b/java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSMatrix.java index f8e60c5555..708009be02 100644 --- a/java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSMatrix.java +++ b/java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSMatrix.java @@ -69,29 +69,33 @@ static CuVSMatrix ofArray(byte[][] vectors) { return CuVSProvider.provider().newMatrixFromArray(vectors); } - interface Builder { + /** + * A builder to construct a new matrix one row at a time + * @param the CuVSMatrix type to build + */ + interface Builder { /** - * Add a single vector to the dataset. + * Adds a single vector to the matrix. * * @param vector A float array of as many elements as the dimensions */ void addVector(float[] vector); /** - * Add a single vector to the dataset. + * Adds a single vector to the matrix. * * @param vector A byte array of as many elements as the dimensions */ void addVector(byte[] vector); /** - * Add a single vector to the dataset. + * Adds a single vector to the matrix. * - * @param vector A int array of as many elements as the dimensions + * @param vector An int array of as many elements as the dimensions */ void addVector(int[] vector); - CuVSMatrix build(); + T build(); } /** @@ -100,10 +104,24 @@ interface Builder { * @param size Number of vectors in the dataset * @param columns Size of each vector in the dataset * @param dataType The data type of the dataset elements - * @return new instance of {@link CuVSMatrix} + * @return a builder for creating a {@link CuVSHostMatrix} + */ + static Builder hostBuilder(long size, long columns, DataType dataType) { + return CuVSProvider.provider().newHostMatrixBuilder(size, columns, dataType); + } + + /** + * Returns a builder to create a new instance of a dataset + * + * @param resources CuVS resources used to allocate the device memory needed + * @param size Number of vectors in the dataset + * @param columns Size of each vector in the dataset + * @param dataType The data type of the dataset elements + * @return a builder for creating a {@link CuVSDeviceMatrix} */ - static CuVSMatrix.Builder builder(int size, int columns, DataType dataType) { - return CuVSProvider.provider().newMatrixBuilder(size, columns, dataType); + static Builder deviceBuilder( + CuVSResources resources, long size, long columns, DataType dataType) { + return CuVSProvider.provider().newDeviceMatrixBuilder(resources, size, columns, dataType); } /** @@ -121,6 +139,13 @@ static CuVSMatrix.Builder builder(int size, int columns, DataType dataType) { */ long columns(); + /** + * Gets the element type + * + * @return a {@link DataType} describing the matrix element type + */ + DataType dataType(); + /** * Get a view (0-copy) of the row data, as a list of integers (32 bit) * diff --git a/java/cuvs-java/src/main/java/com/nvidia/cuvs/RowView.java b/java/cuvs-java/src/main/java/com/nvidia/cuvs/RowView.java index cbace82ae3..676a652427 100644 --- a/java/cuvs-java/src/main/java/com/nvidia/cuvs/RowView.java +++ b/java/cuvs-java/src/main/java/com/nvidia/cuvs/RowView.java @@ -16,7 +16,7 @@ package com.nvidia.cuvs; /** - * Represent a contiguous list of integers (32-bit) backed by off-heap memory. + * Represent a contiguous list of elements backed by off-heap memory. * * @since 25.08 */ 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 e528beff1a..2ee3fb373a 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 @@ -15,13 +15,7 @@ */ package com.nvidia.cuvs.spi; -import com.nvidia.cuvs.BruteForceIndex; -import com.nvidia.cuvs.CagraIndex; -import com.nvidia.cuvs.CagraMergeParams; -import com.nvidia.cuvs.CuVSMatrix; -import com.nvidia.cuvs.CuVSResources; -import com.nvidia.cuvs.HnswIndex; -import com.nvidia.cuvs.TieredIndex; +import com.nvidia.cuvs.*; import java.lang.invoke.MethodHandle; import java.lang.invoke.MethodType; import java.nio.file.Path; @@ -52,23 +46,37 @@ default Path nativeLibraryPath() { /** Creates a new CuVSResources. */ CuVSResources newCuVSResources(Path tempDirectory) throws Throwable; - /** Create a {@link CuVSMatrix.Builder} instance **/ - CuVSMatrix.Builder newMatrixBuilder(int size, int dimensions, CuVSMatrix.DataType dataType); + /** Create a {@link CuVSMatrix.Builder} instance for a host memory matrix **/ + CuVSMatrix.Builder newHostMatrixBuilder( + long size, long dimensions, CuVSMatrix.DataType dataType); + + /** Create a {@link CuVSMatrix.Builder} instance for a device memory matrix **/ + CuVSMatrix.Builder newDeviceMatrixBuilder( + CuVSResources cuVSResources, long size, long dimensions, CuVSMatrix.DataType dataType); + + /** Create a {@link CuVSMatrix.Builder} instance for a device memory matrix **/ + CuVSMatrix.Builder newDeviceMatrixBuilder( + CuVSResources cuVSResources, + long size, + long dimensions, + int rowStride, + int columnStride, + CuVSMatrix.DataType dataType); /** - * Returns the factory method used to build a Dataset from native memory. + * Returns the factory method used to build a CuVSMatrix from native memory. * The factory method will have this signature: - * {@code Dataset createNativeDataset(memorySegment, size, dimensions, dataType)}, + * {@code CuVSMatrix createNativeMatrix(memorySegment, size, dimensions, dataType)}, * where {@code memorySegment} is a {@code java.lang.foreign.MemorySegment} containing {@code int size} vectors of * {@code int dimensions} length of type {@link CuVSMatrix.DataType}. *

* In order to expose this factory in a way that is compatible with Java 21, the factory method is returned as a * {@link MethodHandle} with {@link MethodType} equal to - * {@code (Dataset.class, MemorySegment.class, int.class, int.class, Dataset.DataType.class)}. + * {@code (CuVSMatrix.class, MemorySegment.class, int.class, int.class, CuVSMatrix.DataType.class)}. * The caller will need to invoke the factory via the {@link MethodHandle#invokeExact} method: - * {@code Dataset dataset = (Dataset)newNativeDatasetBuilder().invokeExact(memorySegment, size, dimensions, dataType)} + * {@code var matrix = (CuVSMatrix)newNativeMatrixBuilder().invokeExact(memorySegment, size, dimensions, dataType)} *

- * @return a MethodHandle which can be invoked to build a Dataset from an external {@code MemorySegment} + * @return a MethodHandle which can be invoked to build a CuVSMatrix from an external {@code MemorySegment} */ MethodHandle newNativeMatrixBuilder(); diff --git a/java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/UnsupportedProvider.java b/java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/UnsupportedProvider.java index 8f65bf7068..499ea22fc5 100644 --- a/java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/UnsupportedProvider.java +++ b/java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/UnsupportedProvider.java @@ -15,12 +15,7 @@ */ package com.nvidia.cuvs.spi; -import com.nvidia.cuvs.BruteForceIndex; -import com.nvidia.cuvs.CagraIndex; -import com.nvidia.cuvs.CuVSMatrix; -import com.nvidia.cuvs.CuVSResources; -import com.nvidia.cuvs.HnswIndex; -import com.nvidia.cuvs.TieredIndex; +import com.nvidia.cuvs.*; import java.lang.invoke.MethodHandle; import java.nio.file.Path; @@ -55,13 +50,30 @@ public TieredIndex.Builder newTieredIndexBuilder(CuVSResources cuVSResources) { } @Override - public CagraIndex mergeCagraIndexes(CagraIndex[] indexes) throws Throwable { + public CagraIndex mergeCagraIndexes(CagraIndex[] indexes) { throw new UnsupportedOperationException(); } @Override - public CuVSMatrix.Builder newMatrixBuilder( - int size, int dimensions, CuVSMatrix.DataType dataType) { + public CuVSMatrix.Builder newHostMatrixBuilder( + long size, long dimensions, CuVSMatrix.DataType dataType) { + throw new UnsupportedOperationException(); + } + + @Override + public CuVSMatrix.Builder newDeviceMatrixBuilder( + CuVSResources cuVSResources, long size, long dimensions, CuVSMatrix.DataType dataType) { + throw new UnsupportedOperationException(); + } + + @Override + public CuVSMatrix.Builder newDeviceMatrixBuilder( + CuVSResources cuVSResources, + long size, + long dimensions, + int rowStride, + int columnStride, + CuVSMatrix.DataType dataType) { throw new UnsupportedOperationException(); } diff --git a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/BruteForceIndexImpl.java b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/BruteForceIndexImpl.java index e5f98ad2dd..bc7d6b0a8b 100644 --- a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/BruteForceIndexImpl.java +++ b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/BruteForceIndexImpl.java @@ -28,15 +28,7 @@ import static com.nvidia.cuvs.internal.common.Util.concatenate; import static com.nvidia.cuvs.internal.common.Util.cudaMemcpy; import static com.nvidia.cuvs.internal.common.Util.prepareTensor; -import static com.nvidia.cuvs.internal.panama.headers_h.cuvsBruteForceBuild; -import static com.nvidia.cuvs.internal.panama.headers_h.cuvsBruteForceDeserialize; -import static com.nvidia.cuvs.internal.panama.headers_h.cuvsBruteForceIndexCreate; -import static com.nvidia.cuvs.internal.panama.headers_h.cuvsBruteForceIndexDestroy; -import static com.nvidia.cuvs.internal.panama.headers_h.cuvsBruteForceIndex_t; -import static com.nvidia.cuvs.internal.panama.headers_h.cuvsBruteForceSearch; -import static com.nvidia.cuvs.internal.panama.headers_h.cuvsBruteForceSerialize; -import static com.nvidia.cuvs.internal.panama.headers_h.cuvsStreamSync; -import static com.nvidia.cuvs.internal.panama.headers_h.omp_set_num_threads; +import static com.nvidia.cuvs.internal.panama.headers_h.*; import com.nvidia.cuvs.BruteForceIndex; import com.nvidia.cuvs.BruteForceIndexParams; @@ -153,7 +145,8 @@ private IndexReference build( long[] datasetShape = {rows, cols}; var tensorDataArena = Arena.ofShared(); MemorySegment datasetTensor = - prepareTensor(tensorDataArena, datasetMemorySegmentP, datasetShape, 2, 32, 2, 1); + prepareTensor( + tensorDataArena, datasetMemorySegmentP, datasetShape, kDLFloat(), 32, kDLCUDA()); var returnValue = cuvsStreamSync(cuvsResources); checkCuVSError(returnValue, "cuvsStreamSync"); @@ -236,13 +229,16 @@ public SearchResults search(BruteForceQuery cuvsQuery) throws Throwable { long[] queriesShape = {numQueries, vectorDimension}; MemorySegment queriesTensor = - prepareTensor(localArena, queriesDP.handle(), queriesShape, 2, 32, 2, 1); + prepareTensor( + localArena, queriesDP.handle(), queriesShape, kDLFloat(), 32, kDLCUDA()); long[] neighborsShape = {numQueries, topk}; MemorySegment neighborsTensor = - prepareTensor(localArena, neighborsDP.handle(), neighborsShape, 0, 64, 2, 1); + prepareTensor( + localArena, neighborsDP.handle(), neighborsShape, kDLInt(), 64, kDLCUDA()); long[] distancesShape = {numQueries, topk}; MemorySegment distancesTensor = - prepareTensor(localArena, distancesDP.handle(), distancesShape, 2, 32, 2, 1); + prepareTensor( + localArena, distancesDP.handle(), distancesShape, kDLFloat(), 32, kDLCUDA()); MemorySegment prefilter = cuvsFilter.allocate(localArena); MemorySegment prefilterTensor; @@ -255,8 +251,10 @@ public SearchResults search(BruteForceQuery cuvsQuery) throws Throwable { cudaMemcpy( prefilterDP.handle(), prefilterDataMemorySegment, prefilterBytes, HOST_TO_DEVICE); + // TODO: check if this is correct: filters is a long[] (64 bits) prefilterTensor = - prepareTensor(localArena, prefilterDP.handle(), prefilterShape, 1, 32, 2, 1); + prepareTensor( + localArena, prefilterDP.handle(), prefilterShape, kDLUInt(), 32, kDLCUDA()); cuvsFilter.type(prefilter, 2); cuvsFilter.addr(prefilter, prefilterTensor.address()); diff --git a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CagraIndexImpl.java b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CagraIndexImpl.java index 0d5fc8d90f..aa082482dc 100644 --- a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CagraIndexImpl.java +++ b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CagraIndexImpl.java @@ -30,30 +30,12 @@ import static com.nvidia.cuvs.internal.common.Util.prepareTensor; import static com.nvidia.cuvs.internal.panama.headers_h.*; -import com.nvidia.cuvs.CagraCompressionParams; -import com.nvidia.cuvs.CagraIndex; -import com.nvidia.cuvs.CagraIndexParams; +import com.nvidia.cuvs.*; import com.nvidia.cuvs.CagraIndexParams.CagraGraphBuildAlgo; -import com.nvidia.cuvs.CagraMergeParams; -import com.nvidia.cuvs.CagraQuery; -import com.nvidia.cuvs.CagraSearchParams; -import com.nvidia.cuvs.CuVSIvfPqIndexParams; -import com.nvidia.cuvs.CuVSIvfPqSearchParams; -import com.nvidia.cuvs.CuVSMatrix; -import com.nvidia.cuvs.CuVSResources; -import com.nvidia.cuvs.SearchResults; import com.nvidia.cuvs.internal.common.CloseableHandle; import com.nvidia.cuvs.internal.common.CloseableRMMAllocation; import com.nvidia.cuvs.internal.common.CompositeCloseableHandle; -import com.nvidia.cuvs.internal.common.Util; -import com.nvidia.cuvs.internal.panama.cuvsCagraCompressionParams; -import com.nvidia.cuvs.internal.panama.cuvsCagraIndexParams; -import com.nvidia.cuvs.internal.panama.cuvsCagraMergeParams; -import com.nvidia.cuvs.internal.panama.cuvsCagraSearchParams; -import com.nvidia.cuvs.internal.panama.cuvsFilter; -import com.nvidia.cuvs.internal.panama.cuvsIvfPqIndexParams; -import com.nvidia.cuvs.internal.panama.cuvsIvfPqParams; -import com.nvidia.cuvs.internal.panama.cuvsIvfPqSearchParams; +import com.nvidia.cuvs.internal.panama.*; import java.io.FileInputStream; import java.io.InputStream; import java.io.OutputStream; @@ -180,7 +162,6 @@ public void close() { */ private IndexReference build(CagraIndexParams indexParameters, CuVSMatrixBaseImpl dataset) { long rows = dataset.size(); - long cols = dataset.columns(); try (var indexParams = segmentFromIndexParams(indexParameters); var localArena = Arena.ofConfined()) { @@ -189,14 +170,7 @@ private IndexReference build(CagraIndexParams indexParameters, CuVSMatrixBaseImp int numWriterThreads = indexParameters != null ? indexParameters.getNumWriterThreads() : 1; omp_set_num_threads(numWriterThreads); - MemorySegment dataSeg = dataset.memorySegment(); - // TODO: type kDLCPU()/kDLCUDA() should be aligned with the CuVSMatrixBaseImpl type (host or - // device?) - - long[] datasetShape = {rows, cols}; - MemorySegment datasetTensor = - prepareTensor(localArena, dataSeg, datasetShape, kDLFloat(), 32, kDLCPU(), 1); - + var datasetTensor = dataset.toTensor(localArena); var index = createCagraIndex(); if (cuvsCagraIndexParams.build_algo(indexParamsMemorySegment) @@ -212,6 +186,7 @@ private IndexReference build(CagraIndexParams indexParameters, CuVSMatrixBaseImp try (var resourcesAccessor = resources.access()) { var cuvsRes = resourcesAccessor.handle(); + // TODO: do we need a stream sync here? var returnValue = cuvsStreamSync(cuvsRes); checkCuVSError(returnValue, "cuvsStreamSync"); @@ -293,15 +268,15 @@ public SearchResults search(CagraQuery query) throws Throwable { long[] queriesShape = {numQueries, vectorDimension}; MemorySegment queriesTensor = prepareTensor( - localArena, queriesDP.handle(), queriesShape, kDLFloat(), 32, kDLCUDA(), 1); + localArena, queriesDP.handle(), queriesShape, kDLFloat(), 32, kDLCUDA()); long[] neighborsShape = {numQueries, topK}; MemorySegment neighborsTensor = prepareTensor( - localArena, neighborsDP.handle(), neighborsShape, kDLUInt(), 32, kDLCUDA(), 1); + localArena, neighborsDP.handle(), neighborsShape, kDLUInt(), 32, kDLCUDA()); long[] distancesShape = {numQueries, topK}; MemorySegment distancesTensor = prepareTensor( - localArena, distancesDP.handle(), distancesShape, kDLFloat(), 32, kDLCUDA(), 1); + localArena, distancesDP.handle(), distancesShape, kDLFloat(), 32, kDLCUDA()); var returnValue = cuvsStreamSync(cuvsRes); checkCuVSError(returnValue, "cuvsStreamSync"); @@ -328,7 +303,7 @@ public SearchResults search(CagraQuery query) throws Throwable { prefilterTensor = prepareTensor( - localArena, prefilterDP.handle(), prefilterShape, kDLUInt(), 32, kDLCUDA(), 1); + localArena, prefilterDP.handle(), prefilterShape, kDLUInt(), 32, kDLCUDA()); cuvsFilter.type(prefilter, 1); cuvsFilter.addr(prefilter, prefilterTensor.address()); @@ -411,46 +386,24 @@ public void serialize(OutputStream outputStream, Path tempFile, int bufferLength } @Override - public CuVSMatrix getGraph() { + public CuVSDeviceMatrix getGraph() { try (var localArena = Arena.ofConfined()) { - var outPtr = localArena.allocate(uint32_t); - checkCuVSError( - cuvsCagraIndexGetGraphDegree(cagraIndexReference.getMemorySegment(), outPtr), - "cuvsCagraIndexGetGraphDegree"); - long graphDegree = Util.dereferenceUnsignedInt(outPtr); + // Use a "device" graph + tensor, avoid (defer) copy + MemorySegment graphDeviceTensor = DLManagedTensor.allocate(localArena); + DLManagedTensor.dl_tensor(graphDeviceTensor, DLTensor.allocate(localArena)); - checkCuVSError( - cuvsCagraIndexGetSize(cagraIndexReference.getMemorySegment(), outPtr), - "cuvsCagraIndexGetSize"); - long size = Util.dereferenceUnsignedInt(outPtr); - - // TODO: use a "device" graph + tensor, avoid (defer) copy - var graph = new CuVSHostMatrixArenaImpl(size, graphDegree, CuVSMatrix.DataType.UINT); - var graphHostTensor = graph.toTensor(localArena); - var graphDeviceTensor = - prepareTensor( - localArena, - MemorySegment.NULL, - new long[] {size, graphDegree}, - kDLUInt(), - 32, - kDLCUDA(), - 1); checkCuVSError( cuvsCagraIndexGetGraph(cagraIndexReference.getMemorySegment(), graphDeviceTensor), "cuvsCagraIndexGetGraph"); - try (var resourceAccess = resources.access()) { - var cuvsRes = resourceAccess.handle(); - checkCuVSError(cuvsStreamSync(cuvsRes), "cuvsStreamSync"); - - checkCuVSError( - cuvsMatrixCopy(cuvsRes, graphDeviceTensor, graphHostTensor), "cuvsMatrixCopy"); + assert DLTensor.ndim(DLManagedTensor.dl_tensor(graphDeviceTensor)) == 2; + assert DLTensor.shape(DLManagedTensor.dl_tensor(graphDeviceTensor)).get(int64_t, 0) > 0; + assert DLTensor.shape(DLManagedTensor.dl_tensor(graphDeviceTensor)).getAtIndex(int64_t, 1) + > 0; - checkCuVSError(cuvsStreamSync(cuvsRes), "cuvsStreamSync"); - } - - return graph; + var graph = CuVSMatrixBaseImpl.fromTensor(graphDeviceTensor, resources); + assert graph instanceof CuVSDeviceMatrix; + return (CuVSDeviceMatrix) graph; } } @@ -459,28 +412,17 @@ private IndexReference fromGraph( CuVSMatrixBaseImpl graph, CuVSMatrixBaseImpl dataset) { try (var localArena = Arena.ofConfined()) { - long rows = dataset.size(); - long cols = dataset.columns(); - var index = createCagraIndex(); try (var resourcesAccess = resources.access()) { long cuvsRes = resourcesAccess.handle(); - long[] datasetShape = {rows, cols}; - MemorySegment datasetTensor = - prepareTensor( - localArena, dataset.memorySegment(), datasetShape, kDLFloat(), 32, kDLCPU(), 1); - - long[] graphShape = {graph.size(), graph.columns()}; - MemorySegment graphTensor = - prepareTensor( - localArena, graph.memorySegment(), graphShape, kDLUInt(), 32, kDLCPU(), 1); + MemorySegment datasetTensor = dataset.toTensor(localArena); + MemorySegment graphTensor = graph.toTensor(localArena); checkCuVSError( cuvsCagraIndexFromArgs(cuvsRes, metric.value, graphTensor, datasetTensor, index), "cuvsCagraIndexFromArgs"); } - return new IndexReference(index, dataset); } } diff --git a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CuVSDeviceMatrixImpl.java b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CuVSDeviceMatrixImpl.java new file mode 100644 index 0000000000..50b914c76b --- /dev/null +++ b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CuVSDeviceMatrixImpl.java @@ -0,0 +1,246 @@ +/* + * Copyright (c) 2025, NVIDIA CORPORATION. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.nvidia.cuvs.internal; + +import static com.nvidia.cuvs.internal.common.LinkerHelper.C_POINTER; +import static com.nvidia.cuvs.internal.common.Util.*; +import static com.nvidia.cuvs.internal.panama.headers_h.*; + +import com.nvidia.cuvs.*; +import com.nvidia.cuvs.internal.panama.DLManagedTensor; +import com.nvidia.cuvs.internal.panama.DLTensor; +import java.lang.foreign.*; + +public class CuVSDeviceMatrixImpl extends CuVSMatrixBaseImpl implements CuVSDeviceMatrix { + + private static final int CHUNK_BYTES = + 8 * 1024 * 1024; // Based on benchmarks, 8MB seems the minimum size to optimize PCIe bandwidth + private final long hostBufferBytes; + + private long bufferedMatrixRowStart = 0; + private long bufferedMatrixRowEnd = 0; + + private final CuVSResources resources; + + private final long rowStride; + private final long columnStride; + + private MemorySegment hostBuffer = MemorySegment.NULL; + + protected CuVSDeviceMatrixImpl( + CuVSResources resources, + MemorySegment deviceMemorySegment, + long size, + long columns, + DataType dataType, + ValueLayout valueLayout) { + this(resources, deviceMemorySegment, size, columns, -1, -1, dataType, valueLayout); + } + + protected CuVSDeviceMatrixImpl( + CuVSResources resources, + MemorySegment deviceMemorySegment, + long size, + long columns, + long rowStride, + long columnStride, + DataType dataType, + ValueLayout valueLayout) { + super(deviceMemorySegment, dataType, valueLayout, size, columns); + this.resources = resources; + this.rowStride = rowStride; + this.columnStride = columnStride; + + long rowBytes = columns * valueLayout.byteSize(); + long matrixBytes = size * rowBytes; + if (matrixBytes < CHUNK_BYTES) { + this.hostBufferBytes = matrixBytes; + } else if (rowBytes > CHUNK_BYTES) { + // We need to buffer at least one row at time + this.hostBufferBytes = rowBytes; + } else { + var rowCount = (CHUNK_BYTES / rowBytes); + this.hostBufferBytes = rowBytes * rowCount; + } + } + + @Override + public MemorySegment toTensor(Arena arena) { + var strides = rowStride >= 0 ? new long[] {rowStride, columnStride} : null; + return prepareTensor( + arena, memorySegment, new long[] {size, columns}, strides, code(), bits(), kDLCUDA()); + } + + private static MemorySegment createPinnedBuffer(long bufferBytes) { + try (var localArena = Arena.ofConfined()) { + MemorySegment pointer = localArena.allocate(C_POINTER); + checkCudaError(cudaMallocHost(pointer, bufferBytes), "cudaMallocHost"); + return pointer.get(C_POINTER, 0); + } + } + + private static void destroyPinnedBuffer(MemorySegment bufferSegment) { + checkCudaError(cudaFreeHost(bufferSegment), "cudaFreeHost"); + } + + private void populateBuffer(long startRow) { + if (hostBuffer == MemorySegment.NULL) { + // System.out.println("Creating a buffer of size " + hostBufferBytes); + hostBuffer = createPinnedBuffer(hostBufferBytes); + } + + try (var localArena = Arena.ofConfined()) { + long rowBytes = columns * valueLayout.byteSize(); + var endRow = Math.min(startRow + (hostBufferBytes / rowBytes), size); + var rowCount = endRow - startRow; + + // System.out.printf( + // Locale.ROOT, "startRow: %d, endRow %d, count: %d\n", startRow, endRow, rowCount); + + MemorySegment sliceManagedTensor = DLManagedTensor.allocate(localArena); + DLManagedTensor.dl_tensor(sliceManagedTensor, DLTensor.allocate(localArena)); + + checkCuVSError( + cuvsMatrixSliceRows(0, toTensor(localArena), startRow, endRow, sliceManagedTensor), + "cuvsMatrixSliceRows"); + assert DLTensor.shape(DLManagedTensor.dl_tensor(sliceManagedTensor)).get(C_LONG, 0) + == rowCount; + assert DLTensor.shape(DLManagedTensor.dl_tensor(sliceManagedTensor)).getAtIndex(C_LONG, 1) + == columns; + + MemorySegment bufferTensor = + prepareTensor( + localArena, hostBuffer, new long[] {rowCount, columns}, code(), bits(), kDLCPU()); + + try (var resourceAccess = resources.access()) { + checkCuVSError( + cuvsMatrixCopy(resourceAccess.handle(), sliceManagedTensor, bufferTensor), + "cuvsMatrixCopy"); + checkCuVSError(cuvsStreamSync(resourceAccess.handle()), "cuvsStreamSync"); + + bufferedMatrixRowStart = startRow; + bufferedMatrixRowEnd = endRow; + } + } + } + + @Override + public RowView getRow(long row) { + if (row < bufferedMatrixRowStart || row >= bufferedMatrixRowEnd) { + populateBuffer(row); + } + var valueByteSize = valueLayout.byteSize(); + var startRow = row - bufferedMatrixRowStart; + + return new SliceRowView( + hostBuffer.asSlice(startRow * columns * valueByteSize, columns * valueByteSize), + columns, + valueLayout, + dataType, + valueByteSize); + } + + @Override + public void toArray(int[][] array) { + assert dataType == DataType.INT || dataType == DataType.UINT; + assert (array.length >= size) : "Input array is not large enough"; + assert (array.length == 0 || array[0].length >= columns) : "Input array is not large enough"; + try (var localArena = Arena.ofConfined()) { + var rowBytes = columns * valueLayout.byteSize(); + var tmpRowSegment = localArena.allocate(rowBytes); + for (int r = 0; r < size; ++r) { + copyRow(array[r], localArena, r, tmpRowSegment); + } + } + } + + @Override + public void toArray(float[][] array) { + assert dataType == DataType.FLOAT; + assert (array.length >= size) : "Input array is not large enough"; + assert (array.length == 0 || array[0].length >= columns) : "Input array is not large enough"; + try (var localArena = Arena.ofConfined()) { + var rowBytes = columns * valueLayout.byteSize(); + var tmpRowSegment = localArena.allocate(rowBytes); + for (int r = 0; r < size; ++r) { + copyRow(array[r], localArena, r, tmpRowSegment); + } + } + } + + @Override + public void toArray(byte[][] array) { + assert dataType == DataType.BYTE; + assert (array.length >= size) : "Input array is not large enough"; + assert (array.length == 0 || array[0].length >= columns) : "Input array is not large enough"; + try (var localArena = Arena.ofConfined()) { + var rowSegmentLayout = MemoryLayout.sequenceLayout(columns, valueLayout); + var tmpRowSegment = localArena.allocate(rowSegmentLayout); + for (int r = 0; r < size; ++r) { + copyRow(array[r], localArena, r, tmpRowSegment); + } + } + } + + private void copyRow(Object array, Arena localArena, int r, MemorySegment tmpRowSegment) { + MemorySegment sliceManagedTensor = DLManagedTensor.allocate(localArena); + DLManagedTensor.dl_tensor(sliceManagedTensor, DLTensor.allocate(localArena)); + + checkCuVSError( + cuvsMatrixSliceRows(0, toTensor(localArena), r, r + 1, sliceManagedTensor), + "cuvsMatrixSliceRows"); + + MemorySegment bufferTensor = + prepareTensor( + localArena, tmpRowSegment, new long[] {1, columns}, code(), bits(), kDLCUDA()); + try (var resourceAccess = resources.access()) { + checkCuVSError( + cuvsMatrixCopy(resourceAccess.handle(), sliceManagedTensor, bufferTensor), + "cuvsMatrixCopy"); + checkCuVSError(cuvsStreamSync(resourceAccess.handle()), "cuvsStreamSync"); + } + MemorySegment.copy(tmpRowSegment, valueLayout, 0L, array, 0, (int) columns); + } + + @Override + public void toHost(CuVSHostMatrix hostMatrix) { + if (hostMatrix.columns() != columns || hostMatrix.size() != size) { + throw new IllegalArgumentException("[hostMatrix] must have the same dimensions"); + } + if (hostMatrix.dataType() != dataType) { + throw new IllegalArgumentException("[hostMatrix] must have the same dataType"); + } + try (var localArena = Arena.ofConfined()) { + var hostMatrixTensor = ((CuVSHostMatrixImpl) hostMatrix).toTensor(localArena); + + try (var resourceAccess = resources.access()) { + var cuvsRes = resourceAccess.handle(); + var deviceMatrixTensor = toTensor(localArena); + checkCuVSError( + cuvsMatrixCopy(cuvsRes, deviceMatrixTensor, hostMatrixTensor), "cuvsMatrixCopy"); + checkCuVSError(cuvsStreamSync(cuvsRes), "cuvsStreamSync"); + } + } + } + + @Override + public void close() { + if (hostBuffer != MemorySegment.NULL) { + destroyPinnedBuffer(hostBuffer); + hostBuffer = MemorySegment.NULL; + } + } +} diff --git a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CuVSDeviceMatrixRMMImpl.java b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CuVSDeviceMatrixRMMImpl.java new file mode 100644 index 0000000000..7ad84a0e1f --- /dev/null +++ b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CuVSDeviceMatrixRMMImpl.java @@ -0,0 +1,97 @@ +/* + * Copyright (c) 2025, NVIDIA CORPORATION. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.nvidia.cuvs.internal; + +import com.nvidia.cuvs.CuVSDeviceMatrix; +import com.nvidia.cuvs.CuVSResources; +import com.nvidia.cuvs.internal.common.CloseableRMMAllocation; +import java.lang.foreign.ValueLayout; + +/** + * A Dataset implementation backed by device (GPU) memory, allocated + * using RMM functions from {@link CuVSResources} pools (via {@link CloseableRMMAllocation}). + */ +public class CuVSDeviceMatrixRMMImpl extends CuVSDeviceMatrixImpl implements CuVSDeviceMatrix { + + private final CloseableRMMAllocation rmmAllocation; + + private CuVSDeviceMatrixRMMImpl( + CuVSResources resources, + CloseableRMMAllocation rmmAllocation, + long size, + long columns, + DataType dataType, + ValueLayout valueLayout) { + super(resources, rmmAllocation.handle(), size, columns, dataType, valueLayout); + this.rmmAllocation = rmmAllocation; + } + + private CuVSDeviceMatrixRMMImpl( + CuVSResources resources, + CloseableRMMAllocation rmmAllocation, + long size, + long columns, + long rowStride, + long columnStride, + DataType dataType, + ValueLayout valueLayout) { + super( + resources, + rmmAllocation.handle(), + size, + columns, + rowStride, + columnStride, + dataType, + valueLayout); + this.rmmAllocation = rmmAllocation; + } + + public static CuVSDeviceMatrixImpl create( + CuVSResources resources, long size, long columns, DataType dataType) { + try (var resourcesAccess = resources.access()) { + var valueLayout = valueLayoutFromType(dataType); + var rmmAllocation = + CloseableRMMAllocation.allocateRMMSegment( + resourcesAccess.handle(), size * columns * valueLayout.byteSize()); + return new CuVSDeviceMatrixRMMImpl( + resources, rmmAllocation, size, columns, dataType, valueLayout); + } + } + + public static CuVSDeviceMatrixImpl create( + CuVSResources resources, + long size, + long columns, + long rowStride, + long columnStride, + DataType dataType) { + try (var resourcesAccess = resources.access()) { + var valueLayout = valueLayoutFromType(dataType); + var rmmAllocation = + CloseableRMMAllocation.allocateRMMSegment( + resourcesAccess.handle(), size * columns * valueLayout.byteSize()); + return new CuVSDeviceMatrixRMMImpl( + resources, rmmAllocation, size, columns, rowStride, columnStride, dataType, valueLayout); + } + } + + @Override + public void close() { + super.close(); + rmmAllocation.close(); + } +} diff --git a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CuVSHostMatrixImpl.java b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CuVSHostMatrixImpl.java index c2895282ca..317b315c77 100644 --- a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CuVSHostMatrixImpl.java +++ b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CuVSHostMatrixImpl.java @@ -15,9 +15,6 @@ */ package com.nvidia.cuvs.internal; -import static com.nvidia.cuvs.internal.common.LinkerHelper.C_CHAR; -import static com.nvidia.cuvs.internal.common.LinkerHelper.C_FLOAT; -import static com.nvidia.cuvs.internal.common.LinkerHelper.C_INT; import static com.nvidia.cuvs.internal.common.Util.prepareTensor; import static com.nvidia.cuvs.internal.panama.headers_h.*; @@ -25,12 +22,12 @@ import com.nvidia.cuvs.RowView; import java.lang.foreign.*; import java.lang.invoke.VarHandle; +import java.util.Locale; /** * A Dataset implementation backed by host (CPU) memory. */ public class CuVSHostMatrixImpl extends CuVSMatrixBaseImpl implements CuVSHostMatrix { - private final ValueLayout valueLayout; protected final VarHandle accessor$vh; public CuVSHostMatrixImpl( @@ -52,30 +49,18 @@ protected CuVSHostMatrixImpl( DataType dataType, ValueLayout valueLayout, MemoryLayout sequenceLayout) { - super(memorySegment, dataType, size, columns); + super(memorySegment, dataType, valueLayout, size, columns); this.accessor$vh = sequenceLayout.varHandle(MemoryLayout.PathElement.sequenceElement()); - this.valueLayout = valueLayout; - } - - protected static ValueLayout valueLayoutFromType(DataType dataType) { - return switch (dataType) { - case FLOAT -> C_FLOAT; - case INT, UINT -> C_INT; - case BYTE -> C_CHAR; - }; - } - - protected static SequenceLayout sequenceLayoutFromType( - long size, long columns, DataType dataType) { - return MemoryLayout.sequenceLayout(size * columns, valueLayoutFromType(dataType)) - .withByteAlignment(32); } @Override - public RowView getRow(long nodeIndex) { + public RowView getRow(long index) { + assert (index < size) + : String.format(Locale.ROOT, "Index out of bound ([%d], size [%d])", index, size); + var valueByteSize = valueLayout.byteSize(); return new SliceRowView( - memorySegment.asSlice(nodeIndex * columns * valueByteSize, columns * valueByteSize), + memorySegment.asSlice(index * columns * valueByteSize, columns * valueByteSize), columns, valueLayout, dataType, @@ -84,9 +69,22 @@ public RowView getRow(long nodeIndex) { @Override public void toArray(int[][] array) { - assert dataType == DataType.INT; - assert (array.length >= size) : "Input array is not large enough"; - assert (array.length == 0 || array[0].length >= columns) : "Input array is not large enough"; + assert (array.length >= size) + : String.format( + Locale.ROOT, + "Input array is not large enough (required: [%d], actual [%d])", + size, + array.length); + assert (array.length == 0 || array[0].length >= columns) + : String.format( + Locale.ROOT, + "Input array is not wide enough (required: [%d], actual [%d])", + columns, + array[0].length); + assert dataType == DataType.INT || dataType == DataType.UINT + : String.format( + Locale.ROOT, "Input array is of the wrong type for dataType [%s]", dataType.toString()); + var valueByteSize = valueLayout.byteSize(); for (int r = 0; r < size; ++r) { MemorySegment.copy( @@ -96,9 +94,22 @@ public void toArray(int[][] array) { @Override public void toArray(float[][] array) { - assert dataType == DataType.FLOAT; - assert (array.length >= size) : "Input array is not large enough"; - assert (array.length == 0 || array[0].length >= columns) : "Input array is not large enough"; + assert (array.length >= size) + : String.format( + Locale.ROOT, + "Input array is not large enough (required: [%d], actual [%d])", + size, + array.length); + assert (array.length == 0 || array[0].length >= columns) + : String.format( + Locale.ROOT, + "Input array is not wide enough (required: [%d], actual [%d])", + columns, + array[0].length); + assert dataType == DataType.FLOAT + : String.format( + Locale.ROOT, "Input array is of the wrong type for dataType [%s]", dataType.toString()); + var valueByteSize = valueLayout.byteSize(); for (int r = 0; r < size; ++r) { MemorySegment.copy( @@ -108,9 +119,22 @@ public void toArray(float[][] array) { @Override public void toArray(byte[][] array) { - assert dataType == DataType.BYTE; - assert (array.length >= size) : "Input array is not large enough"; - assert (array.length == 0 || array[0].length >= columns) : "Input array is not large enough"; + assert (array.length >= size) + : String.format( + Locale.ROOT, + "Input array is not large enough (required: [%d], actual [%d])", + size, + array.length); + assert (array.length == 0 || array[0].length >= columns) + : String.format( + Locale.ROOT, + "Input array is not wide enough (required: [%d], actual [%d])", + columns, + array[0].length); + assert dataType == DataType.BYTE + : String.format( + Locale.ROOT, "Input array is of the wrong type for dataType [%s]", dataType.toString()); + var valueByteSize = valueLayout.byteSize(); for (int r = 0; r < size; ++r) { MemorySegment.copy( @@ -126,78 +150,9 @@ public int get(int row, int col) { return (int) accessor$vh.get(memorySegment, 0L, (long) row * columns + col); } - public ValueLayout valueLayout() { - return valueLayout; - } - @Override public MemorySegment toTensor(Arena arena) { return prepareTensor( - arena, memorySegment, new long[] {size, columns}, code(), bits(), kDLCPU(), 1); - } - - private static class SliceRowView implements RowView { - private final MemorySegment memorySegment; - private final long size; - private final ValueLayout valueLayout; - private final DataType dataType; - private final long valueByteSize; - - SliceRowView( - MemorySegment slice, - long size, - ValueLayout valueLayout, - DataType dataType, - long valueByteSize) { - this.memorySegment = slice; - this.size = size; - this.valueLayout = valueLayout; - this.dataType = dataType; - this.valueByteSize = valueByteSize; - } - - @Override - public long size() { - return size; - } - - @Override - public float getAsFloat(long index) { - assert dataType == DataType.FLOAT; - return memorySegment.get((ValueLayout.OfFloat) valueLayout, index * valueByteSize); - } - - @Override - public byte getAsByte(long index) { - assert dataType == DataType.BYTE; - return memorySegment.get((ValueLayout.OfByte) valueLayout, index * valueByteSize); - } - - @Override - public int getAsInt(long index) { - assert dataType == DataType.INT; - return memorySegment.get((ValueLayout.OfInt) valueLayout, index * valueByteSize); - } - - @Override - public void toArray(int[] array) { - assert (array.length >= size) : "Input array is not large enough"; - assert dataType == DataType.INT; - MemorySegment.copy(memorySegment, valueLayout, 0, array, 0, (int) size); - } - - @Override - public void toArray(float[] array) { - assert (array.length >= size) : "Input array is not large enough"; - assert dataType == DataType.FLOAT; - MemorySegment.copy(memorySegment, valueLayout, 0, array, 0, (int) size); - } - - @Override - public void toArray(byte[] array) { - assert (array.length >= size) : "Input array is not large enough"; - assert dataType == DataType.BYTE; - MemorySegment.copy(memorySegment, valueLayout, 0, array, 0, (int) size); - } + arena, memorySegment, new long[] {size, columns}, code(), bits(), kDLCPU()); } } diff --git a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CuVSMatrixBaseImpl.java b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CuVSMatrixBaseImpl.java index 120960cfc4..d629645e61 100644 --- a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CuVSMatrixBaseImpl.java +++ b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CuVSMatrixBaseImpl.java @@ -15,22 +15,36 @@ */ package com.nvidia.cuvs.internal; +import static com.nvidia.cuvs.internal.common.LinkerHelper.C_CHAR; +import static com.nvidia.cuvs.internal.common.LinkerHelper.C_FLOAT; +import static com.nvidia.cuvs.internal.common.LinkerHelper.C_INT; import static com.nvidia.cuvs.internal.panama.headers_h.*; import com.nvidia.cuvs.CuVSMatrix; -import java.lang.foreign.Arena; -import java.lang.foreign.MemorySegment; +import com.nvidia.cuvs.CuVSResources; +import com.nvidia.cuvs.internal.panama.DLDataType; +import com.nvidia.cuvs.internal.panama.DLDevice; +import com.nvidia.cuvs.internal.panama.DLManagedTensor; +import com.nvidia.cuvs.internal.panama.DLTensor; +import java.lang.foreign.*; +import java.util.Locale; public abstract class CuVSMatrixBaseImpl implements CuVSMatrix { protected final MemorySegment memorySegment; protected final DataType dataType; + protected final ValueLayout valueLayout; protected final long size; protected final long columns; protected CuVSMatrixBaseImpl( - MemorySegment memorySegment, DataType dataType, long size, long columns) { + MemorySegment memorySegment, + DataType dataType, + ValueLayout valueLayout, + long size, + long columns) { this.memorySegment = memorySegment; this.dataType = dataType; + this.valueLayout = valueLayout; this.size = size; this.columns = columns; } @@ -45,17 +59,29 @@ public long columns() { return columns; } + @Override + public DataType dataType() { + return dataType; + } + public MemorySegment memorySegment() { return memorySegment; } + public ValueLayout valueLayout() { + return valueLayout; + } + + /** + * Size (in bits) for the element type of this matrix + */ protected int bits() { - return switch (dataType) { - case FLOAT, INT, UINT -> 32; - case BYTE -> 8; - }; + return (int) (valueLayout.byteSize() * 8); } + /** + * DLTensor data type {@code code} for the element type of this matrix + */ protected int code() { return switch (dataType) { case FLOAT -> kDLFloat(); @@ -64,5 +90,104 @@ protected int code() { }; } + protected static ValueLayout valueLayoutFromType(DataType dataType) { + return switch (dataType) { + case FLOAT -> C_FLOAT; + case INT, UINT -> C_INT; + case BYTE -> C_CHAR; + }; + } + + protected static SequenceLayout sequenceLayoutFromType( + long size, long columns, DataType dataType) { + return MemoryLayout.sequenceLayout(size * columns, valueLayoutFromType(dataType)) + .withByteAlignment(32); + } + + /** + * Creates a {@link DLManagedTensor} representing the matrix data and shape, to be + * passed to the CuVS C API. + * @param arena The Arena to use to allocate DL data structures + * @return a {@link MemorySegment} for the newly allocated DLManagedTensor + */ public abstract MemorySegment toTensor(Arena arena); + + /** + * Creates a {@link CuVSMatrix} from data and infos from a {@link DLManagedTensor} + * + * @param dlManagedTensor a {@link MemorySegment} representing the source DLManagedTensor + * @param resources {@link CuVSResources} to allocate the resulting matrix + * @return a {@link CuVSMatrix} encapsulating the same data as the input {@link DLManagedTensor} + */ + public static CuVSMatrix fromTensor(MemorySegment dlManagedTensor, CuVSResources resources) { + var dlTensor = DLManagedTensor.dl_tensor(dlManagedTensor); + var dlDevice = DLTensor.device(dlTensor); + + var deviceType = DLDevice.device_type(dlDevice); + + var data = DLTensor.data(dlTensor); + if (data.equals(MemorySegment.NULL)) { + throw new IllegalArgumentException("[data] must not be NULL"); + } + + var ndim = DLTensor.ndim(dlTensor); + if (ndim != 2) { + throw new IllegalArgumentException("CuVSMatrix only supports 2D data"); + } + + var dtype = DLTensor.dtype(dlTensor); + var code = DLDataType.code(dtype); + var bits = DLDataType.bits(dtype); + + final DataType dataType = dataTypeFromTensor(code, bits); + + var shape = DLTensor.shape(dlTensor); + if (shape.equals(MemorySegment.NULL)) { + throw new IllegalArgumentException("[shape] must not be NULL"); + } + + var rows = shape.get(int64_t, 0); + var cols = shape.getAtIndex(int64_t, 1); + + if (deviceType == kDLCUDA()) { + var strides = DLTensor.strides(dlTensor); + if (strides.equals(MemorySegment.NULL)) { + return new CuVSDeviceMatrixImpl( + resources, data, rows, cols, dataType, valueLayoutFromType(dataType)); + } else { + var rowStride = strides.get(int64_t, 0); + var colStride = strides.getAtIndex(int64_t, 1); + return new CuVSDeviceMatrixImpl( + resources, + data, + rows, + cols, + rowStride, + colStride, + dataType, + valueLayoutFromType(dataType)); + } + } else if (deviceType == kDLCPU()) { + return new CuVSHostMatrixImpl(data, rows, cols, dataType); + } else { + throw new IllegalArgumentException("Unsupported device type: " + deviceType); + } + } + + private static DataType dataTypeFromTensor(byte code, byte bits) { + final DataType dataType; + if (code == kDLUInt() && bits == 32) { + dataType = DataType.UINT; + } else if (code == kDLInt() && bits == 32) { + dataType = DataType.INT; + } else if (code == kDLFloat() && bits == 32) { + dataType = DataType.FLOAT; + } else if ((code == kDLInt() || code == kDLUInt()) && bits == 8) { + dataType = DataType.BYTE; + } else { + throw new IllegalArgumentException( + String.format(Locale.ROOT, "Unsupported data type (code=%d, bits=%d)", code, bits)); + } + return dataType; + } } diff --git a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/HnswIndexImpl.java b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/HnswIndexImpl.java index 736b22f567..9428d2eb6f 100644 --- a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/HnswIndexImpl.java +++ b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/HnswIndexImpl.java @@ -98,18 +98,23 @@ public SearchResults search(HnswQuery query) throws Throwable { SequenceLayout neighborsSequenceLayout = MemoryLayout.sequenceLayout(numBlocks, C_LONG); SequenceLayout distancesSequenceLayout = MemoryLayout.sequenceLayout(numBlocks, C_FLOAT); + // TODO: these could be CuVSHostMatrix MemorySegment neighborsMemorySegment = localArena.allocate(neighborsSequenceLayout); MemorySegment distancesMemorySegment = localArena.allocate(distancesSequenceLayout); MemorySegment querySeg = buildMemorySegment(localArena, queryVectors); long[] queriesShape = {numQueries, vectorDimension}; - MemorySegment queriesTensor = prepareTensor(localArena, querySeg, queriesShape, 2, 32, 1, 1); + MemorySegment queriesTensor = + prepareTensor(localArena, querySeg, queriesShape, kDLFloat(), 32, kDLCPU()); long[] neighborsShape = {numQueries, topK}; + // TODO: check type code and bits across all implementations -- they are inconsistent MemorySegment neighborsTensor = - prepareTensor(localArena, neighborsMemorySegment, neighborsShape, 1, 64, 1, 1); + prepareTensor( + localArena, neighborsMemorySegment, neighborsShape, kDLUInt(), 64, kDLCPU()); long[] distancesShape = {numQueries, topK}; MemorySegment distancesTensor = - prepareTensor(localArena, distancesMemorySegment, distancesShape, 2, 32, 1, 1); + prepareTensor( + localArena, distancesMemorySegment, distancesShape, kDLFloat(), 32, kDLCPU()); try (var resourcesAccessor = query.getResources().access()) { var cuvsRes = resourcesAccessor.handle(); diff --git a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/SliceRowView.java b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/SliceRowView.java new file mode 100644 index 0000000000..3620527449 --- /dev/null +++ b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/SliceRowView.java @@ -0,0 +1,126 @@ +/* + * Copyright (c) 2025, NVIDIA CORPORATION. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.nvidia.cuvs.internal; + +import com.nvidia.cuvs.CuVSMatrix; +import com.nvidia.cuvs.RowView; +import java.lang.foreign.MemorySegment; +import java.lang.foreign.ValueLayout; +import java.util.Locale; + +/** + * A {@link RowView} over a {@link MemorySegment} slice. + */ +class SliceRowView implements RowView { + private final MemorySegment memorySegment; + private final long size; + private final ValueLayout valueLayout; + private final CuVSMatrix.DataType dataType; + private final long valueByteSize; + + SliceRowView( + MemorySegment slice, + long size, + ValueLayout valueLayout, + CuVSMatrix.DataType dataType, + long valueByteSize) { + this.memorySegment = slice; + this.size = size; + this.valueLayout = valueLayout; + this.dataType = dataType; + this.valueByteSize = valueByteSize; + } + + @Override + public long size() { + return size; + } + + @Override + public float getAsFloat(long index) { + assert (index < size) + : String.format(Locale.ROOT, "Index out of bound ([%d], size [%d])", index, size); + assert dataType == CuVSMatrix.DataType.FLOAT + : String.format( + Locale.ROOT, "Input array is of the wrong type for dataType [%s]", dataType.toString()); + + return memorySegment.get((ValueLayout.OfFloat) valueLayout, index * valueByteSize); + } + + @Override + public byte getAsByte(long index) { + assert (index < size) + : String.format(Locale.ROOT, "Index out of bound ([%d], size [%d])", index, size); + assert dataType == CuVSMatrix.DataType.BYTE + : String.format( + Locale.ROOT, "Input array is of the wrong type for dataType [%s]", dataType.toString()); + + return memorySegment.get((ValueLayout.OfByte) valueLayout, index * valueByteSize); + } + + @Override + public int getAsInt(long index) { + assert (index < size) + : String.format(Locale.ROOT, "Index out of bound ([%d], size [%d])", index, size); + assert dataType == CuVSMatrix.DataType.INT || dataType == CuVSMatrix.DataType.UINT + : String.format( + Locale.ROOT, "Input array is of the wrong type for dataType [%s]", dataType.toString()); + + return memorySegment.get((ValueLayout.OfInt) valueLayout, index * valueByteSize); + } + + @Override + public void toArray(int[] array) { + assert (array.length >= size) + : String.format( + Locale.ROOT, + "Input array is not large enough (required: [%d], actual [%d])", + size, + array.length); + assert dataType == CuVSMatrix.DataType.INT || dataType == CuVSMatrix.DataType.UINT + : String.format( + Locale.ROOT, "Input array is of the wrong type for dataType [%s]", dataType.toString()); + MemorySegment.copy(memorySegment, valueLayout, 0, array, 0, (int) size); + } + + @Override + public void toArray(float[] array) { + assert (array.length >= size) + : String.format( + Locale.ROOT, + "Input array is not large enough (required: [%d], actual [%d])", + size, + array.length); + assert dataType == CuVSMatrix.DataType.FLOAT + : String.format( + Locale.ROOT, "Input array is of the wrong type for dataType [%s]", dataType.toString()); + MemorySegment.copy(memorySegment, valueLayout, 0, array, 0, (int) size); + } + + @Override + public void toArray(byte[] array) { + assert (array.length >= size) + : String.format( + Locale.ROOT, + "Input array is not large enough (required: [%d], actual [%d])", + size, + array.length); + assert dataType == CuVSMatrix.DataType.BYTE + : String.format( + Locale.ROOT, "Input array is of the wrong type for dataType [%s]", dataType.toString()); + MemorySegment.copy(memorySegment, valueLayout, 0, array, 0, (int) size); + } +} diff --git a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/TieredIndexImpl.java b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/TieredIndexImpl.java index c507045326..5edbcabd71 100644 --- a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/TieredIndexImpl.java +++ b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/TieredIndexImpl.java @@ -143,7 +143,7 @@ private IndexReference build() { long[] datasetShape = {rows, cols}; MemorySegment datasetTensor = prepareTensor( - localArena, datasetDP.handle(), datasetShape, kDLFloat(), 32, kDLCUDA(), 1); + localArena, datasetDP.handle(), datasetShape, kDLFloat(), 32, kDLCUDA()); MemorySegment index = localArena.allocate(cuvsTieredIndex_t); var returnValue = cuvsTieredIndexCreate(index); @@ -205,6 +205,7 @@ public SearchResults search(TieredIndexQuery query) throws Throwable { hasPreFilter ? new BitSet[] {query.getPrefilter()} : EMPTY_PREFILTER_BITSET; final long prefilterDataLength = hasPreFilter ? query.getNumDocs() * prefilters.length : 0; final long prefilterLen = hasPreFilter ? (prefilterDataLength + 31) / 32 : 0; + // TODO: is this correct? prefilters is a LONG array final long prefilterBytes = C_INT_BYTE_SIZE * prefilterLen; try (var resourceAccess = query.getResources().access()) { @@ -229,7 +230,7 @@ public SearchResults search(TieredIndexQuery query) throws Throwable { long[] queriesShape = {numQueries, vectorDimension}; MemorySegment queriesTensor = prepareTensor( - localArena, queriesDP.handle(), queriesShape, kDLFloat(), 32, kDLCUDA(), 1); + localArena, queriesDP.handle(), queriesShape, kDLFloat(), 32, kDLCUDA()); long[] neighborsShape = {numQueries, topK}; MemorySegment neighborsTensor = prepareTensor( @@ -238,12 +239,11 @@ public SearchResults search(TieredIndexQuery query) throws Throwable { neighborsShape, kDLInt(), 64, - kDLCUDA(), - 1); // 64-bit int + kDLCUDA()); // 64-bit int long[] distancesShape = {numQueries, topK}; MemorySegment distancesTensor = prepareTensor( - localArena, distancesDP.handle(), distancesShape, kDLFloat(), 32, kDLCUDA(), 1); + localArena, distancesDP.handle(), distancesShape, kDLFloat(), 32, kDLCUDA()); // Sync before prefilter setup returnValue = cuvsStreamSync(cuvsRes); @@ -270,7 +270,7 @@ public SearchResults search(TieredIndexQuery query) throws Throwable { MemorySegment prefilterTensor = prepareTensor( - localArena, prefilterDP.handle(), prefilterShape, kDLUInt(), 32, kDLCUDA(), 1); + localArena, prefilterDP.handle(), prefilterShape, kDLUInt(), 32, kDLCUDA()); cuvsFilter.type(prefilter, 1); // BITSET cuvsFilter.addr(prefilter, prefilterTensor.address()); @@ -347,7 +347,7 @@ private void performExtend(CuVSMatrix extendDataset) { long[] datasetShape = {rows, cols}; MemorySegment datasetTensor = prepareTensor( - localArena, datasetDP.handle(), datasetShape, kDLFloat(), 32, kDLCUDA(), 1); + localArena, datasetDP.handle(), datasetShape, kDLFloat(), 32, kDLCUDA()); checkCuVSError(cuvsStreamSync(cuvsRes), "cuvsStreamSync"); diff --git a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/common/Util.java b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/common/Util.java index 8125497078..d62450ece0 100644 --- a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/common/Util.java +++ b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/common/Util.java @@ -20,7 +20,9 @@ import static com.nvidia.cuvs.internal.common.LinkerHelper.C_INT; import static com.nvidia.cuvs.internal.common.LinkerHelper.C_LONG; import static com.nvidia.cuvs.internal.panama.headers_h.*; +import static com.nvidia.cuvs.internal.panama.headers_h_1.cudaStream_t; +import com.nvidia.cuvs.CuVSResources; import com.nvidia.cuvs.GPUInfo; import com.nvidia.cuvs.internal.panama.DLDataType; import com.nvidia.cuvs.internal.panama.DLDevice; @@ -29,9 +31,11 @@ import com.nvidia.cuvs.internal.panama.cudaDeviceProp; import com.nvidia.cuvs.internal.panama.headers_h; import java.lang.foreign.Arena; +import java.lang.foreign.Linker; import java.lang.foreign.MemoryLayout; import java.lang.foreign.MemoryLayout.PathElement; import java.lang.foreign.MemorySegment; +import java.lang.invoke.MethodHandle; import java.lang.invoke.VarHandle; import java.util.ArrayList; import java.util.BitSet; @@ -39,11 +43,25 @@ public class Util { - public static final int CUVS_SUCCESS = 1; + public static final int CUVS_SUCCESS = headers_h.CUVS_SUCCESS(); public static final int CUDA_SUCCESS = 0; private Util() {} + private static final Linker LINKER = Linker.nativeLinker(); + + /** + * Bindings for {@code cudaMemcpyAsync}; differently from the {@code headers_h} bindings (which are + * automatically generated by {@code jextract}), these bindings specify the {@code critical} linker option, + * which allows to directly access heap-based {@link MemorySegment}s. + * {@snippet lang=c : + * extern cudaError_t cudaMemcpyAsync(void *dst, const void *src, size_t count, enum cudaMemcpyKind kind, cudaStream_t stream) + * } + */ + private static final MethodHandle cudaMemcpyAsync$mh = + LINKER.downcallHandle( + cudaMemcpyAsync$address(), cudaMemcpyAsync$descriptor(), Linker.Option.critical(true)); + /** * Checks the result value of a (CuVS) native method handle call. * @@ -81,10 +99,10 @@ public static long dereferenceUnsignedInt(MemorySegment ptr) { * @see CUDA Runtime API */ public enum CudaMemcpyKind { - HOST_TO_HOST(0), - HOST_TO_DEVICE(1), - DEVICE_TO_HOST(2), - DEVICE_TO_DEVICE(3), + HOST_TO_HOST(cudaMemcpyHostToHost()), + HOST_TO_DEVICE(cudaMemcpyHostToDevice()), + DEVICE_TO_HOST(cudaMemcpyDeviceToHost()), + DEVICE_TO_DEVICE(cudaMemcpyDeviceToDevice()), INFER_DIRECTION(4); CudaMemcpyKind(int k) { @@ -120,6 +138,41 @@ public static void cudaMemcpy(MemorySegment dest, MemorySegment src, long numByt Util.cudaMemcpy(dest, src, numBytes, CudaMemcpyKind.INFER_DIRECTION); } + /** + * Helper to invoke cudaMemcpyAsync CUDA runtime function to copy data between host/device memory. + * + * @param dst Destination address for data copy + * @param src Source address for data copy + * @param numBytes Number of bytes to be copied + * @param kind "Direction" of data copy (Host->Device, Device->Host, etc.) + * @throws RuntimeException on failure of copy + */ + public static void cudaMemcpyAsync( + MemorySegment dst, + MemorySegment src, + long numBytes, + CudaMemcpyKind kind, + MemorySegment stream) { + try { + int returnValue = (int) cudaMemcpyAsync$mh.invokeExact(dst, src, numBytes, kind.kind, stream); + checkCudaError(returnValue, "cudaMemcpyAsync"); + } catch (Throwable ex$) { + throw new AssertionError("should not reach here", ex$); + } + } + + /** + * Helper to get the CUDA stream associated with a {@link CuVSResources} + */ + public static MemorySegment getStream(CuVSResources resources) { + try (var resourcesAccess = resources.access(); + var localArena = Arena.ofConfined()) { + var streamPointer = localArena.allocate(cudaStream_t); + checkCuVSError(cuvsStreamGet(resourcesAccess.handle(), streamPointer), "cuvsStreamGet"); + return streamPointer.get(cudaStream_t, 0); + } + } + static final long MAX_ERROR_TEXT = 1_000_000L; static String getLastErrorText() { @@ -309,47 +362,71 @@ public static BitSet concatenate(BitSet[] arr, int maxSizeOfEachBitSet) { } /** - * @brief Helper function for creating DLManagedTensor instance + * Helper function for creating a DLManagedTensor instance * - * @param[in] data the data pointer points to the allocated data - * @param[in] shape the shape of the tensor - * @param[in] code the type code of base types - * @param[in] bits the shape of the tensor + * @param arena the Arena to use to allocate native memory for this data structure + * @param data the data pointer points to the allocated data + * @param shape the shape of the tensor + * @param code the type code of the matrix elements (e.g. kDLFloat()) + * @param bits the size in bits of the matrix elements + * @param deviceType the device where {@code data} is held (kDLCPU() or kDLCUDA()) + * @return DLManagedTensor + */ + public static MemorySegment prepareTensor( + Arena arena, MemorySegment data, long[] shape, int code, int bits, int deviceType) { + return prepareTensor(arena, data, shape, null, code, bits, deviceType); + } + + /** + * Helper function for creating a DLManagedTensor instance + * + * @param arena the Arena to use to allocate native memory for this data structure + * @param data the data pointer points to the allocated data + * @param shape the shape of the tensor + * @param strides the strides for each dimension (can be null) + * @param code the type code of the matrix elements (e.g. kDLFloat()) + * @param bits the size in bits of the matrix elements + * @param deviceType the device where {@code data} is held (kDLCPU() or kDLCUDA()) * @return DLManagedTensor */ public static MemorySegment prepareTensor( Arena arena, MemorySegment data, long[] shape, + long[] strides, int code, int bits, - int deviceType, - int lanes) { + int deviceType) { - MemorySegment tensor = DLManagedTensor.allocate(arena); - MemorySegment dlTensor = DLTensor.allocate(arena); + MemorySegment managedTensor = DLManagedTensor.allocate(arena); + MemorySegment tensor = DLTensor.allocate(arena); - DLTensor.data(dlTensor, data); + DLTensor.data(tensor, data); MemorySegment dlDevice = DLDevice.allocate(arena); DLDevice.device_type(dlDevice, deviceType); - DLTensor.device(dlTensor, dlDevice); - - var ndim = shape.length; - DLTensor.ndim(dlTensor, ndim); + DLTensor.device(tensor, dlDevice); MemorySegment dtype = DLDataType.allocate(arena); DLDataType.code(dtype, (byte) code); DLDataType.bits(dtype, (byte) bits); - DLDataType.lanes(dtype, (short) lanes); - DLTensor.dtype(dlTensor, dtype); - - DLTensor.shape(dlTensor, Util.buildMemorySegment(arena, shape)); - - DLTensor.strides(dlTensor, MemorySegment.NULL); + // Number of lanes for vectorized types (1 for scalar types) + DLDataType.lanes(dtype, (short) 1); + DLTensor.dtype(tensor, dtype); + + DLTensor.ndim(tensor, shape.length); + DLTensor.shape(tensor, Util.buildMemorySegment(arena, shape)); + if (strides != null) { + assert shape.length == strides.length; + DLTensor.strides(tensor, Util.buildMemorySegment(arena, strides)); + } else { + DLTensor.strides(tensor, MemorySegment.NULL); + } - DLManagedTensor.dl_tensor(tensor, dlTensor); + // Copy tensor information into the DLManagedTensor struct + DLManagedTensor.dl_tensor(managedTensor, tensor); - return tensor; + assert bits == DLDataType.bits(DLTensor.dtype(DLManagedTensor.dl_tensor(managedTensor))); + return managedTensor; } } diff --git a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/spi/JDKProvider.java b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/spi/JDKProvider.java index 668befe236..01eb72ed69 100644 --- a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/spi/JDKProvider.java +++ b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/spi/JDKProvider.java @@ -15,6 +15,8 @@ */ package com.nvidia.cuvs.spi; +import static com.nvidia.cuvs.internal.common.Util.*; + import com.nvidia.cuvs.*; import com.nvidia.cuvs.internal.*; import com.nvidia.cuvs.internal.common.Util; @@ -24,6 +26,7 @@ import java.lang.invoke.MethodType; import java.nio.file.Files; import java.nio.file.Path; +import java.util.Locale; import java.util.Objects; final class JDKProvider implements CuVSProvider { @@ -84,7 +87,7 @@ public TieredIndex.Builder newTieredIndexBuilder(CuVSResources cuVSResources) { } @Override - public CagraIndex mergeCagraIndexes(CagraIndex[] indexes) throws Throwable { + public CagraIndex mergeCagraIndexes(CagraIndex[] indexes) { if (indexes == null || indexes.length == 0) { throw new IllegalArgumentException("At least one index must be provided for merging"); } @@ -92,8 +95,7 @@ public CagraIndex mergeCagraIndexes(CagraIndex[] indexes) throws Throwable { } @Override - public CagraIndex mergeCagraIndexes(CagraIndex[] indexes, CagraMergeParams mergeParams) - throws Throwable { + public CagraIndex mergeCagraIndexes(CagraIndex[] indexes, CagraMergeParams mergeParams) { if (indexes == null || indexes.length == 0) { throw new IllegalArgumentException("At least one index must be provided for merging"); } @@ -101,26 +103,40 @@ public CagraIndex mergeCagraIndexes(CagraIndex[] indexes, CagraMergeParams merge } @Override - public CuVSMatrix.Builder newMatrixBuilder(int size, int dimensions, CuVSMatrix.DataType dataType) - throws UnsupportedOperationException { + public CuVSMatrix.Builder newHostMatrixBuilder( + long size, long columns, CuVSMatrix.DataType dataType) throws UnsupportedOperationException { - var dataset = new CuVSHostMatrixArenaImpl(size, dimensions, dataType); - - return new CuVSMatrix.Builder() { + return new CuVSMatrix.Builder<>() { + final CuVSHostMatrixArenaImpl matrix = new CuVSHostMatrixArenaImpl(size, columns, dataType); int current = 0; @Override public void addVector(float[] vector) { + if (vector.length != columns) { + throw new IllegalArgumentException( + String.format( + Locale.ROOT, "Expected a vector of size [%d], got [%d]", columns, vector.length)); + } internalAddVector(vector); } @Override public void addVector(byte[] vector) { + if (vector.length != columns) { + throw new IllegalArgumentException( + String.format( + Locale.ROOT, "Expected a vector of size [%d], got [%d]", columns, vector.length)); + } internalAddVector(vector); } @Override public void addVector(int[] vector) { + if (vector.length != columns) { + throw new IllegalArgumentException( + String.format( + Locale.ROOT, "Expected a vector of size [%d], got [%d]", columns, vector.length)); + } internalAddVector(vector); } @@ -129,19 +145,37 @@ private void internalAddVector(Object vector) { MemorySegment.copy( vector, 0, - dataset.memorySegment(), - dataset.valueLayout(), - ((current++) * dimensions * dataset.valueLayout().byteSize()), - dimensions); + matrix.memorySegment(), + matrix.valueLayout(), + ((current++) * columns * matrix.valueLayout().byteSize()), + (int) columns); } @Override - public CuVSMatrix build() { - return dataset; + public CuVSHostMatrix build() { + return matrix; } }; } + @Override + public CuVSMatrix.Builder newDeviceMatrixBuilder( + CuVSResources resources, long size, long columns, CuVSMatrix.DataType dataType) + throws UnsupportedOperationException { + return new HeapSegmentBuilder(resources, size, columns, dataType); + } + + @Override + public CuVSMatrix.Builder newDeviceMatrixBuilder( + CuVSResources resources, + long size, + long columns, + int rowStride, + int columnStride, + CuVSMatrix.DataType dataType) { + return new HeapSegmentBuilder(resources, size, columns, rowStride, columnStride, dataType); + } + @Override public MethodHandle newNativeMatrixBuilder() { return createNativeDataset$mh; @@ -188,4 +222,91 @@ public CuVSMatrix newMatrixFromArray(byte[][] vectors) { Util.copy(dataset.memorySegment(), vectors); return dataset; } + + /** + * This {@link CuVSDeviceMatrix} builder implementation returns a {@link CuVSDeviceMatrix} backed by managed RMM + * device memory. It uses a non-native {@link MemorySegment} created directly from on-heap java arrays to avoid + * an intermediate allocation and copy to a native (off-heap) segment. + * It requires the copy function ({@code cudaMemcpyAsync}) to have the {@code Critical} linker option in order + * to allow the access to on-heap memory (see {@link Util#cudaMemcpyAsync}). + */ + private static class HeapSegmentBuilder implements CuVSMatrix.Builder { + private final long columns; + private final long size; + private final CuVSDeviceMatrixImpl matrix; + private final MemorySegment stream; + private int current; + + private HeapSegmentBuilder( + CuVSResources resources, long size, long columns, CuVSMatrix.DataType dataType) { + this.columns = columns; + this.size = size; + this.matrix = CuVSDeviceMatrixRMMImpl.create(resources, size, columns, dataType); + this.stream = Util.getStream(resources); + this.current = 0; + } + + private HeapSegmentBuilder( + CuVSResources resources, + long size, + long columns, + int rowStride, + int columnStride, + CuVSMatrix.DataType dataType) { + this.columns = columns; + this.size = size; + this.matrix = + CuVSDeviceMatrixRMMImpl.create( + resources, size, columns, rowStride, columnStride, dataType); + this.stream = Util.getStream(resources); + this.current = 0; + } + + @Override + public void addVector(float[] vector) { + if (vector.length != columns) { + throw new IllegalArgumentException( + String.format( + Locale.ROOT, "Expected a vector of size [%d], got [%d]", columns, vector.length)); + } + internalAddVector(MemorySegment.ofArray(vector)); + } + + @Override + public void addVector(byte[] vector) { + if (vector.length != columns) { + throw new IllegalArgumentException( + String.format( + Locale.ROOT, "Expected a vector of size [%d], got [%d]", columns, vector.length)); + } + internalAddVector(MemorySegment.ofArray(vector)); + } + + @Override + public void addVector(int[] vector) { + if (vector.length != columns) { + throw new IllegalArgumentException( + String.format( + Locale.ROOT, "Expected a vector of size [%d], got [%d]", columns, vector.length)); + } + internalAddVector(MemorySegment.ofArray(vector)); + } + + private void internalAddVector(MemorySegment vector) { + if (current >= size) { + throw new ArrayIndexOutOfBoundsException(); + } + + long rowBytes = columns * matrix.valueLayout().byteSize(); + + var dstOffset = ((current++) * rowBytes); + var dst = matrix.memorySegment().asSlice(dstOffset); + cudaMemcpyAsync(dst, vector, rowBytes, CudaMemcpyKind.HOST_TO_DEVICE, stream); + } + + @Override + public CuVSDeviceMatrix build() { + return matrix; + } + } } 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 6bfbf4d2c7..f65680d631 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 @@ -121,7 +121,7 @@ private void tmpResultsTopKWithRandomValues(boolean useNativeMemoryDataset) thro BruteForceIndex index; if (useNativeMemoryDataset) { var datasetBuilder = - CuVSMatrix.builder(vectors.length, vectors[0].length, CuVSMatrix.DataType.FLOAT); + CuVSMatrix.hostBuilder(vectors.length, vectors[0].length, CuVSMatrix.DataType.FLOAT); for (float[] v : vectors) { datasetBuilder.addVector(v); } 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 a4133c4dc5..c0c0c8a37a 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 @@ -31,7 +31,6 @@ import java.lang.foreign.MemoryLayout; import java.lang.foreign.MemorySegment; import java.lang.foreign.ValueLayout; -import java.lang.invoke.MethodHandles; import java.nio.file.Files; import java.nio.file.Path; import java.util.Arrays; @@ -56,7 +55,7 @@ @RunWith(RandomizedRunner.class) public class CagraBuildAndSearchIT extends CuVSTestCase { - private static final Logger log = LoggerFactory.getLogger(MethodHandles.lookup().lookupClass()); + private static final Logger log = LoggerFactory.getLogger(CagraBuildAndSearchIT.class); @Before public void setup() { 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 144090866c..6e02793478 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 @@ -122,7 +122,7 @@ private void tmpResultsTopKWithRandomValues(boolean useNativeMemoryDataset) thro CagraIndex index; if (useNativeMemoryDataset) { var datasetBuilder = - CuVSMatrix.builder(vectors.length, vectors[0].length, CuVSMatrix.DataType.FLOAT); + CuVSMatrix.hostBuilder(vectors.length, vectors[0].length, CuVSMatrix.DataType.FLOAT); for (float[] v : vectors) { datasetBuilder.addVector(v); } 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 3518b8acd6..0eaf666d15 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 @@ -23,10 +23,14 @@ import org.junit.Before; import org.junit.Test; import org.junit.runner.RunWith; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; @RunWith(RandomizedRunner.class) public class CuVSMatrixIT extends CuVSTestCase { + private static final Logger log = LoggerFactory.getLogger(CuVSMatrixIT.class); + @Before public void setup() { assumeTrue("not supported on " + System.getProperty("os.name"), isLinuxAmd64()); @@ -61,6 +65,10 @@ private float[][] createFloatMatrix() { int rows = randomIntBetween(1, 32); int cols = randomIntBetween(1, 100); + return createFloatMatrix(rows, cols); + } + + private float[][] createFloatMatrix(int rows, int cols) { float[][] result = new float[rows][cols]; for (int r = 0; r < rows; ++r) { @@ -71,42 +79,99 @@ private float[][] createFloatMatrix() { return result; } + private void testByteDatasetRowGetAccess(CuVSMatrix dataset) { + for (int n = 0; n < dataset.size(); ++n) { + var row = dataset.getRow(n); + assertEquals(dataset.columns(), row.size()); + for (int i = 0; i < dataset.columns(); ++i) { + assertEquals(byteData[n][i], row.getAsByte(i)); + } + } + } + @Test - public void testByteDatasetRowGetAccess() { - try (var dataset = CuVSMatrix.ofArray(byteData)) { - for (int n = 0; n < dataset.size(); ++n) { - var row = dataset.getRow(n); - assertEquals(dataset.columns(), row.size()); - for (int i = 0; i < dataset.columns(); ++i) { - assertEquals(byteData[n][i], row.getAsByte(i)); - } + public void testByteHostDatasetRowGetAccess() { + try (var matrix = CuVSMatrix.ofArray(byteData)) { + testByteDatasetRowGetAccess(matrix); + } + } + + @Test + public void testByteDeviceDatasetRowGetAccess() throws Throwable { + try (var resources = CheckedCuVSResources.create()) { + var builder = + CuVSMatrix.deviceBuilder( + resources, byteData.length, byteData[0].length, CuVSMatrix.DataType.BYTE); + for (int i = 0; i < byteData.length; i++) { + builder.addVector(byteData[i]); + } + try (var matrix = builder.build()) { + testByteDatasetRowGetAccess(matrix); } } } + private void testByteDatasetRowCopy(CuVSMatrix dataset) { + for (int n = 0; n < dataset.size(); ++n) { + var row = dataset.getRow(n); + assertEquals(dataset.columns(), row.size()); + + var rowCopy = new byte[(int) row.size()]; + row.toArray(rowCopy); + assertArrayEquals(byteData[n], rowCopy); + } + } + @Test - public void testByteDatasetRowCopy() { + public void testByteHostDatasetRowCopy() { try (var dataset = CuVSMatrix.ofArray(byteData)) { - for (int n = 0; n < dataset.size(); ++n) { - var row = dataset.getRow(n); - assertEquals(dataset.columns(), row.size()); + testByteDatasetRowCopy(dataset); + } + } - var rowCopy = new byte[(int) row.size()]; - row.toArray(rowCopy); - assertArrayEquals(byteData[n], rowCopy); + @Test + public void testByteDeviceDatasetRowCopy() throws Throwable { + try (var resources = CheckedCuVSResources.create()) { + var builder = + CuVSMatrix.deviceBuilder( + resources, byteData.length, byteData[0].length, CuVSMatrix.DataType.BYTE); + for (int i = 0; i < byteData.length; i++) { + builder.addVector(byteData[i]); + } + try (var matrix = builder.build()) { + testByteDatasetRowCopy(matrix); + } + } + } + + private void testByteDatasetCopy(CuVSMatrix dataset) { + var dataCopy = new byte[(int) dataset.size()][(int) dataset.columns()]; + dataset.toArray(dataCopy); + for (int n = 0; n < dataset.size(); ++n) { + for (int i = 0; i < dataset.columns(); ++i) { + assertEquals(byteData[n][i], dataCopy[n][i]); } } } @Test - public void testByteDatasetCopy() { + public void testByteHostDatasetCopy() { try (var dataset = CuVSMatrix.ofArray(byteData)) { - var dataCopy = new byte[(int) dataset.size()][(int) dataset.columns()]; - dataset.toArray(dataCopy); - for (int n = 0; n < dataset.size(); ++n) { - for (int i = 0; i < dataset.columns(); ++i) { - assertEquals(byteData[n][i], dataCopy[n][i]); - } + testByteDatasetCopy(dataset); + } + } + + @Test + public void testByteDeviceDatasetCopy() throws Throwable { + try (var resources = CheckedCuVSResources.create()) { + var builder = + CuVSMatrix.deviceBuilder( + resources, byteData.length, byteData[0].length, CuVSMatrix.DataType.BYTE); + for (int i = 0; i < byteData.length; i++) { + builder.addVector(byteData[i]); + } + try (var matrix = builder.build()) { + testByteDatasetCopy(matrix); } } } @@ -197,18 +262,15 @@ public void testFloatDatasetCopy() { } } - public void testFloatDatasetBuilder() { - int rows = randomIntBetween(1, 32); - int cols = randomIntBetween(1, 100); + private void testFloatDatasetBuilder(int rows, int cols, CuVSMatrix.Builder builder) { float[][] data = new float[rows][cols]; for (int r = 0; r < rows; ++r) { for (int c = 0; c < cols; ++c) { - data[c][r] = randomFloat(); + data[r][c] = randomFloat(); } } - var builder = CuVSMatrix.builder(rows, cols, CuVSMatrix.DataType.FLOAT); for (int r = 0; r < rows; ++r) { builder.addVector(data[r]); } @@ -226,18 +288,33 @@ public void testFloatDatasetBuilder() { } } - public void testIntDatasetBuilder() { + @Test + public void testFloatDatasetHostBuilder() { int rows = randomIntBetween(1, 32); int cols = randomIntBetween(1, 100); + testFloatDatasetBuilder( + rows, cols, CuVSMatrix.hostBuilder(rows, cols, CuVSMatrix.DataType.FLOAT)); + } + + @Test + public void testFloatDatasetDeviceBuilder() throws Throwable { + int rows = randomIntBetween(1, 1024); + int cols = randomIntBetween(1, 2048); + try (var resources = CheckedCuVSResources.create()) { + testFloatDatasetBuilder( + rows, cols, CuVSMatrix.deviceBuilder(resources, rows, cols, CuVSMatrix.DataType.FLOAT)); + } + } + + private void testIntDatasetBuilder(int rows, int cols, CuVSMatrix.Builder builder) { var data = new int[rows][cols]; for (int r = 0; r < rows; ++r) { for (int c = 0; c < cols; ++c) { - data[c][r] = randomInt(); + data[r][c] = randomInt(); } } - var builder = CuVSMatrix.builder(rows, cols, CuVSMatrix.DataType.INT); for (int r = 0; r < rows; ++r) { builder.addVector(data[r]); } @@ -255,18 +332,32 @@ public void testIntDatasetBuilder() { } } - public void testByteDatasetBuilder() { + @Test + public void testIntDatasetHostBuilder() { int rows = randomIntBetween(1, 32); int cols = randomIntBetween(1, 100); + testIntDatasetBuilder(rows, cols, CuVSMatrix.hostBuilder(rows, cols, CuVSMatrix.DataType.INT)); + } + + @Test + public void testIntDatasetDeviceBuilder() throws Throwable { + int rows = randomIntBetween(1, 1024); + int cols = randomIntBetween(1, 2048); + try (var resources = CheckedCuVSResources.create()) { + testIntDatasetBuilder( + rows, cols, CuVSMatrix.deviceBuilder(resources, rows, cols, CuVSMatrix.DataType.INT)); + } + } + + private void testByteDatasetBuilder(int rows, int cols, CuVSMatrix.Builder builder) { var data = new byte[rows][cols]; for (int r = 0; r < rows; ++r) { for (int c = 0; c < cols; ++c) { - data[c][r] = randomByte(); + data[r][c] = randomByte(); } } - var builder = CuVSMatrix.builder(rows, cols, CuVSMatrix.DataType.BYTE); for (int r = 0; r < rows; ++r) { builder.addVector(data[r]); } @@ -283,4 +374,59 @@ public void testByteDatasetBuilder() { } } } + + @Test + public void testByteDatasetHostBuilder() { + int rows = randomIntBetween(1, 32); + int cols = randomIntBetween(1, 100); + testByteDatasetBuilder( + rows, cols, CuVSMatrix.hostBuilder(rows, cols, CuVSMatrix.DataType.BYTE)); + } + + @Test + public void testByteDatasetDeviceBuilder() throws Throwable { + int rows = randomIntBetween(1, 1024); + int cols = randomIntBetween(1, 2048); + try (var resources = CheckedCuVSResources.create()) { + testByteDatasetBuilder( + rows, cols, CuVSMatrix.deviceBuilder(resources, rows, cols, CuVSMatrix.DataType.BYTE)); + } + } + + @Test + public void testDeviceToHost() throws Throwable { + + final int size = 16 * 1024; + final int columns = 2048; + final float[][] data = createFloatMatrix(size, columns); + + try (var resources = CuVSResources.create()) { + + var builder = CuVSMatrix.deviceBuilder(resources, size, columns, CuVSMatrix.DataType.FLOAT); + for (int i = 0; i < size; ++i) { + var array = data[i]; + builder.addVector(array); + } + + try (var deviceMatrix = (CuVSDeviceMatrix) builder.build(); + var hostMatrix = deviceMatrix.toHost()) { + + assertEquals(data.length, deviceMatrix.size()); + assertEquals(data[0].length, deviceMatrix.columns()); + + assertEquals(deviceMatrix.size(), hostMatrix.size()); + assertEquals(deviceMatrix.columns(), hostMatrix.columns()); + + var roundTripData = new float[size][columns]; + + hostMatrix.toArray(roundTripData); + + for (int n = 0; n < hostMatrix.size(); ++n) { + for (int i = 0; i < hostMatrix.columns(); ++i) { + assertEquals(data[n][i], roundTripData[n][i], 1e-9); + } + } + } + } + } } 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 02ac3c3a14..ac4a491cff 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 @@ -114,7 +114,7 @@ private void tmpResultsTopKWithRandomValues(boolean useNativeMemoryDataset) thro final CagraIndex index; if (useNativeMemoryDataset) { var datasetBuilder = - CuVSMatrix.builder(vectors.length, vectors[0].length, CuVSMatrix.DataType.FLOAT); + CuVSMatrix.hostBuilder(vectors.length, vectors[0].length, CuVSMatrix.DataType.FLOAT); for (float[] v : vectors) { datasetBuilder.addVector(v); } From 0b1b1ec47ff5364e40148ae9ae4ff375242c05be Mon Sep 17 00:00:00 2001 From: tsuki <12711693+enp1s0@users.noreply.github.com> Date: Tue, 26 Aug 2025 01:57:10 +0900 Subject: [PATCH 43/93] Fix balanced kmeans API (#1268) This PR updates cluster::balanced_kmeans to: 1. use `int64_t` as the `mdspan` index type to support a large dataset 2. fix the `fit` function for `data_t` = `int8_t` 3. test `fit` and `predict` in the unit test. ## 1. Use `int64_t` as the `mdspan` index type to support a large dataset In the current implementation, the mdspan index type is `int32_t`. Even if `int64_t` is specified, it is cast to `int32_t`. e.g. ```cpp auto dataset = raft::make_device_matrix(handle, n, d); cuvs::cluster::kmeans::fit(handle, params, mkconst(dataset.view()), centroids_view); // Here, the index type of the dataset view is cast to int32_t since there is only an API for int32_t. ``` This cast results in an unexpected error in this check. `IndexT` is `int32_t` even if the input index type is `int64_t`. https://github.com/rapidsai/cuvs/blob/d0a83d43f7be934f0c4dcef585302a7fad287df0/cpp/src/cluster/kmeans_balanced.cuh#L86-L88 This PR changes the index type to `int64_t`. If `int32_t` is specified, it is cast to `int64_t`. Therefore, this change is not breaking. ## 2. Fix the `fit` function for `data_t` = `int8_t` The centroid data type of the `predict` and `fit_predict` functions is `float`, but `int8_t` for the `fit` function (`fit_predict` is a function that sequentially calls the `fit` and `predict` functions internally). This PR changes it to `float` since this data type must be the same as the one of `predict`. ## 3. Test `fit` and `predict` in the unit test. In the current unit test, only the `fit_predict` function is tested. This PR adds tests for the `fit` and `predict` functions. ## Discussion There is no consistency in the label data types for `predict` and `fit_predict`: `predict` is `int32_t`, while `fit_predict` is `uint32_t`. Would it be better to: 1. Support both types in both functions (e.g., by reinterpret_cast-ing int32_t* to uint32*), 2. Unify on one type, or 3. Keep the inconsistency? Authors: - tsuki (https://github.com/enp1s0) Approvers: - Corey J. Nolet (https://github.com/cjnolet) URL: https://github.com/rapidsai/cuvs/pull/1268 --- cpp/include/cuvs/cluster/kmeans.hpp | 113 ++++++++++++------ cpp/src/cluster/kmeans_balanced_fit_float.cu | 6 +- cpp/src/cluster/kmeans_balanced_fit_int8.cu | 6 +- .../kmeans_balanced_fit_predict_float.cu | 8 +- .../kmeans_balanced_fit_predict_int8.cu | 18 ++- .../cluster/kmeans_balanced_predict_float.cu | 8 +- .../cluster/kmeans_balanced_predict_int8.cu | 18 ++- cpp/tests/cluster/kmeans_balanced.cu | 53 +++++--- 8 files changed, 159 insertions(+), 71 deletions(-) diff --git a/cpp/include/cuvs/cluster/kmeans.hpp b/cpp/include/cuvs/cluster/kmeans.hpp index ad08cd1394..43d2db5a5c 100644 --- a/cpp/include/cuvs/cluster/kmeans.hpp +++ b/cpp/include/cuvs/cluster/kmeans.hpp @@ -1,5 +1,5 @@ /* - * Copyright (c) 2022, NVIDIA CORPORATION. + * Copyright (c) 2022-2025, NVIDIA CORPORATION. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -407,9 +407,9 @@ void fit(raft::resources const& handle, * ... * raft::resources handle; * cuvs::cluster::kmeans::balanced_params params; - * int n_features = 15; - * int n_clusters = 8; - * auto centroids = raft::make_device_matrix(handle, n_clusters, n_features); + * int64_t n_features = 15; + * int64_t n_clusters = 8; + * auto centroids = raft::make_device_matrix(handle, n_clusters, n_features); * * kmeans::fit(handle, * params, @@ -429,8 +429,8 @@ void fit(raft::resources const& handle, */ void fit(const raft::resources& handle, cuvs::cluster::kmeans::balanced_params const& params, - raft::device_matrix_view X, - raft::device_matrix_view centroids); + raft::device_matrix_view X, + raft::device_matrix_view centroids); /** * @brief Find balanced clusters with k-means algorithm. @@ -442,8 +442,8 @@ void fit(const raft::resources& handle, * ... * raft::resources handle; * cuvs::cluster::kmeans::balanced_params params; - * int n_features = 15, n_clusters = 8; - * auto centroids = raft::make_device_matrix(handle, n_clusters, n_features); + * int64_t n_features = 15, n_clusters = 8; + * auto centroids = raft::make_device_matrix(handle, n_clusters, n_features); * * kmeans::fit(handle, * params, @@ -463,8 +463,9 @@ void fit(const raft::resources& handle, */ void fit(const raft::resources& handle, cuvs::cluster::kmeans::balanced_params const& params, - raft::device_matrix_view X, - raft::device_matrix_view centroids); + raft::device_matrix_view X, + raft::device_matrix_view centroids); + /** * @brief Predict the closest cluster each sample in X belongs to. * @@ -707,15 +708,15 @@ void predict(raft::resources const& handle, * ... * raft::resources handle; * cuvs::cluster::kmeans::balanced_params params; - * int n_features = 15, n_clusters = 8; - * auto centroids = raft::make_device_matrix(handle, n_clusters, n_features); + * int64_t n_features = 15, n_clusters = 8; + * auto centroids = raft::make_device_matrix(handle, n_clusters, n_features); * * kmeans::fit(handle, * params, * X, * centroids.view()); * ... - * auto labels = raft::make_device_vector(handle, X.extent(0)); + * auto labels = raft::make_device_vector(handle, X.extent(0)); * * kmeans::predict(handle, * params, @@ -737,9 +738,9 @@ void predict(raft::resources const& handle, */ void predict(const raft::resources& handle, cuvs::cluster::kmeans::balanced_params const& params, - raft::device_matrix_view X, - raft::device_matrix_view centroids, - raft::device_vector_view labels); + raft::device_matrix_view X, + raft::device_matrix_view centroids, + raft::device_vector_view labels); /** * @brief Predict the closest cluster each sample in X belongs to. @@ -751,15 +752,15 @@ void predict(const raft::resources& handle, * ... * raft::resources handle; * cuvs::cluster::kmeans::balanced_params params; - * int n_features = 15, n_clusters = 8; - * auto centroids = raft::make_device_matrix(handle, n_clusters, n_features); + * int64_t n_features = 15, n_clusters = 8; + * auto centroids = raft::make_device_matrix(handle, n_clusters, n_features); * * kmeans::fit(handle, * params, * X, * centroids.view()); * ... - * auto labels = raft::make_device_vector(handle, X.extent(0)); + * auto labels = raft::make_device_vector(handle, X.extent(0)); * * kmeans::predict(handle, * params, @@ -781,9 +782,53 @@ void predict(const raft::resources& handle, */ void predict(const raft::resources& handle, cuvs::cluster::kmeans::balanced_params const& params, - raft::device_matrix_view X, - raft::device_matrix_view centroids, - raft::device_vector_view labels); + raft::device_matrix_view X, + raft::device_matrix_view centroids, + raft::device_vector_view labels); + +/** + * @brief Predict the closest cluster each sample in X belongs to. + * + * @code{.cpp} + * #include + * #include + * using namespace cuvs::cluster; + * ... + * raft::resources handle; + * cuvs::cluster::kmeans::balanced_params params; + * int64_t n_features = 15, n_clusters = 8; + * auto centroids = raft::make_device_matrix(handle, n_clusters, n_features); + * + * kmeans::fit(handle, + * params, + * X, + * centroids.view()); + * ... + * auto labels = raft::make_device_vector(handle, X.extent(0)); + * + * kmeans::predict(handle, + * params, + * X, + * centroids.view(), + * labels.view()); + * @endcode + * + * @param[in] handle The raft handle. + * @param[in] params Parameters for KMeans model. + * @param[in] X New data to predict. + * [dim = n_samples x n_features] + * @param[in] centroids Cluster centroids. The data must be in + * row-major format. + * [dim = n_clusters x n_features] + * @param[out] labels Index of the cluster each sample in X + * belongs to. + * [len = n_samples] + */ +void predict(const raft::resources& handle, + cuvs::cluster::kmeans::balanced_params const& params, + raft::device_matrix_view X, + raft::device_matrix_view centroids, + raft::device_vector_view labels); /** * @brief Compute k-means clustering and predicts cluster index for each sample @@ -1016,9 +1061,9 @@ void fit_predict(raft::resources const& handle, * ... * raft::resources handle; * cuvs::cluster::kmeans::balanced_params params; - * int n_features = 15, n_clusters = 8; - * auto centroids = raft::make_device_matrix(handle, n_clusters, n_features); - * auto labels = raft::make_device_vector(handle, X.extent(0)); + * int64_t n_features = 15, n_clusters = 8; + * auto centroids = raft::make_device_matrix(handle, n_clusters, n_features); + * auto labels = raft::make_device_vector(handle, X.extent(0)); * * kmeans::fit_predict(handle, * params, @@ -1045,9 +1090,9 @@ void fit_predict(raft::resources const& handle, */ void fit_predict(const raft::resources& handle, cuvs::cluster::kmeans::balanced_params const& params, - raft::device_matrix_view X, - raft::device_matrix_view centroids, - raft::device_vector_view labels); + raft::device_matrix_view X, + raft::device_matrix_view centroids, + raft::device_vector_view labels); /** * @brief Compute balanced k-means clustering and predicts cluster index for each sample @@ -1060,9 +1105,9 @@ void fit_predict(const raft::resources& handle, * ... * raft::resources handle; * cuvs::cluster::kmeans::balanced_params params; - * int n_features = 15, n_clusters = 8; - * auto centroids = raft::make_device_matrix(handle, n_clusters, n_features); - * auto labels = raft::make_device_vector(handle, X.extent(0)); + * int64_t n_features = 15, n_clusters = 8; + * auto centroids = raft::make_device_matrix(handle, n_clusters, n_features); + * auto labels = raft::make_device_vector(handle, X.extent(0)); * * kmeans::fit_predict(handle, * params, @@ -1089,9 +1134,9 @@ void fit_predict(const raft::resources& handle, */ void fit_predict(const raft::resources& handle, cuvs::cluster::kmeans::balanced_params const& params, - raft::device_matrix_view X, - raft::device_matrix_view centroids, - raft::device_vector_view labels); + raft::device_matrix_view X, + raft::device_matrix_view centroids, + raft::device_vector_view labels); /** * @brief Transform X to a cluster-distance space. diff --git a/cpp/src/cluster/kmeans_balanced_fit_float.cu b/cpp/src/cluster/kmeans_balanced_fit_float.cu index ac5a664929..41e13d9f93 100644 --- a/cpp/src/cluster/kmeans_balanced_fit_float.cu +++ b/cpp/src/cluster/kmeans_balanced_fit_float.cu @@ -1,5 +1,5 @@ /* - * Copyright (c) 2022-2023, NVIDIA CORPORATION. + * Copyright (c) 2022-2025, NVIDIA CORPORATION. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -24,8 +24,8 @@ namespace cuvs::cluster::kmeans { void fit(const raft::resources& handle, cuvs::cluster::kmeans::balanced_params const& params, - raft::device_matrix_view X, - raft::device_matrix_view centroids) + raft::device_matrix_view X, + raft::device_matrix_view centroids) { cuvs::cluster::kmeans_balanced::fit( handle, params, X, centroids, cuvs::spatial::knn::detail::utils::mapping{}); diff --git a/cpp/src/cluster/kmeans_balanced_fit_int8.cu b/cpp/src/cluster/kmeans_balanced_fit_int8.cu index ad9c118684..fd74c2dc69 100644 --- a/cpp/src/cluster/kmeans_balanced_fit_int8.cu +++ b/cpp/src/cluster/kmeans_balanced_fit_int8.cu @@ -1,5 +1,5 @@ /* - * Copyright (c) 2022-2023, NVIDIA CORPORATION. + * Copyright (c) 2022-2025, NVIDIA CORPORATION. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -24,8 +24,8 @@ namespace cuvs::cluster::kmeans { void fit(const raft::resources& handle, cuvs::cluster::kmeans::balanced_params const& params, - raft::device_matrix_view X, - raft::device_matrix_view centroids) + raft::device_matrix_view X, + raft::device_matrix_view centroids) { cuvs::cluster::kmeans_balanced::fit( handle, params, X, centroids, cuvs::spatial::knn::detail::utils::mapping{}); diff --git a/cpp/src/cluster/kmeans_balanced_fit_predict_float.cu b/cpp/src/cluster/kmeans_balanced_fit_predict_float.cu index f6f177121f..5c0109ffb4 100644 --- a/cpp/src/cluster/kmeans_balanced_fit_predict_float.cu +++ b/cpp/src/cluster/kmeans_balanced_fit_predict_float.cu @@ -1,5 +1,5 @@ /* - * Copyright (c) 2022-2023, NVIDIA CORPORATION. + * Copyright (c) 2022-2025, NVIDIA CORPORATION. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -24,9 +24,9 @@ namespace cuvs::cluster::kmeans { void fit_predict(const raft::resources& handle, cuvs::cluster::kmeans::balanced_params const& params, - raft::device_matrix_view X, - raft::device_matrix_view centroids, - raft::device_vector_view labels) + raft::device_matrix_view X, + raft::device_matrix_view centroids, + raft::device_vector_view labels) { cuvs::cluster::kmeans_balanced::fit_predict( handle, params, X, centroids, labels, cuvs::spatial::knn::detail::utils::mapping{}); diff --git a/cpp/src/cluster/kmeans_balanced_fit_predict_int8.cu b/cpp/src/cluster/kmeans_balanced_fit_predict_int8.cu index b6171ece5a..007f50df9b 100644 --- a/cpp/src/cluster/kmeans_balanced_fit_predict_int8.cu +++ b/cpp/src/cluster/kmeans_balanced_fit_predict_int8.cu @@ -1,5 +1,5 @@ /* - * Copyright (c) 2022-2023, NVIDIA CORPORATION. + * Copyright (c) 2022-2025, NVIDIA CORPORATION. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -24,9 +24,19 @@ namespace cuvs::cluster::kmeans { void fit_predict(const raft::resources& handle, cuvs::cluster::kmeans::balanced_params const& params, - raft::device_matrix_view X, - raft::device_matrix_view centroids, - raft::device_vector_view labels) + raft::device_matrix_view X, + raft::device_matrix_view centroids, + raft::device_vector_view labels) +{ + cuvs::cluster::kmeans_balanced::fit_predict( + handle, params, X, centroids, labels, cuvs::spatial::knn::detail::utils::mapping{}); +} + +void fit_predict(const raft::resources& handle, + cuvs::cluster::kmeans::balanced_params const& params, + raft::device_matrix_view X, + raft::device_matrix_view centroids, + raft::device_vector_view labels) { cuvs::cluster::kmeans_balanced::fit_predict( handle, params, X, centroids, labels, cuvs::spatial::knn::detail::utils::mapping{}); diff --git a/cpp/src/cluster/kmeans_balanced_predict_float.cu b/cpp/src/cluster/kmeans_balanced_predict_float.cu index e951c20748..2b4086bcfc 100644 --- a/cpp/src/cluster/kmeans_balanced_predict_float.cu +++ b/cpp/src/cluster/kmeans_balanced_predict_float.cu @@ -1,5 +1,5 @@ /* - * Copyright (c) 2022-2023, NVIDIA CORPORATION. + * Copyright (c) 2022-2025, NVIDIA CORPORATION. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -24,9 +24,9 @@ namespace cuvs::cluster::kmeans { void predict(const raft::resources& handle, cuvs::cluster::kmeans::balanced_params const& params, - raft::device_matrix_view X, - raft::device_matrix_view centroids, - raft::device_vector_view labels) + raft::device_matrix_view X, + raft::device_matrix_view centroids, + raft::device_vector_view labels) { cuvs::cluster::kmeans_balanced::predict( handle, params, X, centroids, labels, cuvs::spatial::knn::detail::utils::mapping{}); diff --git a/cpp/src/cluster/kmeans_balanced_predict_int8.cu b/cpp/src/cluster/kmeans_balanced_predict_int8.cu index e08d8fc87e..c91c6f35ba 100644 --- a/cpp/src/cluster/kmeans_balanced_predict_int8.cu +++ b/cpp/src/cluster/kmeans_balanced_predict_int8.cu @@ -1,5 +1,5 @@ /* - * Copyright (c) 2022-2023, NVIDIA CORPORATION. + * Copyright (c) 2022-2025, NVIDIA CORPORATION. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -24,9 +24,19 @@ namespace cuvs::cluster::kmeans { void predict(const raft::resources& handle, cuvs::cluster::kmeans::balanced_params const& params, - raft::device_matrix_view X, - raft::device_matrix_view centroids, - raft::device_vector_view labels) + raft::device_matrix_view X, + raft::device_matrix_view centroids, + raft::device_vector_view labels) +{ + cuvs::cluster::kmeans_balanced::predict( + handle, params, X, centroids, labels, cuvs::spatial::knn::detail::utils::mapping{}); +} + +void predict(const raft::resources& handle, + cuvs::cluster::kmeans::balanced_params const& params, + raft::device_matrix_view X, + raft::device_matrix_view centroids, + raft::device_vector_view labels) { cuvs::cluster::kmeans_balanced::predict( handle, params, X, centroids, labels, cuvs::spatial::knn::detail::utils::mapping{}); diff --git a/cpp/tests/cluster/kmeans_balanced.cu b/cpp/tests/cluster/kmeans_balanced.cu index f273dd92fa..339f8d9124 100644 --- a/cpp/tests/cluster/kmeans_balanced.cu +++ b/cpp/tests/cluster/kmeans_balanced.cu @@ -1,5 +1,5 @@ /* - * Copyright (c) 2022-2024, NVIDIA CORPORATION. + * Copyright (c) 2022-2025, NVIDIA CORPORATION. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -60,7 +60,12 @@ template return os; } -template +template class KmeansBalancedTest : public ::testing::TestWithParam> { protected: KmeansBalancedTest() @@ -123,8 +128,14 @@ class KmeansBalancedTest : public ::testing::TestWithParam(d_centroids.data(), p.n_clusters, p.n_cols); auto d_labels_view = raft::make_device_vector_view(d_labels.data(), p.n_rows); - cuvs::cluster::kmeans::fit_predict( - handle, p.kb_params, X_view, d_centroids_view, d_labels_view); + if constexpr (SeparateFitPredict) { + cuvs::cluster::kmeans::fit(handle, p.kb_params, X_view, d_centroids_view); + cuvs::cluster::kmeans::predict( + handle, p.kb_params, X_view, raft::make_const_mdspan(d_centroids_view), d_labels_view); + } else { + cuvs::cluster::kmeans::fit_predict( + handle, p.kb_params, X_view, d_centroids_view, d_labels_view); + } raft::resource::sync_stream(handle, stream); @@ -161,12 +172,15 @@ std::vector> get_kmeans_balanced_inputs() p.kb_params.n_iters = 20; p.kb_params.metric = cuvs::distance::DistanceType::L2Expanded; p.tol = MathT{0.0001}; - std::vector> row_cols_k = {{1000, 32, 5}, - {1000, 100, 20}, - {10000, 32, 10}, - {10000, 100, 50}, - {10000, 500, 100}, - {1000000, 128, 10}}; + std::vector> row_cols_k = { + {1000, 32, 5}, + {1000, 100, 20}, + {10000, 32, 10}, + {10000, 100, 50}, + {10000, 500, 100}, + {1000000, 128, 10}, + {10000000, 128, 10}, + }; for (auto& rck : row_cols_k) { p.n_rows = static_cast(std::get<0>(rck)); p.n_cols = static_cast(std::get<1>(rck)); @@ -190,15 +204,21 @@ const auto inputsf_i64 = get_kmeans_balanced_inputs(); * First set of tests: no conversion */ -KB_TEST((KmeansBalancedTest), +KB_TEST((KmeansBalancedTest), KmeansBalancedTestFFU32I32, inputsf_i32); +KB_TEST((KmeansBalancedTest), + KmeansBalancedTestFFI32I32_SEP, + inputsf_i32); // KB_TEST((KmeansBalancedTest), // KmeansBalancedTestDDU32I32, // inputsd_i32); -// KB_TEST((KmeansBalancedTest), -// KmeansBalancedTestFFU32I64, -// inputsf_i64); +KB_TEST((KmeansBalancedTest), + KmeansBalancedTestFFU32I64, + inputsf_i64); +KB_TEST((KmeansBalancedTest), + KmeansBalancedTestFFI32I64_SEP, + inputsf_i64); // KB_TEST((KmeansBalancedTest), // KmeansBalancedTestDDU32I64, // inputsd_i64); @@ -231,8 +251,11 @@ struct i2f_scaler { RAFT_INLINE_FUNCTION auto operator()(const DataT& x) const { return op(x); }; }; -KB_TEST((KmeansBalancedTest>), +KB_TEST((KmeansBalancedTest, false>), KmeansBalancedTestFI8U32I32, inputsf_i32); +KB_TEST((KmeansBalancedTest, true>), + KmeansBalancedTestFI8I32I32_SEP, + inputsf_i32); } // namespace cuvs From 73672634c3841df358d15da5dea647cf6472b199 Mon Sep 17 00:00:00 2001 From: Kyle Edwards Date: Tue, 26 Aug 2025 16:23:51 -0400 Subject: [PATCH 44/93] Update rapids-dependency-file-generator (#1285) This PR updates the rapids-dependency-file-generator hook to get https://github.com/rapidsai/dependency-file-generator/pull/163. Authors: - Kyle Edwards (https://github.com/KyleFromNVIDIA) Approvers: - James Lamb (https://github.com/jameslamb) URL: https://github.com/rapidsai/cuvs/pull/1285 --- .pre-commit-config.yaml | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 05489ee94d..c135ff4ff4 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -126,10 +126,10 @@ repos: - id: verify-codeowners args: [--fix, --project-prefix=cuvs] - repo: https://github.com/rapidsai/dependency-file-generator - rev: v1.19.0 + rev: v1.20.0 hooks: - id: rapids-dependency-file-generator - args: ["--clean"] + args: ["--clean", "--warn-all", "--strict"] - repo: https://github.com/shellcheck-py/shellcheck-py rev: v0.10.0.1 hooks: From 6e0f85997d86e0e9be0744dbb3674b2f58484c08 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Lorenzo=20Dematt=C3=A9?= Date: Wed, 27 Aug 2025 08:56:51 +0200 Subject: [PATCH 45/93] [Java][C] Expose GPUInfo (#1267) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit cuvs-java already contains a public `GPUInfo` class, but methods to retrieve the information and fill it are internal. This PR exposes them through and interface, `GPUInfoProvider`. It also separates immutable data related to a GPU (which is kept in `GPUInfo`) from transient resources-related data and counters (at the moment, only the amount of free memory, which is kept in the new `CuVSResourcesInfo`). The change let you query transient data at a later moment; to do this, we need to find the device ID associated with a `CuVSResource` object. The change to the C API exposes the raft function that does it. Authors: - Lorenzo Dematté (https://github.com/ldematte) - MithunR (https://github.com/mythrocks) Approvers: - Corey J. Nolet (https://github.com/cjnolet) - MithunR (https://github.com/mythrocks) URL: https://github.com/rapidsai/cuvs/pull/1267 --- cpp/include/cuvs/core/c_api.h | 11 +- cpp/src/core/c_api.cpp | 10 ++ .../java/com/nvidia/cuvs/CuVSResources.java | 7 + .../com/nvidia/cuvs/CuVSResourcesInfo.java | 26 ++++ .../main/java/com/nvidia/cuvs/GPUInfo.java | 22 +++- .../java/com/nvidia/cuvs/GPUInfoProvider.java | 48 +++++++ .../cuvs/SynchronizedCuVSResources.java | 5 + .../com/nvidia/cuvs/spi/CuVSProvider.java | 3 + .../nvidia/cuvs/spi/UnsupportedProvider.java | 5 + .../cuvs/internal/CuVSResourcesImpl.java | 17 ++- .../cuvs/internal/GPUInfoProviderImpl.java | 121 ++++++++++++++++++ .../com/nvidia/cuvs/internal/common/Util.java | 83 ------------ .../com/nvidia/cuvs/spi/JDKProvider.java | 5 + .../com/nvidia/cuvs/CheckedCuVSResources.java | 5 + .../java/com/nvidia/cuvs/CuVSResourcesIT.java | 3 +- .../test/java/com/nvidia/cuvs/GPUInfoIT.java | 78 +++++++++++ 16 files changed, 352 insertions(+), 97 deletions(-) create mode 100644 java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSResourcesInfo.java create mode 100644 java/cuvs-java/src/main/java/com/nvidia/cuvs/GPUInfoProvider.java create mode 100644 java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/GPUInfoProviderImpl.java create mode 100644 java/cuvs-java/src/test/java/com/nvidia/cuvs/GPUInfoIT.java diff --git a/cpp/include/cuvs/core/c_api.h b/cpp/include/cuvs/core/c_api.h index d7b6673990..e3cd9bd40b 100644 --- a/cpp/include/cuvs/core/c_api.h +++ b/cpp/include/cuvs/core/c_api.h @@ -86,7 +86,7 @@ cuvsError_t cuvsResourcesDestroy(cuvsResources_t res); cuvsError_t cuvsStreamSet(cuvsResources_t res, cudaStream_t stream); /** - * @brief Get the cudaStream_t from a cuvsResources_t t + * @brief Get the cudaStream_t from a cuvsResources_t * * @param[in] res cuvsResources_t opaque C handle * @param[out] stream cudaStream_t stream to queue CUDA kernels @@ -101,6 +101,15 @@ cuvsError_t cuvsStreamGet(cuvsResources_t res, cudaStream_t* stream); * @return cuvsError_t */ cuvsError_t cuvsStreamSync(cuvsResources_t res); + +/** + * @brief Get the id of the device associated with this cuvsResources_t + * + * @param[in] res cuvsResources_t opaque C handle + * @param[out] device_id int the id of the device associated with res + * @return cuvsError_t + */ +cuvsError_t cuvsDeviceIdGet(cuvsResources_t res, int* device_id); /** @} */ /** diff --git a/cpp/src/core/c_api.cpp b/cpp/src/core/c_api.cpp index 8a9116f4cb..2f66adf416 100644 --- a/cpp/src/core/c_api.cpp +++ b/cpp/src/core/c_api.cpp @@ -19,6 +19,8 @@ #include #include +#include +#include #include #include #include @@ -73,6 +75,14 @@ extern "C" cuvsError_t cuvsStreamSync(cuvsResources_t res) }); } +extern "C" cuvsError_t cuvsDeviceIdGet(cuvsResources_t res, int* device_id) +{ + return cuvs::core::translate_exceptions([=] { + auto res_ptr = reinterpret_cast(res); + *device_id = raft::resource::get_device_id(*res_ptr); + }); +} + extern "C" cuvsError_t cuvsRMMAlloc(cuvsResources_t res, void** ptr, size_t bytes) { return cuvs::core::translate_exceptions([=] { diff --git a/java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSResources.java b/java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSResources.java index 2d362b6c4c..b44e50e890 100644 --- a/java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSResources.java +++ b/java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSResources.java @@ -49,6 +49,13 @@ interface ScopedAccess extends AutoCloseable { */ ScopedAccess access(); + /** + * Get the logical id of the device associated with this resources object. + * Information about the device id is immutable, so it is safe to expose it without getting {@link ScopedAccess} + * to the enclosing resources. + */ + int deviceId(); + /** * Closes this CuVSResources object and releases any resources associated with it. */ diff --git a/java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSResourcesInfo.java b/java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSResourcesInfo.java new file mode 100644 index 0000000000..a100091f12 --- /dev/null +++ b/java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSResourcesInfo.java @@ -0,0 +1,26 @@ +/* + * Copyright (c) 2025, NVIDIA CORPORATION. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.nvidia.cuvs; + +/** + * Contains performance-related information associated to a {@link CuVSResources} and its GPU. + * Can be extended to report different types of GPU memory linked to the resources, + * e.g. the type and capacity of the underlying RMM {@code device_memory_resource} + * + * @param freeDeviceMemoryInBytes free memory in bytes, as reported by the device driver + * @param totalDeviceMemoryInBytes total device memory in bytes + */ +public record CuVSResourcesInfo(long freeDeviceMemoryInBytes, long totalDeviceMemoryInBytes) {} diff --git a/java/cuvs-java/src/main/java/com/nvidia/cuvs/GPUInfo.java b/java/cuvs-java/src/main/java/com/nvidia/cuvs/GPUInfo.java index c0ccaf99c3..8d24df9aae 100644 --- a/java/cuvs-java/src/main/java/com/nvidia/cuvs/GPUInfo.java +++ b/java/cuvs-java/src/main/java/com/nvidia/cuvs/GPUInfo.java @@ -18,11 +18,21 @@ /** * Contains GPU information * - * @param gpuId id of the GPU starting from 0 - * @param name ASCII string identifying device - * @param freeMemory returned free memory in bytes - * @param totalMemory returned total memory in bytes - * @param computeCapability the compute capability of the device + * @param gpuId id of the GPU starting from 0 + * @param name ASCII string identifying device + * @param totalDeviceMemoryInBytes total device memory in bytes + * @param computeCapabilityMajor the compute capability of the device (major) + * @param computeCapabilityMinor the compute capability of the device (minor) + * @param supportsConcurrentCopy whether the device can concurrently copy memory between host and device while + * executing a kernel + * @param supportsConcurrentKernels whether the device supports executing multiple kernels within the same context + * simultaneously */ public record GPUInfo( - int gpuId, String name, long freeMemory, long totalMemory, float computeCapability) {} + int gpuId, + String name, + long totalDeviceMemoryInBytes, + int computeCapabilityMajor, + int computeCapabilityMinor, + boolean supportsConcurrentCopy, + boolean supportsConcurrentKernels) {} diff --git a/java/cuvs-java/src/main/java/com/nvidia/cuvs/GPUInfoProvider.java b/java/cuvs-java/src/main/java/com/nvidia/cuvs/GPUInfoProvider.java new file mode 100644 index 0000000000..fe03934df0 --- /dev/null +++ b/java/cuvs-java/src/main/java/com/nvidia/cuvs/GPUInfoProvider.java @@ -0,0 +1,48 @@ +/* + * Copyright (c) 2025, NVIDIA CORPORATION. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.nvidia.cuvs; + +import java.util.List; + +public interface GPUInfoProvider { + + int MIN_COMPUTE_CAPABILITY_MAJOR = 7; + int MIN_COMPUTE_CAPABILITY_MINOR = 0; + + int MIN_DEVICE_MEMORY_IN_MB = 8192; + + /** + * Gets all the available GPUs + * + * @return a list of {@link GPUInfo} objects with GPU details + */ + List availableGPUs(); + + /** + * Get the list of compatible GPUs based on compute capability >= 7.0 and total + * memory >= 8GB + * + * @return a list of compatible GPUs. See {@link GPUInfo} + */ + List compatibleGPUs(); + + /** + * Gets memory information relative to a {@link CuVSResources} + * @param resources from which to obtain memory information + * @return a {@link CuVSResourcesInfo} record containing the memory information + */ + CuVSResourcesInfo getCurrentInfo(CuVSResources resources); +} diff --git a/java/cuvs-java/src/main/java/com/nvidia/cuvs/SynchronizedCuVSResources.java b/java/cuvs-java/src/main/java/com/nvidia/cuvs/SynchronizedCuVSResources.java index bbba262bb0..6f17d27317 100644 --- a/java/cuvs-java/src/main/java/com/nvidia/cuvs/SynchronizedCuVSResources.java +++ b/java/cuvs-java/src/main/java/com/nvidia/cuvs/SynchronizedCuVSResources.java @@ -51,6 +51,11 @@ public void close() { }; } + @Override + public int deviceId() { + return inner.deviceId(); + } + @Override public void close() { inner.close(); 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 2ee3fb373a..f12389d226 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 @@ -128,6 +128,9 @@ default CagraIndex mergeCagraIndexes(CagraIndex[] indexes, CagraMergeParams merg return mergeCagraIndexes(indexes); } + /** Returns a {@link GPUInfoProvider} to query the system for GPU related information */ + GPUInfoProvider gpuInfoProvider(); + /** Retrieves the system-wide provider. */ static CuVSProvider provider() { return CuVSServiceProvider.Holder.INSTANCE; diff --git a/java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/UnsupportedProvider.java b/java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/UnsupportedProvider.java index 499ea22fc5..60509004c4 100644 --- a/java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/UnsupportedProvider.java +++ b/java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/UnsupportedProvider.java @@ -66,6 +66,11 @@ public CuVSMatrix.Builder newDeviceMatrixBuilder( throw new UnsupportedOperationException(); } + @Override + public GPUInfoProvider gpuInfoProvider() { + throw new UnsupportedOperationException(); + } + @Override public CuVSMatrix.Builder newDeviceMatrixBuilder( CuVSResources cuVSResources, diff --git a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CuVSResourcesImpl.java b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CuVSResourcesImpl.java index 8f9e1d94d0..ef1b9cdab7 100644 --- a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CuVSResourcesImpl.java +++ b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CuVSResourcesImpl.java @@ -16,9 +16,8 @@ package com.nvidia.cuvs.internal; import static com.nvidia.cuvs.internal.common.Util.checkCuVSError; -import static com.nvidia.cuvs.internal.panama.headers_h.cuvsResourcesCreate; -import static com.nvidia.cuvs.internal.panama.headers_h.cuvsResourcesDestroy; -import static com.nvidia.cuvs.internal.panama.headers_h.cuvsResources_t; +import static com.nvidia.cuvs.internal.panama.headers_h.*; +import static com.nvidia.cuvs.internal.panama.headers_h_1.C_INT; import com.nvidia.cuvs.CuVSResources; import java.lang.foreign.Arena; @@ -34,6 +33,7 @@ public class CuVSResourcesImpl implements CuVSResources { private final Path tempDirectory; private final long resourceHandle; private final ScopedAccess access; + private final int deviceId; /** * Constructor that allocates the resources needed for cuVS @@ -43,9 +43,11 @@ public CuVSResourcesImpl(Path tempDirectory) { this.tempDirectory = tempDirectory; try (var localArena = Arena.ofConfined()) { var resourcesMemorySegment = localArena.allocate(cuvsResources_t); - int returnValue = cuvsResourcesCreate(resourcesMemorySegment); - checkCuVSError(returnValue, "cuvsResourcesCreate"); + checkCuVSError(cuvsResourcesCreate(resourcesMemorySegment), "cuvsResourcesCreate"); this.resourceHandle = resourcesMemorySegment.get(cuvsResources_t, 0); + var deviceIdPtr = localArena.allocate(C_INT); + checkCuVSError(cuvsDeviceIdGet(resourceHandle, deviceIdPtr), "cuvsDeviceIdGet"); + this.deviceId = deviceIdPtr.get(C_INT, 0); this.access = new ScopedAccess() { @Override @@ -64,6 +66,11 @@ public ScopedAccess access() { return this.access; } + @Override + public int deviceId() { + return this.deviceId; + } + @Override public void close() { synchronized (this) { diff --git a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/GPUInfoProviderImpl.java b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/GPUInfoProviderImpl.java new file mode 100644 index 0000000000..c2422318ac --- /dev/null +++ b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/GPUInfoProviderImpl.java @@ -0,0 +1,121 @@ +/* + * Copyright (c) 2025, NVIDIA CORPORATION. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.nvidia.cuvs.internal; + +import static com.nvidia.cuvs.internal.common.LinkerHelper.C_INT; +import static com.nvidia.cuvs.internal.common.Util.checkCudaError; +import static com.nvidia.cuvs.internal.panama.headers_h.cudaMemGetInfo; +import static com.nvidia.cuvs.internal.panama.headers_h_1.*; + +import com.nvidia.cuvs.CuVSResources; +import com.nvidia.cuvs.CuVSResourcesInfo; +import com.nvidia.cuvs.GPUInfo; +import com.nvidia.cuvs.GPUInfoProvider; +import com.nvidia.cuvs.internal.panama.cudaDeviceProp; +import java.lang.foreign.Arena; +import java.lang.foreign.MemorySegment; +import java.util.ArrayList; +import java.util.List; + +public class GPUInfoProviderImpl implements GPUInfoProvider { + + // Lazy initialization for list of available GPUs. + private static class AvailableGpuInitializer { + + // Available GPUs are initialized only once when first accessed. + // This is assumed to be invariant for the lifetime of the program. + static final List AVAILABLE_GPUS = getAvailableGpusInfo(); + + private static List getAvailableGpusInfo() { + try (var localArena = Arena.ofConfined()) { + + MemorySegment numGpus = localArena.allocate(C_INT); + int returnValue = cudaGetDeviceCount(numGpus); + checkCudaError(returnValue, "cudaGetDeviceCount"); + + int numGpuCount = numGpus.get(C_INT, 0); + List gpuInfoArr = new ArrayList(); + + MemorySegment deviceProp = cudaDeviceProp.allocate(localArena); + + for (int i = 0; i < numGpuCount; i++) { + returnValue = cudaGetDeviceProperties_v2(deviceProp, i); + checkCudaError(returnValue, "cudaGetDeviceProperties_v2"); + + GPUInfo gpuInfo = + new GPUInfo( + i, + cudaDeviceProp.name(deviceProp).getString(0), + cudaDeviceProp.totalGlobalMem(deviceProp), + cudaDeviceProp.major(deviceProp), + cudaDeviceProp.minor(deviceProp), + cudaDeviceProp.asyncEngineCount(deviceProp) > 0, + cudaDeviceProp.concurrentKernels(deviceProp) > 0); + + gpuInfoArr.add(gpuInfo); + } + return gpuInfoArr; + } + } + } + + private static boolean hasMinimumCapability(GPUInfo gpuInfo) { + return gpuInfo.computeCapabilityMajor() > GPUInfoProvider.MIN_COMPUTE_CAPABILITY_MAJOR + || (gpuInfo.computeCapabilityMajor() == GPUInfoProvider.MIN_COMPUTE_CAPABILITY_MAJOR + && gpuInfo.computeCapabilityMinor() >= GPUInfoProvider.MIN_COMPUTE_CAPABILITY_MINOR); + } + + @Override + public List availableGPUs() { + return AvailableGpuInitializer.AVAILABLE_GPUS; + } + + @Override + public List compatibleGPUs() { + List compatibleGPUs = new ArrayList<>(); + long minDeviceMemoryInBytes = 1024L * 1024L * GPUInfoProvider.MIN_DEVICE_MEMORY_IN_MB; + for (GPUInfo gpuInfo : AvailableGpuInitializer.AVAILABLE_GPUS) { + if (hasMinimumCapability(gpuInfo) + && gpuInfo.totalDeviceMemoryInBytes() >= minDeviceMemoryInBytes) { + compatibleGPUs.add(gpuInfo); + } + } + return compatibleGPUs; + } + + @Override + public CuVSResourcesInfo getCurrentInfo(CuVSResources resources) { + try (var localArena = Arena.ofConfined()) { + var deviceIdPtr = localArena.allocate(C_INT); + checkCudaError(cudaGetDevice(deviceIdPtr), "cudaGetDevice"); + var currentDeviceId = deviceIdPtr.get(C_INT, 0); + + if (resources.deviceId() != currentDeviceId) { + checkCudaError(cudaSetDevice(resources.deviceId()), "cudaSetDevice"); + } + + MemorySegment freeMemoryPtr = localArena.allocate(size_t); + MemorySegment totalMemoryPtr = localArena.allocate(size_t); + checkCudaError(cudaMemGetInfo(freeMemoryPtr, totalMemoryPtr), "cudaMemGetInfo"); + + if (resources.deviceId() != currentDeviceId) { + checkCudaError(cudaSetDevice(currentDeviceId), "cudaSetDevice"); + } + + return new CuVSResourcesInfo(freeMemoryPtr.get(size_t, 0), totalMemoryPtr.get(size_t, 0)); + } + } +} diff --git a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/common/Util.java b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/common/Util.java index d62450ece0..09adf59309 100644 --- a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/common/Util.java +++ b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/common/Util.java @@ -23,12 +23,10 @@ import static com.nvidia.cuvs.internal.panama.headers_h_1.cudaStream_t; import com.nvidia.cuvs.CuVSResources; -import com.nvidia.cuvs.GPUInfo; import com.nvidia.cuvs.internal.panama.DLDataType; import com.nvidia.cuvs.internal.panama.DLDevice; import com.nvidia.cuvs.internal.panama.DLManagedTensor; import com.nvidia.cuvs.internal.panama.DLTensor; -import com.nvidia.cuvs.internal.panama.cudaDeviceProp; import com.nvidia.cuvs.internal.panama.headers_h; import java.lang.foreign.Arena; import java.lang.foreign.Linker; @@ -37,9 +35,7 @@ import java.lang.foreign.MemorySegment; import java.lang.invoke.MethodHandle; import java.lang.invoke.VarHandle; -import java.util.ArrayList; import java.util.BitSet; -import java.util.List; public class Util { @@ -187,85 +183,6 @@ static String getLastErrorText() { } } - /** - * Get the list of compatible GPUs based on compute capability >= 7.0 and total - * memory >= 8GB - * - * @return a list of compatible GPUs. See {@link GPUInfo} - */ - public static List compatibleGPUs() throws Throwable { - return compatibleGPUs(7.0, 8192); - } - - /** - * Get the list of compatible GPUs based on given compute capability and total - * memory - * - * @param minComputeCapability the minimum compute capability - * @param minDeviceMemoryMB the minimum total available memory in MB - * @return a list of compatible GPUs. See {@link GPUInfo} - */ - public static List compatibleGPUs(double minComputeCapability, int minDeviceMemoryMB) - throws Throwable { - List compatibleGPUs = new ArrayList(); - double minDeviceMemoryB = Math.pow(2, 20) * minDeviceMemoryMB; - for (GPUInfo gpuInfo : availableGPUs()) { - if (gpuInfo.computeCapability() >= minComputeCapability - && gpuInfo.totalMemory() >= minDeviceMemoryB) { - compatibleGPUs.add(gpuInfo); - } - } - return compatibleGPUs; - } - - /** - * Gets all the available GPUs - * - * @return a list of {@link GPUInfo} objects with GPU details - */ - public static List availableGPUs() throws Throwable { - try (var localArena = Arena.ofConfined()) { - - MemorySegment numGpus = localArena.allocate(C_INT); - int returnValue = cudaGetDeviceCount(numGpus); - checkCudaError(returnValue, "cudaGetDeviceCount"); - - int numGpuCount = numGpus.get(C_INT, 0); - List gpuInfoArr = new ArrayList(); - - MemorySegment free = localArena.allocate(size_t); - MemorySegment total = localArena.allocate(size_t); - MemorySegment deviceProp = cudaDeviceProp.allocate(localArena); - - for (int i = 0; i < numGpuCount; i++) { - - returnValue = cudaSetDevice(i); - checkCudaError(returnValue, "cudaSetDevice"); - - returnValue = cudaGetDeviceProperties_v2(deviceProp, i); - checkCudaError(returnValue, "cudaGetDeviceProperties_v2"); - - returnValue = cudaMemGetInfo(free, total); - checkCudaError(returnValue, "cudaMemGetInfo"); - - float computeCapability = - Float.parseFloat( - cudaDeviceProp.major(deviceProp) + "." + cudaDeviceProp.minor(deviceProp)); - - GPUInfo gpuInfo = - new GPUInfo( - i, - cudaDeviceProp.name(deviceProp).getString(0), - free.get(C_LONG, 0), - total.get(C_LONG, 0), - computeCapability); - - gpuInfoArr.add(gpuInfo); - } - return gpuInfoArr; - } - } - /** * A utility method for getting an instance of {@link MemorySegment} for a * {@link String}. diff --git a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/spi/JDKProvider.java b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/spi/JDKProvider.java index 01eb72ed69..0ee0cf1d10 100644 --- a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/spi/JDKProvider.java +++ b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/spi/JDKProvider.java @@ -102,6 +102,11 @@ public CagraIndex mergeCagraIndexes(CagraIndex[] indexes, CagraMergeParams merge return CagraIndexImpl.merge(indexes, mergeParams); } + @Override + public GPUInfoProvider gpuInfoProvider() { + return new GPUInfoProviderImpl(); + } + @Override public CuVSMatrix.Builder newHostMatrixBuilder( long size, long columns, CuVSMatrix.DataType dataType) throws UnsupportedOperationException { diff --git a/java/cuvs-java/src/test/java/com/nvidia/cuvs/CheckedCuVSResources.java b/java/cuvs-java/src/test/java/com/nvidia/cuvs/CheckedCuVSResources.java index a58d9c04e7..581e29973c 100644 --- a/java/cuvs-java/src/test/java/com/nvidia/cuvs/CheckedCuVSResources.java +++ b/java/cuvs-java/src/test/java/com/nvidia/cuvs/CheckedCuVSResources.java @@ -61,6 +61,11 @@ public void close() { }; } + @Override + public int deviceId() { + return inner.deviceId(); + } + @Override public void close() { destroyed = true; 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 e479a6ebda..bc012f6a0e 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 @@ -18,7 +18,6 @@ import static com.carrotsearch.randomizedtesting.RandomizedTest.assumeTrue; import static org.junit.Assert.*; -import java.lang.invoke.MethodHandles; import java.util.concurrent.ExecutionException; import java.util.concurrent.Executors; import org.junit.Before; @@ -28,7 +27,7 @@ public class CuVSResourcesIT extends CuVSTestCase { - private static final Logger log = LoggerFactory.getLogger(MethodHandles.lookup().lookupClass()); + private static final Logger log = LoggerFactory.getLogger(CuVSResourcesIT.class); @Before public void setup() { 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 new file mode 100644 index 0000000000..1284df2d31 --- /dev/null +++ b/java/cuvs-java/src/test/java/com/nvidia/cuvs/GPUInfoIT.java @@ -0,0 +1,78 @@ +/* + * Copyright (c) 2025, NVIDIA CORPORATION. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.nvidia.cuvs; + +import static com.carrotsearch.randomizedtesting.RandomizedTest.assumeTrue; +import static org.junit.Assert.*; + +import com.nvidia.cuvs.spi.CuVSProvider; +import org.junit.Before; +import org.junit.Test; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +public class GPUInfoIT extends CuVSTestCase { + + private static final Logger log = LoggerFactory.getLogger(GPUInfoIT.class); + + @Before + public void setup() { + assumeTrue("not supported on " + System.getProperty("os.name"), isLinuxAmd64()); + } + + @Test + public void testAvailableAndCompatibleGpus() throws Throwable { + var gpuInfoProvider = CuVSProvider.provider().gpuInfoProvider(); + var availableGpus = gpuInfoProvider.availableGPUs(); + var compatibleGpus = gpuInfoProvider.compatibleGPUs(); + assertFalse(availableGpus.isEmpty()); + assertTrue(availableGpus.get(0).gpuId() >= 0); + for (var gpuInfo : availableGpus) { + log.trace( + "Available GPU with name [{}], memory [{}MB], compute [{}.{}]", + gpuInfo.name(), + gpuInfo.totalDeviceMemoryInBytes() / (1024L * 1024L), + gpuInfo.computeCapabilityMajor(), + gpuInfo.computeCapabilityMinor()); + } + + assertTrue(availableGpus.size() >= compatibleGpus.size()); + log.trace("Compatible GPUs: [{}]", compatibleGpus.size()); + for (var gpuInfo : compatibleGpus) { + log.trace( + "Compatible GPU with name [{}], memory [{}MB], compute [{}.{}]", + gpuInfo.name(), + gpuInfo.totalDeviceMemoryInBytes() / (1024L * 1024L), + gpuInfo.computeCapabilityMajor(), + gpuInfo.computeCapabilityMinor()); + assertTrue(gpuInfo.computeCapabilityMajor() >= GPUInfoProvider.MIN_COMPUTE_CAPABILITY_MAJOR); + assertTrue( + gpuInfo.totalDeviceMemoryInBytes() + >= GPUInfoProvider.MIN_DEVICE_MEMORY_IN_MB * 1024L * 1024L); + } + } + + @Test + public void testMemoryInfo() throws Throwable { + try (var resources = CheckedCuVSResources.create()) { + var gpuInfoProvider = CuVSProvider.provider().gpuInfoProvider(); + var memoryInfo = gpuInfoProvider.getCurrentInfo(resources); + assertNotNull(memoryInfo); + log.trace("Free memory: {}", memoryInfo.freeDeviceMemoryInBytes()); + assertTrue(memoryInfo.freeDeviceMemoryInBytes() > 0); + } + } +} From 4f710ecfcd5a054179151655721aa6b47ec711aa Mon Sep 17 00:00:00 2001 From: Jinsol Park Date: Wed, 27 Aug 2025 12:40:49 -0700 Subject: [PATCH 46/93] Removing filter templates in ivfpq for binary size reduction (#1211) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ### Summary This PR removed filters from templates in ivfpq. Only applied to `compute_similarity_kernel` in this PR, but can be applied elsewhere (e.g. `ivfflat_interleaved_scan`) which will be a follow-up PR to this one. ### Perf check Ran the default setting with bitset flter 20 times to check if this slows anything down. All numbers are averaged across the 20 runs. Profiing done with nsys. `compute_similarity` kernel runtime Screenshot 2025-08-04 at 1 10 31 PM Think the extra time is coming from the `switch` statements inside `ivf_to_sample_filter_dev`'s `operator()`. But not too much of an overhead. ### Binary size reduction From CI reports: 1107.27 MB ->1040.44 MB Authors: - Jinsol Park (https://github.com/jinsolp) - Divye Gala (https://github.com/divyegala) Approvers: - Micka (https://github.com/lowener) - Divye Gala (https://github.com/divyegala) - Artem M. Chirkin (https://github.com/achirkin) URL: https://github.com/rapidsai/cuvs/pull/1211 --- cpp/CMakeLists.txt | 7 - cpp/include/cuvs/neighbors/common.hpp | 14 +- .../ivf_pq_compute_similarity_float_float.cu | 83 ++++----- ...compute_similarity_float_float_bitset64.cu | 73 -------- ...f_pq_compute_similarity_float_fp8_false.cu | 82 ++++----- ...ute_similarity_float_fp8_false_bitset64.cu | 73 -------- ...vf_pq_compute_similarity_float_fp8_true.cu | 82 ++++----- ...pute_similarity_float_fp8_true_bitset64.cu | 73 -------- .../ivf_pq_compute_similarity_float_half.cu | 83 ++++----- ..._compute_similarity_float_half_bitset64.cu | 73 -------- ...vf_pq_compute_similarity_half_fp8_false.cu | 82 ++++----- ...pute_similarity_half_fp8_false_bitset64.cu | 73 -------- ...ivf_pq_compute_similarity_half_fp8_true.cu | 82 ++++----- ...mpute_similarity_half_fp8_true_bitset64.cu | 73 -------- .../ivf_pq_compute_similarity_half_half.cu | 83 ++++----- ...q_compute_similarity_half_half_bitset64.cu | 73 -------- .../ivf_pq/ivf_pq_compute_similarity.cuh | 114 +++--------- .../ivf_pq/ivf_pq_compute_similarity_impl.cuh | 169 +++++++++--------- cpp/src/neighbors/ivf_pq/ivf_pq_search.cuh | 29 ++- cpp/src/neighbors/sample_filter.cuh | 55 +++++- 20 files changed, 449 insertions(+), 1027 deletions(-) delete mode 100644 cpp/src/neighbors/ivf_pq/detail/ivf_pq_compute_similarity_float_float_bitset64.cu delete mode 100644 cpp/src/neighbors/ivf_pq/detail/ivf_pq_compute_similarity_float_fp8_false_bitset64.cu delete mode 100644 cpp/src/neighbors/ivf_pq/detail/ivf_pq_compute_similarity_float_fp8_true_bitset64.cu delete mode 100644 cpp/src/neighbors/ivf_pq/detail/ivf_pq_compute_similarity_float_half_bitset64.cu delete mode 100644 cpp/src/neighbors/ivf_pq/detail/ivf_pq_compute_similarity_half_fp8_false_bitset64.cu delete mode 100644 cpp/src/neighbors/ivf_pq/detail/ivf_pq_compute_similarity_half_fp8_true_bitset64.cu delete mode 100644 cpp/src/neighbors/ivf_pq/detail/ivf_pq_compute_similarity_half_half_bitset64.cu diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt index 50ee1a0ce2..5606ff85f0 100644 --- a/cpp/CMakeLists.txt +++ b/cpp/CMakeLists.txt @@ -458,13 +458,6 @@ if(BUILD_SHARED_LIBS) src/neighbors/ivf_pq/detail/ivf_pq_compute_similarity_float_float.cu src/neighbors/ivf_pq/detail/ivf_pq_compute_similarity_float_fp8_false.cu src/neighbors/ivf_pq/detail/ivf_pq_compute_similarity_float_fp8_true.cu - src/neighbors/ivf_pq/detail/ivf_pq_compute_similarity_half_fp8_false_bitset64.cu - src/neighbors/ivf_pq/detail/ivf_pq_compute_similarity_half_fp8_true_bitset64.cu - src/neighbors/ivf_pq/detail/ivf_pq_compute_similarity_half_half_bitset64.cu - src/neighbors/ivf_pq/detail/ivf_pq_compute_similarity_float_half_bitset64.cu - src/neighbors/ivf_pq/detail/ivf_pq_compute_similarity_float_float_bitset64.cu - src/neighbors/ivf_pq/detail/ivf_pq_compute_similarity_float_fp8_false_bitset64.cu - src/neighbors/ivf_pq/detail/ivf_pq_compute_similarity_float_fp8_true_bitset64.cu src/neighbors/ivf_pq/detail/ivf_pq_contiguous_list_data.cu src/neighbors/ivf_pq/detail/ivf_pq_process_and_fill_codes.cu src/neighbors/ivf_pq/detail/ivf_pq_search_float_int64_t.cu diff --git a/cpp/include/cuvs/neighbors/common.hpp b/cpp/include/cuvs/neighbors/common.hpp index 2e91f3c15f..70da22a09d 100644 --- a/cpp/include/cuvs/neighbors/common.hpp +++ b/cpp/include/cuvs/neighbors/common.hpp @@ -485,13 +485,13 @@ namespace filtering { enum class FilterType { None, Bitmap, Bitset }; struct base_filter { - ~base_filter() = default; + virtual ~base_filter() = default; virtual FilterType get_filter_type() const = 0; }; /* A filter that filters nothing. This is the default behavior. */ struct none_sample_filter : public base_filter { - inline _RAFT_HOST_DEVICE bool operator()( + constexpr __forceinline__ _RAFT_HOST_DEVICE bool operator()( // query index const uint32_t query_ix, // the current inverted list index @@ -499,7 +499,7 @@ struct none_sample_filter : public base_filter { // the index of the current sample inside the current inverted list const uint32_t sample_ix) const; - inline _RAFT_HOST_DEVICE bool operator()( + constexpr __forceinline__ _RAFT_HOST_DEVICE bool operator()( // query index const uint32_t query_ix, // the index of the current sample @@ -517,7 +517,7 @@ struct none_sample_filter : public base_filter { * @tparam filter_t */ template -struct ivf_to_sample_filter { +struct ivf_to_sample_filter : public base_filter { const index_t* const* inds_ptrs_; const filter_t next_filter_; @@ -534,6 +534,8 @@ struct ivf_to_sample_filter { const uint32_t cluster_ix, // the index of the current sample inside the current inverted list const uint32_t sample_ix) const; + + FilterType get_filter_type() const override { return next_filter_.get_filter_type(); } }; /** @@ -577,8 +579,8 @@ struct bitset_filter : public base_filter { // View of the bitset to use as a filter const view_t bitset_view_; - bitset_filter(const view_t bitset_for_filtering); - inline _RAFT_HOST_DEVICE bool operator()( + _RAFT_HOST_DEVICE bitset_filter(const view_t bitset_for_filtering); + constexpr __forceinline__ _RAFT_HOST_DEVICE bool operator()( // query index const uint32_t query_ix, // the index of the current sample diff --git a/cpp/src/neighbors/ivf_pq/detail/ivf_pq_compute_similarity_float_float.cu b/cpp/src/neighbors/ivf_pq/detail/ivf_pq_compute_similarity_float_float.cu index 33e7067a70..a5c726b6ee 100644 --- a/cpp/src/neighbors/ivf_pq/detail/ivf_pq_compute_similarity_float_float.cu +++ b/cpp/src/neighbors/ivf_pq/detail/ivf_pq_compute_similarity_float_float.cu @@ -1,5 +1,5 @@ /* - * Copyright (c) 2024, NVIDIA CORPORATION. + * Copyright (c) 2024-2025, NVIDIA CORPORATION. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -24,50 +24,43 @@ #include "../ivf_pq_compute_similarity_impl.cuh" #include "../ivf_pq_fp_8bit.cuh" -#define instantiate_cuvs_neighbors_ivf_pq_detail_compute_similarity_select( \ - OutT, LutT, IvfSampleFilterT) \ - template auto \ - cuvs::neighbors::ivf_pq::detail::compute_similarity_select( \ - const cudaDeviceProp& dev_props, \ - bool manage_local_topk, \ - int locality_hint, \ - double preferred_shmem_carveout, \ - uint32_t pq_bits, \ - uint32_t pq_dim, \ - uint32_t precomp_data_count, \ - uint32_t n_queries, \ - uint32_t n_probes, \ - uint32_t topk) -> cuvs::neighbors::ivf_pq::detail::selected; \ - \ - template void \ - cuvs::neighbors::ivf_pq::detail::compute_similarity_run( \ - cuvs::neighbors::ivf_pq::detail::selected s, \ - rmm::cuda_stream_view stream, \ - uint32_t dim, \ - uint32_t n_probes, \ - uint32_t pq_dim, \ - uint32_t n_queries, \ - uint32_t queries_offset, \ - cuvs::distance::DistanceType metric, \ - cuvs::neighbors::ivf_pq::codebook_gen codebook_kind, \ - uint32_t topk, \ - uint32_t max_samples, \ - const float* cluster_centers, \ - const float* pq_centers, \ - const uint8_t* const* pq_dataset, \ - const uint32_t* cluster_labels, \ - const uint32_t* _chunk_indices, \ - const float* queries, \ - const uint32_t* index_list, \ - float* query_kths, \ - IvfSampleFilterT sample_filter, \ - LutT* lut_scores, \ - OutT* _out_scores, \ +#define instantiate_cuvs_neighbors_ivf_pq_detail_compute_similarity_select(OutT, LutT) \ + template auto cuvs::neighbors::ivf_pq::detail::compute_similarity_select( \ + const cudaDeviceProp& dev_props, \ + bool manage_local_topk, \ + int locality_hint, \ + double preferred_shmem_carveout, \ + uint32_t pq_bits, \ + uint32_t pq_dim, \ + uint32_t precomp_data_count, \ + uint32_t n_queries, \ + uint32_t n_probes, \ + uint32_t topk) -> cuvs::neighbors::ivf_pq::detail::selected; \ + \ + template void cuvs::neighbors::ivf_pq::detail::compute_similarity_run( \ + cuvs::neighbors::ivf_pq::detail::selected s, \ + rmm::cuda_stream_view stream, \ + uint32_t dim, \ + uint32_t n_probes, \ + uint32_t pq_dim, \ + uint32_t n_queries, \ + uint32_t queries_offset, \ + cuvs::distance::DistanceType metric, \ + cuvs::neighbors::ivf_pq::codebook_gen codebook_kind, \ + uint32_t topk, \ + uint32_t max_samples, \ + const float* cluster_centers, \ + const float* pq_centers, \ + const uint8_t* const* pq_dataset, \ + const uint32_t* cluster_labels, \ + const uint32_t* _chunk_indices, \ + const float* queries, \ + const uint32_t* index_list, \ + float* query_kths, \ + const cuvs::neighbors::filtering::base_filter& sample_filter, \ + LutT* lut_scores, \ + OutT* _out_scores, \ uint32_t* _out_indices); #define COMMA , -instantiate_cuvs_neighbors_ivf_pq_detail_compute_similarity_select( - float, - float, - cuvs::neighbors::filtering::ivf_to_sample_filter< - int64_t COMMA cuvs::neighbors::filtering::none_sample_filter>); +instantiate_cuvs_neighbors_ivf_pq_detail_compute_similarity_select(float, float); diff --git a/cpp/src/neighbors/ivf_pq/detail/ivf_pq_compute_similarity_float_float_bitset64.cu b/cpp/src/neighbors/ivf_pq/detail/ivf_pq_compute_similarity_float_float_bitset64.cu deleted file mode 100644 index 8c28e2c5b9..0000000000 --- a/cpp/src/neighbors/ivf_pq/detail/ivf_pq_compute_similarity_float_float_bitset64.cu +++ /dev/null @@ -1,73 +0,0 @@ -/* - * Copyright (c) 2024, NVIDIA CORPORATION. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -/* - * NOTE: this file is generated by generate_ivf_pq_compute_similarity.py - * Make changes there and run in this directory: - * > python generate_ivf_pq_compute_similarity.py - */ - -#include "../../sample_filter.cuh" -#include "../ivf_pq_compute_similarity_impl.cuh" -#include "../ivf_pq_fp_8bit.cuh" - -#define instantiate_cuvs_neighbors_ivf_pq_detail_compute_similarity_select( \ - OutT, LutT, IvfSampleFilterT) \ - template auto \ - cuvs::neighbors::ivf_pq::detail::compute_similarity_select( \ - const cudaDeviceProp& dev_props, \ - bool manage_local_topk, \ - int locality_hint, \ - double preferred_shmem_carveout, \ - uint32_t pq_bits, \ - uint32_t pq_dim, \ - uint32_t precomp_data_count, \ - uint32_t n_queries, \ - uint32_t n_probes, \ - uint32_t topk) -> cuvs::neighbors::ivf_pq::detail::selected; \ - \ - template void \ - cuvs::neighbors::ivf_pq::detail::compute_similarity_run( \ - cuvs::neighbors::ivf_pq::detail::selected s, \ - rmm::cuda_stream_view stream, \ - uint32_t dim, \ - uint32_t n_probes, \ - uint32_t pq_dim, \ - uint32_t n_queries, \ - uint32_t queries_offset, \ - cuvs::distance::DistanceType metric, \ - cuvs::neighbors::ivf_pq::codebook_gen codebook_kind, \ - uint32_t topk, \ - uint32_t max_samples, \ - const float* cluster_centers, \ - const float* pq_centers, \ - const uint8_t* const* pq_dataset, \ - const uint32_t* cluster_labels, \ - const uint32_t* _chunk_indices, \ - const float* queries, \ - const uint32_t* index_list, \ - float* query_kths, \ - IvfSampleFilterT sample_filter, \ - LutT* lut_scores, \ - OutT* _out_scores, \ - uint32_t* _out_indices); - -#define COMMA , -instantiate_cuvs_neighbors_ivf_pq_detail_compute_similarity_select( - float, - float, - cuvs::neighbors::filtering::ivf_to_sample_filter< - int64_t COMMA cuvs::neighbors::filtering::bitset_filter>); diff --git a/cpp/src/neighbors/ivf_pq/detail/ivf_pq_compute_similarity_float_fp8_false.cu b/cpp/src/neighbors/ivf_pq/detail/ivf_pq_compute_similarity_float_fp8_false.cu index 4cb3424cb5..b65b26fc8c 100644 --- a/cpp/src/neighbors/ivf_pq/detail/ivf_pq_compute_similarity_float_fp8_false.cu +++ b/cpp/src/neighbors/ivf_pq/detail/ivf_pq_compute_similarity_float_fp8_false.cu @@ -1,5 +1,5 @@ /* - * Copyright (c) 2024, NVIDIA CORPORATION. + * Copyright (c) 2024-2025, NVIDIA CORPORATION. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -24,50 +24,44 @@ #include "../ivf_pq_compute_similarity_impl.cuh" #include "../ivf_pq_fp_8bit.cuh" -#define instantiate_cuvs_neighbors_ivf_pq_detail_compute_similarity_select( \ - OutT, LutT, IvfSampleFilterT) \ - template auto \ - cuvs::neighbors::ivf_pq::detail::compute_similarity_select( \ - const cudaDeviceProp& dev_props, \ - bool manage_local_topk, \ - int locality_hint, \ - double preferred_shmem_carveout, \ - uint32_t pq_bits, \ - uint32_t pq_dim, \ - uint32_t precomp_data_count, \ - uint32_t n_queries, \ - uint32_t n_probes, \ - uint32_t topk) -> cuvs::neighbors::ivf_pq::detail::selected; \ - \ - template void \ - cuvs::neighbors::ivf_pq::detail::compute_similarity_run( \ - cuvs::neighbors::ivf_pq::detail::selected s, \ - rmm::cuda_stream_view stream, \ - uint32_t dim, \ - uint32_t n_probes, \ - uint32_t pq_dim, \ - uint32_t n_queries, \ - uint32_t queries_offset, \ - cuvs::distance::DistanceType metric, \ - cuvs::neighbors::ivf_pq::codebook_gen codebook_kind, \ - uint32_t topk, \ - uint32_t max_samples, \ - const float* cluster_centers, \ - const float* pq_centers, \ - const uint8_t* const* pq_dataset, \ - const uint32_t* cluster_labels, \ - const uint32_t* _chunk_indices, \ - const float* queries, \ - const uint32_t* index_list, \ - float* query_kths, \ - IvfSampleFilterT sample_filter, \ - LutT* lut_scores, \ - OutT* _out_scores, \ +#define instantiate_cuvs_neighbors_ivf_pq_detail_compute_similarity_select(OutT, LutT) \ + template auto cuvs::neighbors::ivf_pq::detail::compute_similarity_select( \ + const cudaDeviceProp& dev_props, \ + bool manage_local_topk, \ + int locality_hint, \ + double preferred_shmem_carveout, \ + uint32_t pq_bits, \ + uint32_t pq_dim, \ + uint32_t precomp_data_count, \ + uint32_t n_queries, \ + uint32_t n_probes, \ + uint32_t topk) -> cuvs::neighbors::ivf_pq::detail::selected; \ + \ + template void cuvs::neighbors::ivf_pq::detail::compute_similarity_run( \ + cuvs::neighbors::ivf_pq::detail::selected s, \ + rmm::cuda_stream_view stream, \ + uint32_t dim, \ + uint32_t n_probes, \ + uint32_t pq_dim, \ + uint32_t n_queries, \ + uint32_t queries_offset, \ + cuvs::distance::DistanceType metric, \ + cuvs::neighbors::ivf_pq::codebook_gen codebook_kind, \ + uint32_t topk, \ + uint32_t max_samples, \ + const float* cluster_centers, \ + const float* pq_centers, \ + const uint8_t* const* pq_dataset, \ + const uint32_t* cluster_labels, \ + const uint32_t* _chunk_indices, \ + const float* queries, \ + const uint32_t* index_list, \ + float* query_kths, \ + const cuvs::neighbors::filtering::base_filter& sample_filter, \ + LutT* lut_scores, \ + OutT* _out_scores, \ uint32_t* _out_indices); #define COMMA , instantiate_cuvs_neighbors_ivf_pq_detail_compute_similarity_select( - float, - cuvs::neighbors::ivf_pq::detail::fp_8bit<5u COMMA false>, - cuvs::neighbors::filtering::ivf_to_sample_filter< - int64_t COMMA cuvs::neighbors::filtering::none_sample_filter>); + float, cuvs::neighbors::ivf_pq::detail::fp_8bit<5u COMMA false>); diff --git a/cpp/src/neighbors/ivf_pq/detail/ivf_pq_compute_similarity_float_fp8_false_bitset64.cu b/cpp/src/neighbors/ivf_pq/detail/ivf_pq_compute_similarity_float_fp8_false_bitset64.cu deleted file mode 100644 index e3f330dccb..0000000000 --- a/cpp/src/neighbors/ivf_pq/detail/ivf_pq_compute_similarity_float_fp8_false_bitset64.cu +++ /dev/null @@ -1,73 +0,0 @@ -/* - * Copyright (c) 2024, NVIDIA CORPORATION. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -/* - * NOTE: this file is generated by generate_ivf_pq_compute_similarity.py - * Make changes there and run in this directory: - * > python generate_ivf_pq_compute_similarity.py - */ - -#include "../../sample_filter.cuh" -#include "../ivf_pq_compute_similarity_impl.cuh" -#include "../ivf_pq_fp_8bit.cuh" - -#define instantiate_cuvs_neighbors_ivf_pq_detail_compute_similarity_select( \ - OutT, LutT, IvfSampleFilterT) \ - template auto \ - cuvs::neighbors::ivf_pq::detail::compute_similarity_select( \ - const cudaDeviceProp& dev_props, \ - bool manage_local_topk, \ - int locality_hint, \ - double preferred_shmem_carveout, \ - uint32_t pq_bits, \ - uint32_t pq_dim, \ - uint32_t precomp_data_count, \ - uint32_t n_queries, \ - uint32_t n_probes, \ - uint32_t topk) -> cuvs::neighbors::ivf_pq::detail::selected; \ - \ - template void \ - cuvs::neighbors::ivf_pq::detail::compute_similarity_run( \ - cuvs::neighbors::ivf_pq::detail::selected s, \ - rmm::cuda_stream_view stream, \ - uint32_t dim, \ - uint32_t n_probes, \ - uint32_t pq_dim, \ - uint32_t n_queries, \ - uint32_t queries_offset, \ - cuvs::distance::DistanceType metric, \ - cuvs::neighbors::ivf_pq::codebook_gen codebook_kind, \ - uint32_t topk, \ - uint32_t max_samples, \ - const float* cluster_centers, \ - const float* pq_centers, \ - const uint8_t* const* pq_dataset, \ - const uint32_t* cluster_labels, \ - const uint32_t* _chunk_indices, \ - const float* queries, \ - const uint32_t* index_list, \ - float* query_kths, \ - IvfSampleFilterT sample_filter, \ - LutT* lut_scores, \ - OutT* _out_scores, \ - uint32_t* _out_indices); - -#define COMMA , -instantiate_cuvs_neighbors_ivf_pq_detail_compute_similarity_select( - float, - cuvs::neighbors::ivf_pq::detail::fp_8bit<5u COMMA false>, - cuvs::neighbors::filtering::ivf_to_sample_filter< - int64_t COMMA cuvs::neighbors::filtering::bitset_filter>); diff --git a/cpp/src/neighbors/ivf_pq/detail/ivf_pq_compute_similarity_float_fp8_true.cu b/cpp/src/neighbors/ivf_pq/detail/ivf_pq_compute_similarity_float_fp8_true.cu index 75f8f2cf8a..30f7eaa30c 100644 --- a/cpp/src/neighbors/ivf_pq/detail/ivf_pq_compute_similarity_float_fp8_true.cu +++ b/cpp/src/neighbors/ivf_pq/detail/ivf_pq_compute_similarity_float_fp8_true.cu @@ -1,5 +1,5 @@ /* - * Copyright (c) 2024, NVIDIA CORPORATION. + * Copyright (c) 2024-2025, NVIDIA CORPORATION. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -24,50 +24,44 @@ #include "../ivf_pq_compute_similarity_impl.cuh" #include "../ivf_pq_fp_8bit.cuh" -#define instantiate_cuvs_neighbors_ivf_pq_detail_compute_similarity_select( \ - OutT, LutT, IvfSampleFilterT) \ - template auto \ - cuvs::neighbors::ivf_pq::detail::compute_similarity_select( \ - const cudaDeviceProp& dev_props, \ - bool manage_local_topk, \ - int locality_hint, \ - double preferred_shmem_carveout, \ - uint32_t pq_bits, \ - uint32_t pq_dim, \ - uint32_t precomp_data_count, \ - uint32_t n_queries, \ - uint32_t n_probes, \ - uint32_t topk) -> cuvs::neighbors::ivf_pq::detail::selected; \ - \ - template void \ - cuvs::neighbors::ivf_pq::detail::compute_similarity_run( \ - cuvs::neighbors::ivf_pq::detail::selected s, \ - rmm::cuda_stream_view stream, \ - uint32_t dim, \ - uint32_t n_probes, \ - uint32_t pq_dim, \ - uint32_t n_queries, \ - uint32_t queries_offset, \ - cuvs::distance::DistanceType metric, \ - cuvs::neighbors::ivf_pq::codebook_gen codebook_kind, \ - uint32_t topk, \ - uint32_t max_samples, \ - const float* cluster_centers, \ - const float* pq_centers, \ - const uint8_t* const* pq_dataset, \ - const uint32_t* cluster_labels, \ - const uint32_t* _chunk_indices, \ - const float* queries, \ - const uint32_t* index_list, \ - float* query_kths, \ - IvfSampleFilterT sample_filter, \ - LutT* lut_scores, \ - OutT* _out_scores, \ +#define instantiate_cuvs_neighbors_ivf_pq_detail_compute_similarity_select(OutT, LutT) \ + template auto cuvs::neighbors::ivf_pq::detail::compute_similarity_select( \ + const cudaDeviceProp& dev_props, \ + bool manage_local_topk, \ + int locality_hint, \ + double preferred_shmem_carveout, \ + uint32_t pq_bits, \ + uint32_t pq_dim, \ + uint32_t precomp_data_count, \ + uint32_t n_queries, \ + uint32_t n_probes, \ + uint32_t topk) -> cuvs::neighbors::ivf_pq::detail::selected; \ + \ + template void cuvs::neighbors::ivf_pq::detail::compute_similarity_run( \ + cuvs::neighbors::ivf_pq::detail::selected s, \ + rmm::cuda_stream_view stream, \ + uint32_t dim, \ + uint32_t n_probes, \ + uint32_t pq_dim, \ + uint32_t n_queries, \ + uint32_t queries_offset, \ + cuvs::distance::DistanceType metric, \ + cuvs::neighbors::ivf_pq::codebook_gen codebook_kind, \ + uint32_t topk, \ + uint32_t max_samples, \ + const float* cluster_centers, \ + const float* pq_centers, \ + const uint8_t* const* pq_dataset, \ + const uint32_t* cluster_labels, \ + const uint32_t* _chunk_indices, \ + const float* queries, \ + const uint32_t* index_list, \ + float* query_kths, \ + const cuvs::neighbors::filtering::base_filter& sample_filter, \ + LutT* lut_scores, \ + OutT* _out_scores, \ uint32_t* _out_indices); #define COMMA , instantiate_cuvs_neighbors_ivf_pq_detail_compute_similarity_select( - float, - cuvs::neighbors::ivf_pq::detail::fp_8bit<5u COMMA true>, - cuvs::neighbors::filtering::ivf_to_sample_filter< - int64_t COMMA cuvs::neighbors::filtering::none_sample_filter>); + float, cuvs::neighbors::ivf_pq::detail::fp_8bit<5u COMMA true>); diff --git a/cpp/src/neighbors/ivf_pq/detail/ivf_pq_compute_similarity_float_fp8_true_bitset64.cu b/cpp/src/neighbors/ivf_pq/detail/ivf_pq_compute_similarity_float_fp8_true_bitset64.cu deleted file mode 100644 index 3d0c610dd2..0000000000 --- a/cpp/src/neighbors/ivf_pq/detail/ivf_pq_compute_similarity_float_fp8_true_bitset64.cu +++ /dev/null @@ -1,73 +0,0 @@ -/* - * Copyright (c) 2024, NVIDIA CORPORATION. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -/* - * NOTE: this file is generated by generate_ivf_pq_compute_similarity.py - * Make changes there and run in this directory: - * > python generate_ivf_pq_compute_similarity.py - */ - -#include "../../sample_filter.cuh" -#include "../ivf_pq_compute_similarity_impl.cuh" -#include "../ivf_pq_fp_8bit.cuh" - -#define instantiate_cuvs_neighbors_ivf_pq_detail_compute_similarity_select( \ - OutT, LutT, IvfSampleFilterT) \ - template auto \ - cuvs::neighbors::ivf_pq::detail::compute_similarity_select( \ - const cudaDeviceProp& dev_props, \ - bool manage_local_topk, \ - int locality_hint, \ - double preferred_shmem_carveout, \ - uint32_t pq_bits, \ - uint32_t pq_dim, \ - uint32_t precomp_data_count, \ - uint32_t n_queries, \ - uint32_t n_probes, \ - uint32_t topk) -> cuvs::neighbors::ivf_pq::detail::selected; \ - \ - template void \ - cuvs::neighbors::ivf_pq::detail::compute_similarity_run( \ - cuvs::neighbors::ivf_pq::detail::selected s, \ - rmm::cuda_stream_view stream, \ - uint32_t dim, \ - uint32_t n_probes, \ - uint32_t pq_dim, \ - uint32_t n_queries, \ - uint32_t queries_offset, \ - cuvs::distance::DistanceType metric, \ - cuvs::neighbors::ivf_pq::codebook_gen codebook_kind, \ - uint32_t topk, \ - uint32_t max_samples, \ - const float* cluster_centers, \ - const float* pq_centers, \ - const uint8_t* const* pq_dataset, \ - const uint32_t* cluster_labels, \ - const uint32_t* _chunk_indices, \ - const float* queries, \ - const uint32_t* index_list, \ - float* query_kths, \ - IvfSampleFilterT sample_filter, \ - LutT* lut_scores, \ - OutT* _out_scores, \ - uint32_t* _out_indices); - -#define COMMA , -instantiate_cuvs_neighbors_ivf_pq_detail_compute_similarity_select( - float, - cuvs::neighbors::ivf_pq::detail::fp_8bit<5u COMMA true>, - cuvs::neighbors::filtering::ivf_to_sample_filter< - int64_t COMMA cuvs::neighbors::filtering::bitset_filter>); diff --git a/cpp/src/neighbors/ivf_pq/detail/ivf_pq_compute_similarity_float_half.cu b/cpp/src/neighbors/ivf_pq/detail/ivf_pq_compute_similarity_float_half.cu index 9a288a1c7c..1545374419 100644 --- a/cpp/src/neighbors/ivf_pq/detail/ivf_pq_compute_similarity_float_half.cu +++ b/cpp/src/neighbors/ivf_pq/detail/ivf_pq_compute_similarity_float_half.cu @@ -1,5 +1,5 @@ /* - * Copyright (c) 2024, NVIDIA CORPORATION. + * Copyright (c) 2024-2025, NVIDIA CORPORATION. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -24,50 +24,43 @@ #include "../ivf_pq_compute_similarity_impl.cuh" #include "../ivf_pq_fp_8bit.cuh" -#define instantiate_cuvs_neighbors_ivf_pq_detail_compute_similarity_select( \ - OutT, LutT, IvfSampleFilterT) \ - template auto \ - cuvs::neighbors::ivf_pq::detail::compute_similarity_select( \ - const cudaDeviceProp& dev_props, \ - bool manage_local_topk, \ - int locality_hint, \ - double preferred_shmem_carveout, \ - uint32_t pq_bits, \ - uint32_t pq_dim, \ - uint32_t precomp_data_count, \ - uint32_t n_queries, \ - uint32_t n_probes, \ - uint32_t topk) -> cuvs::neighbors::ivf_pq::detail::selected; \ - \ - template void \ - cuvs::neighbors::ivf_pq::detail::compute_similarity_run( \ - cuvs::neighbors::ivf_pq::detail::selected s, \ - rmm::cuda_stream_view stream, \ - uint32_t dim, \ - uint32_t n_probes, \ - uint32_t pq_dim, \ - uint32_t n_queries, \ - uint32_t queries_offset, \ - cuvs::distance::DistanceType metric, \ - cuvs::neighbors::ivf_pq::codebook_gen codebook_kind, \ - uint32_t topk, \ - uint32_t max_samples, \ - const float* cluster_centers, \ - const float* pq_centers, \ - const uint8_t* const* pq_dataset, \ - const uint32_t* cluster_labels, \ - const uint32_t* _chunk_indices, \ - const float* queries, \ - const uint32_t* index_list, \ - float* query_kths, \ - IvfSampleFilterT sample_filter, \ - LutT* lut_scores, \ - OutT* _out_scores, \ +#define instantiate_cuvs_neighbors_ivf_pq_detail_compute_similarity_select(OutT, LutT) \ + template auto cuvs::neighbors::ivf_pq::detail::compute_similarity_select( \ + const cudaDeviceProp& dev_props, \ + bool manage_local_topk, \ + int locality_hint, \ + double preferred_shmem_carveout, \ + uint32_t pq_bits, \ + uint32_t pq_dim, \ + uint32_t precomp_data_count, \ + uint32_t n_queries, \ + uint32_t n_probes, \ + uint32_t topk) -> cuvs::neighbors::ivf_pq::detail::selected; \ + \ + template void cuvs::neighbors::ivf_pq::detail::compute_similarity_run( \ + cuvs::neighbors::ivf_pq::detail::selected s, \ + rmm::cuda_stream_view stream, \ + uint32_t dim, \ + uint32_t n_probes, \ + uint32_t pq_dim, \ + uint32_t n_queries, \ + uint32_t queries_offset, \ + cuvs::distance::DistanceType metric, \ + cuvs::neighbors::ivf_pq::codebook_gen codebook_kind, \ + uint32_t topk, \ + uint32_t max_samples, \ + const float* cluster_centers, \ + const float* pq_centers, \ + const uint8_t* const* pq_dataset, \ + const uint32_t* cluster_labels, \ + const uint32_t* _chunk_indices, \ + const float* queries, \ + const uint32_t* index_list, \ + float* query_kths, \ + const cuvs::neighbors::filtering::base_filter& sample_filter, \ + LutT* lut_scores, \ + OutT* _out_scores, \ uint32_t* _out_indices); #define COMMA , -instantiate_cuvs_neighbors_ivf_pq_detail_compute_similarity_select( - float, - half, - cuvs::neighbors::filtering::ivf_to_sample_filter< - int64_t COMMA cuvs::neighbors::filtering::none_sample_filter>); +instantiate_cuvs_neighbors_ivf_pq_detail_compute_similarity_select(float, half); diff --git a/cpp/src/neighbors/ivf_pq/detail/ivf_pq_compute_similarity_float_half_bitset64.cu b/cpp/src/neighbors/ivf_pq/detail/ivf_pq_compute_similarity_float_half_bitset64.cu deleted file mode 100644 index 07af7815e2..0000000000 --- a/cpp/src/neighbors/ivf_pq/detail/ivf_pq_compute_similarity_float_half_bitset64.cu +++ /dev/null @@ -1,73 +0,0 @@ -/* - * Copyright (c) 2024, NVIDIA CORPORATION. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -/* - * NOTE: this file is generated by generate_ivf_pq_compute_similarity.py - * Make changes there and run in this directory: - * > python generate_ivf_pq_compute_similarity.py - */ - -#include "../../sample_filter.cuh" -#include "../ivf_pq_compute_similarity_impl.cuh" -#include "../ivf_pq_fp_8bit.cuh" - -#define instantiate_cuvs_neighbors_ivf_pq_detail_compute_similarity_select( \ - OutT, LutT, IvfSampleFilterT) \ - template auto \ - cuvs::neighbors::ivf_pq::detail::compute_similarity_select( \ - const cudaDeviceProp& dev_props, \ - bool manage_local_topk, \ - int locality_hint, \ - double preferred_shmem_carveout, \ - uint32_t pq_bits, \ - uint32_t pq_dim, \ - uint32_t precomp_data_count, \ - uint32_t n_queries, \ - uint32_t n_probes, \ - uint32_t topk) -> cuvs::neighbors::ivf_pq::detail::selected; \ - \ - template void \ - cuvs::neighbors::ivf_pq::detail::compute_similarity_run( \ - cuvs::neighbors::ivf_pq::detail::selected s, \ - rmm::cuda_stream_view stream, \ - uint32_t dim, \ - uint32_t n_probes, \ - uint32_t pq_dim, \ - uint32_t n_queries, \ - uint32_t queries_offset, \ - cuvs::distance::DistanceType metric, \ - cuvs::neighbors::ivf_pq::codebook_gen codebook_kind, \ - uint32_t topk, \ - uint32_t max_samples, \ - const float* cluster_centers, \ - const float* pq_centers, \ - const uint8_t* const* pq_dataset, \ - const uint32_t* cluster_labels, \ - const uint32_t* _chunk_indices, \ - const float* queries, \ - const uint32_t* index_list, \ - float* query_kths, \ - IvfSampleFilterT sample_filter, \ - LutT* lut_scores, \ - OutT* _out_scores, \ - uint32_t* _out_indices); - -#define COMMA , -instantiate_cuvs_neighbors_ivf_pq_detail_compute_similarity_select( - float, - half, - cuvs::neighbors::filtering::ivf_to_sample_filter< - int64_t COMMA cuvs::neighbors::filtering::bitset_filter>); diff --git a/cpp/src/neighbors/ivf_pq/detail/ivf_pq_compute_similarity_half_fp8_false.cu b/cpp/src/neighbors/ivf_pq/detail/ivf_pq_compute_similarity_half_fp8_false.cu index 8886b01bef..00707e5ecc 100644 --- a/cpp/src/neighbors/ivf_pq/detail/ivf_pq_compute_similarity_half_fp8_false.cu +++ b/cpp/src/neighbors/ivf_pq/detail/ivf_pq_compute_similarity_half_fp8_false.cu @@ -1,5 +1,5 @@ /* - * Copyright (c) 2024, NVIDIA CORPORATION. + * Copyright (c) 2024-2025, NVIDIA CORPORATION. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -24,50 +24,44 @@ #include "../ivf_pq_compute_similarity_impl.cuh" #include "../ivf_pq_fp_8bit.cuh" -#define instantiate_cuvs_neighbors_ivf_pq_detail_compute_similarity_select( \ - OutT, LutT, IvfSampleFilterT) \ - template auto \ - cuvs::neighbors::ivf_pq::detail::compute_similarity_select( \ - const cudaDeviceProp& dev_props, \ - bool manage_local_topk, \ - int locality_hint, \ - double preferred_shmem_carveout, \ - uint32_t pq_bits, \ - uint32_t pq_dim, \ - uint32_t precomp_data_count, \ - uint32_t n_queries, \ - uint32_t n_probes, \ - uint32_t topk) -> cuvs::neighbors::ivf_pq::detail::selected; \ - \ - template void \ - cuvs::neighbors::ivf_pq::detail::compute_similarity_run( \ - cuvs::neighbors::ivf_pq::detail::selected s, \ - rmm::cuda_stream_view stream, \ - uint32_t dim, \ - uint32_t n_probes, \ - uint32_t pq_dim, \ - uint32_t n_queries, \ - uint32_t queries_offset, \ - cuvs::distance::DistanceType metric, \ - cuvs::neighbors::ivf_pq::codebook_gen codebook_kind, \ - uint32_t topk, \ - uint32_t max_samples, \ - const float* cluster_centers, \ - const float* pq_centers, \ - const uint8_t* const* pq_dataset, \ - const uint32_t* cluster_labels, \ - const uint32_t* _chunk_indices, \ - const float* queries, \ - const uint32_t* index_list, \ - float* query_kths, \ - IvfSampleFilterT sample_filter, \ - LutT* lut_scores, \ - OutT* _out_scores, \ +#define instantiate_cuvs_neighbors_ivf_pq_detail_compute_similarity_select(OutT, LutT) \ + template auto cuvs::neighbors::ivf_pq::detail::compute_similarity_select( \ + const cudaDeviceProp& dev_props, \ + bool manage_local_topk, \ + int locality_hint, \ + double preferred_shmem_carveout, \ + uint32_t pq_bits, \ + uint32_t pq_dim, \ + uint32_t precomp_data_count, \ + uint32_t n_queries, \ + uint32_t n_probes, \ + uint32_t topk) -> cuvs::neighbors::ivf_pq::detail::selected; \ + \ + template void cuvs::neighbors::ivf_pq::detail::compute_similarity_run( \ + cuvs::neighbors::ivf_pq::detail::selected s, \ + rmm::cuda_stream_view stream, \ + uint32_t dim, \ + uint32_t n_probes, \ + uint32_t pq_dim, \ + uint32_t n_queries, \ + uint32_t queries_offset, \ + cuvs::distance::DistanceType metric, \ + cuvs::neighbors::ivf_pq::codebook_gen codebook_kind, \ + uint32_t topk, \ + uint32_t max_samples, \ + const float* cluster_centers, \ + const float* pq_centers, \ + const uint8_t* const* pq_dataset, \ + const uint32_t* cluster_labels, \ + const uint32_t* _chunk_indices, \ + const float* queries, \ + const uint32_t* index_list, \ + float* query_kths, \ + const cuvs::neighbors::filtering::base_filter& sample_filter, \ + LutT* lut_scores, \ + OutT* _out_scores, \ uint32_t* _out_indices); #define COMMA , instantiate_cuvs_neighbors_ivf_pq_detail_compute_similarity_select( - half, - cuvs::neighbors::ivf_pq::detail::fp_8bit<5u COMMA false>, - cuvs::neighbors::filtering::ivf_to_sample_filter< - int64_t COMMA cuvs::neighbors::filtering::none_sample_filter>); + half, cuvs::neighbors::ivf_pq::detail::fp_8bit<5u COMMA false>); diff --git a/cpp/src/neighbors/ivf_pq/detail/ivf_pq_compute_similarity_half_fp8_false_bitset64.cu b/cpp/src/neighbors/ivf_pq/detail/ivf_pq_compute_similarity_half_fp8_false_bitset64.cu deleted file mode 100644 index a921545f80..0000000000 --- a/cpp/src/neighbors/ivf_pq/detail/ivf_pq_compute_similarity_half_fp8_false_bitset64.cu +++ /dev/null @@ -1,73 +0,0 @@ -/* - * Copyright (c) 2024, NVIDIA CORPORATION. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -/* - * NOTE: this file is generated by generate_ivf_pq_compute_similarity.py - * Make changes there and run in this directory: - * > python generate_ivf_pq_compute_similarity.py - */ - -#include "../../sample_filter.cuh" -#include "../ivf_pq_compute_similarity_impl.cuh" -#include "../ivf_pq_fp_8bit.cuh" - -#define instantiate_cuvs_neighbors_ivf_pq_detail_compute_similarity_select( \ - OutT, LutT, IvfSampleFilterT) \ - template auto \ - cuvs::neighbors::ivf_pq::detail::compute_similarity_select( \ - const cudaDeviceProp& dev_props, \ - bool manage_local_topk, \ - int locality_hint, \ - double preferred_shmem_carveout, \ - uint32_t pq_bits, \ - uint32_t pq_dim, \ - uint32_t precomp_data_count, \ - uint32_t n_queries, \ - uint32_t n_probes, \ - uint32_t topk) -> cuvs::neighbors::ivf_pq::detail::selected; \ - \ - template void \ - cuvs::neighbors::ivf_pq::detail::compute_similarity_run( \ - cuvs::neighbors::ivf_pq::detail::selected s, \ - rmm::cuda_stream_view stream, \ - uint32_t dim, \ - uint32_t n_probes, \ - uint32_t pq_dim, \ - uint32_t n_queries, \ - uint32_t queries_offset, \ - cuvs::distance::DistanceType metric, \ - cuvs::neighbors::ivf_pq::codebook_gen codebook_kind, \ - uint32_t topk, \ - uint32_t max_samples, \ - const float* cluster_centers, \ - const float* pq_centers, \ - const uint8_t* const* pq_dataset, \ - const uint32_t* cluster_labels, \ - const uint32_t* _chunk_indices, \ - const float* queries, \ - const uint32_t* index_list, \ - float* query_kths, \ - IvfSampleFilterT sample_filter, \ - LutT* lut_scores, \ - OutT* _out_scores, \ - uint32_t* _out_indices); - -#define COMMA , -instantiate_cuvs_neighbors_ivf_pq_detail_compute_similarity_select( - half, - cuvs::neighbors::ivf_pq::detail::fp_8bit<5u COMMA false>, - cuvs::neighbors::filtering::ivf_to_sample_filter< - int64_t COMMA cuvs::neighbors::filtering::bitset_filter>); diff --git a/cpp/src/neighbors/ivf_pq/detail/ivf_pq_compute_similarity_half_fp8_true.cu b/cpp/src/neighbors/ivf_pq/detail/ivf_pq_compute_similarity_half_fp8_true.cu index 05241135b5..7de8719957 100644 --- a/cpp/src/neighbors/ivf_pq/detail/ivf_pq_compute_similarity_half_fp8_true.cu +++ b/cpp/src/neighbors/ivf_pq/detail/ivf_pq_compute_similarity_half_fp8_true.cu @@ -1,5 +1,5 @@ /* - * Copyright (c) 2024, NVIDIA CORPORATION. + * Copyright (c) 2024-2025, NVIDIA CORPORATION. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -24,50 +24,44 @@ #include "../ivf_pq_compute_similarity_impl.cuh" #include "../ivf_pq_fp_8bit.cuh" -#define instantiate_cuvs_neighbors_ivf_pq_detail_compute_similarity_select( \ - OutT, LutT, IvfSampleFilterT) \ - template auto \ - cuvs::neighbors::ivf_pq::detail::compute_similarity_select( \ - const cudaDeviceProp& dev_props, \ - bool manage_local_topk, \ - int locality_hint, \ - double preferred_shmem_carveout, \ - uint32_t pq_bits, \ - uint32_t pq_dim, \ - uint32_t precomp_data_count, \ - uint32_t n_queries, \ - uint32_t n_probes, \ - uint32_t topk) -> cuvs::neighbors::ivf_pq::detail::selected; \ - \ - template void \ - cuvs::neighbors::ivf_pq::detail::compute_similarity_run( \ - cuvs::neighbors::ivf_pq::detail::selected s, \ - rmm::cuda_stream_view stream, \ - uint32_t dim, \ - uint32_t n_probes, \ - uint32_t pq_dim, \ - uint32_t n_queries, \ - uint32_t queries_offset, \ - cuvs::distance::DistanceType metric, \ - cuvs::neighbors::ivf_pq::codebook_gen codebook_kind, \ - uint32_t topk, \ - uint32_t max_samples, \ - const float* cluster_centers, \ - const float* pq_centers, \ - const uint8_t* const* pq_dataset, \ - const uint32_t* cluster_labels, \ - const uint32_t* _chunk_indices, \ - const float* queries, \ - const uint32_t* index_list, \ - float* query_kths, \ - IvfSampleFilterT sample_filter, \ - LutT* lut_scores, \ - OutT* _out_scores, \ +#define instantiate_cuvs_neighbors_ivf_pq_detail_compute_similarity_select(OutT, LutT) \ + template auto cuvs::neighbors::ivf_pq::detail::compute_similarity_select( \ + const cudaDeviceProp& dev_props, \ + bool manage_local_topk, \ + int locality_hint, \ + double preferred_shmem_carveout, \ + uint32_t pq_bits, \ + uint32_t pq_dim, \ + uint32_t precomp_data_count, \ + uint32_t n_queries, \ + uint32_t n_probes, \ + uint32_t topk) -> cuvs::neighbors::ivf_pq::detail::selected; \ + \ + template void cuvs::neighbors::ivf_pq::detail::compute_similarity_run( \ + cuvs::neighbors::ivf_pq::detail::selected s, \ + rmm::cuda_stream_view stream, \ + uint32_t dim, \ + uint32_t n_probes, \ + uint32_t pq_dim, \ + uint32_t n_queries, \ + uint32_t queries_offset, \ + cuvs::distance::DistanceType metric, \ + cuvs::neighbors::ivf_pq::codebook_gen codebook_kind, \ + uint32_t topk, \ + uint32_t max_samples, \ + const float* cluster_centers, \ + const float* pq_centers, \ + const uint8_t* const* pq_dataset, \ + const uint32_t* cluster_labels, \ + const uint32_t* _chunk_indices, \ + const float* queries, \ + const uint32_t* index_list, \ + float* query_kths, \ + const cuvs::neighbors::filtering::base_filter& sample_filter, \ + LutT* lut_scores, \ + OutT* _out_scores, \ uint32_t* _out_indices); #define COMMA , instantiate_cuvs_neighbors_ivf_pq_detail_compute_similarity_select( - half, - cuvs::neighbors::ivf_pq::detail::fp_8bit<5u COMMA true>, - cuvs::neighbors::filtering::ivf_to_sample_filter< - int64_t COMMA cuvs::neighbors::filtering::none_sample_filter>); + half, cuvs::neighbors::ivf_pq::detail::fp_8bit<5u COMMA true>); diff --git a/cpp/src/neighbors/ivf_pq/detail/ivf_pq_compute_similarity_half_fp8_true_bitset64.cu b/cpp/src/neighbors/ivf_pq/detail/ivf_pq_compute_similarity_half_fp8_true_bitset64.cu deleted file mode 100644 index 2ebd3d5270..0000000000 --- a/cpp/src/neighbors/ivf_pq/detail/ivf_pq_compute_similarity_half_fp8_true_bitset64.cu +++ /dev/null @@ -1,73 +0,0 @@ -/* - * Copyright (c) 2024, NVIDIA CORPORATION. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -/* - * NOTE: this file is generated by generate_ivf_pq_compute_similarity.py - * Make changes there and run in this directory: - * > python generate_ivf_pq_compute_similarity.py - */ - -#include "../../sample_filter.cuh" -#include "../ivf_pq_compute_similarity_impl.cuh" -#include "../ivf_pq_fp_8bit.cuh" - -#define instantiate_cuvs_neighbors_ivf_pq_detail_compute_similarity_select( \ - OutT, LutT, IvfSampleFilterT) \ - template auto \ - cuvs::neighbors::ivf_pq::detail::compute_similarity_select( \ - const cudaDeviceProp& dev_props, \ - bool manage_local_topk, \ - int locality_hint, \ - double preferred_shmem_carveout, \ - uint32_t pq_bits, \ - uint32_t pq_dim, \ - uint32_t precomp_data_count, \ - uint32_t n_queries, \ - uint32_t n_probes, \ - uint32_t topk) -> cuvs::neighbors::ivf_pq::detail::selected; \ - \ - template void \ - cuvs::neighbors::ivf_pq::detail::compute_similarity_run( \ - cuvs::neighbors::ivf_pq::detail::selected s, \ - rmm::cuda_stream_view stream, \ - uint32_t dim, \ - uint32_t n_probes, \ - uint32_t pq_dim, \ - uint32_t n_queries, \ - uint32_t queries_offset, \ - cuvs::distance::DistanceType metric, \ - cuvs::neighbors::ivf_pq::codebook_gen codebook_kind, \ - uint32_t topk, \ - uint32_t max_samples, \ - const float* cluster_centers, \ - const float* pq_centers, \ - const uint8_t* const* pq_dataset, \ - const uint32_t* cluster_labels, \ - const uint32_t* _chunk_indices, \ - const float* queries, \ - const uint32_t* index_list, \ - float* query_kths, \ - IvfSampleFilterT sample_filter, \ - LutT* lut_scores, \ - OutT* _out_scores, \ - uint32_t* _out_indices); - -#define COMMA , -instantiate_cuvs_neighbors_ivf_pq_detail_compute_similarity_select( - half, - cuvs::neighbors::ivf_pq::detail::fp_8bit<5u COMMA true>, - cuvs::neighbors::filtering::ivf_to_sample_filter< - int64_t COMMA cuvs::neighbors::filtering::bitset_filter>); diff --git a/cpp/src/neighbors/ivf_pq/detail/ivf_pq_compute_similarity_half_half.cu b/cpp/src/neighbors/ivf_pq/detail/ivf_pq_compute_similarity_half_half.cu index 95a6f0bedd..94406c844a 100644 --- a/cpp/src/neighbors/ivf_pq/detail/ivf_pq_compute_similarity_half_half.cu +++ b/cpp/src/neighbors/ivf_pq/detail/ivf_pq_compute_similarity_half_half.cu @@ -1,5 +1,5 @@ /* - * Copyright (c) 2024, NVIDIA CORPORATION. + * Copyright (c) 2024-2025, NVIDIA CORPORATION. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -24,50 +24,43 @@ #include "../ivf_pq_compute_similarity_impl.cuh" #include "../ivf_pq_fp_8bit.cuh" -#define instantiate_cuvs_neighbors_ivf_pq_detail_compute_similarity_select( \ - OutT, LutT, IvfSampleFilterT) \ - template auto \ - cuvs::neighbors::ivf_pq::detail::compute_similarity_select( \ - const cudaDeviceProp& dev_props, \ - bool manage_local_topk, \ - int locality_hint, \ - double preferred_shmem_carveout, \ - uint32_t pq_bits, \ - uint32_t pq_dim, \ - uint32_t precomp_data_count, \ - uint32_t n_queries, \ - uint32_t n_probes, \ - uint32_t topk) -> cuvs::neighbors::ivf_pq::detail::selected; \ - \ - template void \ - cuvs::neighbors::ivf_pq::detail::compute_similarity_run( \ - cuvs::neighbors::ivf_pq::detail::selected s, \ - rmm::cuda_stream_view stream, \ - uint32_t dim, \ - uint32_t n_probes, \ - uint32_t pq_dim, \ - uint32_t n_queries, \ - uint32_t queries_offset, \ - cuvs::distance::DistanceType metric, \ - cuvs::neighbors::ivf_pq::codebook_gen codebook_kind, \ - uint32_t topk, \ - uint32_t max_samples, \ - const float* cluster_centers, \ - const float* pq_centers, \ - const uint8_t* const* pq_dataset, \ - const uint32_t* cluster_labels, \ - const uint32_t* _chunk_indices, \ - const float* queries, \ - const uint32_t* index_list, \ - float* query_kths, \ - IvfSampleFilterT sample_filter, \ - LutT* lut_scores, \ - OutT* _out_scores, \ +#define instantiate_cuvs_neighbors_ivf_pq_detail_compute_similarity_select(OutT, LutT) \ + template auto cuvs::neighbors::ivf_pq::detail::compute_similarity_select( \ + const cudaDeviceProp& dev_props, \ + bool manage_local_topk, \ + int locality_hint, \ + double preferred_shmem_carveout, \ + uint32_t pq_bits, \ + uint32_t pq_dim, \ + uint32_t precomp_data_count, \ + uint32_t n_queries, \ + uint32_t n_probes, \ + uint32_t topk) -> cuvs::neighbors::ivf_pq::detail::selected; \ + \ + template void cuvs::neighbors::ivf_pq::detail::compute_similarity_run( \ + cuvs::neighbors::ivf_pq::detail::selected s, \ + rmm::cuda_stream_view stream, \ + uint32_t dim, \ + uint32_t n_probes, \ + uint32_t pq_dim, \ + uint32_t n_queries, \ + uint32_t queries_offset, \ + cuvs::distance::DistanceType metric, \ + cuvs::neighbors::ivf_pq::codebook_gen codebook_kind, \ + uint32_t topk, \ + uint32_t max_samples, \ + const float* cluster_centers, \ + const float* pq_centers, \ + const uint8_t* const* pq_dataset, \ + const uint32_t* cluster_labels, \ + const uint32_t* _chunk_indices, \ + const float* queries, \ + const uint32_t* index_list, \ + float* query_kths, \ + const cuvs::neighbors::filtering::base_filter& sample_filter, \ + LutT* lut_scores, \ + OutT* _out_scores, \ uint32_t* _out_indices); #define COMMA , -instantiate_cuvs_neighbors_ivf_pq_detail_compute_similarity_select( - half, - half, - cuvs::neighbors::filtering::ivf_to_sample_filter< - int64_t COMMA cuvs::neighbors::filtering::none_sample_filter>); +instantiate_cuvs_neighbors_ivf_pq_detail_compute_similarity_select(half, half); diff --git a/cpp/src/neighbors/ivf_pq/detail/ivf_pq_compute_similarity_half_half_bitset64.cu b/cpp/src/neighbors/ivf_pq/detail/ivf_pq_compute_similarity_half_half_bitset64.cu deleted file mode 100644 index a0be0b392f..0000000000 --- a/cpp/src/neighbors/ivf_pq/detail/ivf_pq_compute_similarity_half_half_bitset64.cu +++ /dev/null @@ -1,73 +0,0 @@ -/* - * Copyright (c) 2024, NVIDIA CORPORATION. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -/* - * NOTE: this file is generated by generate_ivf_pq_compute_similarity.py - * Make changes there and run in this directory: - * > python generate_ivf_pq_compute_similarity.py - */ - -#include "../../sample_filter.cuh" -#include "../ivf_pq_compute_similarity_impl.cuh" -#include "../ivf_pq_fp_8bit.cuh" - -#define instantiate_cuvs_neighbors_ivf_pq_detail_compute_similarity_select( \ - OutT, LutT, IvfSampleFilterT) \ - template auto \ - cuvs::neighbors::ivf_pq::detail::compute_similarity_select( \ - const cudaDeviceProp& dev_props, \ - bool manage_local_topk, \ - int locality_hint, \ - double preferred_shmem_carveout, \ - uint32_t pq_bits, \ - uint32_t pq_dim, \ - uint32_t precomp_data_count, \ - uint32_t n_queries, \ - uint32_t n_probes, \ - uint32_t topk) -> cuvs::neighbors::ivf_pq::detail::selected; \ - \ - template void \ - cuvs::neighbors::ivf_pq::detail::compute_similarity_run( \ - cuvs::neighbors::ivf_pq::detail::selected s, \ - rmm::cuda_stream_view stream, \ - uint32_t dim, \ - uint32_t n_probes, \ - uint32_t pq_dim, \ - uint32_t n_queries, \ - uint32_t queries_offset, \ - cuvs::distance::DistanceType metric, \ - cuvs::neighbors::ivf_pq::codebook_gen codebook_kind, \ - uint32_t topk, \ - uint32_t max_samples, \ - const float* cluster_centers, \ - const float* pq_centers, \ - const uint8_t* const* pq_dataset, \ - const uint32_t* cluster_labels, \ - const uint32_t* _chunk_indices, \ - const float* queries, \ - const uint32_t* index_list, \ - float* query_kths, \ - IvfSampleFilterT sample_filter, \ - LutT* lut_scores, \ - OutT* _out_scores, \ - uint32_t* _out_indices); - -#define COMMA , -instantiate_cuvs_neighbors_ivf_pq_detail_compute_similarity_select( - half, - half, - cuvs::neighbors::filtering::ivf_to_sample_filter< - int64_t COMMA cuvs::neighbors::filtering::bitset_filter>); 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 235f0aacb3..6f9d97dbf1 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 @@ /* - * Copyright (c) 2022-2024, NVIDIA CORPORATION. + * Copyright (c) 2022-2025, NVIDIA CORPORATION. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -19,7 +19,8 @@ #include "../sample_filter.cuh" // none_sample_filter #include "ivf_pq_fp_8bit.cuh" // cuvs::neighbors::ivf_pq::detail::fp_8bit -#include // cuvs::distance::DistanceType +#include // cuvs::distance::DistanceType +#include #include // cuvs::neighbors::ivf_pq::codebook_gen #include // RAFT_WEAK_FUNCTION #include // rmm::cuda_stream_view @@ -36,7 +37,6 @@ auto RAFT_WEAK_FUNCTION is_local_topk_feasible(uint32_t k, uint32_t n_probes, ui template +template using compute_similarity_kernel_t = - decltype(&compute_similarity_kernel); + decltype(&compute_similarity_kernel); -template +template struct selected { - compute_similarity_kernel_t kernel; + compute_similarity_kernel_t kernel; dim3 grid_dim; dim3 block_dim; size_t smem_size; size_t device_lut_size; }; -template -void compute_similarity_run(selected s, +template +void compute_similarity_run(selected s, rmm::cuda_stream_view stream, uint32_t dim, uint32_t n_probes, @@ -97,7 +97,7 @@ void compute_similarity_run(selected s, const float* queries, const uint32_t* index_list, float* query_kths, - IvfSampleFilterT sample_filter, + const filtering::base_filter& sample_filter, LutT* lut_scores, OutT* _out_scores, uint32_t* _out_indices); @@ -116,7 +116,7 @@ void compute_similarity_run(selected s, * beyond this limit do not consider increasing the number of active blocks per SM * would improve locality anymore. */ -template +template auto compute_similarity_select(const cudaDeviceProp& dev_props, bool manage_local_topk, int locality_hint, @@ -126,14 +126,12 @@ auto compute_similarity_select(const cudaDeviceProp& dev_props, uint32_t precomp_data_count, uint32_t n_queries, uint32_t n_probes, - uint32_t topk) -> selected; + uint32_t topk) -> selected; } // namespace cuvs::neighbors::ivf_pq::detail -#define instantiate_cuvs_neighbors_ivf_pq_detail_compute_similarity_select( \ - OutT, LutT, IvfSampleFilterT) \ - extern template auto \ - cuvs::neighbors::ivf_pq::detail::compute_similarity_select( \ +#define instantiate_cuvs_neighbors_ivf_pq_detail_compute_similarity_select(OutT, LutT) \ + extern template auto cuvs::neighbors::ivf_pq::detail::compute_similarity_select( \ const cudaDeviceProp& dev_props, \ bool manage_local_topk, \ int locality_hint, \ @@ -143,11 +141,10 @@ auto compute_similarity_select(const cudaDeviceProp& dev_props, uint32_t precomp_data_count, \ uint32_t n_queries, \ uint32_t n_probes, \ - uint32_t topk) -> cuvs::neighbors::ivf_pq::detail::selected; \ + uint32_t topk) -> cuvs::neighbors::ivf_pq::detail::selected; \ \ - extern template void \ - cuvs::neighbors::ivf_pq::detail::compute_similarity_run( \ - cuvs::neighbors::ivf_pq::detail::selected s, \ + extern template void cuvs::neighbors::ivf_pq::detail::compute_similarity_run( \ + cuvs::neighbors::ivf_pq::detail::selected s, \ rmm::cuda_stream_view stream, \ uint32_t dim, \ uint32_t n_probes, \ @@ -166,82 +163,23 @@ auto compute_similarity_select(const cudaDeviceProp& dev_props, const float* queries, \ const uint32_t* index_list, \ float* query_kths, \ - IvfSampleFilterT sample_filter, \ + const filtering::base_filter& sample_filter, \ LutT* lut_scores, \ OutT* _out_scores, \ uint32_t* _out_indices); #define COMMA , instantiate_cuvs_neighbors_ivf_pq_detail_compute_similarity_select( - half, - cuvs::neighbors::ivf_pq::detail::fp_8bit<5u COMMA false>, - cuvs::neighbors::filtering::ivf_to_sample_filter< - int64_t COMMA cuvs::neighbors::filtering::none_sample_filter>); + half, cuvs::neighbors::ivf_pq::detail::fp_8bit<5u COMMA false>); instantiate_cuvs_neighbors_ivf_pq_detail_compute_similarity_select( - half, - cuvs::neighbors::ivf_pq::detail::fp_8bit<5u COMMA true>, - cuvs::neighbors::filtering::ivf_to_sample_filter< - int64_t COMMA cuvs::neighbors::filtering::none_sample_filter>); + half, cuvs::neighbors::ivf_pq::detail::fp_8bit<5u COMMA true>); +instantiate_cuvs_neighbors_ivf_pq_detail_compute_similarity_select(half, half); +instantiate_cuvs_neighbors_ivf_pq_detail_compute_similarity_select(float, half); +instantiate_cuvs_neighbors_ivf_pq_detail_compute_similarity_select(float, float); instantiate_cuvs_neighbors_ivf_pq_detail_compute_similarity_select( - half, - half, - cuvs::neighbors::filtering::ivf_to_sample_filter< - int64_t COMMA cuvs::neighbors::filtering::none_sample_filter>); + float, cuvs::neighbors::ivf_pq::detail::fp_8bit<5u COMMA false>); instantiate_cuvs_neighbors_ivf_pq_detail_compute_similarity_select( - float, - half, - cuvs::neighbors::filtering::ivf_to_sample_filter< - int64_t COMMA cuvs::neighbors::filtering::none_sample_filter>); -instantiate_cuvs_neighbors_ivf_pq_detail_compute_similarity_select( - float, - float, - cuvs::neighbors::filtering::ivf_to_sample_filter< - int64_t COMMA cuvs::neighbors::filtering::none_sample_filter>); -instantiate_cuvs_neighbors_ivf_pq_detail_compute_similarity_select( - float, - cuvs::neighbors::ivf_pq::detail::fp_8bit<5u COMMA false>, - cuvs::neighbors::filtering::ivf_to_sample_filter< - int64_t COMMA cuvs::neighbors::filtering::none_sample_filter>); -instantiate_cuvs_neighbors_ivf_pq_detail_compute_similarity_select( - float, - cuvs::neighbors::ivf_pq::detail::fp_8bit<5u COMMA true>, - cuvs::neighbors::filtering::ivf_to_sample_filter< - int64_t COMMA cuvs::neighbors::filtering::none_sample_filter>); -instantiate_cuvs_neighbors_ivf_pq_detail_compute_similarity_select( - half, - cuvs::neighbors::ivf_pq::detail::fp_8bit<5u COMMA false>, - cuvs::neighbors::filtering::ivf_to_sample_filter< - int64_t COMMA cuvs::neighbors::filtering::bitset_filter>); -instantiate_cuvs_neighbors_ivf_pq_detail_compute_similarity_select( - half, - cuvs::neighbors::ivf_pq::detail::fp_8bit<5u COMMA true>, - cuvs::neighbors::filtering::ivf_to_sample_filter< - int64_t COMMA cuvs::neighbors::filtering::bitset_filter>); -instantiate_cuvs_neighbors_ivf_pq_detail_compute_similarity_select( - half, - half, - cuvs::neighbors::filtering::ivf_to_sample_filter< - int64_t COMMA cuvs::neighbors::filtering::bitset_filter>); -instantiate_cuvs_neighbors_ivf_pq_detail_compute_similarity_select( - float, - half, - cuvs::neighbors::filtering::ivf_to_sample_filter< - int64_t COMMA cuvs::neighbors::filtering::bitset_filter>); -instantiate_cuvs_neighbors_ivf_pq_detail_compute_similarity_select( - float, - float, - cuvs::neighbors::filtering::ivf_to_sample_filter< - int64_t COMMA cuvs::neighbors::filtering::bitset_filter>); -instantiate_cuvs_neighbors_ivf_pq_detail_compute_similarity_select( - float, - cuvs::neighbors::ivf_pq::detail::fp_8bit<5u COMMA false>, - cuvs::neighbors::filtering::ivf_to_sample_filter< - int64_t COMMA cuvs::neighbors::filtering::bitset_filter>); -instantiate_cuvs_neighbors_ivf_pq_detail_compute_similarity_select( - float, - cuvs::neighbors::ivf_pq::detail::fp_8bit<5u COMMA true>, - cuvs::neighbors::filtering::ivf_to_sample_filter< - int64_t COMMA cuvs::neighbors::filtering::bitset_filter>); + float, cuvs::neighbors::ivf_pq::detail::fp_8bit<5u COMMA true>); #undef COMMA 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 fbbdd06c24..2bb5b342e8 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 @@ -1,5 +1,5 @@ /* - * Copyright (c) 2022-2024, NVIDIA CORPORATION. + * Copyright (c) 2022-2025, NVIDIA CORPORATION. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -16,9 +16,10 @@ #pragma once -#include "../ivf_common.cuh" // dummy_block_sort_t -#include "../sample_filter.cuh" // none_sample_filter -#include // cuvs::distance::DistanceType +#include "../ivf_common.cuh" // dummy_block_sort_t +#include "../sample_filter.cuh" // none_sample_filter +#include // cuvs::distance::DistanceType +#include #include // codebook_gen #include // matrix::detail::select::warpsort::warp_sort_distributed #include // RAFT_CUDA_TRY @@ -247,8 +248,7 @@ __device__ auto ivfpq_compute_score(uint32_t pq_dim, * query_kths keep the current state of the filtering - atomically updated distances to the * k-th closest neighbors for each query [n_queries]. * @param sample_filter - * A filter that selects samples for a given query. Use an instance of none_sample_filter to - * provide a green light for every sample. + * A filter that selects samples for a given query. * @param lut_scores * The device pointer for storing the lookup table globally [gridDim.x, pq_dim << PqBits]. * Ignored when `EnableSMemLut == true`. @@ -264,7 +264,6 @@ __device__ auto ivfpq_compute_score(uint32_t pq_dim, */ template +template using compute_similarity_kernel_t = - decltype(&compute_similarity_kernel); + decltype(&compute_similarity_kernel); // The config struct lifts the runtime parameters to the template parameters -template +template struct compute_similarity_kernel_config { public: - static auto get(uint32_t pq_bits, uint32_t k_max) - -> compute_similarity_kernel_t + static auto get(uint32_t pq_bits, uint32_t k_max) -> compute_similarity_kernel_t { return kernel_choose_bits(pq_bits, k_max); } private: static auto kernel_choose_bits(uint32_t pq_bits, uint32_t k_max) - -> compute_similarity_kernel_t + -> compute_similarity_kernel_t { switch (pq_bits) { case 4: return kernel_try_capacity<4, kMaxCapacity>(k_max); @@ -549,8 +541,7 @@ struct compute_similarity_kernel_config { } template - static auto kernel_try_capacity(uint32_t k_max) - -> compute_similarity_kernel_t + static auto kernel_try_capacity(uint32_t k_max) -> compute_similarity_kernel_t { if constexpr (Capacity > 0) { if (k_max == 0 || k_max > Capacity) { return kernel_try_capacity(k_max); } @@ -558,36 +549,23 @@ struct compute_similarity_kernel_config { if constexpr (Capacity > 1) { if (k_max * 2 <= Capacity) { return kernel_try_capacity(k_max); } } - return compute_similarity_kernel; + return compute_similarity_kernel; } }; // A standalone accessor function was necessary to make sure template // instantiation work correctly. This accessor function is not used anymore and // may be removed. -template +template auto get_compute_similarity_kernel(uint32_t pq_bits, uint32_t k_max) - -> compute_similarity_kernel_t + -> compute_similarity_kernel_t { - return compute_similarity_kernel_config::get(pq_bits, k_max); + return compute_similarity_kernel_config::get(pq_bits, + k_max); } /** Estimate the occupancy for the given kernel on the given device. */ -template +template struct occupancy_t { using shmem_unit = raft::Pow2<128>; @@ -598,7 +576,7 @@ struct occupancy_t { inline occupancy_t() = default; inline occupancy_t(size_t smem, uint32_t n_threads, - compute_similarity_kernel_t kernel, + compute_similarity_kernel_t kernel, const cudaDeviceProp& dev_props) { RAFT_CUDA_TRY( @@ -609,19 +587,17 @@ struct occupancy_t { } }; -template +template struct selected { - compute_similarity_kernel_t kernel; + compute_similarity_kernel_t kernel; dim3 grid_dim; dim3 block_dim; size_t smem_size; size_t device_lut_size; }; -template -void compute_similarity_run(selected s, +template +void compute_similarity_run(selected s, rmm::cuda_stream_view stream, uint32_t dim, uint32_t n_probes, @@ -640,33 +616,63 @@ void compute_similarity_run(selected s, const float* queries, const uint32_t* index_list, float* query_kths, - IvfSampleFilterT sample_filter, + const filtering::base_filter& sample_filter_ref, LutT* lut_scores, OutT* _out_scores, uint32_t* _out_indices) { - s.kernel<<>>(dim, - n_probes, - pq_dim, - n_queries, - queries_offset, - metric, - codebook_kind, - topk, - max_samples, - cluster_centers, - pq_centers, - pq_dataset, - cluster_labels, - _chunk_indices, - queries, - index_list, - query_kths, - sample_filter, - lut_scores, - _out_scores, - _out_indices); - RAFT_CHECK_CUDA(stream); + auto launch_kernel = [&](filtering::ivf_filter_dev sample_filter) { + s.kernel<<>>(dim, + n_probes, + pq_dim, + n_queries, + queries_offset, + metric, + codebook_kind, + topk, + max_samples, + cluster_centers, + pq_centers, + pq_dataset, + cluster_labels, + _chunk_indices, + queries, + index_list, + query_kths, + sample_filter, + lut_scores, + _out_scores, + _out_indices); + RAFT_CHECK_CUDA(stream); + }; + + switch (sample_filter_ref.get_filter_type()) { + case filtering::FilterType::None: { + try { + auto& typed_sample_filter = dynamic_cast< + const filtering::ivf_to_sample_filter&>( + sample_filter_ref); + filtering::ivf_filter_dev sample_filter{filtering::none_filter_args_t{}}; + launch_kernel(sample_filter); + } catch (const std::bad_cast& e) { + } + break; + } + case filtering::FilterType::Bitset: { + try { + auto& typed_sample_filter = dynamic_cast< + const filtering::ivf_to_sample_filter>&>( + sample_filter_ref); + filtering::ivf_filter_dev sample_filter{filtering::bitset_filter_args_t{ + typed_sample_filter.inds_ptrs_, typed_sample_filter.next_filter_.bitset_view_}}; + launch_kernel(sample_filter); + } catch (const std::bad_cast& e) { + } + break; + } + default: RAFT_FAIL("Unsupported filter type"); + } } /** @@ -683,9 +689,7 @@ void compute_similarity_run(selected s, * beyond this limit do not consider increasing the number of active blocks per SM * would improve locality anymore. */ -template +template auto compute_similarity_select(const cudaDeviceProp& dev_props, bool manage_local_topk, int locality_hint, @@ -695,7 +699,7 @@ auto compute_similarity_select(const cudaDeviceProp& dev_props, uint32_t precomp_data_count, uint32_t n_queries, uint32_t n_probes, - uint32_t topk) -> selected + uint32_t topk) -> selected { // Shared memory for storing the lookup table size_t lut_mem = sizeof(LutT) * (pq_dim << pq_bits); @@ -796,9 +800,9 @@ auto compute_similarity_select(const cudaDeviceProp& dev_props, the minimum number of blocks (just one, really). Then, we tweak the `n_threads` to further optimize occupancy and data locality for the L1 cache. */ - auto conf_fast = get_compute_similarity_kernel; - auto conf_no_basediff = get_compute_similarity_kernel; - auto conf_no_smem_lut = get_compute_similarity_kernel; + auto conf_fast = get_compute_similarity_kernel; + auto conf_no_basediff = get_compute_similarity_kernel; + auto conf_no_smem_lut = get_compute_similarity_kernel; auto topk_or_zero = manage_local_topk ? topk : 0u; std::array candidates{ std::make_tuple(conf_fast(pq_bits, topk_or_zero), @@ -814,8 +818,8 @@ auto compute_similarity_select(const cudaDeviceProp& dev_props, // we may allow slightly lower than 100% occupancy; constexpr double kTargetOccupancy = 0.75; // This struct is used to select the better candidate - occupancy_t selected_perf{}; - selected selected_config; + occupancy_t selected_perf{}; + selected selected_config; for (auto [kernel, smem_size_f, lut_is_in_shmem] : candidates) { if (smem_size_f(raft::WarpSize) > dev_props.sharedMemPerBlockOptin) { // Even a single block cannot fit into an SM due to shmem requirements. Skip the candidate. @@ -852,7 +856,7 @@ auto compute_similarity_select(const cudaDeviceProp& dev_props, continue; } - occupancy_t cur(smem_size, n_threads, kernel, dev_props); + occupancy_t cur(smem_size, n_threads, kernel, dev_props); if (cur.blocks_per_sm <= 0) { // For some reason, we still cannot make this kernel run. Skip the candidate. continue; @@ -867,8 +871,7 @@ auto compute_similarity_select(const cudaDeviceProp& dev_props, if (n_threads_tmp < n_threads) { while (n_threads_tmp >= n_threads_min) { auto smem_size_tmp = smem_size_f(n_threads_tmp); - occupancy_t tmp( - smem_size_tmp, n_threads_tmp, kernel, dev_props); + occupancy_t tmp(smem_size_tmp, n_threads_tmp, kernel, dev_props); bool select_it = false; if (lut_is_in_shmem && locality_hint >= tmp.blocks_per_sm) { // Normally, the smaller the block the better for L1 cache hit rate. diff --git a/cpp/src/neighbors/ivf_pq/ivf_pq_search.cuh b/cpp/src/neighbors/ivf_pq/ivf_pq_search.cuh index 8d9e606671..15ab026cee 100644 --- a/cpp/src/neighbors/ivf_pq/ivf_pq_search.cuh +++ b/cpp/src/neighbors/ivf_pq/ivf_pq_search.cuh @@ -417,10 +417,7 @@ template struct search_kernel_cache { /** Number of matmul invocations to cache. */ static constexpr size_t kDefaultSize = 100; - raft::cache::lru, - selected> + raft::cache::lru, selected> value{kDefaultSize}; }; @@ -557,7 +554,7 @@ void ivfpq_search_worker(raft::resources const& handle, } break; } - selected search_instance; + selected search_instance; search_kernel_key search_key{manage_local_topk, coresidency, preferred_shmem_carveout, @@ -571,17 +568,17 @@ void ivfpq_search_worker(raft::resources const& handle, raft::resource::get_custom_resource>(handle) ->value; if (!cache.get(search_key, &search_instance)) { - search_instance = compute_similarity_select( - raft::resource::get_device_properties(handle), - manage_local_topk, - coresidency, - preferred_shmem_carveout, - index.pq_bits(), - index.pq_dim(), - precomp_data_count, - n_queries, - n_probes, - topK); + search_instance = + compute_similarity_select(raft::resource::get_device_properties(handle), + manage_local_topk, + coresidency, + preferred_shmem_carveout, + index.pq_bits(), + index.pq_dim(), + precomp_data_count, + n_queries, + n_probes, + topK); cache.set(search_key, search_instance); } diff --git a/cpp/src/neighbors/sample_filter.cuh b/cpp/src/neighbors/sample_filter.cuh index b0c61f924b..4df4794206 100644 --- a/cpp/src/neighbors/sample_filter.cuh +++ b/cpp/src/neighbors/sample_filter.cuh @@ -1,5 +1,5 @@ /* - * Copyright (c) 2023-2024, NVIDIA CORPORATION. + * Copyright (c) 2023-2025, NVIDIA CORPORATION. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -28,7 +28,7 @@ namespace cuvs::neighbors::filtering { /* A filter that filters nothing. This is the default behavior. */ -inline _RAFT_HOST_DEVICE bool none_sample_filter::operator()( +constexpr __forceinline__ _RAFT_HOST_DEVICE bool none_sample_filter::operator()( // query index const uint32_t query_ix, // the current inverted list index @@ -40,7 +40,7 @@ inline _RAFT_HOST_DEVICE bool none_sample_filter::operator()( } /* A filter that filters nothing. This is the default behavior. */ -inline _RAFT_HOST_DEVICE bool none_sample_filter::operator()( +constexpr __forceinline__ _RAFT_HOST_DEVICE bool none_sample_filter::operator()( // query index const uint32_t query_ix, // the index of the current sample @@ -93,14 +93,14 @@ inline _RAFT_HOST_DEVICE bool ivf_to_sample_filter::operator( } template -bitset_filter::bitset_filter( +_RAFT_HOST_DEVICE bitset_filter::bitset_filter( const cuvs::core::bitset_view bitset_for_filtering) : bitset_view_{bitset_for_filtering} { } template -inline _RAFT_HOST_DEVICE bool bitset_filter::operator()( +constexpr __forceinline__ _RAFT_HOST_DEVICE bool bitset_filter::operator()( // query index const uint32_t query_ix, // the index of the current sample @@ -140,4 +140,49 @@ void bitmap_filter::to_csr(raft::resources const& handle, csr raft::sparse::convert::bitmap_to_csr(handle, bitmap_view_, csr); } +struct none_filter_args_t {}; +using bitset_filter_args_t = + std::tuple>; + +struct ivf_filter_dev { + filtering::FilterType tag_; + + union ivf_filter_dev_args_variant { + none_filter_args_t none_filter_args; + bitset_filter_args_t bitset_filter_args; + + _RAFT_HOST_DEVICE explicit ivf_filter_dev_args_variant(const none_filter_args_t& args) + : none_filter_args(args) + { + } + + _RAFT_HOST_DEVICE explicit ivf_filter_dev_args_variant(const bitset_filter_args_t& args) + : bitset_filter_args(args) + { + } + } args_; + + _RAFT_HOST_DEVICE ivf_filter_dev(none_filter_args_t args = {}) + : tag_(FilterType::None), args_(args) {}; + + _RAFT_HOST_DEVICE ivf_filter_dev(bitset_filter_args_t args) + : tag_(FilterType::Bitset), args_(args) {}; + + constexpr __forceinline__ _RAFT_HOST_DEVICE bool operator()(const uint32_t query_ix, + const uint32_t cluster_ix, + const uint32_t sample_ix) + { + switch (tag_) { + case FilterType::None: + return filtering::none_sample_filter{}(query_ix, cluster_ix, sample_ix); + case FilterType::Bitset: { + auto& [inds_ptrs_, bitset_view_] = args_.bitset_filter_args; + return filtering::bitset_filter(bitset_view_)( + query_ix, inds_ptrs_[cluster_ix][sample_ix]); + } + default: return true; + } + } +}; + } // namespace cuvs::neighbors::filtering From e3a2d0c1b6967c6b990af7dfb36b165745a3c2b2 Mon Sep 17 00:00:00 2001 From: Anupam <54245698+aamijar@users.noreply.github.com> Date: Wed, 27 Aug 2025 19:06:21 -0700 Subject: [PATCH 47/93] Port `raft::neighbors::epsilon_neighborhood` to cuvs (#1294) Resolves https://github.com/rapidsai/cuvs/issues/1284 Authors: - Anupam (https://github.com/aamijar) Approvers: - Divye Gala (https://github.com/divyegala) URL: https://github.com/rapidsai/cuvs/pull/1294 --- cpp/CMakeLists.txt | 1 + .../cuvs/neighbors/epsilon_neighborhood.hpp | 80 ++++ .../neighbors/detail/epsilon_neighborhood.cuh | 237 ++++++++++ cpp/src/neighbors/epsilon_neighborhood.cu | 99 ++++ cpp/tests/CMakeLists.txt | 5 + cpp/tests/neighbors/epsilon_neighborhood.cu | 433 ++++++++++++++++++ docs/source/cpp_api/neighbors.rst | 1 + .../neighbors_epsilon_neighborhood.rst | 20 + 8 files changed, 876 insertions(+) create mode 100644 cpp/include/cuvs/neighbors/epsilon_neighborhood.hpp create mode 100644 cpp/src/neighbors/detail/epsilon_neighborhood.cuh create mode 100644 cpp/src/neighbors/epsilon_neighborhood.cu create mode 100644 cpp/tests/neighbors/epsilon_neighborhood.cu create mode 100644 docs/source/cpp_api/neighbors_epsilon_neighborhood.rst diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt index 5606ff85f0..7830004c8b 100644 --- a/cpp/CMakeLists.txt +++ b/cpp/CMakeLists.txt @@ -383,6 +383,7 @@ if(BUILD_SHARED_LIBS) src/embed/spectral.cu src/neighbors/all_neighbors/all_neighbors.cu src/neighbors/ball_cover.cu + src/neighbors/epsilon_neighborhood.cu src/neighbors/ball_cover/detail/ball_cover/registers_eps_pass_euclidean.cu src/neighbors/ball_cover/detail/ball_cover/registers_pass_one.cu src/neighbors/ball_cover/detail/ball_cover/registers_pass_two.cu diff --git a/cpp/include/cuvs/neighbors/epsilon_neighborhood.hpp b/cpp/include/cuvs/neighbors/epsilon_neighborhood.hpp new file mode 100644 index 0000000000..ef884193e4 --- /dev/null +++ b/cpp/include/cuvs/neighbors/epsilon_neighborhood.hpp @@ -0,0 +1,80 @@ +/* + * Copyright (c) 2025, NVIDIA CORPORATION. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include +#include +#include + +namespace cuvs::neighbors::epsilon_neighborhood { + +/** + * @defgroup epsilon_neighborhood_cpp_l2 Epsilon Neighborhood L2 Operations + * @{ + */ + +/** + * @brief Computes epsilon neighborhood for the given distance metric and ball size. + * The epsilon neighbors is represented by a dense boolean adjacency matrix of size m * n and + * an array of degrees for each vertex, which can be used as a compressed sparse row (CSR) + * indptr array. + * + * Currently, only L2Unexpanded (L2-squared) distance metric is supported. Other metrics will + * throw an exception. + * + * @code{.cpp} + * #include + * #include + * #include + * using namespace cuvs::neighbors; + * raft::resources handle; + * ... + * auto x = raft::make_device_matrix(handle, m, k); + * auto y = raft::make_device_matrix(handle, n, k); + * auto adj = raft::make_device_matrix(handle, m, n); + * auto vd = raft::make_device_vector(handle, m + 1); + * epsilon_neighborhood::compute(handle, x.view(), y.view(), adj.view(), vd.view(), + * eps, cuvs::distance::DistanceType::L2Unexpanded); + * @endcode + * + * @tparam value_t IO and math type + * @tparam idx_t Index type + * @tparam matrix_idx_t matrix indexing type + * + * @param[in] handle raft handle to manage library resources + * @param[in] x first matrix [row-major] [on device] [dim = m x k] + * @param[in] y second matrix [row-major] [on device] [dim = n x k] + * @param[out] adj adjacency matrix [row-major] [on device] [dim = m x n] + * @param[out] vd vertex degree array [on device] [len = m + 1] + * `vd + m` stores the total number of edges in the adjacency + * matrix. Pass a nullptr if you don't need this info. + * @param[in] eps defines epsilon neighborhood radius (should be passed as + * squared when using L2Unexpanded metric) + * @param[in] metric distance metric to use. Currently only L2Unexpanded is supported. + */ +template +void compute(raft::resources const& handle, + raft::device_matrix_view x, + raft::device_matrix_view y, + raft::device_matrix_view adj, + raft::device_vector_view vd, + value_t eps, + cuvs::distance::DistanceType metric = cuvs::distance::DistanceType::L2Unexpanded); + +/** @} */ // end group epsilon_neighborhood_cpp_l2 + +} // namespace cuvs::neighbors::epsilon_neighborhood diff --git a/cpp/src/neighbors/detail/epsilon_neighborhood.cuh b/cpp/src/neighbors/detail/epsilon_neighborhood.cuh new file mode 100644 index 0000000000..d1ba86e1b2 --- /dev/null +++ b/cpp/src/neighbors/detail/epsilon_neighborhood.cuh @@ -0,0 +1,237 @@ +/* + * Copyright (c) 2025, NVIDIA CORPORATION. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include +#include +#include + +namespace cuvs::neighbors::epsilon_neighborhood::detail { + +template > +struct EpsUnexpL2SqNeighborhood : public BaseClass { + private: + typedef Policy P; + + bool* adj; + DataT eps; + IdxT* vd; + + char* smem; // for final reductions + + DataT acc[P::AccRowsPerTh][P::AccColsPerTh]; + + public: + DI EpsUnexpL2SqNeighborhood(bool* _adj, + IdxT* _vd, + const DataT* _x, + const DataT* _y, + IdxT _m, + IdxT _n, + IdxT _k, + DataT _eps, + char* _smem) + : BaseClass(_x, _y, _m, _n, _k, _smem), adj(_adj), eps(_eps), vd(_vd), smem(_smem) + { + } + + DI void run() + { + prolog(); + loop(); + epilog(); + } + + private: + DI void prolog() + { + this->ldgXY(IdxT(blockIdx.x) * P::Mblk, IdxT(blockIdx.y) * P::Nblk, 0); +#pragma unroll + for (int i = 0; i < P::AccRowsPerTh; ++i) { +#pragma unroll + for (int j = 0; j < P::AccColsPerTh; ++j) { + acc[i][j] = BaseClass::Zero(); + } + } + this->stsXY(); + __syncthreads(); + this->switch_write_buffer(); + } + + DI void loop() + { + for (int kidx = P::Kblk; kidx < this->k; kidx += P::Kblk) { + this->ldgXY(IdxT(blockIdx.x) * P::Mblk, IdxT(blockIdx.y) * P::Nblk, kidx); + accumulate(); // on the previous k-block + this->stsXY(); + __syncthreads(); + this->switch_write_buffer(); + this->switch_read_buffer(); + } + accumulate(); // last iteration + } + + DI void epilog() + { + IdxT startx = blockIdx.x * P::Mblk + this->accrowid; + IdxT starty = blockIdx.y * P::Nblk + this->acccolid; + auto lid = raft::laneId(); + IdxT sums[P::AccRowsPerTh]; +#pragma unroll + for (int i = 0; i < P::AccRowsPerTh; ++i) { + sums[i] = 0; + auto xid = startx + i * P::AccThRows; +#pragma unroll + for (int j = 0; j < P::AccColsPerTh; ++j) { + auto yid = starty + j * P::AccThCols; + auto is_neigh = acc[i][j] <= eps; + ///@todo: fix uncoalesced writes using shared mem + if (xid < this->m && yid < this->n) { + adj[xid * this->n + yid] = is_neigh; + sums[i] += is_neigh; + } + } + } + // perform reduction of adjacency values to compute vertex degrees + if (vd != nullptr) { updateVertexDegree(sums); } + } + + DI void accumulate() + { +#pragma unroll + for (int ki = 0; ki < P::Kblk; ki += P::Veclen) { + this->ldsXY(ki); +#pragma unroll + for (int i = 0; i < P::AccRowsPerTh; ++i) { +#pragma unroll + for (int j = 0; j < P::AccColsPerTh; ++j) { +#pragma unroll + for (int v = 0; v < P::Veclen; ++v) { + auto diff = this->regx[i][v] - this->regy[j][v]; + acc[i][j] += diff * diff; + } + } + } + } + } + + DI void updateVertexDegree(IdxT (&sums)[P::AccRowsPerTh]) + { + __syncthreads(); // so that we can safely reuse smem + int gid = this->accrowid; + int lid = this->acccolid; + auto cidx = IdxT(blockIdx.x) * P::Mblk + gid; + IdxT totalSum = 0; + // update the individual vertex degrees +#pragma unroll + for (int i = 0; i < P::AccRowsPerTh; ++i) { + // P::AccThCols neighboring threads need to reduce + // -> we have P::Nblk/P::AccThCols individual reductions + auto cid = cidx + i * P::AccThRows; + sums[i] = raft::logicalWarpReduce(sums[i], raft::add_op()); + if (lid == 0 && cid < this->m) { + atomicUpdate(cid, sums[i]); + totalSum += sums[i]; + } + __syncthreads(); // for safe smem reuse + } + // update the total edge count + totalSum = raft::blockReduce(totalSum, smem); + if (threadIdx.x == 0) { atomicUpdate(this->m, totalSum); } + } + + DI void atomicUpdate(IdxT addrId, IdxT val) + { + if (sizeof(IdxT) == 4) { + raft::myAtomicAdd((unsigned*)(vd + addrId), val); + } else if (sizeof(IdxT) == 8) { + raft::myAtomicAdd((unsigned long long*)(vd + addrId), val); + } + } +}; // struct EpsUnexpL2SqNeighborhood + +template +__launch_bounds__(Policy::Nthreads, 2) RAFT_KERNEL epsUnexpL2SqNeighKernel( + bool* adj, IdxT* vd, const DataT* x, const DataT* y, IdxT m, IdxT n, IdxT k, DataT eps) +{ + extern __shared__ char smem[]; + EpsUnexpL2SqNeighborhood obj(adj, vd, x, y, m, n, k, eps, smem); + obj.run(); +} + +template +void epsUnexpL2SqNeighImpl(bool* adj, + IdxT* vd, + const DataT* x, + const DataT* y, + IdxT m, + IdxT n, + IdxT k, + DataT eps, + cudaStream_t stream) +{ + typedef typename raft::linalg::Policy4x4::Policy Policy; + dim3 grid(raft::ceildiv(m, Policy::Mblk), raft::ceildiv(n, Policy::Nblk)); + dim3 blk(Policy::Nthreads); + epsUnexpL2SqNeighKernel + <<>>(adj, vd, x, y, m, n, k, eps); + RAFT_CUDA_TRY(cudaGetLastError()); +} + +/** + * @brief Computes epsilon neighborhood for the L2-Squared distance metric + * + * @tparam DataT IO and math type + * @tparam IdxT Index type + * + * @param[out] adj adjacency matrix [row-major] [on device] [dim = m x n] + * @param[out] vd vertex degree array [on device] [len = m + 1] + * `vd + m` stores the total number of edges in the adjacency + * matrix. Pass a nullptr if you don't need this info. + * @param[in] x first matrix [row-major] [on device] [dim = m x k] + * @param[in] y second matrix [row-major] [on device] [dim = n x k] + * @param[in] eps defines epsilon neighborhood radius (should be passed as + * squared as we compute L2-squared distance in this method) + * @param[in] fop device lambda to do any other custom functions + * @param[in] stream cuda stream + */ +template +void epsUnexpL2SqNeighborhood(bool* adj, + IdxT* vd, + const DataT* x, + const DataT* y, + IdxT m, + IdxT n, + IdxT k, + DataT eps, + cudaStream_t stream) +{ + if (vd != nullptr) { RAFT_CUDA_TRY(cudaMemsetAsync(vd, 0, (m + 1) * sizeof(IdxT), stream)); } + + size_t bytes = sizeof(DataT) * k; + if (16 % sizeof(DataT) == 0 && bytes % 16 == 0) { + epsUnexpL2SqNeighImpl(adj, vd, x, y, m, n, k, eps, stream); + } else if (8 % sizeof(DataT) == 0 && bytes % 8 == 0) { + epsUnexpL2SqNeighImpl(adj, vd, x, y, m, n, k, eps, stream); + } else { + epsUnexpL2SqNeighImpl(adj, vd, x, y, m, n, k, eps, stream); + } +} +} // namespace cuvs::neighbors::epsilon_neighborhood::detail diff --git a/cpp/src/neighbors/epsilon_neighborhood.cu b/cpp/src/neighbors/epsilon_neighborhood.cu new file mode 100644 index 0000000000..bfe2c93000 --- /dev/null +++ b/cpp/src/neighbors/epsilon_neighborhood.cu @@ -0,0 +1,99 @@ +/* + * Copyright (c) 2025, NVIDIA CORPORATION. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "detail/epsilon_neighborhood.cuh" +#include +#include + +namespace cuvs::neighbors::epsilon_neighborhood { + +template +void epsUnexpL2SqNeighborhood(bool* adj, + idx_t* vd, + const value_t* x, + const value_t* y, + idx_t m, + idx_t n, + idx_t k, + value_t eps, + cudaStream_t stream) +{ + detail::epsUnexpL2SqNeighborhood(adj, vd, x, y, m, n, k, eps, stream); +} + +template +void compute(raft::resources const& handle, + raft::device_matrix_view x, + raft::device_matrix_view y, + raft::device_matrix_view adj, + raft::device_vector_view vd, + value_t eps, + cuvs::distance::DistanceType metric) +{ + // Currently only L2Unexpanded metric is supported + RAFT_EXPECTS(metric == cuvs::distance::DistanceType::L2Unexpanded, + "Currently only L2Unexpanded distance metric is supported. " + "Other metrics will be supported in future versions."); + + epsUnexpL2SqNeighborhood(adj.data_handle(), + vd.data_handle(), + x.data_handle(), + y.data_handle(), + x.extent(0), + y.extent(0), + x.extent(1), + eps, + raft::resource::get_cuda_stream(handle)); +} + +// Explicit template instantiations +template void compute( + raft::resources const& handle, + raft::device_matrix_view x, + raft::device_matrix_view y, + raft::device_matrix_view adj, + raft::device_vector_view vd, + float eps, + cuvs::distance::DistanceType metric); + +template void compute( + raft::resources const& handle, + raft::device_matrix_view x, + raft::device_matrix_view y, + raft::device_matrix_view adj, + raft::device_vector_view vd, + float eps, + cuvs::distance::DistanceType metric); + +template void compute( + raft::resources const& handle, + raft::device_matrix_view x, + raft::device_matrix_view y, + raft::device_matrix_view adj, + raft::device_vector_view vd, + double eps, + cuvs::distance::DistanceType metric); + +template void compute( + raft::resources const& handle, + raft::device_matrix_view x, + raft::device_matrix_view y, + raft::device_matrix_view adj, + raft::device_vector_view vd, + double eps, + cuvs::distance::DistanceType metric); + +} // namespace cuvs::neighbors::epsilon_neighborhood diff --git a/cpp/tests/CMakeLists.txt b/cpp/tests/CMakeLists.txt index a266b577ed..ac9db70b17 100644 --- a/cpp/tests/CMakeLists.txt +++ b/cpp/tests/CMakeLists.txt @@ -215,6 +215,11 @@ if(BUILD_TESTS) ConfigureTest(NAME NEIGHBORS_BALL_COVER_TEST PATH neighbors/ball_cover.cu GPUS 1 PERCENT 100) + ConfigureTest( + NAME NEIGHBORS_EPSILON_NEIGHBORHOOD_TEST PATH neighbors/epsilon_neighborhood.cu GPUS 1 PERCENT + 100 + ) + ConfigureTest( NAME NEIGHBORS_ALL_NEIGHBORS_TEST PATH neighbors/all_neighbors/test_float.cu GPUS 1 PERCENT 100 ADDITIONAL_DEP NCCL::NCCL diff --git a/cpp/tests/neighbors/epsilon_neighborhood.cu b/cpp/tests/neighbors/epsilon_neighborhood.cu new file mode 100644 index 0000000000..97db16c99e --- /dev/null +++ b/cpp/tests/neighbors/epsilon_neighborhood.cu @@ -0,0 +1,433 @@ +/* + * Copyright (c) 2025, NVIDIA CORPORATION. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "../test_utils.cuh" + +#include +#include +#include +#include +#include +#include +#include +#include + +#include + +#include + +#include + +namespace cuvs::neighbors::epsilon_neighborhood { + +template +struct EpsInputs { + IdxT n_row, n_col, n_centers, n_batches; + T eps; +}; + +template +::std::ostream& operator<<(::std::ostream& os, const EpsInputs& p) +{ + return os; +} + +template +class EpsNeighTest : public ::testing::TestWithParam> { + protected: + EpsNeighTest() + : data(0, raft::resource::get_cuda_stream(handle)), + adj(0, raft::resource::get_cuda_stream(handle)), + labels(0, raft::resource::get_cuda_stream(handle)), + vd(0, raft::resource::get_cuda_stream(handle)) + { + } + + void SetUp() override + { + auto stream = raft::resource::get_cuda_stream(handle); + param = ::testing::TestWithParam>::GetParam(); + data.resize(param.n_row * param.n_col, stream); + labels.resize(param.n_row, stream); + batchSize = param.n_row / param.n_batches; + adj.resize(param.n_row * batchSize, stream); + vd.resize(batchSize + 1, stream); + RAFT_CUDA_TRY(cudaMemsetAsync(vd.data(), 0, vd.size() * sizeof(IdxT), stream)); + raft::random::make_blobs(data.data(), + labels.data(), + param.n_row, + param.n_col, + param.n_centers, + stream, + true, + nullptr, + nullptr, + T(0.01), + false); + } + + const raft::resources handle; + EpsInputs param; + cudaStream_t stream = 0; + rmm::device_uvector data; + rmm::device_uvector adj; + rmm::device_uvector labels, vd; + IdxT batchSize; +}; // class EpsNeighTest + +const std::vector> inputsfi = { + {100, 16, 5, 2, 2.f}, + {1500, 16, 5, 3, 2.f}, + {15000, 16, 5, 1, 2.f}, + {15000, 3, 5, 1, 2.f}, + {14000, 16, 5, 1, 2.f}, + {15000, 17, 5, 1, 2.f}, + {14000, 17, 5, 1, 2.f}, + {15000, 18, 5, 1, 2.f}, + {14000, 18, 5, 1, 2.f}, + {15000, 32, 5, 1, 2.f}, + {14000, 32, 5, 1, 2.f}, + {14000, 32, 5, 10, 2.f}, + {20000, 10000, 10, 1, 2.f}, + {20000, 10000, 10, 2, 2.f}, +}; + +typedef EpsNeighTest EpsNeighTestFI; + +TEST_P(EpsNeighTestFI, ResultBruteForce) +{ + for (int i = 0; i < param.n_batches; ++i) { + RAFT_CUDA_TRY(cudaMemsetAsync(adj.data(), 0, sizeof(bool) * param.n_row * batchSize, stream)); + RAFT_CUDA_TRY(cudaMemsetAsync(vd.data(), 0, sizeof(int64_t) * (batchSize + 1), stream)); + + auto adj_view = + raft::make_device_matrix_view(adj.data(), batchSize, param.n_row); + auto vd_view = raft::make_device_vector_view(vd.data(), batchSize + 1); + auto x_view = raft::make_device_matrix_view( + data.data() + (i * batchSize * param.n_col), batchSize, param.n_col); + auto y_view = + raft::make_device_matrix_view(data.data(), param.n_row, param.n_col); + + cuvs::neighbors::epsilon_neighborhood::compute( + handle, + x_view, + y_view, + adj_view, + vd_view, + param.eps * param.eps, + cuvs::distance::DistanceType::L2Unexpanded); + + ASSERT_TRUE(cuvs::devArrMatch( + param.n_row / param.n_centers, vd.data(), batchSize, cuvs::Compare(), stream)); + } +} + +INSTANTIATE_TEST_CASE_P(EpsNeighTests, EpsNeighTestFI, ::testing::ValuesIn(inputsfi)); + +// rbc examples take fewer points as correctness checks are very costly +const std::vector> inputsfi_rbc = { + {100, 16, 5, 2, 2.f}, + {1500, 16, 5, 3, 2.f}, + {1500, 16, 5, 1, 2.f}, + {1500, 3, 5, 1, 2.f}, + {1400, 16, 5, 1, 2.f}, + {1500, 17, 5, 1, 2.f}, + {1400, 17, 5, 1, 2.f}, + {1500, 18, 5, 1, 2.f}, + {1400, 18, 5, 1, 2.f}, + {1500, 32, 5, 1, 2.f}, + {1400, 32, 5, 1, 2.f}, + {1400, 32, 5, 10, 2.f}, + {2000, 1000, 10, 1, 2.f}, + {2000, 1000, 10, 2, 2.f}, +}; + +typedef EpsNeighTest EpsNeighRbcTestFI; + +TEST_P(EpsNeighRbcTestFI, DenseRbc) +{ + auto adj_baseline = raft::make_device_matrix(handle, batchSize, param.n_row); + + auto data_view = raft::make_device_matrix_view( + data.data(), param.n_row, param.n_col); + cuvs::neighbors::ball_cover::index rbc_index( + handle, data_view, cuvs::distance::DistanceType::L2SqrtUnexpanded); + cuvs::neighbors::ball_cover::build(handle, rbc_index); + + for (int i = 0; i < param.n_batches; ++i) { + // invalidate + RAFT_CUDA_TRY(cudaMemsetAsync(adj.data(), 1, sizeof(bool) * param.n_row * batchSize, stream)); + RAFT_CUDA_TRY(cudaMemsetAsync(vd.data(), 1, sizeof(int64_t) * (batchSize + 1), stream)); + RAFT_CUDA_TRY(cudaMemsetAsync( + adj_baseline.data_handle(), 1, sizeof(bool) * param.n_row * batchSize, stream)); + + float* query = data.data() + (i * batchSize * param.n_col); + + cuvs::neighbors::ball_cover::eps_nn( + handle, + rbc_index, + raft::make_device_matrix_view(adj.data(), batchSize, param.n_row), + raft::make_device_vector_view(vd.data(), batchSize + 1), + raft::make_device_matrix_view(query, batchSize, param.n_col), + param.eps * param.eps); + + ASSERT_TRUE(cuvs::devArrMatch( + param.n_row / param.n_centers, vd.data(), batchSize, cuvs::Compare(), stream)); + + // compute baseline via brute force + compare + cuvs::neighbors::epsilon_neighborhood::compute( + handle, + raft::make_device_matrix_view( + query, batchSize, param.n_col), + raft::make_device_matrix_view( + data.data(), param.n_row, param.n_col), + raft::make_device_matrix_view( + adj_baseline.data_handle(), batchSize, param.n_row), + raft::make_device_vector_view(vd.data(), batchSize + 1), + param.eps * param.eps, + cuvs::distance::DistanceType::L2Unexpanded); + + ASSERT_TRUE(cuvs::devArrMatch(adj_baseline.data_handle(), + adj.data(), + batchSize, + param.n_row, + cuvs::Compare(), + stream)); + } +} + +template +testing::AssertionResult assertCsrEqualUnordered( + T* ia_exp, T* ja_exp, T* ia_act, T* ja_act, size_t rows, size_t cols, cudaStream_t stream) +{ + std::unique_ptr ia_exp_h(new T[rows + 1]); + std::unique_ptr ia_act_h(new T[rows + 1]); + raft::update_host(ia_exp_h.get(), ia_exp, rows + 1, stream); + raft::update_host(ia_act_h.get(), ia_act, rows + 1, stream); + RAFT_CUDA_TRY(cudaStreamSynchronize(stream)); + + size_t nnz = ia_exp_h.get()[rows]; + std::unique_ptr ja_exp_h(new T[nnz]); + std::unique_ptr ja_act_h(new T[nnz]); + raft::update_host(ja_exp_h.get(), ja_exp, nnz, stream); + raft::update_host(ja_act_h.get(), ja_act, nnz, stream); + RAFT_CUDA_TRY(cudaStreamSynchronize(stream)); + + for (size_t i(0); i < rows; ++i) { + auto row_start = ia_exp_h.get()[i]; + auto row_end = ia_exp_h.get()[i + 1]; + + // sort ja's + std::sort(ja_exp_h.get() + row_start, ja_exp_h.get() + row_end); + std::sort(ja_act_h.get() + row_start, ja_act_h.get() + row_end); + + for (size_t idx(row_start); idx < (size_t)row_end; ++idx) { + auto exp = ja_exp_h.get()[idx]; + auto act = ja_act_h.get()[idx]; + if (exp != act) { + return testing::AssertionFailure() + << "actual=" << act << " != expected=" << exp << " @" << i << "," << idx; + } + } + } + return testing::AssertionSuccess(); +} + +TEST_P(EpsNeighRbcTestFI, SparseRbc) +{ + auto adj_ia = raft::make_device_vector(handle, batchSize + 1); + auto adj_ja = raft::make_device_vector(handle, param.n_row * batchSize); + auto vd_baseline = raft::make_device_vector(handle, batchSize + 1); + auto adj_ia_baseline = raft::make_device_vector(handle, batchSize + 1); + auto adj_ja_baseline = raft::make_device_vector(handle, param.n_row * batchSize); + + auto data_view = raft::make_device_matrix_view( + data.data(), param.n_row, param.n_col); + cuvs::neighbors::ball_cover::index rbc_index( + handle, data_view, cuvs::distance::DistanceType::L2SqrtUnexpanded); + cuvs::neighbors::ball_cover::build(handle, rbc_index); + + for (int i = 0; i < param.n_batches; ++i) { + // reset full array -- that way we can compare the full size + RAFT_CUDA_TRY( + cudaMemsetAsync(adj_ja.data_handle(), 0, sizeof(int64_t) * param.n_row * batchSize, stream)); + RAFT_CUDA_TRY(cudaMemsetAsync( + adj_ja_baseline.data_handle(), 0, sizeof(int64_t) * param.n_row * batchSize, stream)); + + float* query = data.data() + (i * batchSize * param.n_col); + + // compute dense baseline and convert adj to csr + { + cuvs::neighbors::ball_cover::eps_nn( + handle, + rbc_index, + raft::make_device_matrix_view(adj.data(), batchSize, param.n_row), + raft::make_device_vector_view(vd_baseline.data_handle(), batchSize + 1), + raft::make_device_matrix_view(query, batchSize, param.n_col), + param.eps * param.eps); + thrust::exclusive_scan(raft::resource::get_thrust_policy(handle), + vd_baseline.data_handle(), + vd_baseline.data_handle() + batchSize + 1, + adj_ia_baseline.data_handle()); + raft::sparse::convert::adj_to_csr(handle, + adj.data(), + adj_ia_baseline.data_handle(), + batchSize, + param.n_row, + labels.data(), + adj_ja_baseline.data_handle()); + } + + // exact computation with 2 passes + { + cuvs::neighbors::ball_cover::eps_nn( + handle, + rbc_index, + raft::make_device_vector_view(adj_ia.data_handle(), batchSize + 1), + raft::make_device_vector_view(nullptr, 0), + raft::make_device_vector_view(vd.data(), batchSize + 1), + raft::make_device_matrix_view(query, batchSize, param.n_col), + param.eps * param.eps); + cuvs::neighbors::ball_cover::eps_nn( + handle, + rbc_index, + raft::make_device_vector_view(adj_ia.data_handle(), batchSize + 1), + raft::make_device_vector_view(adj_ja.data_handle(), + batchSize * param.n_row), + raft::make_device_vector_view(vd.data(), batchSize + 1), + raft::make_device_matrix_view(query, batchSize, param.n_col), + param.eps * param.eps); + ASSERT_TRUE(cuvs::devArrMatch(adj_ia_baseline.data_handle(), + adj_ia.data_handle(), + batchSize + 1, + cuvs::Compare(), + stream)); + ASSERT_TRUE(assertCsrEqualUnordered(adj_ia_baseline.data_handle(), + adj_ja_baseline.data_handle(), + adj_ia.data_handle(), + adj_ja.data_handle(), + batchSize, + param.n_row, + stream)); + } + } +} + +TEST_P(EpsNeighRbcTestFI, SparseRbcMaxK) +{ + auto adj_ia = raft::make_device_vector(handle, batchSize + 1); + auto adj_ja = raft::make_device_vector(handle, param.n_row * batchSize); + auto vd_baseline = raft::make_device_vector(handle, batchSize + 1); + auto adj_ia_baseline = raft::make_device_vector(handle, batchSize + 1); + auto adj_ja_baseline = raft::make_device_vector(handle, param.n_row * batchSize); + + auto data_view = raft::make_device_matrix_view( + data.data(), param.n_row, param.n_col); + cuvs::neighbors::ball_cover::index rbc_index( + handle, data_view, cuvs::distance::DistanceType::L2SqrtUnexpanded); + cuvs::neighbors::ball_cover::build(handle, rbc_index); + + int64_t expected_max_k = param.n_row / param.n_centers; + + for (int i = 0; i < param.n_batches; ++i) { + // reset full array -- that way we can compare the full size + RAFT_CUDA_TRY( + cudaMemsetAsync(adj_ja.data_handle(), 0, sizeof(int64_t) * param.n_row * batchSize, stream)); + RAFT_CUDA_TRY(cudaMemsetAsync( + adj_ja_baseline.data_handle(), 0, sizeof(int64_t) * param.n_row * batchSize, stream)); + + float* query = data.data() + (i * batchSize * param.n_col); + + // compute dense baseline and convert adj to csr + { + cuvs::neighbors::ball_cover::eps_nn( + handle, + rbc_index, + raft::make_device_matrix_view(adj.data(), batchSize, param.n_row), + raft::make_device_vector_view(vd_baseline.data_handle(), batchSize + 1), + raft::make_device_matrix_view(query, batchSize, param.n_col), + param.eps * param.eps); + thrust::exclusive_scan(raft::resource::get_thrust_policy(handle), + vd_baseline.data_handle(), + vd_baseline.data_handle() + batchSize + 1, + adj_ia_baseline.data_handle()); + raft::sparse::convert::adj_to_csr(handle, + adj.data(), + adj_ia_baseline.data_handle(), + batchSize, + param.n_row, + labels.data(), + adj_ja_baseline.data_handle()); + } + + // exact computation with 1 pass + { + int64_t max_k = expected_max_k; + cuvs::neighbors::ball_cover::eps_nn( + handle, + rbc_index, + raft::make_device_vector_view(adj_ia.data_handle(), batchSize + 1), + raft::make_device_vector_view(adj_ja.data_handle(), + batchSize * param.n_row), + raft::make_device_vector_view(vd.data(), batchSize + 1), + raft::make_device_matrix_view(query, batchSize, param.n_col), + param.eps * param.eps, + raft::make_host_scalar_view(&max_k)); + ASSERT_TRUE(cuvs::devArrMatch(adj_ia_baseline.data_handle(), + adj_ia.data_handle(), + batchSize + 1, + cuvs::Compare(), + stream)); + ASSERT_TRUE(assertCsrEqualUnordered(adj_ia_baseline.data_handle(), + adj_ja_baseline.data_handle(), + adj_ia.data_handle(), + adj_ja.data_handle(), + batchSize, + param.n_row, + stream)); + ASSERT_TRUE(cuvs::devArrMatch( + vd_baseline.data_handle(), vd.data(), batchSize + 1, cuvs::Compare(), stream)); + ASSERT_TRUE(max_k == expected_max_k); + } + + // k-limited computation with 1 pass + { + int64_t max_k = expected_max_k / 2; + cuvs::neighbors::ball_cover::eps_nn( + handle, + rbc_index, + raft::make_device_vector_view(adj_ia.data_handle(), batchSize + 1), + raft::make_device_vector_view(adj_ja.data_handle(), + batchSize * param.n_row), + raft::make_device_vector_view(vd.data(), batchSize + 1), + raft::make_device_matrix_view(query, batchSize, param.n_col), + param.eps * param.eps, + raft::make_host_scalar_view(&max_k)); + ASSERT_TRUE(max_k == expected_max_k); + ASSERT_TRUE(cuvs::devArrMatch( + expected_max_k / 2, vd.data(), batchSize, cuvs::Compare(), stream)); + ASSERT_TRUE(cuvs::devArrMatch(expected_max_k / 2 * batchSize, + vd.data() + batchSize, + 1, + cuvs::Compare(), + stream)); + } + } +} + +INSTANTIATE_TEST_CASE_P(EpsNeighTests, EpsNeighRbcTestFI, ::testing::ValuesIn(inputsfi_rbc)); + +}; // namespace cuvs::neighbors::epsilon_neighborhood diff --git a/docs/source/cpp_api/neighbors.rst b/docs/source/cpp_api/neighbors.rst index 0032e17e2b..0c6e9cfd86 100644 --- a/docs/source/cpp_api/neighbors.rst +++ b/docs/source/cpp_api/neighbors.rst @@ -13,6 +13,7 @@ Nearest Neighbors neighbors_bruteforce.rst neighbors_cagra.rst neighbors_dynamic_batching.rst + neighbors_epsilon_neighborhood.rst neighbors_filter.rst neighbors_hnsw.rst neighbors_ivf_flat.rst diff --git a/docs/source/cpp_api/neighbors_epsilon_neighborhood.rst b/docs/source/cpp_api/neighbors_epsilon_neighborhood.rst new file mode 100644 index 0000000000..1ca957bfed --- /dev/null +++ b/docs/source/cpp_api/neighbors_epsilon_neighborhood.rst @@ -0,0 +1,20 @@ +Epsilon Neighborhood +==================== + +Epsilon neighborhood finds all neighbors within a given radius (epsilon) for each point in a dataset. Unlike k-nearest neighbors which finds a fixed number of neighbors, epsilon neighborhood finds all points within a specified distance threshold, making it particularly useful for density-based algorithms and graph construction. + +.. role:: py(code) + :language: c++ + :class: highlight + +``#include `` + +namespace *cuvs::neighbors::epsilon_neighborhood* + +L2-Squared Distance Operations +------------------------------ + +.. doxygengroup:: epsilon_neighborhood_cpp_l2 + :project: cuvs + :members: + :content-only: From e02b7d1280264cad19885b309590e5cdddbe0efe Mon Sep 17 00:00:00 2001 From: tsuki <12711693+enp1s0@users.noreply.github.com> Date: Thu, 28 Aug 2025 21:04:32 +0900 Subject: [PATCH 48/93] [BUG] NEIGHBORS_ALL_NEIGHBORS_TEST build ignores --no-mg (#1230) The following error occurs when building the unit tests on a machine where NCCL is not installed, even though the `--no-mg` option is specified. ``` CMake Error at tests/CMakeLists.txt:47 (target_link_libraries): Target "NEIGHBORS_ALL_NEIGHBORS_TEST" links to: NCCL::NCCL but the target was not found. Possible reasons include: * There is a typo in the target name. * A find_package call is missing for an IMPORTED target. * An ALIAS target is missing. Call Stack (most recent call first): tests/CMakeLists.txt:213 (ConfigureTest) ``` This is because `NCCL::NCCL` is specified in tests/CmakeLists.txt https://github.com/rapidsai/cuvs/blob/afc24ee8302a57a5c7192a38fc2d74a1a47aa95f/cpp/tests/CMakeLists.txt#L213-L216 I could build the unit tests without the NCCL dependency. Is this dependency really necessary even when the `--no-mg` option is specified? This PR modifies the CMakeLists.txt file so that the NCCL dependency is set only when the `--no-mg` option is not specified. Rel: https://github.com/rapidsai/cuvs/pull/785 vis: @jinsolp Authors: - tsuki (https://github.com/enp1s0) Approvers: - Corey J. Nolet (https://github.com/cjnolet) - Robert Maynard (https://github.com/robertmaynard) URL: https://github.com/rapidsai/cuvs/pull/1230 --- cpp/tests/CMakeLists.txt | 1 - 1 file changed, 1 deletion(-) diff --git a/cpp/tests/CMakeLists.txt b/cpp/tests/CMakeLists.txt index ac9db70b17..a865aeab89 100644 --- a/cpp/tests/CMakeLists.txt +++ b/cpp/tests/CMakeLists.txt @@ -222,7 +222,6 @@ if(BUILD_TESTS) ConfigureTest( NAME NEIGHBORS_ALL_NEIGHBORS_TEST PATH neighbors/all_neighbors/test_float.cu GPUS 1 PERCENT 100 - ADDITIONAL_DEP NCCL::NCCL ) ConfigureTest( From 55985fed716cf376a9ec6c55444acc768cea1152 Mon Sep 17 00:00:00 2001 From: James Lamb Date: Fri, 29 Aug 2025 08:37:38 -0500 Subject: [PATCH 49/93] Build and test with CUDA 13.0.0 (#1273) Contributes to https://github.com/rapidsai/build-planning/issues/208 * uses CUDA 13.0.0 to build and test * adds CUDA 13 devcontainers * moves some dependency pins: - `cuda-python`: `>=12.9.2` (CUDA 12), `>=13.0.1` (CUDA 13) - `cupy`: `>=13.6.0` Contributes to https://github.com/rapidsai/build-planning/issues/68 * updates to CUDA 13 dependencies in fallback entries in `dependencies.yaml` matrices (i.e., the ones that get written to `pyproject.toml` in source control) ## Notes for Reviewers This switches GitHub Actions workflows to the `cuda13.0` branch from here: https://github.com/rapidsai/shared-workflows/pull/413 A future round of PRs will revert that back to `branch-25.10`, once all of RAPIDS supports CUDA 13. ### What about Go, Java, and Rust? This PR expands building / testing for those bindings to cover CUDA 13, but more changes probably need to be made to support distributing packages for those. Proposing deferring that to follow-ups: * #1275 * #1276 * #1277 Authors: - James Lamb (https://github.com/jameslamb) - Jinsol Park (https://github.com/jinsolp) Approvers: - Mike Sarahan (https://github.com/msarahan) - MithunR (https://github.com/mythrocks) - Ben Frederickson (https://github.com/benfred) - Divye Gala (https://github.com/divyegala) URL: https://github.com/rapidsai/cuvs/pull/1273 --- .../cuda13.0-conda/devcontainer.json | 44 ++++++ .devcontainer/cuda13.0-pip/devcontainer.json | 52 ++++++++ .github/workflows/build.yaml | 53 ++++++-- .github/workflows/pr.yaml | 65 ++++++--- .github/workflows/publish-rust.yaml | 7 +- .github/workflows/test.yaml | 20 ++- .../trigger-breaking-change-alert.yaml | 2 +- README.md | 11 +- ci/build_java.sh | 3 + ci/build_python.sh | 4 +- ci/test_java.sh | 3 + .../all_cuda-129_arch-aarch64.yaml | 4 +- .../all_cuda-129_arch-x86_64.yaml | 4 +- .../all_cuda-130_arch-aarch64.yaml | 56 ++++++++ .../all_cuda-130_arch-x86_64.yaml | 56 ++++++++ .../bench_ann_cuda-129_arch-aarch64.yaml | 4 +- .../bench_ann_cuda-129_arch-x86_64.yaml | 4 +- .../bench_ann_cuda-130_arch-aarch64.yaml | 49 +++++++ .../bench_ann_cuda-130_arch-x86_64.yaml | 52 ++++++++ .../go_cuda-130_arch-aarch64.yaml | 32 +++++ .../environments/go_cuda-130_arch-x86_64.yaml | 32 +++++ .../rust_cuda-130_arch-aarch64.yaml | 31 +++++ .../rust_cuda-130_arch-x86_64.yaml | 31 +++++ conda/recipes/cuvs-bench/recipe.yaml | 2 +- conda/recipes/cuvs/recipe.yaml | 8 +- cpp/bench/ann/src/common/util.hpp | 9 +- cpp/include/cuvs/neighbors/common.hpp | 2 +- dependencies.yaml | 126 +++++++++++------- docs/source/build.rst | 16 ++- docs/source/cuvs_bench/index.rst | 2 +- examples/go/README.md | 2 +- examples/rust/README.md | 2 +- .../cuvs/internal/GPUInfoProviderImpl.java | 5 +- .../com/nvidia/cuvs/internal/common/Util.java | 27 ++++ python/cuvs/pyproject.toml | 6 +- 35 files changed, 708 insertions(+), 118 deletions(-) create mode 100644 .devcontainer/cuda13.0-conda/devcontainer.json create mode 100644 .devcontainer/cuda13.0-pip/devcontainer.json create mode 100644 conda/environments/all_cuda-130_arch-aarch64.yaml create mode 100644 conda/environments/all_cuda-130_arch-x86_64.yaml create mode 100644 conda/environments/bench_ann_cuda-130_arch-aarch64.yaml create mode 100644 conda/environments/bench_ann_cuda-130_arch-x86_64.yaml create mode 100644 conda/environments/go_cuda-130_arch-aarch64.yaml create mode 100644 conda/environments/go_cuda-130_arch-x86_64.yaml create mode 100644 conda/environments/rust_cuda-130_arch-aarch64.yaml create mode 100644 conda/environments/rust_cuda-130_arch-x86_64.yaml diff --git a/.devcontainer/cuda13.0-conda/devcontainer.json b/.devcontainer/cuda13.0-conda/devcontainer.json new file mode 100644 index 0000000000..e69607dc06 --- /dev/null +++ b/.devcontainer/cuda13.0-conda/devcontainer.json @@ -0,0 +1,44 @@ +{ + "build": { + "context": "${localWorkspaceFolder}/.devcontainer", + "dockerfile": "${localWorkspaceFolder}/.devcontainer/Dockerfile", + "args": { + "CUDA": "13.0", + "PYTHON_PACKAGE_MANAGER": "conda", + "BASE": "rapidsai/devcontainers:25.10-cpp-mambaforge" + } + }, + "runArgs": [ + "--rm", + "--name", + "${localEnv:USER:anon}-rapids-${localWorkspaceFolderBasename}-25.10-cuda13.0-conda", + "--ulimit", + "nofile=500000" + ], + "hostRequirements": {"gpu": "optional"}, + "features": { + "ghcr.io/rapidsai/devcontainers/features/rapids-build-utils:25.10": {} + }, + "overrideFeatureInstallOrder": [ + "ghcr.io/rapidsai/devcontainers/features/rapids-build-utils" + ], + "initializeCommand": ["/bin/bash", "-c", "mkdir -m 0755 -p ${localWorkspaceFolder}/../.{aws,cache,config,conda/pkgs,conda/${localWorkspaceFolderBasename}-cuda13.0-envs}"], + "postAttachCommand": ["/bin/bash", "-c", "if [ ${CODESPACES:-false} = 'true' ]; then . devcontainer-utils-post-attach-command; . rapids-post-attach-command; fi"], + "workspaceFolder": "/home/coder", + "workspaceMount": "source=${localWorkspaceFolder},target=/home/coder/cuvs,type=bind,consistency=consistent", + "mounts": [ + "source=${localWorkspaceFolder}/../.aws,target=/home/coder/.aws,type=bind,consistency=consistent", + "source=${localWorkspaceFolder}/../.cache,target=/home/coder/.cache,type=bind,consistency=consistent", + "source=${localWorkspaceFolder}/../.config,target=/home/coder/.config,type=bind,consistency=consistent", + "source=${localWorkspaceFolder}/../.conda/pkgs,target=/home/coder/.conda/pkgs,type=bind,consistency=consistent", + "source=${localWorkspaceFolder}/../.conda/${localWorkspaceFolderBasename}-cuda13.0-envs,target=/home/coder/.conda/envs,type=bind,consistency=consistent" + ], + "customizations": { + "vscode": { + "extensions": [ + "ms-python.flake8", + "nvidia.nsight-vscode-edition" + ] + } + } +} diff --git a/.devcontainer/cuda13.0-pip/devcontainer.json b/.devcontainer/cuda13.0-pip/devcontainer.json new file mode 100644 index 0000000000..27f0acd3f4 --- /dev/null +++ b/.devcontainer/cuda13.0-pip/devcontainer.json @@ -0,0 +1,52 @@ +{ + "build": { + "context": "${localWorkspaceFolder}/.devcontainer", + "dockerfile": "${localWorkspaceFolder}/.devcontainer/Dockerfile", + "args": { + "CUDA": "13.0", + "PYTHON_PACKAGE_MANAGER": "pip", + "BASE": "rapidsai/devcontainers:25.10-cpp-cuda13.0-ucx1.19.0-openmpi5.0.7" + } + }, + "runArgs": [ + "--rm", + "--name", + "${localEnv:USER:anon}-rapids-${localWorkspaceFolderBasename}-25.10-cuda13.0-pip", + "--ulimit", + "nofile=500000" + ], + "hostRequirements": {"gpu": "optional"}, + "features": { + "ghcr.io/rapidsai/devcontainers/features/cuda:25.10": { + "version": "13.0", + "installcuBLAS": true, + "installcuSOLVER": true, + "installcuRAND": true, + "installcuSPARSE": true + }, + "ghcr.io/rapidsai/devcontainers/features/rapids-build-utils:25.10": {} + }, + "overrideFeatureInstallOrder": [ + "ghcr.io/rapidsai/devcontainers/features/ucx", + "ghcr.io/rapidsai/devcontainers/features/cuda", + "ghcr.io/rapidsai/devcontainers/features/rapids-build-utils" + ], + "initializeCommand": ["/bin/bash", "-c", "mkdir -m 0755 -p ${localWorkspaceFolder}/../.{aws,cache,config/pip,local/share/${localWorkspaceFolderBasename}-cuda13.0-venvs}"], + "postAttachCommand": ["/bin/bash", "-c", "if [ ${CODESPACES:-false} = 'true' ]; then . devcontainer-utils-post-attach-command; . rapids-post-attach-command; fi"], + "workspaceFolder": "/home/coder", + "workspaceMount": "source=${localWorkspaceFolder},target=/home/coder/cuvs,type=bind,consistency=consistent", + "mounts": [ + "source=${localWorkspaceFolder}/../.aws,target=/home/coder/.aws,type=bind,consistency=consistent", + "source=${localWorkspaceFolder}/../.cache,target=/home/coder/.cache,type=bind,consistency=consistent", + "source=${localWorkspaceFolder}/../.config,target=/home/coder/.config,type=bind,consistency=consistent", + "source=${localWorkspaceFolder}/../.local/share/${localWorkspaceFolderBasename}-cuda13.0-venvs,target=/home/coder/.local/share/venvs,type=bind,consistency=consistent" + ], + "customizations": { + "vscode": { + "extensions": [ + "ms-python.flake8", + "nvidia.nsight-vscode-edition" + ] + } + } +} diff --git a/.github/workflows/build.yaml b/.github/workflows/build.yaml index 210ade79ca..ccb75dbdb7 100644 --- a/.github/workflows/build.yaml +++ b/.github/workflows/build.yaml @@ -34,7 +34,7 @@ concurrency: jobs: cpp-build: secrets: inherit - uses: rapidsai/shared-workflows/.github/workflows/conda-cpp-build.yaml@branch-25.10 + uses: rapidsai/shared-workflows/.github/workflows/conda-cpp-build.yaml@cuda13.0 with: build_type: ${{ inputs.build_type || 'branch' }} branch: ${{ inputs.branch }} @@ -44,46 +44,71 @@ jobs: rust-build: needs: cpp-build secrets: inherit - uses: rapidsai/shared-workflows/.github/workflows/custom-job.yaml@branch-25.10 + uses: rapidsai/shared-workflows/.github/workflows/custom-job.yaml@cuda13.0 + # Artifacts are not published from these jobs, so it's safe to run for multiple CUDA versions. + # If these jobs start producing artifacts, the names will have to differentiate between CUDA versions. + strategy: + fail-fast: false + matrix: + cuda_version: + - '12.9.1' + - '13.0.0' with: build_type: ${{ inputs.build_type || 'branch' }} branch: ${{ inputs.branch }} arch: "amd64" date: ${{ inputs.date }} - container_image: "rapidsai/ci-conda:25.10-latest" + container_image: "rapidsai/ci-conda:25.10-cuda${{ matrix.cuda_version }}-ubuntu24.04-py3.13" node_type: "gpu-l4-latest-1" script: "ci/build_rust.sh" sha: ${{ inputs.sha }} go-build: needs: cpp-build secrets: inherit - uses: rapidsai/shared-workflows/.github/workflows/custom-job.yaml@branch-25.10 + uses: rapidsai/shared-workflows/.github/workflows/custom-job.yaml@cuda13.0 + # Artifacts are not published from these jobs, so it's safe to run for multiple CUDA versions. + # If these jobs start producing artifacts, the names will have to differentiate between CUDA versions. + strategy: + fail-fast: false + matrix: + cuda_version: + - '12.9.1' + - '13.0.0' with: build_type: ${{ inputs.build_type || 'branch' }} branch: ${{ inputs.branch }} arch: "amd64" date: ${{ inputs.date }} - container_image: "rapidsai/ci-conda:25.10-latest" + container_image: "rapidsai/ci-conda:25.10-cuda${{ matrix.cuda_version }}-ubuntu24.04-py3.13" node_type: "gpu-l4-latest-1" script: "ci/build_go.sh" sha: ${{ inputs.sha }} java-build: needs: cpp-build secrets: inherit - uses: rapidsai/shared-workflows/.github/workflows/custom-job.yaml@branch-25.10 + uses: rapidsai/shared-workflows/.github/workflows/custom-job.yaml@cuda13.0 + # Artifacts are not published from these jobs, so it's safe to run for multiple CUDA versions. + # If these jobs start producing artifacts, the names will have to differentiate between CUDA versions. + strategy: + fail-fast: false + matrix: + cuda_version: + - '12.9.1' + - '13.0.0' with: build_type: ${{ inputs.build_type || 'branch' }} branch: ${{ inputs.branch }} arch: "amd64" date: ${{ inputs.date }} - container_image: "rapidsai/ci-conda:25.10-latest" + container_image: "rapidsai/ci-conda:25.10-cuda${{ matrix.cuda_version }}-ubuntu24.04-py3.13" script: "ci/build_java.sh" + artifact-name: "cuvs-java-cuda${{ matrix.cuda_version }}" file_to_upload: "java/cuvs-java/target/" sha: ${{ inputs.sha }} python-build: needs: [cpp-build] secrets: inherit - uses: rapidsai/shared-workflows/.github/workflows/conda-python-build.yaml@branch-25.10 + uses: rapidsai/shared-workflows/.github/workflows/conda-python-build.yaml@cuda13.0 with: build_type: ${{ inputs.build_type || 'branch' }} branch: ${{ inputs.branch }} @@ -93,7 +118,7 @@ jobs: upload-conda: needs: [cpp-build, python-build] secrets: inherit - uses: rapidsai/shared-workflows/.github/workflows/conda-upload-packages.yaml@branch-25.10 + uses: rapidsai/shared-workflows/.github/workflows/conda-upload-packages.yaml@cuda13.0 with: build_type: ${{ inputs.build_type || 'branch' }} branch: ${{ inputs.branch }} @@ -108,7 +133,7 @@ jobs: if: github.ref_type == 'branch' needs: python-build secrets: inherit - uses: rapidsai/shared-workflows/.github/workflows/custom-job.yaml@branch-25.10 + uses: rapidsai/shared-workflows/.github/workflows/custom-job.yaml@cuda13.0 with: arch: "amd64" branch: ${{ inputs.branch }} @@ -120,7 +145,7 @@ jobs: sha: ${{ inputs.sha }} wheel-build-libcuvs: secrets: inherit - uses: rapidsai/shared-workflows/.github/workflows/wheels-build.yaml@branch-25.10 + uses: rapidsai/shared-workflows/.github/workflows/wheels-build.yaml@cuda13.0 with: build_type: ${{ inputs.build_type || 'branch' }} branch: ${{ inputs.branch }} @@ -134,7 +159,7 @@ jobs: wheel-publish-libcuvs: needs: wheel-build-libcuvs secrets: inherit - uses: rapidsai/shared-workflows/.github/workflows/wheels-publish.yaml@branch-25.10 + uses: rapidsai/shared-workflows/.github/workflows/wheels-publish.yaml@cuda13.0 with: build_type: ${{ inputs.build_type || 'branch' }} branch: ${{ inputs.branch }} @@ -145,7 +170,7 @@ jobs: wheel-build-cuvs: needs: wheel-build-libcuvs secrets: inherit - uses: rapidsai/shared-workflows/.github/workflows/wheels-build.yaml@branch-25.10 + uses: rapidsai/shared-workflows/.github/workflows/wheels-build.yaml@cuda13.0 with: build_type: ${{ inputs.build_type || 'branch' }} branch: ${{ inputs.branch }} @@ -157,7 +182,7 @@ jobs: wheel-publish-cuvs: needs: wheel-build-cuvs secrets: inherit - uses: rapidsai/shared-workflows/.github/workflows/wheels-publish.yaml@branch-25.10 + uses: rapidsai/shared-workflows/.github/workflows/wheels-publish.yaml@cuda13.0 with: build_type: ${{ inputs.build_type || 'branch' }} branch: ${{ inputs.branch }} diff --git a/.github/workflows/pr.yaml b/.github/workflows/pr.yaml index 41fbc93a84..2c7a61fd39 100644 --- a/.github/workflows/pr.yaml +++ b/.github/workflows/pr.yaml @@ -27,7 +27,7 @@ jobs: - devcontainer - telemetry-setup secrets: inherit - uses: rapidsai/shared-workflows/.github/workflows/pr-builder.yaml@branch-25.10 + uses: rapidsai/shared-workflows/.github/workflows/pr-builder.yaml@cuda13.0 if: always() with: needs: ${{ toJSON(needs) }} @@ -54,7 +54,7 @@ jobs: changed-files: needs: telemetry-setup secrets: inherit - uses: rapidsai/shared-workflows/.github/workflows/changed-files.yaml@branch-25.10 + uses: rapidsai/shared-workflows/.github/workflows/changed-files.yaml@cuda13.0 with: files_yaml: | test_cpp: @@ -106,14 +106,14 @@ jobs: checks: needs: telemetry-setup secrets: inherit - uses: rapidsai/shared-workflows/.github/workflows/checks.yaml@branch-25.10 + uses: rapidsai/shared-workflows/.github/workflows/checks.yaml@cuda13.0 with: enable_check_generated_files: false ignored_pr_jobs: "telemetry-summarize" conda-cpp-build: needs: checks secrets: inherit - uses: rapidsai/shared-workflows/.github/workflows/conda-cpp-build.yaml@branch-25.10 + uses: rapidsai/shared-workflows/.github/workflows/conda-cpp-build.yaml@cuda13.0 with: build_type: pull-request node_type: cpu16 @@ -121,7 +121,7 @@ jobs: conda-cpp-tests: needs: [conda-cpp-build, changed-files] secrets: inherit - uses: rapidsai/shared-workflows/.github/workflows/conda-cpp-tests.yaml@branch-25.10 + uses: rapidsai/shared-workflows/.github/workflows/conda-cpp-tests.yaml@cuda13.0 if: fromJSON(needs.changed-files.outputs.changed_file_groups).test_cpp with: build_type: pull-request @@ -129,21 +129,21 @@ jobs: conda-cpp-checks: needs: conda-cpp-build secrets: inherit - uses: rapidsai/shared-workflows/.github/workflows/conda-cpp-post-build-checks.yaml@branch-25.10 + uses: rapidsai/shared-workflows/.github/workflows/conda-cpp-post-build-checks.yaml@cuda13.0 with: build_type: pull-request symbol_exclusions: (void (thrust::|cub::)) conda-python-build: needs: conda-cpp-build secrets: inherit - uses: rapidsai/shared-workflows/.github/workflows/conda-python-build.yaml@branch-25.10 + uses: rapidsai/shared-workflows/.github/workflows/conda-python-build.yaml@cuda13.0 with: build_type: pull-request script: ci/build_python.sh conda-python-tests: needs: [conda-python-build, changed-files] secrets: inherit - uses: rapidsai/shared-workflows/.github/workflows/conda-python-tests.yaml@branch-25.10 + uses: rapidsai/shared-workflows/.github/workflows/conda-python-tests.yaml@cuda13.0 if: fromJSON(needs.changed-files.outputs.changed_file_groups).test_python with: build_type: pull-request @@ -151,19 +151,28 @@ jobs: conda-java-tests: needs: [conda-cpp-build, changed-files] secrets: inherit - uses: rapidsai/shared-workflows/.github/workflows/custom-job.yaml@branch-25.10 + uses: rapidsai/shared-workflows/.github/workflows/custom-job.yaml@cuda13.0 if: fromJSON(needs.changed-files.outputs.changed_file_groups).test_java + # Artifacts are not published from these jobs, so it's safe to run for multiple CUDA versions. + # If these jobs start producing artifacts, the names will have to differentiate between CUDA versions. + strategy: + fail-fast: false + matrix: + cuda_version: + - '12.9.1' + - '13.0.0' with: build_type: pull-request node_type: "gpu-l4-latest-1" arch: "amd64" - container_image: "rapidsai/ci-conda:25.10-latest" + container_image: "rapidsai/ci-conda:25.10-cuda${{ matrix.cuda_version }}-ubuntu24.04-py3.13" script: "ci/test_java.sh" + artifact-name: "cuvs-java-cuda${{ matrix.cuda_version }}" file_to_upload: "java/cuvs-java/target/" docs-build: needs: conda-python-build secrets: inherit - uses: rapidsai/shared-workflows/.github/workflows/custom-job.yaml@branch-25.10 + uses: rapidsai/shared-workflows/.github/workflows/custom-job.yaml@cuda13.0 with: build_type: pull-request node_type: "gpu-l4-latest-1" @@ -173,27 +182,43 @@ jobs: rust-build: needs: conda-cpp-build secrets: inherit - uses: rapidsai/shared-workflows/.github/workflows/custom-job.yaml@branch-25.10 + uses: rapidsai/shared-workflows/.github/workflows/custom-job.yaml@cuda13.0 + # Artifacts are not published from these jobs, so it's safe to run for multiple CUDA versions. + # If these jobs start producing artifacts, the names will have to differentiate between CUDA versions. + strategy: + fail-fast: false + matrix: + cuda_version: + - '12.9.1' + - '13.0.0' with: build_type: pull-request node_type: "gpu-l4-latest-1" arch: "amd64" - container_image: "rapidsai/ci-conda:25.10-latest" + container_image: "rapidsai/ci-conda:25.10-cuda${{ matrix.cuda_version }}-ubuntu24.04-py3.13" script: "ci/build_rust.sh" go-build: needs: conda-cpp-build secrets: inherit - uses: rapidsai/shared-workflows/.github/workflows/custom-job.yaml@branch-25.10 + uses: rapidsai/shared-workflows/.github/workflows/custom-job.yaml@cuda13.0 + # Artifacts are not published from these jobs, so it's safe to run for multiple CUDA versions. + # If these jobs start producing artifacts, the names will have to differentiate between CUDA versions. + strategy: + fail-fast: false + matrix: + cuda_version: + - '12.9.1' + - '13.0.0' with: build_type: pull-request node_type: "gpu-l4-latest-1" arch: "amd64" - container_image: "rapidsai/ci-conda:25.10-latest" + container_image: "rapidsai/ci-conda:25.10-cuda${{ matrix.cuda_version }}-ubuntu24.04-py3.13" script: "ci/build_go.sh" wheel-build-libcuvs: needs: checks secrets: inherit - uses: rapidsai/shared-workflows/.github/workflows/wheels-build.yaml@branch-25.10 + uses: rapidsai/shared-workflows/.github/workflows/wheels-build.yaml@cuda13.0 with: build_type: pull-request script: ci/build_wheel_libcuvs.sh @@ -204,7 +229,7 @@ jobs: wheel-build-cuvs: needs: wheel-build-libcuvs secrets: inherit - uses: rapidsai/shared-workflows/.github/workflows/wheels-build.yaml@branch-25.10 + uses: rapidsai/shared-workflows/.github/workflows/wheels-build.yaml@cuda13.0 with: build_type: pull-request script: ci/build_wheel_cuvs.sh @@ -213,7 +238,7 @@ jobs: wheel-tests-cuvs: needs: [wheel-build-cuvs, changed-files] secrets: inherit - uses: rapidsai/shared-workflows/.github/workflows/wheels-test.yaml@branch-25.10 + uses: rapidsai/shared-workflows/.github/workflows/wheels-test.yaml@cuda13.0 if: fromJSON(needs.changed-files.outputs.changed_file_groups).test_python with: build_type: pull-request @@ -221,10 +246,10 @@ jobs: devcontainer: secrets: inherit needs: telemetry-setup - uses: rapidsai/shared-workflows/.github/workflows/build-in-devcontainer.yaml@branch-25.10 + uses: rapidsai/shared-workflows/.github/workflows/build-in-devcontainer.yaml@cuda13.0 with: arch: '["amd64", "arm64"]' - cuda: '["12.9"]' + cuda: '["13.0"]' node_type: "cpu8" rapids-aux-secret-1: GIST_REPO_READ_ORG_GITHUB_TOKEN env: | diff --git a/.github/workflows/publish-rust.yaml b/.github/workflows/publish-rust.yaml index 7834e6bc96..c74af7e7f1 100644 --- a/.github/workflows/publish-rust.yaml +++ b/.github/workflows/publish-rust.yaml @@ -10,8 +10,13 @@ defaults: jobs: rust-publish: runs-on: ubuntu-latest + strategy: + fail-fast: false + matrix: + cuda_version: + - '12.9.1' container: - image: rapidsai/ci-conda:25.10-latest + image: "rapidsai/ci-conda:25.10-cuda${{ matrix.cuda_version }}-ubuntu24.04-py3.13" steps: - uses: actions/checkout@v4 - name: Check if release build diff --git a/.github/workflows/test.yaml b/.github/workflows/test.yaml index 8479ce38a4..51418cbd22 100644 --- a/.github/workflows/test.yaml +++ b/.github/workflows/test.yaml @@ -25,7 +25,7 @@ on: jobs: conda-cpp-checks: secrets: inherit - uses: rapidsai/shared-workflows/.github/workflows/conda-cpp-post-build-checks.yaml@branch-25.10 + uses: rapidsai/shared-workflows/.github/workflows/conda-cpp-post-build-checks.yaml@cuda13.0 with: build_type: ${{ inputs.build_type }} branch: ${{ inputs.branch }} @@ -34,7 +34,7 @@ jobs: symbol_exclusions: (void (thrust::|cub::)) conda-cpp-tests: secrets: inherit - uses: rapidsai/shared-workflows/.github/workflows/conda-cpp-tests.yaml@branch-25.10 + uses: rapidsai/shared-workflows/.github/workflows/conda-cpp-tests.yaml@cuda13.0 with: build_type: ${{ inputs.build_type }} branch: ${{ inputs.branch }} @@ -43,7 +43,7 @@ jobs: sha: ${{ inputs.sha }} conda-python-tests: secrets: inherit - uses: rapidsai/shared-workflows/.github/workflows/conda-python-tests.yaml@branch-25.10 + uses: rapidsai/shared-workflows/.github/workflows/conda-python-tests.yaml@cuda13.0 with: build_type: ${{ inputs.build_type }} branch: ${{ inputs.branch }} @@ -52,7 +52,15 @@ jobs: sha: ${{ inputs.sha }} conda-java-tests: secrets: inherit - uses: rapidsai/shared-workflows/.github/workflows/custom-job.yaml@branch-25.10 + uses: rapidsai/shared-workflows/.github/workflows/custom-job.yaml@cuda13.0 + # Artifacts are not published from these jobs, so it's safe to run for multiple CUDA versions. + # If these jobs start producing artifacts, the names will have to differentiate between CUDA versions. + strategy: + fail-fast: false + matrix: + cuda_version: + - '12.9.1' + - '13.0.0' with: build_type: ${{ inputs.build_type }} branch: ${{ inputs.branch }} @@ -60,11 +68,11 @@ jobs: sha: ${{ inputs.sha }} node_type: "gpu-l4-latest-1" arch: "amd64" - container_image: "rapidsai/ci-conda:25.10-latest" + container_image: "rapidsai/ci-conda:25.10-cuda${{ matrix.cuda_version }}-ubuntu24.04-py3.13" script: "ci/test_java.sh" wheel-tests-cuvs: secrets: inherit - uses: rapidsai/shared-workflows/.github/workflows/wheels-test.yaml@branch-25.10 + uses: rapidsai/shared-workflows/.github/workflows/wheels-test.yaml@cuda13.0 with: build_type: ${{ inputs.build_type }} branch: ${{ inputs.branch }} diff --git a/.github/workflows/trigger-breaking-change-alert.yaml b/.github/workflows/trigger-breaking-change-alert.yaml index 48bf37afc4..72751d071b 100644 --- a/.github/workflows/trigger-breaking-change-alert.yaml +++ b/.github/workflows/trigger-breaking-change-alert.yaml @@ -12,7 +12,7 @@ jobs: trigger-notifier: if: contains(github.event.pull_request.labels.*.name, 'breaking') secrets: inherit - uses: rapidsai/shared-workflows/.github/workflows/breaking-change-alert.yaml@branch-25.10 + uses: rapidsai/shared-workflows/.github/workflows/breaking-change-alert.yaml@cuda13.0 with: sender_login: ${{ github.event.sender.login }} sender_avatar: ${{ github.event.sender.avatar_url }} diff --git a/README.md b/README.md index 1604428492..a81dd1b64f 100755 --- a/README.md +++ b/README.md @@ -95,8 +95,11 @@ conda install -c rapidsai -c conda-forge cuvs The cuVS Python package can also be `installed through pip `_. -For CUDA 12 packages: ```bash +# CUDA 13 +pip install cuvs-cu13 --extra-index-url=https://pypi.nvidia.com + +# CUDA 12 pip install cuvs-cu12 --extra-index-url=https://pypi.nvidia.com ``` @@ -104,7 +107,11 @@ pip install cuvs-cu12 --extra-index-url=https://pypi.nvidia.com If installing a version that has not yet been released, the `rapidsai` channel can be replaced with `rapidsai-nightly`: ```bash -conda install -c rapidsai-nightly -c conda-forge cuvs=25.10 +# CUDA 13 +conda install -c rapidsai-nightly -c conda-forge cuvs=25.10 cuda-version=13.0 + +# CUDA 12 +conda install -c rapidsai-nightly -c conda-forge cuvs=25.10 cuda-version=12.9 ``` cuVS also has `pip` wheel packages that can be installed. 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. diff --git a/ci/build_java.sh b/ci/build_java.sh index 175cc9a559..b692bbbffd 100755 --- a/ci/build_java.sh +++ b/ci/build_java.sh @@ -41,6 +41,9 @@ set +e rapids-logger "Run Java build" +RAPIDS_CUDA_MAJOR="${RAPIDS_CUDA_VERSION%%.*}" +export RAPIDS_CUDA_MAJOR + bash ./build.sh java "${EXTRA_BUILD_ARGS[@]}" rapids-logger "Test script exiting with value: $EXITCODE" diff --git a/ci/build_python.sh b/ci/build_python.sh index 8ee818341c..d4956cce78 100755 --- a/ci/build_python.sh +++ b/ci/build_python.sh @@ -1,5 +1,5 @@ #!/bin/bash -# Copyright (c) 2022-2024, NVIDIA CORPORATION. +# Copyright (c) 2022-2025, NVIDIA CORPORATION. set -euo pipefail @@ -54,7 +54,7 @@ sccache --zero-stats # Build cuvs-bench-cpu only in one CUDA major version since it only depends on # python version RAPIDS_CUDA_MAJOR="${RAPIDS_CUDA_VERSION%%.*}" -if [[ ${RAPIDS_CUDA_MAJOR} == "12" ]]; then +if [[ ${RAPIDS_CUDA_MAJOR} == "13" ]]; then rattler-build build --recipe conda/recipes/cuvs-bench-cpu \ "${RATTLER_ARGS[@]}" \ "${RATTLER_CHANNELS[@]}" diff --git a/ci/test_java.sh b/ci/test_java.sh index 7c7d77860a..20c19b6dee 100755 --- a/ci/test_java.sh +++ b/ci/test_java.sh @@ -12,6 +12,9 @@ nvidia-smi rapids-logger "Run Java build and tests" +RAPIDS_CUDA_MAJOR="${RAPIDS_CUDA_VERSION%%.*}" +export RAPIDS_CUDA_MAJOR + # TODO: switch to installing pre-built artifacts instead of rebuilding in test jobs # ref: https://github.com/rapidsai/cuvs/issues/868 ci/build_java.sh --run-java-tests diff --git a/conda/environments/all_cuda-129_arch-aarch64.yaml b/conda/environments/all_cuda-129_arch-aarch64.yaml index 45bb00e8df..a877f09e23 100644 --- a/conda/environments/all_cuda-129_arch-aarch64.yaml +++ b/conda/environments/all_cuda-129_arch-aarch64.yaml @@ -15,9 +15,9 @@ dependencies: - cuda-nvcc - cuda-nvtx-dev - cuda-profiler-api -- cuda-python>=12.6.2,<13.0a0 +- cuda-python>=12.9.2,<13.0a0 - cuda-version=12.9 -- cupy>=12.0.0 +- cupy>=13.6.0 - cxx-compiler - cython>=3.0.0 - dlpack>=0.8,<1.0 diff --git a/conda/environments/all_cuda-129_arch-x86_64.yaml b/conda/environments/all_cuda-129_arch-x86_64.yaml index 46a0f234c7..588bfa8d64 100644 --- a/conda/environments/all_cuda-129_arch-x86_64.yaml +++ b/conda/environments/all_cuda-129_arch-x86_64.yaml @@ -15,9 +15,9 @@ dependencies: - cuda-nvcc - cuda-nvtx-dev - cuda-profiler-api -- cuda-python>=12.6.2,<13.0a0 +- cuda-python>=12.9.2,<13.0a0 - cuda-version=12.9 -- cupy>=12.0.0 +- cupy>=13.6.0 - cxx-compiler - cython>=3.0.0 - dlpack>=0.8,<1.0 diff --git a/conda/environments/all_cuda-130_arch-aarch64.yaml b/conda/environments/all_cuda-130_arch-aarch64.yaml new file mode 100644 index 0000000000..ebf59a3f99 --- /dev/null +++ b/conda/environments/all_cuda-130_arch-aarch64.yaml @@ -0,0 +1,56 @@ +# This file is generated by `rapids-dependency-file-generator`. +# To make changes, edit ../../dependencies.yaml and run `rapids-dependency-file-generator`. +channels: +- rapidsai +- rapidsai-nightly +- conda-forge +dependencies: +- _go_select *=cgo +- breathe>=4.35.0 +- c-compiler +- clang-tools==20.1.4 +- clang==20.1.4 +- cmake>=3.30.4 +- cuda-cudart-dev +- cuda-nvcc +- cuda-nvtx-dev +- cuda-profiler-api +- cuda-python>=13.0.1,<14.0a0 +- cuda-version=13.0 +- cupy>=13.6.0 +- cxx-compiler +- cython>=3.0.0 +- dlpack>=0.8,<1.0 +- doxygen>=1.8.20 +- gcc_linux-aarch64=14.* +- go +- graphviz +- ipython +- libclang==20.1.4 +- libcublas-dev +- libcurand-dev +- libcusolver-dev +- libcusparse-dev +- librmm==25.10.*,>=0.0.0a0 +- make +- nccl>=2.19 +- ninja +- numpy>=1.23,<3.0a0 +- numpydoc +- openblas +- pre-commit +- pylibraft==25.10.*,>=0.0.0a0 +- pytest +- pytest-cov +- rapids-build-backend>=0.4.0,<0.5.0.dev0 +- recommonmark +- rust +- scikit-build-core>=0.10.0 +- scikit-learn +- sphinx-copybutton +- sphinx-markdown-tables +- sphinx>=8.0.0 +- sysroot_linux-aarch64==2.28 +- pip: + - nvidia-sphinx-theme +name: all_cuda-130_arch-aarch64 diff --git a/conda/environments/all_cuda-130_arch-x86_64.yaml b/conda/environments/all_cuda-130_arch-x86_64.yaml new file mode 100644 index 0000000000..8433f4adff --- /dev/null +++ b/conda/environments/all_cuda-130_arch-x86_64.yaml @@ -0,0 +1,56 @@ +# This file is generated by `rapids-dependency-file-generator`. +# To make changes, edit ../../dependencies.yaml and run `rapids-dependency-file-generator`. +channels: +- rapidsai +- rapidsai-nightly +- conda-forge +dependencies: +- _go_select *=cgo +- breathe>=4.35.0 +- c-compiler +- clang-tools==20.1.4 +- clang==20.1.4 +- cmake>=3.30.4 +- cuda-cudart-dev +- cuda-nvcc +- cuda-nvtx-dev +- cuda-profiler-api +- cuda-python>=13.0.1,<14.0a0 +- cuda-version=13.0 +- cupy>=13.6.0 +- cxx-compiler +- cython>=3.0.0 +- dlpack>=0.8,<1.0 +- doxygen>=1.8.20 +- gcc_linux-64=14.* +- go +- graphviz +- ipython +- libclang==20.1.4 +- libcublas-dev +- libcurand-dev +- libcusolver-dev +- libcusparse-dev +- librmm==25.10.*,>=0.0.0a0 +- make +- nccl>=2.19 +- ninja +- numpy>=1.23,<3.0a0 +- numpydoc +- openblas +- pre-commit +- pylibraft==25.10.*,>=0.0.0a0 +- pytest +- pytest-cov +- rapids-build-backend>=0.4.0,<0.5.0.dev0 +- recommonmark +- rust +- scikit-build-core>=0.10.0 +- scikit-learn +- sphinx-copybutton +- sphinx-markdown-tables +- sphinx>=8.0.0 +- sysroot_linux-64==2.28 +- pip: + - nvidia-sphinx-theme +name: all_cuda-130_arch-x86_64 diff --git a/conda/environments/bench_ann_cuda-129_arch-aarch64.yaml b/conda/environments/bench_ann_cuda-129_arch-aarch64.yaml index 7ffb09777b..7735923261 100644 --- a/conda/environments/bench_ann_cuda-129_arch-aarch64.yaml +++ b/conda/environments/bench_ann_cuda-129_arch-aarch64.yaml @@ -15,9 +15,9 @@ dependencies: - cuda-nvcc - cuda-nvtx-dev - cuda-profiler-api -- cuda-python>=12.6.2,<13.0a0 +- cuda-python>=12.9.2,<13.0a0 - cuda-version=12.9 -- cupy>=12.0.0 +- cupy>=13.6.0 - cuvs==25.10.*,>=0.0.0a0 - cxx-compiler - cython>=3.0.0 diff --git a/conda/environments/bench_ann_cuda-129_arch-x86_64.yaml b/conda/environments/bench_ann_cuda-129_arch-x86_64.yaml index fcd0f89053..96e2c34d5d 100644 --- a/conda/environments/bench_ann_cuda-129_arch-x86_64.yaml +++ b/conda/environments/bench_ann_cuda-129_arch-x86_64.yaml @@ -15,9 +15,9 @@ dependencies: - cuda-nvcc - cuda-nvtx-dev - cuda-profiler-api -- cuda-python>=12.6.2,<13.0a0 +- cuda-python>=12.9.2,<13.0a0 - cuda-version=12.9 -- cupy>=12.0.0 +- cupy>=13.6.0 - cuvs==25.10.*,>=0.0.0a0 - cxx-compiler - cython>=3.0.0 diff --git a/conda/environments/bench_ann_cuda-130_arch-aarch64.yaml b/conda/environments/bench_ann_cuda-130_arch-aarch64.yaml new file mode 100644 index 0000000000..939a24b445 --- /dev/null +++ b/conda/environments/bench_ann_cuda-130_arch-aarch64.yaml @@ -0,0 +1,49 @@ +# This file is generated by `rapids-dependency-file-generator`. +# To make changes, edit ../../dependencies.yaml and run `rapids-dependency-file-generator`. +channels: +- rapidsai +- rapidsai-nightly +- conda-forge +dependencies: +- benchmark>=1.8.2 +- c-compiler +- clang-tools==20.1.4 +- clang==20.1.4 +- click +- cmake>=3.30.4 +- cuda-cudart-dev +- cuda-nvcc +- cuda-nvtx-dev +- cuda-profiler-api +- cuda-python>=13.0.1,<14.0a0 +- cuda-version=13.0 +- cupy>=13.6.0 +- cuvs==25.10.*,>=0.0.0a0 +- cxx-compiler +- cython>=3.0.0 +- dlpack>=0.8,<1.0 +- gcc_linux-aarch64=14.* +- glog>=0.6.0 +- h5py>=3.8.0 +- libclang==20.1.4 +- libcublas-dev +- libcurand-dev +- libcusolver-dev +- libcusparse-dev +- libcuvs==25.10.*,>=0.0.0a0 +- librmm==25.10.*,>=0.0.0a0 +- matplotlib-base>=3.9 +- nccl>=2.19 +- ninja +- nlohmann_json>=3.11.2 +- openblas +- pandas +- pylibraft==25.10.*,>=0.0.0a0 +- pyyaml +- rapids-build-backend>=0.4.0,<0.5.0.dev0 +- requests +- scikit-learn +- setuptools +- sysroot_linux-aarch64==2.28 +- wheel +name: bench_ann_cuda-130_arch-aarch64 diff --git a/conda/environments/bench_ann_cuda-130_arch-x86_64.yaml b/conda/environments/bench_ann_cuda-130_arch-x86_64.yaml new file mode 100644 index 0000000000..296c1fc03a --- /dev/null +++ b/conda/environments/bench_ann_cuda-130_arch-x86_64.yaml @@ -0,0 +1,52 @@ +# This file is generated by `rapids-dependency-file-generator`. +# To make changes, edit ../../dependencies.yaml and run `rapids-dependency-file-generator`. +channels: +- rapidsai +- rapidsai-nightly +- conda-forge +dependencies: +- benchmark>=1.8.2 +- c-compiler +- clang-tools==20.1.4 +- clang==20.1.4 +- click +- cmake>=3.30.4 +- cuda-cudart-dev +- cuda-nvcc +- cuda-nvtx-dev +- cuda-profiler-api +- cuda-python>=13.0.1,<14.0a0 +- cuda-version=13.0 +- cupy>=13.6.0 +- cuvs==25.10.*,>=0.0.0a0 +- cxx-compiler +- cython>=3.0.0 +- dlpack>=0.8,<1.0 +- gcc_linux-64=14.* +- glog>=0.6.0 +- h5py>=3.8.0 +- libaio +- libboost-devel=1.87 +- libclang==20.1.4 +- libcublas-dev +- libcurand-dev +- libcusolver-dev +- libcusparse-dev +- libcuvs==25.10.*,>=0.0.0a0 +- librmm==25.10.*,>=0.0.0a0 +- matplotlib-base>=3.9 +- mkl-devel=2023 +- nccl>=2.19 +- ninja +- nlohmann_json>=3.11.2 +- openblas +- pandas +- pylibraft==25.10.*,>=0.0.0a0 +- pyyaml +- rapids-build-backend>=0.4.0,<0.5.0.dev0 +- requests +- scikit-learn +- setuptools +- sysroot_linux-64==2.28 +- wheel +name: bench_ann_cuda-130_arch-x86_64 diff --git a/conda/environments/go_cuda-130_arch-aarch64.yaml b/conda/environments/go_cuda-130_arch-aarch64.yaml new file mode 100644 index 0000000000..3359dd084a --- /dev/null +++ b/conda/environments/go_cuda-130_arch-aarch64.yaml @@ -0,0 +1,32 @@ +# This file is generated by `rapids-dependency-file-generator`. +# To make changes, edit ../../dependencies.yaml and run `rapids-dependency-file-generator`. +channels: +- rapidsai +- rapidsai-nightly +- conda-forge +dependencies: +- _go_select *=cgo +- c-compiler +- clang-tools==20.1.4 +- clang==20.1.4 +- cmake>=3.30.4 +- cuda-cudart-dev +- cuda-nvcc +- cuda-nvtx-dev +- cuda-profiler-api +- cuda-version=13.0 +- cxx-compiler +- dlpack>=0.8,<1.0 +- gcc_linux-aarch64=14.* +- go +- libclang==20.1.4 +- libcublas-dev +- libcurand-dev +- libcusolver-dev +- libcusparse-dev +- libcuvs==25.10.*,>=0.0.0a0 +- libraft==25.10.*,>=0.0.0a0 +- nccl>=2.19 +- ninja +- sysroot_linux-aarch64==2.28 +name: go_cuda-130_arch-aarch64 diff --git a/conda/environments/go_cuda-130_arch-x86_64.yaml b/conda/environments/go_cuda-130_arch-x86_64.yaml new file mode 100644 index 0000000000..a9a6a90cf1 --- /dev/null +++ b/conda/environments/go_cuda-130_arch-x86_64.yaml @@ -0,0 +1,32 @@ +# This file is generated by `rapids-dependency-file-generator`. +# To make changes, edit ../../dependencies.yaml and run `rapids-dependency-file-generator`. +channels: +- rapidsai +- rapidsai-nightly +- conda-forge +dependencies: +- _go_select *=cgo +- c-compiler +- clang-tools==20.1.4 +- clang==20.1.4 +- cmake>=3.30.4 +- cuda-cudart-dev +- cuda-nvcc +- cuda-nvtx-dev +- cuda-profiler-api +- cuda-version=13.0 +- cxx-compiler +- dlpack>=0.8,<1.0 +- gcc_linux-64=14.* +- go +- libclang==20.1.4 +- libcublas-dev +- libcurand-dev +- libcusolver-dev +- libcusparse-dev +- libcuvs==25.10.*,>=0.0.0a0 +- libraft==25.10.*,>=0.0.0a0 +- nccl>=2.19 +- ninja +- sysroot_linux-64==2.28 +name: go_cuda-130_arch-x86_64 diff --git a/conda/environments/rust_cuda-130_arch-aarch64.yaml b/conda/environments/rust_cuda-130_arch-aarch64.yaml new file mode 100644 index 0000000000..751e44f36b --- /dev/null +++ b/conda/environments/rust_cuda-130_arch-aarch64.yaml @@ -0,0 +1,31 @@ +# This file is generated by `rapids-dependency-file-generator`. +# To make changes, edit ../../dependencies.yaml and run `rapids-dependency-file-generator`. +channels: +- rapidsai +- rapidsai-nightly +- conda-forge +dependencies: +- c-compiler +- clang-tools==20.1.4 +- clang==20.1.4 +- cmake>=3.30.4 +- cuda-cudart-dev +- cuda-nvcc +- cuda-nvtx-dev +- cuda-profiler-api +- cuda-version=13.0 +- cxx-compiler +- gcc_linux-aarch64=14.* +- libclang==20.1.4 +- libcublas-dev +- libcurand-dev +- libcusolver-dev +- libcusparse-dev +- libcuvs==25.10.*,>=0.0.0a0 +- libraft==25.10.*,>=0.0.0a0 +- make +- nccl>=2.19 +- ninja +- rust +- sysroot_linux-aarch64==2.28 +name: rust_cuda-130_arch-aarch64 diff --git a/conda/environments/rust_cuda-130_arch-x86_64.yaml b/conda/environments/rust_cuda-130_arch-x86_64.yaml new file mode 100644 index 0000000000..56f7c9830e --- /dev/null +++ b/conda/environments/rust_cuda-130_arch-x86_64.yaml @@ -0,0 +1,31 @@ +# This file is generated by `rapids-dependency-file-generator`. +# To make changes, edit ../../dependencies.yaml and run `rapids-dependency-file-generator`. +channels: +- rapidsai +- rapidsai-nightly +- conda-forge +dependencies: +- c-compiler +- clang-tools==20.1.4 +- clang==20.1.4 +- cmake>=3.30.4 +- cuda-cudart-dev +- cuda-nvcc +- cuda-nvtx-dev +- cuda-profiler-api +- cuda-version=13.0 +- cxx-compiler +- gcc_linux-64=14.* +- libclang==20.1.4 +- libcublas-dev +- libcurand-dev +- libcusolver-dev +- libcusparse-dev +- libcuvs==25.10.*,>=0.0.0a0 +- libraft==25.10.*,>=0.0.0a0 +- make +- nccl>=2.19 +- ninja +- rust +- sysroot_linux-64==2.28 +name: rust_cuda-130_arch-x86_64 diff --git a/conda/recipes/cuvs-bench/recipe.yaml b/conda/recipes/cuvs-bench/recipe.yaml index a823877926..fe12bea9ae 100644 --- a/conda/recipes/cuvs-bench/recipe.yaml +++ b/conda/recipes/cuvs-bench/recipe.yaml @@ -40,7 +40,7 @@ requirements: - benchmark - click - cuda-cudart - - cupy >=12.0.0 + - cupy >=13.6.0 - cuvs =${{ version }} - glog ${{ glog_version }} - h5py ${{ h5py_version }} diff --git a/conda/recipes/cuvs/recipe.yaml b/conda/recipes/cuvs/recipe.yaml index e7b9f6de99..f3fd05f8f3 100644 --- a/conda/recipes/cuvs/recipe.yaml +++ b/conda/recipes/cuvs/recipe.yaml @@ -58,7 +58,9 @@ requirements: - python =${{ py_version }} - rapids-build-backend >=0.4.0,<0.5.0.dev0 - scikit-build-core >=0.10.0 - - cuda-python >=12.6.2,<13.0a0 + - if: cuda_major == "12" + then: cuda-python >=12.9.2,<13.0a0 + else: cuda-python >=13.0.1,<14.0a0 - cuda-cudart-dev run: - ${{ pin_compatible("cuda-version", upper_bound="x", lower_bound="x") }} @@ -66,7 +68,9 @@ requirements: - pylibraft =${{ minor_version }} - python - numpy >=1.23,<3.0a0 - - cuda-python >=12.6.2,<13.0a0 + - if: cuda_major == "12" + then: cuda-python >=12.9.2,<13.0a0 + else: cuda-python >=13.0.1,<14.0a0 - cuda-cudart ignore_run_exports: by_name: diff --git a/cpp/bench/ann/src/common/util.hpp b/cpp/bench/ann/src/common/util.hpp index b3bcc77950..11361c88fe 100644 --- a/cpp/bench/ann/src/common/util.hpp +++ b/cpp/bench/ann/src/common/util.hpp @@ -446,10 +446,15 @@ inline auto cuda_info() throw std::runtime_error{"cuda_info: call to cudaGetDeviceProperties failed with code " + std::to_string(err_code)}; } + int clockRate = 0; + int memoryClockRate = 0; + err_code = cudaDeviceGetAttribute(&clockRate, cudaDevAttrClockRate, dev); + err_code = cudaDeviceGetAttribute(&memoryClockRate, cudaDevAttrMemoryClockRate, dev); + props.emplace_back("gpu_name", std::string(device_prop.name)); props.emplace_back("gpu_sm_count", std::to_string(device_prop.multiProcessorCount)); - props.emplace_back("gpu_sm_freq", std::to_string(device_prop.clockRate * 1e3)); - props.emplace_back("gpu_mem_freq", std::to_string(device_prop.memoryClockRate * 1e3)); + props.emplace_back("gpu_sm_freq", std::to_string(clockRate * 1e3)); + props.emplace_back("gpu_mem_freq", std::to_string(memoryClockRate * 1e3)); props.emplace_back("gpu_mem_bus_width", std::to_string(device_prop.memoryBusWidth)); props.emplace_back("gpu_mem_global_size", std::to_string(device_prop.totalGlobalMem)); props.emplace_back("gpu_mem_shared_size", std::to_string(device_prop.sharedMemPerMultiprocessor)); diff --git a/cpp/include/cuvs/neighbors/common.hpp b/cpp/include/cuvs/neighbors/common.hpp index 70da22a09d..b49b7a4e8a 100644 --- a/cpp/include/cuvs/neighbors/common.hpp +++ b/cpp/include/cuvs/neighbors/common.hpp @@ -485,7 +485,7 @@ namespace filtering { enum class FilterType { None, Bitmap, Bitset }; struct base_filter { - virtual ~base_filter() = default; + ~base_filter() = default; virtual FilterType get_filter_type() const = 0; }; diff --git a/dependencies.yaml b/dependencies.yaml index 7570ce1973..c370e90f88 100644 --- a/dependencies.yaml +++ b/dependencies.yaml @@ -3,7 +3,7 @@ files: all: output: conda matrix: - cuda: ["12.9"] + cuda: ["12.9", "13.0"] arch: [x86_64, aarch64] includes: - build @@ -14,6 +14,7 @@ files: - clang - cuda - cuda_version + - depends_on_cuda_python - depends_on_cupy - depends_on_librmm - depends_on_pylibraft @@ -29,7 +30,7 @@ files: bench_ann: output: conda matrix: - cuda: ["12.9"] + cuda: ["12.9", "13.0"] arch: [x86_64, aarch64] includes: - bench @@ -39,6 +40,7 @@ files: - clang - cuda - cuda_version + - depends_on_cuda_python - depends_on_cupy - depends_on_pylibraft - depends_on_libcuvs @@ -85,7 +87,7 @@ files: rust: output: conda matrix: - cuda: ["12.9"] + cuda: ["12.9", "13.0"] arch: [x86_64, aarch64] includes: # clang/libclang only needed for bindgen support @@ -100,7 +102,7 @@ files: go: output: conda matrix: - cuda: ["12.9"] + cuda: ["12.9", "13.0"] arch: [x86_64, aarch64] includes: - clang @@ -165,6 +167,7 @@ files: includes: - build_cython - build_py_cuvs + - depends_on_cuda_python - depends_on_libcuvs - depends_on_libraft - depends_on_librmm @@ -175,6 +178,7 @@ files: extras: table: project includes: + - depends_on_cuda_python - depends_on_libcuvs - depends_on_pylibraft - run_py_cuvs @@ -239,26 +243,21 @@ dependencies: - output_types: [conda] packages: - c-compiler + - cuda-nvcc - cxx-compiler specific: - output_types: conda matrices: - matrix: arch: x86_64 - cuda: "12.*" packages: - gcc_linux-64=14.* - sysroot_linux-64==2.28 - matrix: arch: aarch64 - cuda: "12.*" packages: - gcc_linux-aarch64=14.* - sysroot_linux-aarch64==2.28 - - output_types: conda - matrices: - - matrix: {cuda: "12.*"} - packages: [cuda-nvcc] rapids_build_setuptools: common: - output_types: [conda, requirements, pyproject] @@ -271,16 +270,6 @@ dependencies: - output_types: [conda] packages: - dlpack>=0.8,<1.0 - specific: - - output_types: [conda, requirements, pyproject] - matrices: - - matrix: - cuda: "12.*" - packages: - - &cuda_python12 cuda-python>=12.6.2,<13.0a0 - - matrix: - packages: - - &cuda_python cuda-python checks: common: - output_types: [conda, requirements] @@ -317,20 +306,21 @@ dependencies: cuda: "12.9" packages: - cuda-version=12.9 - cuda: - specific: - - output_types: conda - matrices: - matrix: - cuda: "12.*" + cuda: "13.0" packages: - - cuda-nvtx-dev - - cuda-cudart-dev - - cuda-profiler-api - - libcublas-dev - - libcurand-dev - - libcusolver-dev - - libcusparse-dev + - cuda-version=13.0 + cuda: + common: + - output_types: [conda] + packages: + - cuda-nvtx-dev + - cuda-cudart-dev + - cuda-profiler-api + - libcublas-dev + - libcurand-dev + - libcusolver-dev + - libcusparse-dev cuda_wheels: specific: - output_types: [requirements, pyproject] @@ -343,6 +333,14 @@ dependencies: - nvidia-curand-cu12 - nvidia-cusolver-cu12 - nvidia-cusparse-cu12 + - matrix: + cuda: "13.*" + use_cuda_wheels: "true" + packages: + - nvidia-cublas-cu13 + - nvidia-curand-cu13 + - nvidia-cusolver-cu13 + - nvidia-cusparse-cu13 - matrix: use_cuda_wheels: "false" packages: @@ -358,14 +356,21 @@ dependencies: common: - output_types: conda packages: - - cupy>=12.0.0 + - cupy>=13.6.0 + # NOTE: This is intentionally not broken into groups by a 'cuda_suffixed' selector like + # other packages with -cu{nn}x suffixes in this file. + # All RAPIDS wheel builds (including in devcontainers) expect cupy to be suffixed. specific: - output_types: [requirements, pyproject] matrices: - - matrix: {cuda: "12.*"} + - matrix: + cuda: "12.*" packages: - - cupy-cuda12x>=12.0.0 - - {matrix: null, packages: [cupy-cuda12x>=12.0.0]} + - cupy-cuda12x>=13.6.0 + # fallback to CUDA 13 versions if 'cuda' is '13.*' or not provided + - matrix: + packages: + - &cupy_cu13 cupy-cuda13x>=13.6.0 test_libcuvs: common: - output_types: [conda] @@ -434,16 +439,6 @@ dependencies: - output_types: [conda, requirements, pyproject] packages: - &numpy numpy>=1.23,<3.0a0 - specific: - - output_types: [conda, requirements, pyproject] - matrices: - - matrix: - cuda: "12.*" - packages: - - *cuda_python12 - - matrix: - packages: - - *cuda_python test_python_common: common: - output_types: [conda, requirements, pyproject] @@ -500,6 +495,18 @@ dependencies: - output_types: [requirements, pyproject] packages: - matplotlib>=3.9 + depends_on_cuda_python: + specific: + - output_types: [conda, requirements, pyproject] + matrices: + - matrix: + cuda: "12.*" + packages: + - cuda-python>=12.9.2,<13.0a0 + # fallback to CUDA 13 versions if 'cuda' is '13.*' or not provided + - matrix: + packages: + - cuda-python>=13.0.1,<14.0a0 depends_on_cuvs: common: - output_types: conda @@ -528,6 +535,11 @@ dependencies: cuda_suffixed: "true" packages: - libcuvs-cu12==25.10.*,>=0.0.0a0 + - matrix: + cuda: "13.*" + cuda_suffixed: "true" + packages: + - libcuvs-cu13==25.10.*,>=0.0.0a0 - {matrix: null, packages: [*libcuvs_unsuffixed]} depends_on_libcuvs_tests: common: @@ -552,6 +564,11 @@ dependencies: cuda_suffixed: "true" packages: - libraft-cu12==25.10.*,>=0.0.0a0 + - matrix: + cuda: "13.*" + cuda_suffixed: "true" + packages: + - libraft-cu13==25.10.*,>=0.0.0a0 - {matrix: null, packages: [*libraft_unsuffixed]} depends_on_librmm: common: @@ -571,6 +588,11 @@ dependencies: cuda_suffixed: "true" packages: - librmm-cu12==25.10.*,>=0.0.0a0 + - matrix: + cuda: "13.*" + cuda_suffixed: "true" + packages: + - librmm-cu13==25.10.*,>=0.0.0a0 - {matrix: null, packages: [*librmm_unsuffixed]} depends_on_pylibraft: common: @@ -590,12 +612,17 @@ dependencies: cuda_suffixed: "true" packages: - pylibraft-cu12==25.10.*,>=0.0.0a0 + - matrix: + cuda: "13.*" + cuda_suffixed: "true" + packages: + - pylibraft-cu13==25.10.*,>=0.0.0a0 - {matrix: null, packages: [*pylibraft_unsuffixed]} depends_on_nccl: common: - output_types: conda packages: - - &nccl_unsuffixed nccl>=2.19 + - nccl>=2.19 specific: - output_types: [pyproject, requirements] matrices: @@ -604,5 +631,10 @@ dependencies: cuda_suffixed: "true" packages: - nvidia-nccl-cu12>=2.19 + - matrix: + cuda: "13.*" + cuda_suffixed: "true" + packages: + - nvidia-nccl-cu13>=2.19 - matrix: packages: diff --git a/docs/source/build.rst b/docs/source/build.rst index 4acf3bea54..26c361dfbc 100644 --- a/docs/source/build.rst +++ b/docs/source/build.rst @@ -43,6 +43,10 @@ C/C++ Package .. code-block:: bash + # CUDA 13 + conda install -c rapidsai -c conda-forge libcuvs cuda-version=13.0 + + # CUDA 12 conda install -c rapidsai -c conda-forge libcuvs cuda-version=12.9 Python Package @@ -50,6 +54,10 @@ Python Package .. code-block:: bash + # CUDA 13 + conda install -c rapidsai -c conda-forge cuvs cuda-version=13.0 + + # CUDA 12 conda install -c rapidsai -c conda-forge cuvs cuda-version=12.9 Python through Pip @@ -57,10 +65,12 @@ Python through Pip The cuVS Python package can also be `installed through pip `_. -For CUDA 12 packages: - .. code-block:: bash + # CUDA 13 + pip install cuvs-cu13 --extra-index-url=https://pypi.nvidia.com + + # CUDA 12 pip install cuvs-cu12 --extra-index-url=https://pypi.nvidia.com Note: these packages statically link the C and C++ libraries so the `libcuvs` and `libcuvs_c` shared libraries won't be readily available to use in your code. @@ -86,7 +96,7 @@ Conda environment scripts are provided for installing the necessary dependencies .. code-block:: bash - conda env create --name cuvs -f conda/environments/all_cuda-128_arch-x86_64.yaml + conda env create --name cuvs -f conda/environments/all_cuda-130_arch-x86_64.yaml conda activate cuvs The recommended way to build and install cuVS from source is to use the `build.sh` script in the root of the repository. This script can build both the C++ and Python artifacts and provides CMake options for building and installing the headers, tests, benchmarks, and the pre-compiled shared library. diff --git a/docs/source/cuvs_bench/index.rst b/docs/source/cuvs_bench/index.rst index 5a6f0f17e2..cfaa30acb2 100644 --- a/docs/source/cuvs_bench/index.rst +++ b/docs/source/cuvs_bench/index.rst @@ -65,7 +65,7 @@ Conda conda activate cuvs_benchmarks # to install GPU package: - conda install -c rapidsai -c conda-forge cuvs-bench= cuda-version=12.9* + conda install -c rapidsai -c conda-forge cuvs-bench= cuda-version=13.0* # to install CPU package for usage in CPU-only systems: conda install -c rapidsai -c conda-forge cuvs-bench-cpu diff --git a/examples/go/README.md b/examples/go/README.md index bcbea09642..093e87d17d 100644 --- a/examples/go/README.md +++ b/examples/go/README.md @@ -7,7 +7,7 @@ This package provides Go bindings for the cuVS (CUDA Vector Search) library. The required dependencies can be installed with a simple command (which creates your build environment): ```bash -conda env create --name go -f conda/environments/go_cuda-128_arch-x86_64.yaml +conda env create --name go -f conda/environments/go_cuda-130_arch-x86_64.yaml conda activate go ``` You may prefer to use `mamba`, as it provides significant speedup over `conda`. diff --git a/examples/rust/README.md b/examples/rust/README.md index 98fe526975..1d2671de34 100644 --- a/examples/rust/README.md +++ b/examples/rust/README.md @@ -9,7 +9,7 @@ Once the minimum requirements are satisfied, this example template application c You may follow these steps to quickly get set up: ```bash -conda env create --name rust -f conda/environments/rust_cuda-128_arch-x86_64.yaml +conda env create --name rust -f conda/environments/rust_cuda-130_arch-x86_64.yaml conda activate rust ``` You may prefer to use `mamba`, as it provides significant speedup over `conda`. diff --git a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/GPUInfoProviderImpl.java b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/GPUInfoProviderImpl.java index c2422318ac..ab4b0dde1d 100644 --- a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/GPUInfoProviderImpl.java +++ b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/GPUInfoProviderImpl.java @@ -17,6 +17,7 @@ import static com.nvidia.cuvs.internal.common.LinkerHelper.C_INT; import static com.nvidia.cuvs.internal.common.Util.checkCudaError; +import static com.nvidia.cuvs.internal.common.Util.cudaGetDeviceProperties; import static com.nvidia.cuvs.internal.panama.headers_h.cudaMemGetInfo; import static com.nvidia.cuvs.internal.panama.headers_h_1.*; @@ -52,8 +53,8 @@ private static List getAvailableGpusInfo() { MemorySegment deviceProp = cudaDeviceProp.allocate(localArena); for (int i = 0; i < numGpuCount; i++) { - returnValue = cudaGetDeviceProperties_v2(deviceProp, i); - checkCudaError(returnValue, "cudaGetDeviceProperties_v2"); + returnValue = cudaGetDeviceProperties(deviceProp, i); + checkCudaError(returnValue, "cudaGetDeviceProperties"); GPUInfo gpuInfo = new GPUInfo( diff --git a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/common/Util.java b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/common/Util.java index 09adf59309..41f6b3736d 100644 --- a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/common/Util.java +++ b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/common/Util.java @@ -29,10 +29,12 @@ import com.nvidia.cuvs.internal.panama.DLTensor; import com.nvidia.cuvs.internal.panama.headers_h; import java.lang.foreign.Arena; +import java.lang.foreign.FunctionDescriptor; import java.lang.foreign.Linker; import java.lang.foreign.MemoryLayout; import java.lang.foreign.MemoryLayout.PathElement; import java.lang.foreign.MemorySegment; +import java.lang.foreign.SymbolLookup; import java.lang.invoke.MethodHandle; import java.lang.invoke.VarHandle; import java.util.BitSet; @@ -46,6 +48,11 @@ private Util() {} private static final Linker LINKER = Linker.nativeLinker(); + static final SymbolLookup SYMBOL_LOOKUP = + SymbolLookup.libraryLookup(System.mapLibraryName("cuvs_c"), Arena.ofAuto()) + .or(SymbolLookup.loaderLookup()) + .or(Linker.nativeLinker().defaultLookup()); + /** * Bindings for {@code cudaMemcpyAsync}; differently from the {@code headers_h} bindings (which are * automatically generated by {@code jextract}), these bindings specify the {@code critical} linker option, @@ -58,6 +65,26 @@ private Util() {} LINKER.downcallHandle( cudaMemcpyAsync$address(), cudaMemcpyAsync$descriptor(), Linker.Option.critical(true)); + private static final String cudaGetDevicePropertiesSymbolName = + "12".equals(System.getenv("RAPIDS_CUDA_MAJOR")) + ? "cudaGetDeviceProperties_v2" + : "cudaGetDeviceProperties"; + + private static final MethodHandle cudaGetDeviceProperties$mh = + LINKER.downcallHandle( + SYMBOL_LOOKUP + .find(cudaGetDevicePropertiesSymbolName) + .orElseThrow(UnsatisfiedLinkError::new), + FunctionDescriptor.of(headers_h.C_INT, headers_h.C_POINTER, headers_h.C_INT)); + + public static int cudaGetDeviceProperties(MemorySegment prop, int device) { + try { + return (int) cudaGetDeviceProperties$mh.invokeExact(prop, device); + } catch (Throwable ex$) { + throw new AssertionError("should not reach here", ex$); + } + } + /** * Checks the result value of a (CuVS) native method handle call. * diff --git a/python/cuvs/pyproject.toml b/python/cuvs/pyproject.toml index 99e31fc620..b5b7dba8cf 100644 --- a/python/cuvs/pyproject.toml +++ b/python/cuvs/pyproject.toml @@ -31,7 +31,7 @@ authors = [ license = { text = "Apache-2.0" } requires-python = ">=3.10" dependencies = [ - "cuda-python", + "cuda-python>=13.0.1,<14.0a0", "libcuvs==25.10.*,>=0.0.0a0", "numpy>=1.23,<3.0a0", "pylibraft==25.10.*,>=0.0.0a0", @@ -47,7 +47,7 @@ classifiers = [ [project.optional-dependencies] test = [ - "cupy-cuda12x>=12.0.0", + "cupy-cuda13x>=13.6.0", "pytest", "pytest-cov", "scikit-learn", @@ -117,7 +117,7 @@ regex = "(?P.*)" [tool.rapids-build-backend] requires = [ "cmake>=3.30.4", - "cuda-python", + "cuda-python>=13.0.1,<14.0a0", "cython>=3.0.0", "libcuvs==25.10.*,>=0.0.0a0", "libraft==25.10.*,>=0.0.0a0", From 6254dd9993bd9e3afc503b82a599147933c7eed5 Mon Sep 17 00:00:00 2001 From: Victor Lafargue Date: Tue, 2 Sep 2025 16:57:11 +0200 Subject: [PATCH 50/93] MG C API (#1160) Authors: - Victor Lafargue (https://github.com/viclafargue) - Corey J. Nolet (https://github.com/cjnolet) Approvers: - Ben Frederickson (https://github.com/benfred) URL: https://github.com/rapidsai/cuvs/pull/1160 --- cpp/CMakeLists.txt | 10 + cpp/include/cuvs/core/c_api.h | 17 + cpp/include/cuvs/neighbors/ivf_flat.hpp | 100 ++- cpp/include/cuvs/neighbors/mg_cagra.h | 285 +++++++++ cpp/include/cuvs/neighbors/mg_common.h | 67 ++ cpp/include/cuvs/neighbors/mg_ivf_flat.h | 285 +++++++++ cpp/include/cuvs/neighbors/mg_ivf_pq.h | 285 +++++++++ cpp/src/core/c_api.cpp | 17 + cpp/src/neighbors/iface/generate_iface.py | 5 +- .../iface/iface_flat_half_int64_t.cu | 96 +++ cpp/src/neighbors/mg/generate_mg.py | 13 +- .../neighbors/mg/mg_cagra_float_uint32_t.cu | 10 +- .../neighbors/mg/mg_cagra_half_uint32_t.cu | 10 +- .../neighbors/mg/mg_cagra_int8_t_uint32_t.cu | 10 +- .../neighbors/mg/mg_cagra_uint8_t_uint32_t.cu | 10 +- cpp/src/neighbors/mg/mg_flat_half_int64_t.cu | 96 +++ cpp/src/neighbors/mg/snmg.cuh | 15 +- cpp/src/neighbors/mg_cagra_c.cpp | 419 +++++++++++++ cpp/src/neighbors/mg_ivf_flat_c.cpp | 416 +++++++++++++ cpp/src/neighbors/mg_ivf_pq_c.cpp | 408 +++++++++++++ cpp/tests/CMakeLists.txt | 2 + cpp/tests/neighbors/ann_mg_c.cu | 341 +++++++++++ cpp/tests/neighbors/run_mg_c.c | 576 ++++++++++++++++++ docs/source/c_api/neighbors.rst | 1 + docs/source/c_api/neighbors_mg.rst | 257 ++++++++ 25 files changed, 3739 insertions(+), 12 deletions(-) create mode 100644 cpp/include/cuvs/neighbors/mg_cagra.h create mode 100644 cpp/include/cuvs/neighbors/mg_common.h create mode 100644 cpp/include/cuvs/neighbors/mg_ivf_flat.h create mode 100644 cpp/include/cuvs/neighbors/mg_ivf_pq.h create mode 100644 cpp/src/neighbors/iface/iface_flat_half_int64_t.cu create mode 100644 cpp/src/neighbors/mg/mg_flat_half_int64_t.cu create mode 100644 cpp/src/neighbors/mg_cagra_c.cpp create mode 100644 cpp/src/neighbors/mg_ivf_flat_c.cpp create mode 100644 cpp/src/neighbors/mg_ivf_pq_c.cpp create mode 100644 cpp/tests/neighbors/ann_mg_c.cu create mode 100644 cpp/tests/neighbors/run_mg_c.c create mode 100644 docs/source/c_api/neighbors_mg.rst diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt index 7830004c8b..04621aaffd 100644 --- a/cpp/CMakeLists.txt +++ b/cpp/CMakeLists.txt @@ -295,6 +295,7 @@ if(BUILD_SHARED_LIBS) if(BUILD_MG_ALGOS) set(CUVS_MG_ALGOS src/neighbors/mg/mg_flat_float_int64_t.cu + src/neighbors/mg/mg_flat_half_int64_t.cu src/neighbors/mg/mg_flat_int8_t_int64_t.cu src/neighbors/mg/mg_flat_uint8_t_int64_t.cu src/neighbors/mg/mg_pq_float_int64_t.cu @@ -410,6 +411,7 @@ if(BUILD_SHARED_LIBS) src/neighbors/iface/iface_cagra_int8_t_uint32_t.cu src/neighbors/iface/iface_cagra_uint8_t_uint32_t.cu src/neighbors/iface/iface_flat_float_int64_t.cu + src/neighbors/iface/iface_flat_half_int64_t.cu src/neighbors/iface/iface_flat_int8_t_int64_t.cu src/neighbors/iface/iface_flat_uint8_t_int64_t.cu src/neighbors/iface/iface_pq_float_int64_t.cu @@ -712,6 +714,13 @@ target_compile_definitions(cuvs::cuvs INTERFACE $<$:NVTX_ENAB # ################################################################################################ # * cuvs_c ------------------------------------------------------------------------------- + + if(BUILD_MG_ALGOS) + set(CUVS_C_MG_ALGOS src/neighbors/mg_ivf_pq_c.cpp src/neighbors/mg_ivf_flat_c.cpp + src/neighbors/mg_cagra_c.cpp + ) + endif() + if(BUILD_C_LIBRARY) add_library( cuvs_c SHARED @@ -728,6 +737,7 @@ target_compile_definitions(cuvs::cuvs INTERFACE $<$:NVTX_ENAB src/preprocessing/quantize/binary_c.cpp src/preprocessing/quantize/scalar_c.cpp src/distance/pairwise_distance_c.cpp + ${CUVS_C_MG_ALGOS} ) if(BUILD_CAGRA_HNSWLIB) diff --git a/cpp/include/cuvs/core/c_api.h b/cpp/include/cuvs/core/c_api.h index e3cd9bd40b..e2fc4ebc94 100644 --- a/cpp/include/cuvs/core/c_api.h +++ b/cpp/include/cuvs/core/c_api.h @@ -75,6 +75,23 @@ cuvsError_t cuvsResourcesCreate(cuvsResources_t* res); */ cuvsError_t cuvsResourcesDestroy(cuvsResources_t res); +/** + * @brief Create an Initialized opaque C handle for C++ type `raft::device_resources_snmg` + * for multi-GPU operations + * + * @param[in] res cuvsResources_t opaque C handle + * @return cuvsError_t + */ +cuvsError_t cuvsMultiGpuResourcesCreate(cuvsResources_t* res); + +/** + * @brief Destroy and de-allocate opaque C handle for C++ type `raft::device_resources_snmg` + * + * @param[in] res cuvsResources_t opaque C handle + * @return cuvsError_t + */ +cuvsError_t cuvsMultiGpuResourcesDestroy(cuvsResources_t res); + /** * @brief Set cudaStream_t on cuvsResources_t to queue CUDA kernels on APIs * that accept a cuvsResources_t handle diff --git a/cpp/include/cuvs/neighbors/ivf_flat.hpp b/cpp/include/cuvs/neighbors/ivf_flat.hpp index c206fefde1..445cd527bc 100644 --- a/cpp/include/cuvs/neighbors/ivf_flat.hpp +++ b/cpp/include/cuvs/neighbors/ivf_flat.hpp @@ -1,5 +1,5 @@ /* - * Copyright (c) 2024, NVIDIA CORPORATION. + * Copyright (c) 2024-2025, NVIDIA CORPORATION. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -2042,6 +2042,28 @@ auto build(const raft::resources& clique, raft::host_matrix_view index_dataset) -> cuvs::neighbors::mg_index, float, int64_t>; +/// \ingroup mg_cpp_index_build +/** + * @brief Builds a multi-GPU index + * + * Usage example: + * @code{.cpp} + * raft::device_resources_snmg clique; + * cuvs::neighbors::mg_index_params index_params; + * auto index = cuvs::neighbors::ivf_flat::build(clique, index_params, index_dataset); + * @endcode + * + * @param[in] clique a `raft::resources` object specifying the NCCL clique configuration + * @param[in] index_params configure the index building + * @param[in] index_dataset a row-major matrix on host [n_rows, dim] + * + * @return the constructed IVF-Flat MG index + */ +auto build(const raft::resources& clique, + const cuvs::neighbors::mg_index_params& index_params, + raft::host_matrix_view index_dataset) + -> cuvs::neighbors::mg_index, half, int64_t>; + /// \ingroup mg_cpp_index_build /** * @brief Builds a multi-GPU index @@ -2112,6 +2134,30 @@ void extend(const raft::resources& clique, raft::host_matrix_view new_vectors, std::optional> new_indices); +/// \ingroup mg_cpp_index_extend +/** + * @brief Extends a multi-GPU index + * + * Usage example: + * @code{.cpp} + * raft::device_resources_snmg clique; + * cuvs::neighbors::mg_index_params index_params; + * auto index = cuvs::neighbors::ivf_flat::build(clique, index_params, index_dataset); + * cuvs::neighbors::ivf_flat::extend(clique, index, new_vectors, std::nullopt); + * @endcode + * + * @param[in] clique a `raft::resources` object specifying the NCCL clique configuration + * @param[in] index the pre-built index + * @param[in] new_vectors a row-major matrix on host [n_rows, dim] + * @param[in] new_indices optional vector on host [n_rows], + * `std::nullopt` means default continuous range `[0...n_rows)` + * + */ +void extend(const raft::resources& clique, + cuvs::neighbors::mg_index, half, int64_t>& index, + raft::host_matrix_view new_vectors, + std::optional> new_indices); + /// \ingroup mg_cpp_index_extend /** * @brief Extends a multi-GPU index @@ -2191,6 +2237,35 @@ void search(const raft::resources& clique, raft::host_matrix_view neighbors, raft::host_matrix_view distances); +/// \ingroup mg_cpp_index_search +/** + * @brief Searches a multi-GPU index + * + * Usage example: + * @code{.cpp} + * raft::device_resources_snmg clique; + * cuvs::neighbors::mg_index_params index_params; + * auto index = cuvs::neighbors::ivf_flat::build(clique, index_params, index_dataset); + * cuvs::neighbors::mg_search_params search_params; + * cuvs::neighbors::ivf_flat::search(clique, index, search_params, queries, neighbors, + * distances); + * @endcode + * + * @param[in] clique a `raft::resources` object specifying the NCCL clique configuration + * @param[in] index the pre-built index + * @param[in] search_params configure the index search + * @param[in] queries a row-major matrix on host [n_rows, dim] + * @param[out] neighbors a row-major matrix on host [n_rows, n_neighbors] + * @param[out] distances a row-major matrix on host [n_rows, n_neighbors] + * + */ +void search(const raft::resources& clique, + const cuvs::neighbors::mg_index, half, int64_t>& index, + const cuvs::neighbors::mg_search_params& search_params, + raft::host_matrix_view queries, + raft::host_matrix_view neighbors, + raft::host_matrix_view distances); + /// \ingroup mg_cpp_index_search /** * @brief Searches a multi-GPU index @@ -2276,6 +2351,29 @@ void serialize( const cuvs::neighbors::mg_index, float, int64_t>& index, const std::string& filename); +/// \ingroup mg_cpp_serialize +/** + * @brief Serializes a multi-GPU index + * + * Usage example: + * @code{.cpp} + * raft::device_resources_snmg clique; + * cuvs::neighbors::mg_index_params index_params; + * auto index = cuvs::neighbors::ivf_flat::build(clique, index_params, index_dataset); + * const std::string filename = "mg_index.cuvs"; + * cuvs::neighbors::ivf_flat::serialize(clique, index, filename); + * @endcode + * + * @param[in] clique a `raft::resources` object specifying the NCCL clique configuration + * @param[in] index the pre-built index + * @param[in] filename path to the file to be serialized + * + */ +void serialize( + const raft::resources& clique, + const cuvs::neighbors::mg_index, half, int64_t>& index, + const std::string& filename); + /// \ingroup mg_cpp_serialize /** * @brief Serializes a multi-GPU index diff --git a/cpp/include/cuvs/neighbors/mg_cagra.h b/cpp/include/cuvs/neighbors/mg_cagra.h new file mode 100644 index 0000000000..8680ba9202 --- /dev/null +++ b/cpp/include/cuvs/neighbors/mg_cagra.h @@ -0,0 +1,285 @@ +/* + * Copyright (c) 2025, NVIDIA CORPORATION. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include +#include +#include +#include +#include + +#ifdef __cplusplus +extern "C" { +#endif + +/** + * @defgroup mg_cagra_c_index_params Multi-GPU CAGRA index build parameters + * @{ + */ + +/** + * @brief Multi-GPU parameters to build CAGRA Index + * + * This structure extends the base CAGRA index parameters with multi-GPU specific settings. + */ +struct cuvsMultiGpuCagraIndexParams { + /** Base CAGRA index parameters */ + cuvsCagraIndexParams_t base_params; + /** Distribution mode for multi-GPU setup */ + cuvsMultiGpuDistributionMode mode; +}; + +typedef struct cuvsMultiGpuCagraIndexParams* cuvsMultiGpuCagraIndexParams_t; + +/** + * @brief Allocate Multi-GPU CAGRA Index params, and populate with default values + * + * @param[in] index_params cuvsMultiGpuCagraIndexParams_t to allocate + * @return cuvsError_t + */ +cuvsError_t cuvsMultiGpuCagraIndexParamsCreate(cuvsMultiGpuCagraIndexParams_t* index_params); + +/** + * @brief De-allocate Multi-GPU CAGRA Index params + * + * @param[in] index_params + * @return cuvsError_t + */ +cuvsError_t cuvsMultiGpuCagraIndexParamsDestroy(cuvsMultiGpuCagraIndexParams_t index_params); + +/** + * @} + */ + +/** + * @defgroup mg_cagra_c_search_params Multi-GPU CAGRA index search parameters + * @{ + */ + +/** + * @brief Multi-GPU parameters to search CAGRA index + * + * This structure extends the base CAGRA search parameters with multi-GPU specific settings. + */ +struct cuvsMultiGpuCagraSearchParams { + /** Base CAGRA search parameters */ + cuvsCagraSearchParams_t base_params; + /** Replicated search mode */ + cuvsMultiGpuReplicatedSearchMode search_mode; + /** Sharded merge mode */ + cuvsMultiGpuShardedMergeMode merge_mode; + /** Number of rows per batch */ + int64_t n_rows_per_batch; +}; + +typedef struct cuvsMultiGpuCagraSearchParams* cuvsMultiGpuCagraSearchParams_t; + +/** + * @brief Allocate Multi-GPU CAGRA search params, and populate with default values + * + * @param[in] params cuvsMultiGpuCagraSearchParams_t to allocate + * @return cuvsError_t + */ +cuvsError_t cuvsMultiGpuCagraSearchParamsCreate(cuvsMultiGpuCagraSearchParams_t* params); + +/** + * @brief De-allocate Multi-GPU CAGRA search params + * + * @param[in] params + * @return cuvsError_t + */ +cuvsError_t cuvsMultiGpuCagraSearchParamsDestroy(cuvsMultiGpuCagraSearchParams_t params); + +/** + * @} + */ + +/** + * @defgroup mg_cagra_c_index Multi-GPU CAGRA index + * @{ + */ + +/** + * @brief Struct to hold address of cuvs::neighbors::mg_index and its active trained + * dtype + */ +typedef struct { + uintptr_t addr; + DLDataType dtype; +} cuvsMultiGpuCagraIndex; + +typedef cuvsMultiGpuCagraIndex* cuvsMultiGpuCagraIndex_t; + +/** + * @brief Allocate Multi-GPU CAGRA index + * + * @param[in] index cuvsMultiGpuCagraIndex_t to allocate + * @return cuvsError_t + */ +cuvsError_t cuvsMultiGpuCagraIndexCreate(cuvsMultiGpuCagraIndex_t* index); + +/** + * @brief De-allocate Multi-GPU CAGRA index + * + * @param[in] index cuvsMultiGpuCagraIndex_t to de-allocate + * @return cuvsError_t + */ +cuvsError_t cuvsMultiGpuCagraIndexDestroy(cuvsMultiGpuCagraIndex_t index); + +/** + * @} + */ + +/** + * @defgroup mg_cagra_c_index_build Multi-GPU CAGRA index build + * @{ + */ + +/** + * @brief Build a Multi-GPU CAGRA index + * + * @param[in] res cuvsResources_t opaque C handle + * @param[in] params Multi-GPU CAGRA index parameters + * @param[in] dataset_tensor DLManagedTensor* training dataset + * @param[out] index Multi-GPU CAGRA index + * @return cuvsError_t + */ +cuvsError_t cuvsMultiGpuCagraBuild(cuvsResources_t res, + cuvsMultiGpuCagraIndexParams_t params, + DLManagedTensor* dataset_tensor, + cuvsMultiGpuCagraIndex_t index); + +/** + * @} + */ + +/** + * @defgroup mg_cagra_c_index_search Multi-GPU CAGRA index search + * @{ + */ + +/** + * @brief Search a Multi-GPU CAGRA index + * + * @param[in] res cuvsResources_t opaque C handle + * @param[in] params Multi-GPU CAGRA search parameters + * @param[in] index Multi-GPU CAGRA index + * @param[in] queries_tensor DLManagedTensor* queries dataset + * @param[out] neighbors_tensor DLManagedTensor* output neighbors + * @param[out] distances_tensor DLManagedTensor* output distances + * @return cuvsError_t + */ +cuvsError_t cuvsMultiGpuCagraSearch(cuvsResources_t res, + cuvsMultiGpuCagraSearchParams_t params, + cuvsMultiGpuCagraIndex_t index, + DLManagedTensor* queries_tensor, + DLManagedTensor* neighbors_tensor, + DLManagedTensor* distances_tensor); + +/** + * @} + */ + +/** + * @defgroup mg_cagra_c_index_extend Multi-GPU CAGRA index extend + * @{ + */ + +/** + * @brief Extend a Multi-GPU CAGRA index + * + * @param[in] res cuvsResources_t opaque C handle + * @param[in,out] index Multi-GPU CAGRA index to extend + * @param[in] new_vectors_tensor DLManagedTensor* new vectors to add + * @param[in] new_indices_tensor DLManagedTensor* new indices (optional, can be NULL) + * @return cuvsError_t + */ +cuvsError_t cuvsMultiGpuCagraExtend(cuvsResources_t res, + cuvsMultiGpuCagraIndex_t index, + DLManagedTensor* new_vectors_tensor, + DLManagedTensor* new_indices_tensor); + +/** + * @} + */ + +/** + * @defgroup mg_cagra_c_index_serialize Multi-GPU CAGRA index serialize + * @{ + */ + +/** + * @brief Serialize a Multi-GPU CAGRA index to file + * + * @param[in] res cuvsResources_t opaque C handle + * @param[in] index Multi-GPU CAGRA index to serialize + * @param[in] filename Path to the output file + * @return cuvsError_t + */ +cuvsError_t cuvsMultiGpuCagraSerialize(cuvsResources_t res, + cuvsMultiGpuCagraIndex_t index, + const char* filename); + +/** + * @} + */ + +/** + * @defgroup mg_cagra_c_index_deserialize Multi-GPU CAGRA index deserialize + * @{ + */ + +/** + * @brief Deserialize a Multi-GPU CAGRA index from file + * + * @param[in] res cuvsResources_t opaque C handle + * @param[in] filename Path to the input file + * @param[out] index Multi-GPU CAGRA index + * @return cuvsError_t + */ +cuvsError_t cuvsMultiGpuCagraDeserialize(cuvsResources_t res, + const char* filename, + cuvsMultiGpuCagraIndex_t index); + +/** + * @} + */ + +/** + * @defgroup mg_cagra_c_index_distribute Multi-GPU CAGRA index distribute + * @{ + */ + +/** + * @brief Distribute a local CAGRA index to create a Multi-GPU index + * + * @param[in] res cuvsResources_t opaque C handle + * @param[in] filename Path to the local index file + * @param[out] index Multi-GPU CAGRA index + * @return cuvsError_t + */ +cuvsError_t cuvsMultiGpuCagraDistribute(cuvsResources_t res, + const char* filename, + cuvsMultiGpuCagraIndex_t index); + +/** + * @} + */ + +#ifdef __cplusplus +} +#endif diff --git a/cpp/include/cuvs/neighbors/mg_common.h b/cpp/include/cuvs/neighbors/mg_common.h new file mode 100644 index 0000000000..06b6771573 --- /dev/null +++ b/cpp/include/cuvs/neighbors/mg_common.h @@ -0,0 +1,67 @@ +/* + * Copyright (c) 2025, NVIDIA CORPORATION. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include +#include + +#ifdef __cplusplus +extern "C" { +#endif + +/** + * @defgroup mg_c_common_types Multi-GPU common types and enums + * @{ + */ + +/** + * @brief Distribution mode for multi-GPU indexes + */ +typedef enum { + /** Index is replicated on each device, favors throughput */ + CUVS_NEIGHBORS_MG_REPLICATED, + /** Index is split on several devices, favors scaling */ + CUVS_NEIGHBORS_MG_SHARDED +} cuvsMultiGpuDistributionMode; + +/** + * @brief Search mode when using a replicated index + */ +typedef enum { + /** Search queries are split to maintain equal load on GPUs */ + CUVS_NEIGHBORS_MG_LOAD_BALANCER, + /** Each search query is processed by a single GPU in a round-robin fashion */ + CUVS_NEIGHBORS_MG_ROUND_ROBIN +} cuvsMultiGpuReplicatedSearchMode; + +/** + * @brief Merge mode when using a sharded index + */ +typedef enum { + /** Search batches are merged on the root rank */ + CUVS_NEIGHBORS_MG_MERGE_ON_ROOT_RANK, + /** Search batches are merged in a tree reduction fashion */ + CUVS_NEIGHBORS_MG_TREE_MERGE +} cuvsMultiGpuShardedMergeMode; + +/** + * @} + */ + +#ifdef __cplusplus +} +#endif diff --git a/cpp/include/cuvs/neighbors/mg_ivf_flat.h b/cpp/include/cuvs/neighbors/mg_ivf_flat.h new file mode 100644 index 0000000000..b6ead7e108 --- /dev/null +++ b/cpp/include/cuvs/neighbors/mg_ivf_flat.h @@ -0,0 +1,285 @@ +/* + * Copyright (c) 2025, NVIDIA CORPORATION. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include +#include +#include +#include +#include + +#ifdef __cplusplus +extern "C" { +#endif + +/** + * @defgroup mg_ivf_flat_c_index_params Multi-GPU IVF-Flat index build parameters + * @{ + */ + +/** + * @brief Multi-GPU parameters to build IVF-Flat Index + * + * This structure extends the base IVF-Flat index parameters with multi-GPU specific settings. + */ +struct cuvsMultiGpuIvfFlatIndexParams { + /** Base IVF-Flat index parameters */ + cuvsIvfFlatIndexParams_t base_params; + /** Distribution mode for multi-GPU setup */ + cuvsMultiGpuDistributionMode mode; +}; + +typedef struct cuvsMultiGpuIvfFlatIndexParams* cuvsMultiGpuIvfFlatIndexParams_t; + +/** + * @brief Allocate Multi-GPU IVF-Flat Index params, and populate with default values + * + * @param[in] index_params cuvsMultiGpuIvfFlatIndexParams_t to allocate + * @return cuvsError_t + */ +cuvsError_t cuvsMultiGpuIvfFlatIndexParamsCreate(cuvsMultiGpuIvfFlatIndexParams_t* index_params); + +/** + * @brief De-allocate Multi-GPU IVF-Flat Index params + * + * @param[in] index_params + * @return cuvsError_t + */ +cuvsError_t cuvsMultiGpuIvfFlatIndexParamsDestroy(cuvsMultiGpuIvfFlatIndexParams_t index_params); + +/** + * @} + */ + +/** + * @defgroup mg_ivf_flat_c_search_params Multi-GPU IVF-Flat index search parameters + * @{ + */ + +/** + * @brief Multi-GPU parameters to search IVF-Flat index + * + * This structure extends the base IVF-Flat search parameters with multi-GPU specific settings. + */ +struct cuvsMultiGpuIvfFlatSearchParams { + /** Base IVF-Flat search parameters */ + cuvsIvfFlatSearchParams_t base_params; + /** Replicated search mode */ + cuvsMultiGpuReplicatedSearchMode search_mode; + /** Sharded merge mode */ + cuvsMultiGpuShardedMergeMode merge_mode; + /** Number of rows per batch */ + int64_t n_rows_per_batch; +}; + +typedef struct cuvsMultiGpuIvfFlatSearchParams* cuvsMultiGpuIvfFlatSearchParams_t; + +/** + * @brief Allocate Multi-GPU IVF-Flat search params, and populate with default values + * + * @param[in] params cuvsMultiGpuIvfFlatSearchParams_t to allocate + * @return cuvsError_t + */ +cuvsError_t cuvsMultiGpuIvfFlatSearchParamsCreate(cuvsMultiGpuIvfFlatSearchParams_t* params); + +/** + * @brief De-allocate Multi-GPU IVF-Flat search params + * + * @param[in] params + * @return cuvsError_t + */ +cuvsError_t cuvsMultiGpuIvfFlatSearchParamsDestroy(cuvsMultiGpuIvfFlatSearchParams_t params); + +/** + * @} + */ + +/** + * @defgroup mg_ivf_flat_c_index Multi-GPU IVF-Flat index + * @{ + */ + +/** + * @brief Struct to hold address of cuvs::neighbors::mg_index and its active + * trained dtype + */ +typedef struct { + uintptr_t addr; + DLDataType dtype; +} cuvsMultiGpuIvfFlatIndex; + +typedef cuvsMultiGpuIvfFlatIndex* cuvsMultiGpuIvfFlatIndex_t; + +/** + * @brief Allocate Multi-GPU IVF-Flat index + * + * @param[in] index cuvsMultiGpuIvfFlatIndex_t to allocate + * @return cuvsError_t + */ +cuvsError_t cuvsMultiGpuIvfFlatIndexCreate(cuvsMultiGpuIvfFlatIndex_t* index); + +/** + * @brief De-allocate Multi-GPU IVF-Flat index + * + * @param[in] index cuvsMultiGpuIvfFlatIndex_t to de-allocate + * @return cuvsError_t + */ +cuvsError_t cuvsMultiGpuIvfFlatIndexDestroy(cuvsMultiGpuIvfFlatIndex_t index); + +/** + * @} + */ + +/** + * @defgroup mg_ivf_flat_c_index_build Multi-GPU IVF-Flat index build + * @{ + */ + +/** + * @brief Build a Multi-GPU IVF-Flat index + * + * @param[in] res cuvsResources_t opaque C handle + * @param[in] params Multi-GPU IVF-Flat index parameters + * @param[in] dataset_tensor DLManagedTensor* training dataset + * @param[out] index Multi-GPU IVF-Flat index + * @return cuvsError_t + */ +cuvsError_t cuvsMultiGpuIvfFlatBuild(cuvsResources_t res, + cuvsMultiGpuIvfFlatIndexParams_t params, + DLManagedTensor* dataset_tensor, + cuvsMultiGpuIvfFlatIndex_t index); + +/** + * @} + */ + +/** + * @defgroup mg_ivf_flat_c_index_search Multi-GPU IVF-Flat index search + * @{ + */ + +/** + * @brief Search a Multi-GPU IVF-Flat index + * + * @param[in] res cuvsResources_t opaque C handle + * @param[in] params Multi-GPU IVF-Flat search parameters + * @param[in] index Multi-GPU IVF-Flat index + * @param[in] queries_tensor DLManagedTensor* queries dataset + * @param[out] neighbors_tensor DLManagedTensor* output neighbors + * @param[out] distances_tensor DLManagedTensor* output distances + * @return cuvsError_t + */ +cuvsError_t cuvsMultiGpuIvfFlatSearch(cuvsResources_t res, + cuvsMultiGpuIvfFlatSearchParams_t params, + cuvsMultiGpuIvfFlatIndex_t index, + DLManagedTensor* queries_tensor, + DLManagedTensor* neighbors_tensor, + DLManagedTensor* distances_tensor); + +/** + * @} + */ + +/** + * @defgroup mg_ivf_flat_c_index_extend Multi-GPU IVF-Flat index extend + * @{ + */ + +/** + * @brief Extend a Multi-GPU IVF-Flat index + * + * @param[in] res cuvsResources_t opaque C handle + * @param[in,out] index Multi-GPU IVF-Flat index to extend + * @param[in] new_vectors_tensor DLManagedTensor* new vectors to add + * @param[in] new_indices_tensor DLManagedTensor* new indices (optional, can be NULL) + * @return cuvsError_t + */ +cuvsError_t cuvsMultiGpuIvfFlatExtend(cuvsResources_t res, + cuvsMultiGpuIvfFlatIndex_t index, + DLManagedTensor* new_vectors_tensor, + DLManagedTensor* new_indices_tensor); + +/** + * @} + */ + +/** + * @defgroup mg_ivf_flat_c_index_serialize Multi-GPU IVF-Flat index serialize + * @{ + */ + +/** + * @brief Serialize a Multi-GPU IVF-Flat index to file + * + * @param[in] res cuvsResources_t opaque C handle + * @param[in] index Multi-GPU IVF-Flat index to serialize + * @param[in] filename Path to the output file + * @return cuvsError_t + */ +cuvsError_t cuvsMultiGpuIvfFlatSerialize(cuvsResources_t res, + cuvsMultiGpuIvfFlatIndex_t index, + const char* filename); + +/** + * @} + */ + +/** + * @defgroup mg_ivf_flat_c_index_deserialize Multi-GPU IVF-Flat index deserialize + * @{ + */ + +/** + * @brief Deserialize a Multi-GPU IVF-Flat index from file + * + * @param[in] res cuvsResources_t opaque C handle + * @param[in] filename Path to the input file + * @param[out] index Multi-GPU IVF-Flat index + * @return cuvsError_t + */ +cuvsError_t cuvsMultiGpuIvfFlatDeserialize(cuvsResources_t res, + const char* filename, + cuvsMultiGpuIvfFlatIndex_t index); + +/** + * @} + */ + +/** + * @defgroup mg_ivf_flat_c_index_distribute Multi-GPU IVF-Flat index distribute + * @{ + */ + +/** + * @brief Distribute a local IVF-Flat index to create a Multi-GPU index + * + * @param[in] res cuvsResources_t opaque C handle + * @param[in] filename Path to the local index file + * @param[out] index Multi-GPU IVF-Flat index + * @return cuvsError_t + */ +cuvsError_t cuvsMultiGpuIvfFlatDistribute(cuvsResources_t res, + const char* filename, + cuvsMultiGpuIvfFlatIndex_t index); + +/** + * @} + */ + +#ifdef __cplusplus +} +#endif diff --git a/cpp/include/cuvs/neighbors/mg_ivf_pq.h b/cpp/include/cuvs/neighbors/mg_ivf_pq.h new file mode 100644 index 0000000000..6654bc3a89 --- /dev/null +++ b/cpp/include/cuvs/neighbors/mg_ivf_pq.h @@ -0,0 +1,285 @@ +/* + * Copyright (c) 2025, NVIDIA CORPORATION. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include +#include +#include +#include +#include + +#ifdef __cplusplus +extern "C" { +#endif + +/** + * @defgroup mg_ivf_pq_c_index_params Multi-GPU IVF-PQ index build parameters + * @{ + */ + +/** + * @brief Multi-GPU parameters to build IVF-PQ Index + * + * This structure extends the base IVF-PQ index parameters with multi-GPU specific settings. + */ +struct cuvsMultiGpuIvfPqIndexParams { + /** Base IVF-PQ index parameters */ + cuvsIvfPqIndexParams_t base_params; + /** Distribution mode for multi-GPU setup */ + cuvsMultiGpuDistributionMode mode; +}; + +typedef struct cuvsMultiGpuIvfPqIndexParams* cuvsMultiGpuIvfPqIndexParams_t; + +/** + * @brief Allocate Multi-GPU IVF-PQ Index params, and populate with default values + * + * @param[in] index_params cuvsMultiGpuIvfPqIndexParams_t to allocate + * @return cuvsError_t + */ +cuvsError_t cuvsMultiGpuIvfPqIndexParamsCreate(cuvsMultiGpuIvfPqIndexParams_t* index_params); + +/** + * @brief De-allocate Multi-GPU IVF-PQ Index params + * + * @param[in] index_params + * @return cuvsError_t + */ +cuvsError_t cuvsMultiGpuIvfPqIndexParamsDestroy(cuvsMultiGpuIvfPqIndexParams_t index_params); + +/** + * @} + */ + +/** + * @defgroup mg_ivf_pq_c_search_params Multi-GPU IVF-PQ index search parameters + * @{ + */ + +/** + * @brief Multi-GPU parameters to search IVF-PQ index + * + * This structure extends the base IVF-PQ search parameters with multi-GPU specific settings. + */ +struct cuvsMultiGpuIvfPqSearchParams { + /** Base IVF-PQ search parameters */ + cuvsIvfPqSearchParams_t base_params; + /** Replicated search mode */ + cuvsMultiGpuReplicatedSearchMode search_mode; + /** Sharded merge mode */ + cuvsMultiGpuShardedMergeMode merge_mode; + /** Number of rows per batch */ + int64_t n_rows_per_batch; +}; + +typedef struct cuvsMultiGpuIvfPqSearchParams* cuvsMultiGpuIvfPqSearchParams_t; + +/** + * @brief Allocate Multi-GPU IVF-PQ search params, and populate with default values + * + * @param[in] params cuvsMultiGpuIvfPqSearchParams_t to allocate + * @return cuvsError_t + */ +cuvsError_t cuvsMultiGpuIvfPqSearchParamsCreate(cuvsMultiGpuIvfPqSearchParams_t* params); + +/** + * @brief De-allocate Multi-GPU IVF-PQ search params + * + * @param[in] params + * @return cuvsError_t + */ +cuvsError_t cuvsMultiGpuIvfPqSearchParamsDestroy(cuvsMultiGpuIvfPqSearchParams_t params); + +/** + * @} + */ + +/** + * @defgroup mg_ivf_pq_c_index Multi-GPU IVF-PQ index + * @{ + */ + +/** + * @brief Struct to hold address of cuvs::neighbors::mg_index and its active trained + * dtype + */ +typedef struct { + uintptr_t addr; + DLDataType dtype; +} cuvsMultiGpuIvfPqIndex; + +typedef cuvsMultiGpuIvfPqIndex* cuvsMultiGpuIvfPqIndex_t; + +/** + * @brief Allocate Multi-GPU IVF-PQ index + * + * @param[in] index cuvsMultiGpuIvfPqIndex_t to allocate + * @return cuvsError_t + */ +cuvsError_t cuvsMultiGpuIvfPqIndexCreate(cuvsMultiGpuIvfPqIndex_t* index); + +/** + * @brief De-allocate Multi-GPU IVF-PQ index + * + * @param[in] index cuvsMultiGpuIvfPqIndex_t to de-allocate + * @return cuvsError_t + */ +cuvsError_t cuvsMultiGpuIvfPqIndexDestroy(cuvsMultiGpuIvfPqIndex_t index); + +/** + * @} + */ + +/** + * @defgroup mg_ivf_pq_c_index_build Multi-GPU IVF-PQ index build + * @{ + */ + +/** + * @brief Build a Multi-GPU IVF-PQ index + * + * @param[in] res cuvsResources_t opaque C handle + * @param[in] params Multi-GPU IVF-PQ index parameters + * @param[in] dataset_tensor DLManagedTensor* training dataset + * @param[out] index Multi-GPU IVF-PQ index + * @return cuvsError_t + */ +cuvsError_t cuvsMultiGpuIvfPqBuild(cuvsResources_t res, + cuvsMultiGpuIvfPqIndexParams_t params, + DLManagedTensor* dataset_tensor, + cuvsMultiGpuIvfPqIndex_t index); + +/** + * @} + */ + +/** + * @defgroup mg_ivf_pq_c_index_search Multi-GPU IVF-PQ index search + * @{ + */ + +/** + * @brief Search a Multi-GPU IVF-PQ index + * + * @param[in] res cuvsResources_t opaque C handle + * @param[in] params Multi-GPU IVF-PQ search parameters + * @param[in] index Multi-GPU IVF-PQ index + * @param[in] queries_tensor DLManagedTensor* queries dataset + * @param[out] neighbors_tensor DLManagedTensor* output neighbors + * @param[out] distances_tensor DLManagedTensor* output distances + * @return cuvsError_t + */ +cuvsError_t cuvsMultiGpuIvfPqSearch(cuvsResources_t res, + cuvsMultiGpuIvfPqSearchParams_t params, + cuvsMultiGpuIvfPqIndex_t index, + DLManagedTensor* queries_tensor, + DLManagedTensor* neighbors_tensor, + DLManagedTensor* distances_tensor); + +/** + * @} + */ + +/** + * @defgroup mg_ivf_pq_c_index_extend Multi-GPU IVF-PQ index extend + * @{ + */ + +/** + * @brief Extend a Multi-GPU IVF-PQ index + * + * @param[in] res cuvsResources_t opaque C handle + * @param[in,out] index Multi-GPU IVF-PQ index to extend + * @param[in] new_vectors_tensor DLManagedTensor* new vectors to add + * @param[in] new_indices_tensor DLManagedTensor* new indices (optional, can be NULL) + * @return cuvsError_t + */ +cuvsError_t cuvsMultiGpuIvfPqExtend(cuvsResources_t res, + cuvsMultiGpuIvfPqIndex_t index, + DLManagedTensor* new_vectors_tensor, + DLManagedTensor* new_indices_tensor); + +/** + * @} + */ + +/** + * @defgroup mg_ivf_pq_c_index_serialize Multi-GPU IVF-PQ index serialize + * @{ + */ + +/** + * @brief Serialize a Multi-GPU IVF-PQ index to file + * + * @param[in] res cuvsResources_t opaque C handle + * @param[in] index Multi-GPU IVF-PQ index to serialize + * @param[in] filename Path to the output file + * @return cuvsError_t + */ +cuvsError_t cuvsMultiGpuIvfPqSerialize(cuvsResources_t res, + cuvsMultiGpuIvfPqIndex_t index, + const char* filename); + +/** + * @} + */ + +/** + * @defgroup mg_ivf_pq_c_index_deserialize Multi-GPU IVF-PQ index deserialize + * @{ + */ + +/** + * @brief Deserialize a Multi-GPU IVF-PQ index from file + * + * @param[in] res cuvsResources_t opaque C handle + * @param[in] filename Path to the input file + * @param[out] index Multi-GPU IVF-PQ index + * @return cuvsError_t + */ +cuvsError_t cuvsMultiGpuIvfPqDeserialize(cuvsResources_t res, + const char* filename, + cuvsMultiGpuIvfPqIndex_t index); + +/** + * @} + */ + +/** + * @defgroup mg_ivf_pq_c_index_distribute Multi-GPU IVF-PQ index distribute + * @{ + */ + +/** + * @brief Distribute a local IVF-PQ index to create a Multi-GPU index + * + * @param[in] res cuvsResources_t opaque C handle + * @param[in] filename Path to the local index file + * @param[out] index Multi-GPU IVF-PQ index + * @return cuvsError_t + */ +cuvsError_t cuvsMultiGpuIvfPqDistribute(cuvsResources_t res, + const char* filename, + cuvsMultiGpuIvfPqIndex_t index); + +/** + * @} + */ + +#ifdef __cplusplus +} +#endif diff --git a/cpp/src/core/c_api.cpp b/cpp/src/core/c_api.cpp index 2f66adf416..a36bb472b6 100644 --- a/cpp/src/core/c_api.cpp +++ b/cpp/src/core/c_api.cpp @@ -18,6 +18,7 @@ #include #include +#include #include #include #include @@ -51,6 +52,22 @@ extern "C" cuvsError_t cuvsResourcesDestroy(cuvsResources_t res) }); } +extern "C" cuvsError_t cuvsMultiGpuResourcesCreate(cuvsResources_t* res) +{ + return cuvs::core::translate_exceptions([=] { + auto res_ptr = new raft::device_resources_snmg{}; + *res = reinterpret_cast(res_ptr); + }); +} + +extern "C" cuvsError_t cuvsMultiGpuResourcesDestroy(cuvsResources_t res) +{ + return cuvs::core::translate_exceptions([=] { + auto res_ptr = reinterpret_cast(res); + delete res_ptr; + }); +} + extern "C" cuvsError_t cuvsStreamSet(cuvsResources_t res, cudaStream_t stream) { return cuvs::core::translate_exceptions([=] { diff --git a/cpp/src/neighbors/iface/generate_iface.py b/cpp/src/neighbors/iface/generate_iface.py index d5fa6f56ca..66209e3f47 100644 --- a/cpp/src/neighbors/iface/generate_iface.py +++ b/cpp/src/neighbors/iface/generate_iface.py @@ -1,4 +1,4 @@ -# Copyright (c) 2024, NVIDIA CORPORATION. +# Copyright (c) 2024-2025, NVIDIA CORPORATION. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -13,7 +13,7 @@ # limitations under the License. header = """/* - * Copyright (c) 2024, NVIDIA CORPORATION. + * Copyright (c) 2024-2025, NVIDIA CORPORATION. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -253,6 +253,7 @@ flat_types = dict( float_int64_t=("float", "int64_t"), + half_int64_t=("half", "int64_t"), int8_t_int64_t=("int8_t", "int64_t"), uint8_t_int64_t=("uint8_t", "int64_t"), ) diff --git a/cpp/src/neighbors/iface/iface_flat_half_int64_t.cu b/cpp/src/neighbors/iface/iface_flat_half_int64_t.cu new file mode 100644 index 0000000000..4907e11063 --- /dev/null +++ b/cpp/src/neighbors/iface/iface_flat_half_int64_t.cu @@ -0,0 +1,96 @@ +/* + * Copyright (c) 2024-2025, NVIDIA CORPORATION. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +/* + * NOTE: this file is generated by generate_iface.py + * + * Make changes there and run in this directory: + * + * > python generate_iface.py + * + */ + +#include "iface.hpp" + +namespace cuvs::neighbors { + +#define CUVS_INST_MG_FLAT(T, IdxT) \ + using T_ha = raft::host_device_accessor, \ + raft::memory_type::device>; \ + using T_da = raft::host_device_accessor, \ + raft::memory_type::host>; \ + using IdxT_ha = raft::host_device_accessor, \ + raft::memory_type::device>; \ + using IdxT_da = raft::host_device_accessor, \ + raft::memory_type::host>; \ + \ + template void build( \ + const raft::resources& handle, \ + cuvs::neighbors::iface, T, IdxT>& interface, \ + const cuvs::neighbors::index_params* index_params, \ + raft::mdspan, row_major, T_ha> index_dataset); \ + \ + template void build( \ + const raft::resources& handle, \ + cuvs::neighbors::iface, T, IdxT>& interface, \ + const cuvs::neighbors::index_params* index_params, \ + raft::mdspan, row_major, T_da> index_dataset); \ + \ + template void extend( \ + const raft::resources& handle, \ + cuvs::neighbors::iface, T, IdxT>& interface, \ + raft::mdspan, row_major, T_ha> new_vectors, \ + std::optional, layout_c_contiguous, IdxT_ha>> \ + new_indices); \ + \ + template void extend( \ + const raft::resources& handle, \ + cuvs::neighbors::iface, T, IdxT>& interface, \ + raft::mdspan, row_major, T_da> new_vectors, \ + std::optional, layout_c_contiguous, IdxT_da>> \ + new_indices); \ + \ + template void search(const raft::resources& handle, \ + const cuvs::neighbors::iface, T, IdxT>& interface, \ + const cuvs::neighbors::search_params* search_params, \ + raft::device_matrix_view queries, \ + raft::device_matrix_view neighbors, \ + raft::device_matrix_view distances); \ + \ + template void search(const raft::resources& handle, \ + const cuvs::neighbors::iface, T, IdxT>& interface, \ + const cuvs::neighbors::search_params* search_params, \ + raft::host_matrix_view h_queries, \ + raft::device_matrix_view d_neighbors, \ + raft::device_matrix_view d_distances); \ + \ + template void serialize( \ + const raft::resources& handle, \ + const cuvs::neighbors::iface, T, IdxT>& interface, \ + std::ostream& os); \ + \ + template void deserialize(const raft::resources& handle, \ + cuvs::neighbors::iface, T, IdxT>& interface, \ + std::istream& is); \ + \ + template void deserialize(const raft::resources& handle, \ + cuvs::neighbors::iface, T, IdxT>& interface, \ + const std::string& filename); +CUVS_INST_MG_FLAT(half, int64_t); + +#undef CUVS_INST_MG_FLAT + +} // namespace cuvs::neighbors diff --git a/cpp/src/neighbors/mg/generate_mg.py b/cpp/src/neighbors/mg/generate_mg.py index f6ee47213e..895c1a31fc 100644 --- a/cpp/src/neighbors/mg/generate_mg.py +++ b/cpp/src/neighbors/mg/generate_mg.py @@ -1,4 +1,4 @@ -# Copyright (c) 2024, NVIDIA CORPORATION. +# Copyright (c) 2024-2025, NVIDIA CORPORATION. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -13,7 +13,7 @@ # limitations under the License. header = """/* - * Copyright (c) 2024, NVIDIA CORPORATION. + * Copyright (c) 2024-2025, NVIDIA CORPORATION. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -190,6 +190,14 @@ return index; \\ } \\ \\ + void extend(const raft::resources& res, \\ + cuvs::neighbors::mg_index, T, IdxT>& index, \\ + raft::host_matrix_view new_vectors, \\ + std::optional> new_indices) \\ + { \\ + cuvs::neighbors::snmg::detail::extend(res, index, new_vectors, new_indices); \\ + } \\ + \\ void search(const raft::resources& res, \\ const cuvs::neighbors::mg_index, T, IdxT>& index, \\ const mg_search_params& search_params, \\ @@ -268,6 +276,7 @@ flat_types = dict( float_int64_t=("float", "int64_t"), + half_int64_t=("half", "int64_t"), int8_t_int64_t=("int8_t", "int64_t"), uint8_t_int64_t=("uint8_t", "int64_t"), ) diff --git a/cpp/src/neighbors/mg/mg_cagra_float_uint32_t.cu b/cpp/src/neighbors/mg/mg_cagra_float_uint32_t.cu index 19d6ccc4f9..cbde5ef929 100644 --- a/cpp/src/neighbors/mg/mg_cagra_float_uint32_t.cu +++ b/cpp/src/neighbors/mg/mg_cagra_float_uint32_t.cu @@ -1,5 +1,5 @@ /* - * Copyright (c) 2024, NVIDIA CORPORATION. + * Copyright (c) 2024-2025, NVIDIA CORPORATION. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -43,6 +43,14 @@ return index; \ } \ \ + void extend(const raft::resources& res, \ + cuvs::neighbors::mg_index, T, IdxT>& index, \ + raft::host_matrix_view new_vectors, \ + std::optional> new_indices) \ + { \ + cuvs::neighbors::snmg::detail::extend(res, index, new_vectors, new_indices); \ + } \ + \ void search(const raft::resources& res, \ const cuvs::neighbors::mg_index, T, IdxT>& index, \ const mg_search_params& search_params, \ diff --git a/cpp/src/neighbors/mg/mg_cagra_half_uint32_t.cu b/cpp/src/neighbors/mg/mg_cagra_half_uint32_t.cu index 8b6b6a5c8c..4483f39178 100644 --- a/cpp/src/neighbors/mg/mg_cagra_half_uint32_t.cu +++ b/cpp/src/neighbors/mg/mg_cagra_half_uint32_t.cu @@ -1,5 +1,5 @@ /* - * Copyright (c) 2024, NVIDIA CORPORATION. + * Copyright (c) 2024-2025, NVIDIA CORPORATION. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -43,6 +43,14 @@ return index; \ } \ \ + void extend(const raft::resources& res, \ + cuvs::neighbors::mg_index, T, IdxT>& index, \ + raft::host_matrix_view new_vectors, \ + std::optional> new_indices) \ + { \ + cuvs::neighbors::snmg::detail::extend(res, index, new_vectors, new_indices); \ + } \ + \ void search(const raft::resources& res, \ const cuvs::neighbors::mg_index, T, IdxT>& index, \ const mg_search_params& search_params, \ diff --git a/cpp/src/neighbors/mg/mg_cagra_int8_t_uint32_t.cu b/cpp/src/neighbors/mg/mg_cagra_int8_t_uint32_t.cu index 4050a013a2..64282e531f 100644 --- a/cpp/src/neighbors/mg/mg_cagra_int8_t_uint32_t.cu +++ b/cpp/src/neighbors/mg/mg_cagra_int8_t_uint32_t.cu @@ -1,5 +1,5 @@ /* - * Copyright (c) 2024, NVIDIA CORPORATION. + * Copyright (c) 2024-2025, NVIDIA CORPORATION. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -43,6 +43,14 @@ return index; \ } \ \ + void extend(const raft::resources& res, \ + cuvs::neighbors::mg_index, T, IdxT>& index, \ + raft::host_matrix_view new_vectors, \ + std::optional> new_indices) \ + { \ + cuvs::neighbors::snmg::detail::extend(res, index, new_vectors, new_indices); \ + } \ + \ void search(const raft::resources& res, \ const cuvs::neighbors::mg_index, T, IdxT>& index, \ const mg_search_params& search_params, \ diff --git a/cpp/src/neighbors/mg/mg_cagra_uint8_t_uint32_t.cu b/cpp/src/neighbors/mg/mg_cagra_uint8_t_uint32_t.cu index 76d6a44aa7..64464034fa 100644 --- a/cpp/src/neighbors/mg/mg_cagra_uint8_t_uint32_t.cu +++ b/cpp/src/neighbors/mg/mg_cagra_uint8_t_uint32_t.cu @@ -1,5 +1,5 @@ /* - * Copyright (c) 2024, NVIDIA CORPORATION. + * Copyright (c) 2024-2025, NVIDIA CORPORATION. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -43,6 +43,14 @@ return index; \ } \ \ + void extend(const raft::resources& res, \ + cuvs::neighbors::mg_index, T, IdxT>& index, \ + raft::host_matrix_view new_vectors, \ + std::optional> new_indices) \ + { \ + cuvs::neighbors::snmg::detail::extend(res, index, new_vectors, new_indices); \ + } \ + \ void search(const raft::resources& res, \ const cuvs::neighbors::mg_index, T, IdxT>& index, \ const mg_search_params& search_params, \ diff --git a/cpp/src/neighbors/mg/mg_flat_half_int64_t.cu b/cpp/src/neighbors/mg/mg_flat_half_int64_t.cu new file mode 100644 index 0000000000..12446aee03 --- /dev/null +++ b/cpp/src/neighbors/mg/mg_flat_half_int64_t.cu @@ -0,0 +1,96 @@ +/* + * Copyright (c) 2024-2025, NVIDIA CORPORATION. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +/* + * NOTE: this file is generated by generate_mg.py + * + * Make changes there and run in this directory: + * + * > python generate_mg.py + * + */ + +#include "snmg.cuh" + +#define CUVS_INST_MG_FLAT(T, IdxT) \ + namespace cuvs::neighbors::ivf_flat { \ + using namespace cuvs::neighbors; \ + \ + cuvs::neighbors::mg_index, T, IdxT> build( \ + const raft::resources& res, \ + const mg_index_params& index_params, \ + raft::host_matrix_view index_dataset) \ + { \ + cuvs::neighbors::mg_index, T, IdxT> index(res, index_params.mode); \ + cuvs::neighbors::snmg::detail::build( \ + res, \ + index, \ + static_cast(&index_params), \ + index_dataset); \ + return index; \ + } \ + \ + void extend(const raft::resources& res, \ + cuvs::neighbors::mg_index, T, IdxT>& index, \ + raft::host_matrix_view new_vectors, \ + std::optional> new_indices) \ + { \ + cuvs::neighbors::snmg::detail::extend(res, index, new_vectors, new_indices); \ + } \ + \ + void search(const raft::resources& res, \ + const cuvs::neighbors::mg_index, T, IdxT>& index, \ + const mg_search_params& search_params, \ + raft::host_matrix_view queries, \ + raft::host_matrix_view neighbors, \ + raft::host_matrix_view distances) \ + { \ + cuvs::neighbors::snmg::detail::search( \ + res, \ + index, \ + static_cast(&search_params), \ + queries, \ + neighbors, \ + distances); \ + } \ + \ + void serialize(const raft::resources& res, \ + const cuvs::neighbors::mg_index, T, IdxT>& index, \ + const std::string& filename) \ + { \ + cuvs::neighbors::snmg::detail::serialize(res, index, filename); \ + } \ + \ + template <> \ + cuvs::neighbors::mg_index, T, IdxT> deserialize( \ + const raft::resources& res, const std::string& filename) \ + { \ + auto idx = cuvs::neighbors::mg_index, T, IdxT>(res, filename); \ + return idx; \ + } \ + \ + template <> \ + cuvs::neighbors::mg_index, T, IdxT> distribute( \ + const raft::resources& res, const std::string& filename) \ + { \ + auto idx = cuvs::neighbors::mg_index, T, IdxT>(res, REPLICATED); \ + cuvs::neighbors::snmg::detail::deserialize_and_distribute(res, idx, filename); \ + return idx; \ + } \ + } // namespace cuvs::neighbors::ivf_flat +CUVS_INST_MG_FLAT(half, int64_t); + +#undef CUVS_INST_MG_FLAT diff --git a/cpp/src/neighbors/mg/snmg.cuh b/cpp/src/neighbors/mg/snmg.cuh index abff1f82b6..b5658825f0 100644 --- a/cpp/src/neighbors/mg/snmg.cuh +++ b/cpp/src/neighbors/mg/snmg.cuh @@ -72,9 +72,13 @@ void deserialize(const raft::resources& clique, std::ifstream is(filename, std::ios::in | std::ios::binary); if (!is) { RAFT_FAIL("Cannot open file %s", filename.c_str()); } + char dtype_string[4]; + is.read(dtype_string, 4); + const auto& handle = raft::resource::set_current_device_to_root_rank(clique); - index.mode_ = (cuvs::neighbors::distribution_mode)deserialize_scalar(handle, is); - index.num_ranks_ = deserialize_scalar(handle, is); + index.mode_ = + static_cast(deserialize_scalar(handle, is)); + index.num_ranks_ = deserialize_scalar(handle, is); if (index.num_ranks_ != raft::resource::get_num_ranks(clique)) { RAFT_FAIL("Serialized index has %d ranks whereas NCCL clique has %d ranks", @@ -665,9 +669,14 @@ void serialize(const raft::resources& clique, std::ofstream of(filename, std::ios::out | std::ios::binary); if (!of) { RAFT_FAIL("Cannot open file %s", filename.c_str()); } + std::string dtype_string = raft::detail::numpy_serializer::get_numpy_dtype().to_string(); + dtype_string.resize(4); + of << dtype_string; + const auto& handle = raft::resource::set_current_device_to_root_rank(clique); + serialize_scalar(handle, of, (int)index.mode_); - serialize_scalar(handle, of, index.num_ranks_); + serialize_scalar(handle, of, (int)index.num_ranks_); for (int rank = 0; rank < index.num_ranks_; rank++) { const raft::resources& dev_res = raft::resource::set_current_device_to_rank(clique, rank); diff --git a/cpp/src/neighbors/mg_cagra_c.cpp b/cpp/src/neighbors/mg_cagra_c.cpp new file mode 100644 index 0000000000..c6d05605d7 --- /dev/null +++ b/cpp/src/neighbors/mg_cagra_c.cpp @@ -0,0 +1,419 @@ +/* + * Copyright (c) 2025, NVIDIA CORPORATION. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "cagra_c.hpp" +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include + +extern "C" cuvsError_t cuvsMultiGpuCagraIndexParamsCreate( + cuvsMultiGpuCagraIndexParams_t* index_params) +{ + return cuvs::core::translate_exceptions([=] { + // Create base CAGRA parameters + cuvsCagraIndexParams_t base_params; + cuvsCagraIndexParamsCreate(&base_params); + + // Create MG wrapper with default values + *index_params = new cuvsMultiGpuCagraIndexParams{ + .base_params = base_params, + .mode = CUVS_NEIGHBORS_MG_SHARDED // Default to sharded mode + }; + }); +} + +extern "C" cuvsError_t cuvsMultiGpuCagraIndexParamsDestroy( + cuvsMultiGpuCagraIndexParams_t index_params) +{ + return cuvs::core::translate_exceptions([=] { + if (index_params) { + cuvsCagraIndexParamsDestroy(index_params->base_params); + delete index_params; + } + }); +} + +extern "C" cuvsError_t cuvsMultiGpuCagraSearchParamsCreate(cuvsMultiGpuCagraSearchParams_t* params) +{ + return cuvs::core::translate_exceptions([=] { + // Create base CAGRA search parameters + cuvsCagraSearchParams_t base_params; + cuvsCagraSearchParamsCreate(&base_params); + + // Create MG wrapper with default values + *params = new cuvsMultiGpuCagraSearchParams{ + .base_params = base_params, + .search_mode = CUVS_NEIGHBORS_MG_LOAD_BALANCER, // Default to load balancer + .merge_mode = CUVS_NEIGHBORS_MG_TREE_MERGE, // Default to tree merge + .n_rows_per_batch = 1LL << 20 // Default to 1M rows per batch + }; + }); +} + +extern "C" cuvsError_t cuvsMultiGpuCagraSearchParamsDestroy(cuvsMultiGpuCagraSearchParams_t params) +{ + return cuvs::core::translate_exceptions([=] { + if (params) { + cuvsCagraSearchParamsDestroy(params->base_params); + delete params; + } + }); +} + +extern "C" cuvsError_t cuvsMultiGpuCagraIndexCreate(cuvsMultiGpuCagraIndex_t* index) +{ + return cuvs::core::translate_exceptions([=] { *index = new cuvsMultiGpuCagraIndex{}; }); +} + +extern "C" cuvsError_t cuvsMultiGpuCagraIndexDestroy(cuvsMultiGpuCagraIndex_t index) +{ + return cuvs::core::translate_exceptions([=] { + if (index) { + // Properly clean up the templated inner object based on dtype, like single GPU API + if (index->dtype.code == kDLFloat && index->dtype.bits == 32) { + auto mg_index_ptr = + reinterpret_cast, + float, + uint32_t>*>(index->addr); + delete mg_index_ptr; + } else if (index->dtype.code == kDLFloat && index->dtype.bits == 16) { + auto mg_index_ptr = + reinterpret_cast, + half, + uint32_t>*>(index->addr); + delete mg_index_ptr; + } else if (index->dtype.code == kDLInt && index->dtype.bits == 8) { + auto mg_index_ptr = reinterpret_cast< + cuvs::neighbors:: + mg_index, int8_t, uint32_t>*>( + index->addr); + delete mg_index_ptr; + } else if (index->dtype.code == kDLUInt && index->dtype.bits == 8) { + auto mg_index_ptr = reinterpret_cast< + cuvs::neighbors:: + mg_index, uint8_t, uint32_t>*>( + index->addr); + delete mg_index_ptr; + } + delete index; + } + }); +} + +namespace cuvs::neighbors::cagra { + +void convert_c_mg_index_params( + cuvsMultiGpuCagraIndexParams params, + int64_t n_rows, + int64_t dim, + cuvs::neighbors::mg_index_params* out) +{ + convert_c_index_params(*params.base_params, n_rows, dim, out); + out->mode = (params.mode == CUVS_NEIGHBORS_MG_SHARDED) + ? cuvs::neighbors::distribution_mode::SHARDED + : cuvs::neighbors::distribution_mode::REPLICATED; +} + +void convert_c_mg_search_params( + cuvsMultiGpuCagraSearchParams params, + cuvs::neighbors::mg_search_params* out) +{ + convert_c_search_params(*params.base_params, out); + out->search_mode = (params.search_mode == CUVS_NEIGHBORS_MG_LOAD_BALANCER) + ? cuvs::neighbors::replicated_search_mode::LOAD_BALANCER + : cuvs::neighbors::replicated_search_mode::ROUND_ROBIN; + out->merge_mode = (params.merge_mode == CUVS_NEIGHBORS_MG_TREE_MERGE) + ? cuvs::neighbors::sharded_merge_mode::TREE_MERGE + : cuvs::neighbors::sharded_merge_mode::MERGE_ON_ROOT_RANK; + out->n_rows_per_batch = params.n_rows_per_batch; +} +} // namespace cuvs::neighbors::cagra + +namespace { + +template +void* _mg_build(cuvsResources_t res, + cuvsMultiGpuCagraIndexParams params, + DLManagedTensor* dataset_tensor) +{ + auto res_ptr = reinterpret_cast(res); + auto dataset = dataset_tensor->dl_tensor; + + auto mg_params = cuvs::neighbors::mg_index_params(); + cuvs::neighbors::cagra::convert_c_mg_index_params( + params, dataset.shape[0], dataset.shape[1], &mg_params); + + using mdspan_type = raft::host_matrix_view; + auto mds = cuvs::core::from_dlpack(dataset_tensor); + + auto mg_index = + new cuvs::neighbors::mg_index, T, uint32_t>( + cuvs::neighbors::cagra::build(*res_ptr, mg_params, mds)); + + return mg_index; +} + +template +void _mg_search(cuvsResources_t res, + cuvsMultiGpuCagraSearchParams params, + cuvsMultiGpuCagraIndex index, + DLManagedTensor* queries_tensor, + DLManagedTensor* neighbors_tensor, + DLManagedTensor* distances_tensor) +{ + auto res_ptr = reinterpret_cast(res); + auto mg_index_ptr = reinterpret_cast< + cuvs::neighbors::mg_index, T, uint32_t>*>( + index.addr); + + auto mg_search_params = + cuvs::neighbors::mg_search_params(); + cuvs::neighbors::cagra::convert_c_mg_search_params(params, &mg_search_params); + + using queries_mdspan_type = raft::host_matrix_view; + using neighbors_mdspan_type = raft::host_matrix_view; + using distances_mdspan_type = raft::host_matrix_view; + + auto queries_mds = cuvs::core::from_dlpack(queries_tensor); + auto neighbors_mds = cuvs::core::from_dlpack(neighbors_tensor); + auto distances_mds = cuvs::core::from_dlpack(distances_tensor); + + cuvs::neighbors::cagra::search( + *res_ptr, *mg_index_ptr, mg_search_params, queries_mds, neighbors_mds, distances_mds); +} + +template +void _mg_extend(cuvsResources_t res, + cuvsMultiGpuCagraIndex index, + DLManagedTensor* new_vectors_tensor, + DLManagedTensor* new_indices_tensor) +{ + auto res_ptr = reinterpret_cast(res); + auto mg_index_ptr = reinterpret_cast< + cuvs::neighbors::mg_index, T, uint32_t>*>( + index.addr); + + using vectors_mdspan_type = raft::host_matrix_view; + auto new_vectors_mds = cuvs::core::from_dlpack(new_vectors_tensor); + + std::optional> new_indices_mds = std::nullopt; + if (new_indices_tensor != nullptr) { + using indices_mdspan_type = raft::host_vector_view; + new_indices_mds = cuvs::core::from_dlpack(new_indices_tensor); + } + + cuvs::neighbors::cagra::extend(*res_ptr, *mg_index_ptr, new_vectors_mds, new_indices_mds); +} + +template +void _mg_serialize(cuvsResources_t res, cuvsMultiGpuCagraIndex index, const char* filename) +{ + auto res_ptr = reinterpret_cast(res); + auto mg_index_ptr = reinterpret_cast< + cuvs::neighbors::mg_index, T, uint32_t>*>( + index.addr); + + cuvs::neighbors::cagra::serialize(*res_ptr, *mg_index_ptr, std::string(filename)); +} + +template +void* _mg_deserialize(cuvsResources_t res, const char* filename) +{ + auto res_ptr = reinterpret_cast(res); + auto mg_index = + new cuvs::neighbors::mg_index, T, uint32_t>( + cuvs::neighbors::cagra::deserialize(*res_ptr, std::string(filename))); + + return mg_index; +} + +template +void* _mg_distribute(cuvsResources_t res, const char* filename) +{ + auto res_ptr = reinterpret_cast(res); + auto mg_index = + new cuvs::neighbors::mg_index, T, uint32_t>( + cuvs::neighbors::cagra::distribute(*res_ptr, std::string(filename))); + + return mg_index; +} + +} // anonymous namespace + +extern "C" cuvsError_t cuvsMultiGpuCagraBuild(cuvsResources_t res, + cuvsMultiGpuCagraIndexParams_t params, + DLManagedTensor* dataset_tensor, + cuvsMultiGpuCagraIndex_t index) +{ + return cuvs::core::translate_exceptions([=] { + auto dataset = dataset_tensor->dl_tensor; + index->dtype.code = dataset.dtype.code; + index->dtype.bits = dataset.dtype.bits; + + if (dataset.dtype.code == kDLFloat && dataset.dtype.bits == 32) { + index->addr = reinterpret_cast(_mg_build(res, *params, dataset_tensor)); + } else if (dataset.dtype.code == kDLFloat && dataset.dtype.bits == 16) { + index->addr = reinterpret_cast(_mg_build(res, *params, dataset_tensor)); + } else if (dataset.dtype.code == kDLInt && dataset.dtype.bits == 8) { + index->addr = reinterpret_cast(_mg_build(res, *params, dataset_tensor)); + } else if (dataset.dtype.code == kDLUInt && dataset.dtype.bits == 8) { + index->addr = reinterpret_cast(_mg_build(res, *params, dataset_tensor)); + } else { + RAFT_FAIL("Unsupported dataset DLtensor dtype: %d and bits: %d", + dataset.dtype.code, + dataset.dtype.bits); + } + }); +} + +extern "C" cuvsError_t cuvsMultiGpuCagraSearch(cuvsResources_t res, + cuvsMultiGpuCagraSearchParams_t params, + cuvsMultiGpuCagraIndex_t index, + DLManagedTensor* queries_tensor, + DLManagedTensor* neighbors_tensor, + DLManagedTensor* distances_tensor) +{ + return cuvs::core::translate_exceptions([=] { + auto queries = queries_tensor->dl_tensor; + + if (queries.dtype.code == kDLFloat && queries.dtype.bits == 32) { + _mg_search(res, *params, *index, queries_tensor, neighbors_tensor, distances_tensor); + } else if (queries.dtype.code == kDLFloat && queries.dtype.bits == 16) { + _mg_search(res, *params, *index, queries_tensor, neighbors_tensor, distances_tensor); + } else if (queries.dtype.code == kDLInt && queries.dtype.bits == 8) { + _mg_search(res, *params, *index, queries_tensor, neighbors_tensor, distances_tensor); + } else if (queries.dtype.code == kDLUInt && queries.dtype.bits == 8) { + _mg_search(res, *params, *index, queries_tensor, neighbors_tensor, distances_tensor); + } else { + RAFT_FAIL("Unsupported queries DLtensor dtype: %d and bits: %d", + queries.dtype.code, + queries.dtype.bits); + } + }); +} + +extern "C" cuvsError_t cuvsMultiGpuCagraExtend(cuvsResources_t res, + cuvsMultiGpuCagraIndex_t index, + DLManagedTensor* new_vectors_tensor, + DLManagedTensor* new_indices_tensor) +{ + return cuvs::core::translate_exceptions([=] { + auto vectors = new_vectors_tensor->dl_tensor; + + if (vectors.dtype.code == kDLFloat && vectors.dtype.bits == 32) { + _mg_extend(res, *index, new_vectors_tensor, new_indices_tensor); + } else if (vectors.dtype.code == kDLFloat && vectors.dtype.bits == 16) { + _mg_extend(res, *index, new_vectors_tensor, new_indices_tensor); + } else if (vectors.dtype.code == kDLInt && vectors.dtype.bits == 8) { + _mg_extend(res, *index, new_vectors_tensor, new_indices_tensor); + } else if (vectors.dtype.code == kDLUInt && vectors.dtype.bits == 8) { + _mg_extend(res, *index, new_vectors_tensor, new_indices_tensor); + } else { + RAFT_FAIL("Unsupported new_vectors DLtensor dtype: %d and bits: %d", + vectors.dtype.code, + vectors.dtype.bits); + } + }); +} + +extern "C" cuvsError_t cuvsMultiGpuCagraSerialize(cuvsResources_t res, + cuvsMultiGpuCagraIndex_t index, + const char* filename) +{ + return cuvs::core::translate_exceptions([=] { + if (index->dtype.code == kDLFloat && index->dtype.bits == 32) { + _mg_serialize(res, *index, filename); + } else if (index->dtype.code == kDLFloat && index->dtype.bits == 16) { + _mg_serialize(res, *index, filename); + } else if (index->dtype.code == kDLInt && index->dtype.bits == 8) { + _mg_serialize(res, *index, filename); + } else if (index->dtype.code == kDLUInt && index->dtype.bits == 8) { + _mg_serialize(res, *index, filename); + } else { + RAFT_FAIL("Unsupported index dtype: %d and bits: %d", index->dtype.code, index->dtype.bits); + } + }); +} + +extern "C" cuvsError_t cuvsMultiGpuCagraDeserialize(cuvsResources_t res, + const char* filename, + cuvsMultiGpuCagraIndex_t index) +{ + return cuvs::core::translate_exceptions([=] { + std::ifstream is(filename, std::ios::in | std::ios::binary); + if (!is) { RAFT_FAIL("Cannot open file %s", filename); } + char dtype_string[4]; + is.read(dtype_string, 4); + auto dtype = raft::detail::numpy_serializer::parse_descr(std::string(dtype_string, 4)); + is.close(); + + index->dtype.bits = dtype.itemsize * 8; + if (dtype.kind == 'f' && dtype.itemsize == 4) { + index->dtype.code = kDLFloat; + index->addr = reinterpret_cast(_mg_deserialize(res, filename)); + } else if (dtype.kind == 'f' && dtype.itemsize == 2) { + index->dtype.code = kDLFloat; + index->addr = reinterpret_cast(_mg_deserialize(res, filename)); + } else if (dtype.kind == 'i' && dtype.itemsize == 1) { + index->dtype.code = kDLInt; + index->addr = reinterpret_cast(_mg_deserialize(res, filename)); + } else if (dtype.kind == 'u' && dtype.itemsize == 1) { + index->dtype.code = kDLUInt; + index->addr = reinterpret_cast(_mg_deserialize(res, filename)); + } else { + RAFT_FAIL("Unsupported index dtype"); + } + }); +} + +extern "C" cuvsError_t cuvsMultiGpuCagraDistribute(cuvsResources_t res, + const char* filename, + cuvsMultiGpuCagraIndex_t index) +{ + return cuvs::core::translate_exceptions([=] { + std::ifstream is(filename, std::ios::in | std::ios::binary); + if (!is) { RAFT_FAIL("Cannot open file %s", filename); } + char dtype_string[4]; + is.read(dtype_string, 4); + auto dtype = raft::detail::numpy_serializer::parse_descr(std::string(dtype_string, 4)); + is.close(); + + index->dtype.bits = dtype.itemsize * 8; + if (dtype.kind == 'f' && dtype.itemsize == 4) { + index->dtype.code = kDLFloat; + index->addr = reinterpret_cast(_mg_distribute(res, filename)); + } else if (dtype.kind == 'f' && dtype.itemsize == 2) { + index->dtype.code = kDLFloat; + index->addr = reinterpret_cast(_mg_distribute(res, filename)); + } else if (dtype.kind == 'i' && dtype.itemsize == 1) { + index->dtype.code = kDLInt; + index->addr = reinterpret_cast(_mg_distribute(res, filename)); + } else if (dtype.kind == 'u' && dtype.itemsize == 1) { + index->dtype.code = kDLUInt; + index->addr = reinterpret_cast(_mg_distribute(res, filename)); + } else { + RAFT_FAIL("Unsupported index dtype"); + } + }); +} diff --git a/cpp/src/neighbors/mg_ivf_flat_c.cpp b/cpp/src/neighbors/mg_ivf_flat_c.cpp new file mode 100644 index 0000000000..c012cb4c7c --- /dev/null +++ b/cpp/src/neighbors/mg_ivf_flat_c.cpp @@ -0,0 +1,416 @@ +/* + * Copyright (c) 2025, NVIDIA CORPORATION. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "ivf_flat_c.hpp" +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include + +extern "C" cuvsError_t cuvsMultiGpuIvfFlatIndexParamsCreate( + cuvsMultiGpuIvfFlatIndexParams_t* index_params) +{ + return cuvs::core::translate_exceptions([=] { + // Create base IVF-Flat parameters + cuvsIvfFlatIndexParams_t base_params; + cuvsIvfFlatIndexParamsCreate(&base_params); + + // Create MG wrapper with default values + *index_params = new cuvsMultiGpuIvfFlatIndexParams{ + .base_params = base_params, + .mode = CUVS_NEIGHBORS_MG_SHARDED // Default to sharded mode + }; + }); +} + +extern "C" cuvsError_t cuvsMultiGpuIvfFlatIndexParamsDestroy( + cuvsMultiGpuIvfFlatIndexParams_t index_params) +{ + return cuvs::core::translate_exceptions([=] { + if (index_params) { + cuvsIvfFlatIndexParamsDestroy(index_params->base_params); + delete index_params; + } + }); +} + +extern "C" cuvsError_t cuvsMultiGpuIvfFlatSearchParamsCreate( + cuvsMultiGpuIvfFlatSearchParams_t* params) +{ + return cuvs::core::translate_exceptions([=] { + // Create base IVF-Flat search parameters + cuvsIvfFlatSearchParams_t base_params; + cuvsIvfFlatSearchParamsCreate(&base_params); + + // Create MG wrapper with default values + *params = new cuvsMultiGpuIvfFlatSearchParams{ + .base_params = base_params, + .search_mode = CUVS_NEIGHBORS_MG_LOAD_BALANCER, // Default to load balancer + .merge_mode = CUVS_NEIGHBORS_MG_TREE_MERGE, // Default to tree merge + .n_rows_per_batch = 1LL << 20 // Default to 1M rows per batch + }; + }); +} + +extern "C" cuvsError_t cuvsMultiGpuIvfFlatSearchParamsDestroy( + cuvsMultiGpuIvfFlatSearchParams_t params) +{ + return cuvs::core::translate_exceptions([=] { + if (params) { + cuvsIvfFlatSearchParamsDestroy(params->base_params); + delete params; + } + }); +} + +extern "C" cuvsError_t cuvsMultiGpuIvfFlatIndexCreate(cuvsMultiGpuIvfFlatIndex_t* index) +{ + return cuvs::core::translate_exceptions([=] { *index = new cuvsMultiGpuIvfFlatIndex{}; }); +} + +extern "C" cuvsError_t cuvsMultiGpuIvfFlatIndexDestroy(cuvsMultiGpuIvfFlatIndex_t index) +{ + return cuvs::core::translate_exceptions([=] { + if (index) { + // Properly clean up the templated inner object based on dtype, like single GPU API + if (index->dtype.code == kDLFloat && index->dtype.bits == 32) { + auto mg_index_ptr = reinterpret_cast< + cuvs::neighbors:: + mg_index, float, int64_t>*>( + index->addr); + delete mg_index_ptr; + } else if (index->dtype.code == kDLFloat && index->dtype.bits == 16) { + auto mg_index_ptr = reinterpret_cast< + cuvs::neighbors:: + mg_index, half, int64_t>*>(index->addr); + delete mg_index_ptr; + } else if (index->dtype.code == kDLInt && index->dtype.bits == 8) { + auto mg_index_ptr = reinterpret_cast< + cuvs::neighbors:: + mg_index, int8_t, int64_t>*>( + index->addr); + delete mg_index_ptr; + } else if (index->dtype.code == kDLUInt && index->dtype.bits == 8) { + auto mg_index_ptr = reinterpret_cast< + cuvs::neighbors:: + mg_index, uint8_t, int64_t>*>( + index->addr); + delete mg_index_ptr; + } + delete index; + } + }); +} + +namespace cuvs::neighbors::ivf_flat { + +void convert_c_mg_index_params( + cuvsMultiGpuIvfFlatIndexParams params, + cuvs::neighbors::mg_index_params* out) +{ + convert_c_index_params(*params.base_params, out); + out->mode = (params.mode == CUVS_NEIGHBORS_MG_SHARDED) + ? cuvs::neighbors::distribution_mode::SHARDED + : cuvs::neighbors::distribution_mode::REPLICATED; +} + +void convert_c_mg_search_params( + cuvsMultiGpuIvfFlatSearchParams params, + cuvs::neighbors::mg_search_params* out) +{ + convert_c_search_params(*params.base_params, out); + out->search_mode = (params.search_mode == CUVS_NEIGHBORS_MG_LOAD_BALANCER) + ? cuvs::neighbors::replicated_search_mode::LOAD_BALANCER + : cuvs::neighbors::replicated_search_mode::ROUND_ROBIN; + out->merge_mode = (params.merge_mode == CUVS_NEIGHBORS_MG_TREE_MERGE) + ? cuvs::neighbors::sharded_merge_mode::TREE_MERGE + : cuvs::neighbors::sharded_merge_mode::MERGE_ON_ROOT_RANK; + out->n_rows_per_batch = params.n_rows_per_batch; +} +} // namespace cuvs::neighbors::ivf_flat + +namespace { + +template +void* _mg_build(cuvsResources_t res, + cuvsMultiGpuIvfFlatIndexParams params, + DLManagedTensor* dataset_tensor) +{ + auto res_ptr = reinterpret_cast(res); + + auto mg_params = cuvs::neighbors::mg_index_params(); + cuvs::neighbors::ivf_flat::convert_c_mg_index_params(params, &mg_params); + + using mdspan_type = raft::host_matrix_view; + auto mds = cuvs::core::from_dlpack(dataset_tensor); + + auto mg_index = + new cuvs::neighbors::mg_index, T, int64_t>( + cuvs::neighbors::ivf_flat::build(*res_ptr, mg_params, mds)); + + return mg_index; +} + +template +void _mg_search(cuvsResources_t res, + cuvsMultiGpuIvfFlatSearchParams params, + cuvsMultiGpuIvfFlatIndex index, + DLManagedTensor* queries_tensor, + DLManagedTensor* neighbors_tensor, + DLManagedTensor* distances_tensor) +{ + auto res_ptr = reinterpret_cast(res); + auto mg_index_ptr = reinterpret_cast< + cuvs::neighbors::mg_index, T, int64_t>*>( + index.addr); + + auto mg_search_params = + cuvs::neighbors::mg_search_params(); + cuvs::neighbors::ivf_flat::convert_c_mg_search_params(params, &mg_search_params); + + using queries_mdspan_type = raft::host_matrix_view; + using neighbors_mdspan_type = raft::host_matrix_view; + using distances_mdspan_type = raft::host_matrix_view; + + auto queries_mds = cuvs::core::from_dlpack(queries_tensor); + auto neighbors_mds = cuvs::core::from_dlpack(neighbors_tensor); + auto distances_mds = cuvs::core::from_dlpack(distances_tensor); + + cuvs::neighbors::ivf_flat::search( + *res_ptr, *mg_index_ptr, mg_search_params, queries_mds, neighbors_mds, distances_mds); +} + +template +void _mg_extend(cuvsResources_t res, + cuvsMultiGpuIvfFlatIndex index, + DLManagedTensor* new_vectors_tensor, + DLManagedTensor* new_indices_tensor) +{ + auto res_ptr = reinterpret_cast(res); + auto mg_index_ptr = reinterpret_cast< + cuvs::neighbors::mg_index, T, int64_t>*>( + index.addr); + + using vectors_mdspan_type = raft::host_matrix_view; + auto new_vectors_mds = cuvs::core::from_dlpack(new_vectors_tensor); + + std::optional> new_indices_mds = std::nullopt; + if (new_indices_tensor != nullptr) { + using indices_mdspan_type = raft::host_vector_view; + new_indices_mds = cuvs::core::from_dlpack(new_indices_tensor); + } + + cuvs::neighbors::ivf_flat::extend(*res_ptr, *mg_index_ptr, new_vectors_mds, new_indices_mds); +} + +template +void _mg_serialize(cuvsResources_t res, cuvsMultiGpuIvfFlatIndex index, const char* filename) +{ + auto res_ptr = reinterpret_cast(res); + auto mg_index_ptr = reinterpret_cast< + cuvs::neighbors::mg_index, T, int64_t>*>( + index.addr); + + cuvs::neighbors::ivf_flat::serialize(*res_ptr, *mg_index_ptr, std::string(filename)); +} + +template +void* _mg_deserialize(cuvsResources_t res, const char* filename) +{ + auto res_ptr = reinterpret_cast(res); + auto mg_index = + new cuvs::neighbors::mg_index, T, int64_t>( + cuvs::neighbors::ivf_flat::deserialize(*res_ptr, std::string(filename))); + + return mg_index; +} + +template +void* _mg_distribute(cuvsResources_t res, const char* filename) +{ + auto res_ptr = reinterpret_cast(res); + auto mg_index = + new cuvs::neighbors::mg_index, T, int64_t>( + cuvs::neighbors::ivf_flat::distribute(*res_ptr, std::string(filename))); + + return mg_index; +} + +} // anonymous namespace + +extern "C" cuvsError_t cuvsMultiGpuIvfFlatBuild(cuvsResources_t res, + cuvsMultiGpuIvfFlatIndexParams_t params, + DLManagedTensor* dataset_tensor, + cuvsMultiGpuIvfFlatIndex_t index) +{ + return cuvs::core::translate_exceptions([=] { + auto dataset = dataset_tensor->dl_tensor; + index->dtype.code = dataset.dtype.code; + index->dtype.bits = dataset.dtype.bits; + + if (dataset.dtype.code == kDLFloat && dataset.dtype.bits == 32) { + index->addr = reinterpret_cast(_mg_build(res, *params, dataset_tensor)); + } else if (dataset.dtype.code == kDLFloat && dataset.dtype.bits == 16) { + index->addr = reinterpret_cast(_mg_build(res, *params, dataset_tensor)); + } else if (dataset.dtype.code == kDLInt && dataset.dtype.bits == 8) { + index->addr = reinterpret_cast(_mg_build(res, *params, dataset_tensor)); + } else if (dataset.dtype.code == kDLUInt && dataset.dtype.bits == 8) { + index->addr = reinterpret_cast(_mg_build(res, *params, dataset_tensor)); + } else { + RAFT_FAIL("Unsupported dataset DLtensor dtype: %d and bits: %d", + dataset.dtype.code, + dataset.dtype.bits); + } + }); +} + +extern "C" cuvsError_t cuvsMultiGpuIvfFlatSearch(cuvsResources_t res, + cuvsMultiGpuIvfFlatSearchParams_t params, + cuvsMultiGpuIvfFlatIndex_t index, + DLManagedTensor* queries_tensor, + DLManagedTensor* neighbors_tensor, + DLManagedTensor* distances_tensor) +{ + return cuvs::core::translate_exceptions([=] { + auto queries = queries_tensor->dl_tensor; + + if (queries.dtype.code == kDLFloat && queries.dtype.bits == 32) { + _mg_search(res, *params, *index, queries_tensor, neighbors_tensor, distances_tensor); + } else if (queries.dtype.code == kDLFloat && queries.dtype.bits == 16) { + _mg_search(res, *params, *index, queries_tensor, neighbors_tensor, distances_tensor); + } else if (queries.dtype.code == kDLInt && queries.dtype.bits == 8) { + _mg_search(res, *params, *index, queries_tensor, neighbors_tensor, distances_tensor); + } else if (queries.dtype.code == kDLUInt && queries.dtype.bits == 8) { + _mg_search(res, *params, *index, queries_tensor, neighbors_tensor, distances_tensor); + } else { + RAFT_FAIL("Unsupported queries DLtensor dtype: %d and bits: %d", + queries.dtype.code, + queries.dtype.bits); + } + }); +} + +extern "C" cuvsError_t cuvsMultiGpuIvfFlatExtend(cuvsResources_t res, + cuvsMultiGpuIvfFlatIndex_t index, + DLManagedTensor* new_vectors_tensor, + DLManagedTensor* new_indices_tensor) +{ + return cuvs::core::translate_exceptions([=] { + auto vectors = new_vectors_tensor->dl_tensor; + + if (vectors.dtype.code == kDLFloat && vectors.dtype.bits == 32) { + _mg_extend(res, *index, new_vectors_tensor, new_indices_tensor); + } else if (vectors.dtype.code == kDLFloat && vectors.dtype.bits == 16) { + _mg_extend(res, *index, new_vectors_tensor, new_indices_tensor); + } else if (vectors.dtype.code == kDLInt && vectors.dtype.bits == 8) { + _mg_extend(res, *index, new_vectors_tensor, new_indices_tensor); + } else if (vectors.dtype.code == kDLUInt && vectors.dtype.bits == 8) { + _mg_extend(res, *index, new_vectors_tensor, new_indices_tensor); + } else { + RAFT_FAIL("Unsupported new_vectors DLtensor dtype: %d and bits: %d", + vectors.dtype.code, + vectors.dtype.bits); + } + }); +} + +extern "C" cuvsError_t cuvsMultiGpuIvfFlatSerialize(cuvsResources_t res, + cuvsMultiGpuIvfFlatIndex_t index, + const char* filename) +{ + return cuvs::core::translate_exceptions([=] { + if (index->dtype.code == kDLFloat && index->dtype.bits == 32) { + _mg_serialize(res, *index, filename); + } else if (index->dtype.code == kDLFloat && index->dtype.bits == 16) { + _mg_serialize(res, *index, filename); + } else if (index->dtype.code == kDLInt && index->dtype.bits == 8) { + _mg_serialize(res, *index, filename); + } else if (index->dtype.code == kDLUInt && index->dtype.bits == 8) { + _mg_serialize(res, *index, filename); + } else { + RAFT_FAIL("Unsupported index dtype: %d and bits: %d", index->dtype.code, index->dtype.bits); + } + }); +} + +extern "C" cuvsError_t cuvsMultiGpuIvfFlatDeserialize(cuvsResources_t res, + const char* filename, + cuvsMultiGpuIvfFlatIndex_t index) +{ + return cuvs::core::translate_exceptions([=] { + std::ifstream is(filename, std::ios::in | std::ios::binary); + if (!is) { RAFT_FAIL("Cannot open file %s", filename); } + char dtype_string[4]; + is.read(dtype_string, 4); + auto dtype = raft::detail::numpy_serializer::parse_descr(std::string(dtype_string, 4)); + is.close(); + + index->dtype.bits = dtype.itemsize * 8; + if (dtype.kind == 'f' && dtype.itemsize == 4) { + index->dtype.code = kDLFloat; + index->addr = reinterpret_cast(_mg_deserialize(res, filename)); + } else if (dtype.kind == 'f' && dtype.itemsize == 2) { + index->dtype.code = kDLFloat; + index->addr = reinterpret_cast(_mg_deserialize(res, filename)); + } else if (dtype.kind == 'i' && dtype.itemsize == 1) { + index->dtype.code = kDLInt; + index->addr = reinterpret_cast(_mg_deserialize(res, filename)); + } else if (dtype.kind == 'u' && dtype.itemsize == 1) { + index->dtype.code = kDLUInt; + index->addr = reinterpret_cast(_mg_deserialize(res, filename)); + } else { + RAFT_FAIL("Unsupported index dtype"); + } + }); +} + +extern "C" cuvsError_t cuvsMultiGpuIvfFlatDistribute(cuvsResources_t res, + const char* filename, + cuvsMultiGpuIvfFlatIndex_t index) +{ + return cuvs::core::translate_exceptions([=] { + std::ifstream is(filename, std::ios::in | std::ios::binary); + if (!is) { RAFT_FAIL("Cannot open file %s", filename); } + char dtype_string[4]; + is.read(dtype_string, 4); + auto dtype = raft::detail::numpy_serializer::parse_descr(std::string(dtype_string, 4)); + is.close(); + + index->dtype.bits = dtype.itemsize * 8; + if (dtype.kind == 'f' && dtype.itemsize == 4) { + index->dtype.code = kDLFloat; + index->addr = reinterpret_cast(_mg_distribute(res, filename)); + } else if (dtype.kind == 'f' && dtype.itemsize == 2) { + index->dtype.code = kDLFloat; + index->addr = reinterpret_cast(_mg_distribute(res, filename)); + } else if (dtype.kind == 'i' && dtype.itemsize == 1) { + index->dtype.code = kDLInt; + index->addr = reinterpret_cast(_mg_distribute(res, filename)); + } else if (dtype.kind == 'u' && dtype.itemsize == 1) { + index->dtype.code = kDLUInt; + index->addr = reinterpret_cast(_mg_distribute(res, filename)); + } else { + RAFT_FAIL("Unsupported index dtype"); + } + }); +} diff --git a/cpp/src/neighbors/mg_ivf_pq_c.cpp b/cpp/src/neighbors/mg_ivf_pq_c.cpp new file mode 100644 index 0000000000..57d11f5264 --- /dev/null +++ b/cpp/src/neighbors/mg_ivf_pq_c.cpp @@ -0,0 +1,408 @@ +/* + * Copyright (c) 2025, NVIDIA CORPORATION. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "ivf_pq_c.hpp" +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include + +extern "C" cuvsError_t cuvsMultiGpuIvfPqIndexParamsCreate( + cuvsMultiGpuIvfPqIndexParams_t* index_params) +{ + return cuvs::core::translate_exceptions([=] { + // Create base IVF-PQ parameters + cuvsIvfPqIndexParams_t base_params; + cuvsIvfPqIndexParamsCreate(&base_params); + + // Create MG wrapper with default values + *index_params = new cuvsMultiGpuIvfPqIndexParams{ + .base_params = base_params, + .mode = CUVS_NEIGHBORS_MG_SHARDED // Default to sharded mode + }; + }); +} + +extern "C" cuvsError_t cuvsMultiGpuIvfPqIndexParamsDestroy( + cuvsMultiGpuIvfPqIndexParams_t index_params) +{ + return cuvs::core::translate_exceptions([=] { + if (index_params) { + cuvsIvfPqIndexParamsDestroy(index_params->base_params); + delete index_params; + } + }); +} + +extern "C" cuvsError_t cuvsMultiGpuIvfPqSearchParamsCreate(cuvsMultiGpuIvfPqSearchParams_t* params) +{ + return cuvs::core::translate_exceptions([=] { + // Create base IVF-PQ search parameters + cuvsIvfPqSearchParams_t base_params; + cuvsIvfPqSearchParamsCreate(&base_params); + + // Create MG wrapper with default values + *params = new cuvsMultiGpuIvfPqSearchParams{ + .base_params = base_params, + .search_mode = CUVS_NEIGHBORS_MG_LOAD_BALANCER, // Default to load balancer + .merge_mode = CUVS_NEIGHBORS_MG_TREE_MERGE, // Default to tree merge + .n_rows_per_batch = 1LL << 20 // Default to 1M rows per batch + }; + }); +} + +extern "C" cuvsError_t cuvsMultiGpuIvfPqSearchParamsDestroy(cuvsMultiGpuIvfPqSearchParams_t params) +{ + return cuvs::core::translate_exceptions([=] { + if (params) { + cuvsIvfPqSearchParamsDestroy(params->base_params); + delete params; + } + }); +} + +extern "C" cuvsError_t cuvsMultiGpuIvfPqIndexCreate(cuvsMultiGpuIvfPqIndex_t* index) +{ + return cuvs::core::translate_exceptions([=] { *index = new cuvsMultiGpuIvfPqIndex{}; }); +} + +extern "C" cuvsError_t cuvsMultiGpuIvfPqIndexDestroy(cuvsMultiGpuIvfPqIndex_t index) +{ + return cuvs::core::translate_exceptions([=] { + if (index) { + // Properly clean up the templated inner object based on dtype, like single GPU API + if (index->dtype.code == kDLFloat && index->dtype.bits == 32) { + auto mg_index_ptr = reinterpret_cast< + cuvs::neighbors::mg_index, float, int64_t>*>( + index->addr); + delete mg_index_ptr; + } else if (index->dtype.code == kDLFloat && index->dtype.bits == 16) { + auto mg_index_ptr = reinterpret_cast< + cuvs::neighbors::mg_index, half, int64_t>*>( + index->addr); + delete mg_index_ptr; + } else if (index->dtype.code == kDLInt && index->dtype.bits == 8) { + auto mg_index_ptr = reinterpret_cast< + cuvs::neighbors::mg_index, int8_t, int64_t>*>( + index->addr); + delete mg_index_ptr; + } else if (index->dtype.code == kDLUInt && index->dtype.bits == 8) { + auto mg_index_ptr = reinterpret_cast< + cuvs::neighbors::mg_index, uint8_t, int64_t>*>( + index->addr); + delete mg_index_ptr; + } + delete index; + } + }); +} + +namespace cuvs::neighbors::ivf_pq { + +void convert_c_mg_index_params( + cuvsMultiGpuIvfPqIndexParams params, + cuvs::neighbors::mg_index_params* out) +{ + convert_c_index_params(*params.base_params, out); + out->mode = (params.mode == CUVS_NEIGHBORS_MG_SHARDED) + ? cuvs::neighbors::distribution_mode::SHARDED + : cuvs::neighbors::distribution_mode::REPLICATED; +} + +void convert_c_mg_search_params( + cuvsMultiGpuIvfPqSearchParams params, + cuvs::neighbors::mg_search_params* out) +{ + convert_c_search_params(*params.base_params, out); + out->search_mode = (params.search_mode == CUVS_NEIGHBORS_MG_LOAD_BALANCER) + ? cuvs::neighbors::replicated_search_mode::LOAD_BALANCER + : cuvs::neighbors::replicated_search_mode::ROUND_ROBIN; + out->merge_mode = (params.merge_mode == CUVS_NEIGHBORS_MG_TREE_MERGE) + ? cuvs::neighbors::sharded_merge_mode::TREE_MERGE + : cuvs::neighbors::sharded_merge_mode::MERGE_ON_ROOT_RANK; + out->n_rows_per_batch = params.n_rows_per_batch; +} +} // namespace cuvs::neighbors::ivf_pq + +namespace { + +template +void* _mg_build(cuvsResources_t res, + cuvsMultiGpuIvfPqIndexParams params, + DLManagedTensor* dataset_tensor) +{ + auto res_ptr = reinterpret_cast(res); + + auto mg_params = cuvs::neighbors::mg_index_params(); + cuvs::neighbors::ivf_pq::convert_c_mg_index_params(params, &mg_params); + + using mdspan_type = raft::host_matrix_view; + auto mds = cuvs::core::from_dlpack(dataset_tensor); + + auto mg_index = + new cuvs::neighbors::mg_index, T, int64_t>( + cuvs::neighbors::ivf_pq::build(*res_ptr, mg_params, mds)); + + return mg_index; +} + +template +void _mg_search(cuvsResources_t res, + cuvsMultiGpuIvfPqSearchParams params, + cuvsMultiGpuIvfPqIndex index, + DLManagedTensor* queries_tensor, + DLManagedTensor* neighbors_tensor, + DLManagedTensor* distances_tensor) +{ + auto res_ptr = reinterpret_cast(res); + auto mg_index_ptr = reinterpret_cast< + cuvs::neighbors::mg_index, T, int64_t>*>(index.addr); + + auto mg_search_params = + cuvs::neighbors::mg_search_params(); + cuvs::neighbors::ivf_pq::convert_c_mg_search_params(params, &mg_search_params); + + using queries_mdspan_type = raft::host_matrix_view; + using neighbors_mdspan_type = raft::host_matrix_view; + using distances_mdspan_type = raft::host_matrix_view; + + auto queries_mds = cuvs::core::from_dlpack(queries_tensor); + auto neighbors_mds = cuvs::core::from_dlpack(neighbors_tensor); + auto distances_mds = cuvs::core::from_dlpack(distances_tensor); + + cuvs::neighbors::ivf_pq::search( + *res_ptr, *mg_index_ptr, mg_search_params, queries_mds, neighbors_mds, distances_mds); +} + +template +void _mg_extend(cuvsResources_t res, + cuvsMultiGpuIvfPqIndex index, + DLManagedTensor* new_vectors_tensor, + DLManagedTensor* new_indices_tensor) +{ + auto res_ptr = reinterpret_cast(res); + auto mg_index_ptr = reinterpret_cast< + cuvs::neighbors::mg_index, T, int64_t>*>(index.addr); + + using vectors_mdspan_type = raft::host_matrix_view; + auto new_vectors_mds = cuvs::core::from_dlpack(new_vectors_tensor); + + std::optional> new_indices_mds = std::nullopt; + if (new_indices_tensor != nullptr) { + using indices_mdspan_type = raft::host_vector_view; + new_indices_mds = cuvs::core::from_dlpack(new_indices_tensor); + } + + cuvs::neighbors::ivf_pq::extend(*res_ptr, *mg_index_ptr, new_vectors_mds, new_indices_mds); +} + +template +void _mg_serialize(cuvsResources_t res, cuvsMultiGpuIvfPqIndex index, const char* filename) +{ + auto res_ptr = reinterpret_cast(res); + auto mg_index_ptr = reinterpret_cast< + cuvs::neighbors::mg_index, T, int64_t>*>(index.addr); + + cuvs::neighbors::ivf_pq::serialize(*res_ptr, *mg_index_ptr, std::string(filename)); +} + +template +void* _mg_deserialize(cuvsResources_t res, const char* filename) +{ + auto res_ptr = reinterpret_cast(res); + auto mg_index = + new cuvs::neighbors::mg_index, T, int64_t>( + cuvs::neighbors::ivf_pq::deserialize(*res_ptr, std::string(filename))); + + return mg_index; +} + +template +void* _mg_distribute(cuvsResources_t res, const char* filename) +{ + auto res_ptr = reinterpret_cast(res); + auto mg_index = + new cuvs::neighbors::mg_index, T, int64_t>( + cuvs::neighbors::ivf_pq::distribute(*res_ptr, std::string(filename))); + + return mg_index; +} + +} // anonymous namespace + +extern "C" cuvsError_t cuvsMultiGpuIvfPqBuild(cuvsResources_t res, + cuvsMultiGpuIvfPqIndexParams_t params, + DLManagedTensor* dataset_tensor, + cuvsMultiGpuIvfPqIndex_t index) +{ + return cuvs::core::translate_exceptions([=] { + auto dataset = dataset_tensor->dl_tensor; + index->dtype.code = dataset.dtype.code; + index->dtype.bits = dataset.dtype.bits; + + if (dataset.dtype.code == kDLFloat && dataset.dtype.bits == 32) { + index->addr = reinterpret_cast(_mg_build(res, *params, dataset_tensor)); + } else if (dataset.dtype.code == kDLFloat && dataset.dtype.bits == 16) { + index->addr = reinterpret_cast(_mg_build(res, *params, dataset_tensor)); + } else if (dataset.dtype.code == kDLInt && dataset.dtype.bits == 8) { + index->addr = reinterpret_cast(_mg_build(res, *params, dataset_tensor)); + } else if (dataset.dtype.code == kDLUInt && dataset.dtype.bits == 8) { + index->addr = reinterpret_cast(_mg_build(res, *params, dataset_tensor)); + } else { + RAFT_FAIL("Unsupported dataset DLtensor dtype: %d and bits: %d", + dataset.dtype.code, + dataset.dtype.bits); + } + }); +} + +extern "C" cuvsError_t cuvsMultiGpuIvfPqSearch(cuvsResources_t res, + cuvsMultiGpuIvfPqSearchParams_t params, + cuvsMultiGpuIvfPqIndex_t index, + DLManagedTensor* queries_tensor, + DLManagedTensor* neighbors_tensor, + DLManagedTensor* distances_tensor) +{ + return cuvs::core::translate_exceptions([=] { + auto queries = queries_tensor->dl_tensor; + + if (queries.dtype.code == kDLFloat && queries.dtype.bits == 32) { + _mg_search(res, *params, *index, queries_tensor, neighbors_tensor, distances_tensor); + } else if (queries.dtype.code == kDLFloat && queries.dtype.bits == 16) { + _mg_search(res, *params, *index, queries_tensor, neighbors_tensor, distances_tensor); + } else if (queries.dtype.code == kDLInt && queries.dtype.bits == 8) { + _mg_search(res, *params, *index, queries_tensor, neighbors_tensor, distances_tensor); + } else if (queries.dtype.code == kDLUInt && queries.dtype.bits == 8) { + _mg_search(res, *params, *index, queries_tensor, neighbors_tensor, distances_tensor); + } else { + RAFT_FAIL("Unsupported queries DLtensor dtype: %d and bits: %d", + queries.dtype.code, + queries.dtype.bits); + } + }); +} + +extern "C" cuvsError_t cuvsMultiGpuIvfPqExtend(cuvsResources_t res, + cuvsMultiGpuIvfPqIndex_t index, + DLManagedTensor* new_vectors_tensor, + DLManagedTensor* new_indices_tensor) +{ + return cuvs::core::translate_exceptions([=] { + auto vectors = new_vectors_tensor->dl_tensor; + + if (vectors.dtype.code == kDLFloat && vectors.dtype.bits == 32) { + _mg_extend(res, *index, new_vectors_tensor, new_indices_tensor); + } else if (vectors.dtype.code == kDLFloat && vectors.dtype.bits == 16) { + _mg_extend(res, *index, new_vectors_tensor, new_indices_tensor); + } else if (vectors.dtype.code == kDLInt && vectors.dtype.bits == 8) { + _mg_extend(res, *index, new_vectors_tensor, new_indices_tensor); + } else if (vectors.dtype.code == kDLUInt && vectors.dtype.bits == 8) { + _mg_extend(res, *index, new_vectors_tensor, new_indices_tensor); + } else { + RAFT_FAIL("Unsupported new_vectors DLtensor dtype: %d and bits: %d", + vectors.dtype.code, + vectors.dtype.bits); + } + }); +} + +extern "C" cuvsError_t cuvsMultiGpuIvfPqSerialize(cuvsResources_t res, + cuvsMultiGpuIvfPqIndex_t index, + const char* filename) +{ + return cuvs::core::translate_exceptions([=] { + if (index->dtype.code == kDLFloat && index->dtype.bits == 32) { + _mg_serialize(res, *index, filename); + } else if (index->dtype.code == kDLFloat && index->dtype.bits == 16) { + _mg_serialize(res, *index, filename); + } else if (index->dtype.code == kDLInt && index->dtype.bits == 8) { + _mg_serialize(res, *index, filename); + } else if (index->dtype.code == kDLUInt && index->dtype.bits == 8) { + _mg_serialize(res, *index, filename); + } else { + RAFT_FAIL("Unsupported index dtype: %d and bits: %d", index->dtype.code, index->dtype.bits); + } + }); +} + +extern "C" cuvsError_t cuvsMultiGpuIvfPqDeserialize(cuvsResources_t res, + const char* filename, + cuvsMultiGpuIvfPqIndex_t index) +{ + return cuvs::core::translate_exceptions([=] { + std::ifstream is(filename, std::ios::in | std::ios::binary); + if (!is) { RAFT_FAIL("Cannot open file %s", filename); } + char dtype_string[4]; + is.read(dtype_string, 4); + auto dtype = raft::detail::numpy_serializer::parse_descr(std::string(dtype_string, 4)); + is.close(); + + index->dtype.bits = dtype.itemsize * 8; + if (dtype.kind == 'f' && dtype.itemsize == 4) { + index->dtype.code = kDLFloat; + index->addr = reinterpret_cast(_mg_deserialize(res, filename)); + } else if (dtype.kind == 'f' && dtype.itemsize == 2) { + index->dtype.code = kDLFloat; + index->addr = reinterpret_cast(_mg_deserialize(res, filename)); + } else if (dtype.kind == 'i' && dtype.itemsize == 1) { + index->dtype.code = kDLInt; + index->addr = reinterpret_cast(_mg_deserialize(res, filename)); + } else if (dtype.kind == 'u' && dtype.itemsize == 1) { + index->dtype.code = kDLUInt; + index->addr = reinterpret_cast(_mg_deserialize(res, filename)); + } else { + RAFT_FAIL("Unsupported index dtype"); + } + }); +} + +extern "C" cuvsError_t cuvsMultiGpuIvfPqDistribute(cuvsResources_t res, + const char* filename, + cuvsMultiGpuIvfPqIndex_t index) +{ + return cuvs::core::translate_exceptions([=] { + std::ifstream is(filename, std::ios::in | std::ios::binary); + if (!is) { RAFT_FAIL("Cannot open file %s", filename); } + char dtype_string[4]; + is.read(dtype_string, 4); + auto dtype = raft::detail::numpy_serializer::parse_descr(std::string(dtype_string, 4)); + is.close(); + + index->dtype.bits = dtype.itemsize * 8; + if (dtype.kind == 'f' && dtype.itemsize == 4) { + index->dtype.code = kDLFloat; + index->addr = reinterpret_cast(_mg_distribute(res, filename)); + } else if (dtype.kind == 'f' && dtype.itemsize == 2) { + index->dtype.code = kDLFloat; + index->addr = reinterpret_cast(_mg_distribute(res, filename)); + } else if (dtype.kind == 'i' && dtype.itemsize == 1) { + index->dtype.code = kDLInt; + index->addr = reinterpret_cast(_mg_distribute(res, filename)); + } else if (dtype.kind == 'u' && dtype.itemsize == 1) { + index->dtype.code = kDLUInt; + index->addr = reinterpret_cast(_mg_distribute(res, filename)); + } else { + RAFT_FAIL("Unsupported index dtype"); + } + }); +} diff --git a/cpp/tests/CMakeLists.txt b/cpp/tests/CMakeLists.txt index a865aeab89..c6bb7ce8e4 100644 --- a/cpp/tests/CMakeLists.txt +++ b/cpp/tests/CMakeLists.txt @@ -315,6 +315,8 @@ if(TARGET cuvs::c_api) ConfigureTest(NAME CAGRA_C_TEST PATH neighbors/ann_cagra_c.cu C_LIB) + ConfigureTest(NAME MG_C_TEST PATH neighbors/run_mg_c.c neighbors/ann_mg_c.cu C_LIB) + if(BUILD_CAGRA_HNSWLIB) ConfigureTest(NAME HNSW_C_TEST PATH neighbors/ann_hnsw_c.cu C_LIB) target_link_libraries(NEIGHBORS_HNSW_TEST PRIVATE hnswlib::hnswlib) diff --git a/cpp/tests/neighbors/ann_mg_c.cu b/cpp/tests/neighbors/ann_mg_c.cu new file mode 100644 index 0000000000..8a82982951 --- /dev/null +++ b/cpp/tests/neighbors/ann_mg_c.cu @@ -0,0 +1,341 @@ +/* + * Copyright (c) 2025, NVIDIA CORPORATION. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include +#include +#include +#include +#include + +#include "ann_utils.cuh" +#include "naive_knn.cuh" +#include +#include +#include +#include + +extern "C" { + +typedef enum { MG_ALGO_IVF_FLAT, MG_ALGO_IVF_PQ, MG_ALGO_CAGRA } mg_algo_t; + +typedef enum { MG_MODE_REPLICATED, MG_MODE_SHARDED, MG_MODE_LOCAL_THEN_DISTRIBUTED } mg_mode_t; + +typedef struct { + int64_t num_queries; + int64_t num_db_vecs; + int64_t dim; + int64_t k; + mg_mode_t mode; + mg_algo_t algo; + int64_t nprobe; + int64_t nlist; + cuvsDistanceType metric; +} mg_test_params; + +int run_mg_ivf_flat_test(mg_test_params params, + float* index_data, + float* query_data, + float* distances_data, + int64_t* neighbors_data, + float* ref_distances_data, + int64_t* ref_neighbors_data); + +int run_mg_ivf_pq_test(mg_test_params params, + float* index_data, + float* query_data, + float* distances_data, + int64_t* neighbors_data, + float* ref_distances_data, + int64_t* ref_neighbors_data); + +int run_mg_cagra_test(mg_test_params params, + float* index_data, + float* query_data, + float* distances_data, + int64_t* neighbors_data, + float* ref_distances_data, + int64_t* ref_neighbors_data); + +int generate_reference_results(mg_test_params params, + float* index_data, + float* query_data, + float* ref_distances_data, + int64_t* ref_neighbors_data); +} + +template +void generate_random_data(T* devPtr, size_t size) +{ + raft::handle_t handle; + raft::random::RngState r(1234ULL); + if constexpr (std::is_same_v) { + raft::random::uniform(handle, r, devPtr, size, T(0.1), T(2.0)); + } else { + raft::random::uniformInt(handle, r, devPtr, size, T(1), T(20)); + } +} + +template +bool eval_recall(T* query_data, + T* index_data, + IdxT* neighbors, + T* distances, + IdxT* ref_neighbors, + T* ref_distances, + size_t n_queries, + size_t n_rows, + size_t n_dim, + size_t n_neighbors, + cuvsDistanceType metric, + double min_recall = 0.9) +{ + raft::handle_t handle; + auto stream = raft::resource::get_cuda_stream(handle); + + size_t size = n_queries * n_neighbors; + std::vector neighbors_h(size); + std::vector distances_h(size); + std::vector neighbors_ref_h(size); + std::vector distances_ref_h(size); + + raft::copy(neighbors_h.data(), neighbors, size, stream); + raft::copy(distances_h.data(), distances, size, stream); + raft::copy(neighbors_ref_h.data(), ref_neighbors, size, stream); + raft::copy(distances_ref_h.data(), ref_distances, size, stream); + + raft::resource::sync_stream(handle); + + // Evaluate recall + return cuvs::neighbors::eval_neighbours(neighbors_ref_h, + neighbors_h, + distances_ref_h, + distances_h, + n_queries, + n_neighbors, + 0.001, + min_recall); +} + +class MgCTest : public ::testing::TestWithParam { + public: + void SetUp() override + { + params = GetParam(); + + raft::handle_t handle; + auto stream = raft::resource::get_cuda_stream(handle); + + // Allocate device memory + index_data.resize(params.num_db_vecs * params.dim, stream); + query_data.resize(params.num_queries * params.dim, stream); + neighbors_data.resize(params.num_queries * params.k, stream); + distances_data.resize(params.num_queries * params.k, stream); + ref_neighbors_data.resize(params.num_queries * params.k, stream); + ref_distances_data.resize(params.num_queries * params.k, stream); + + // Generate random test data + generate_random_data(index_data.data(), params.num_db_vecs * params.dim); + generate_random_data(query_data.data(), params.num_queries * params.dim); + + raft::resource::sync_stream(handle); + + // Allocate host memory for multi-GPU tests + index_data_host.resize(params.num_db_vecs * params.dim); + query_data_host.resize(params.num_queries * params.dim); + neighbors_data_host.resize(params.num_queries * params.k); + distances_data_host.resize(params.num_queries * params.k); + + // Copy data from device to host for multi-GPU functions + raft::copy(index_data_host.data(), index_data.data(), index_data.size(), stream); + raft::copy(query_data_host.data(), query_data.data(), query_data.size(), stream); + + raft::resource::sync_stream(handle); + } + + void TearDown() override {} + + protected: + mg_test_params params; + rmm::device_uvector index_data{0, rmm::cuda_stream_default}; + rmm::device_uvector query_data{0, rmm::cuda_stream_default}; + rmm::device_uvector neighbors_data{0, rmm::cuda_stream_default}; + rmm::device_uvector distances_data{0, rmm::cuda_stream_default}; + rmm::device_uvector ref_neighbors_data{0, rmm::cuda_stream_default}; + rmm::device_uvector ref_distances_data{0, rmm::cuda_stream_default}; + + // Host memory for multi-GPU tests + std::vector index_data_host; + std::vector query_data_host; + std::vector neighbors_data_host; + std::vector distances_data_host; +}; + +TEST_P(MgCTest, MgIvfFlatTest) +{ + auto params = GetParam(); + if (params.algo != MG_ALGO_IVF_FLAT) return; + + // Generate reference results using brute force + int ref_result = generate_reference_results(params, + index_data.data(), + query_data.data(), + ref_distances_data.data(), + ref_neighbors_data.data()); + ASSERT_EQ(ref_result, 0) << "Failed to generate reference results"; + + // Run MG IVF-Flat test (use host data for multi-GPU) + int test_result = run_mg_ivf_flat_test(params, + index_data_host.data(), + query_data_host.data(), + distances_data_host.data(), + neighbors_data_host.data(), + ref_distances_data.data(), + ref_neighbors_data.data()); + ASSERT_EQ(test_result, 0) << "MG IVF-Flat test failed"; + + // Copy results back from host to device for evaluation + raft::handle_t handle; + auto stream = raft::resource::get_cuda_stream(handle); + raft::copy(neighbors_data.data(), neighbors_data_host.data(), neighbors_data_host.size(), stream); + raft::copy(distances_data.data(), distances_data_host.data(), distances_data_host.size(), stream); + raft::resource::sync_stream(handle); + + // Evaluate recall compared to reference + double min_recall = static_cast(params.nprobe) / static_cast(params.nlist); + bool recall_ok = eval_recall(query_data.data(), + index_data.data(), + neighbors_data.data(), + distances_data.data(), + ref_neighbors_data.data(), + ref_distances_data.data(), + params.num_queries, + params.num_db_vecs, + params.dim, + params.k, + params.metric, + min_recall); + ASSERT_TRUE(recall_ok) << "Recall evaluation failed for MG IVF-Flat"; +} + +TEST_P(MgCTest, MgIvfPqTest) +{ + auto params = GetParam(); + if (params.algo != MG_ALGO_IVF_PQ) return; + + // Generate reference results using brute force + int ref_result = generate_reference_results(params, + index_data.data(), + query_data.data(), + ref_distances_data.data(), + ref_neighbors_data.data()); + ASSERT_EQ(ref_result, 0) << "Failed to generate reference results"; + + // Run MG IVF-PQ test (use host data for multi-GPU) + int test_result = run_mg_ivf_pq_test(params, + index_data_host.data(), + query_data_host.data(), + distances_data_host.data(), + neighbors_data_host.data(), + ref_distances_data.data(), + ref_neighbors_data.data()); + ASSERT_EQ(test_result, 0) << "MG IVF-PQ test failed"; + + // Copy results back from host to device for evaluation + raft::handle_t handle; + auto stream = raft::resource::get_cuda_stream(handle); + raft::copy(neighbors_data.data(), neighbors_data_host.data(), neighbors_data_host.size(), stream); + raft::copy(distances_data.data(), distances_data_host.data(), distances_data_host.size(), stream); + raft::resource::sync_stream(handle); + + // Evaluate recall compared to reference + double min_recall = static_cast(params.nprobe) / static_cast(params.nlist); + bool recall_ok = eval_recall(query_data.data(), + index_data.data(), + neighbors_data.data(), + distances_data.data(), + ref_neighbors_data.data(), + ref_distances_data.data(), + params.num_queries, + params.num_db_vecs, + params.dim, + params.k, + params.metric, + min_recall); + ASSERT_TRUE(recall_ok) << "Recall evaluation failed for MG IVF-PQ"; +} + +TEST_P(MgCTest, MgCagraTest) +{ + auto params = GetParam(); + if (params.algo != MG_ALGO_CAGRA) return; + + // Generate reference results using brute force + int ref_result = generate_reference_results(params, + index_data.data(), + query_data.data(), + ref_distances_data.data(), + ref_neighbors_data.data()); + ASSERT_EQ(ref_result, 0) << "Failed to generate reference results"; + + // Run MG CAGRA test (use host data for multi-GPU) + int test_result = run_mg_cagra_test(params, + index_data_host.data(), + query_data_host.data(), + distances_data_host.data(), + neighbors_data_host.data(), + ref_distances_data.data(), + ref_neighbors_data.data()); + ASSERT_EQ(test_result, 0) << "MG CAGRA test failed"; + + // Copy results back from host to device for evaluation + raft::handle_t handle; + auto stream = raft::resource::get_cuda_stream(handle); + raft::copy(neighbors_data.data(), neighbors_data_host.data(), neighbors_data_host.size(), stream); + raft::copy(distances_data.data(), distances_data_host.data(), distances_data_host.size(), stream); + raft::resource::sync_stream(handle); + + // Evaluate recall compared to reference + bool recall_ok = eval_recall(query_data.data(), + index_data.data(), + neighbors_data.data(), + distances_data.data(), + ref_neighbors_data.data(), + ref_distances_data.data(), + params.num_queries, + params.num_db_vecs, + params.dim, + params.k, + params.metric, + 0.9); + ASSERT_TRUE(recall_ok) << "Recall evaluation failed for MG CAGRA"; +} + +// Test parameters that mirror the C++ test cases +const std::vector test_inputs = { + // IVF-Flat tests + {1000, 5000, 8, 16, MG_MODE_REPLICATED, MG_ALGO_IVF_FLAT, 40, 256, L2Expanded}, + {1000, 5000, 8, 16, MG_MODE_SHARDED, MG_ALGO_IVF_FLAT, 40, 256, L2Expanded}, + + // IVF-PQ tests + {1000, 5000, 8, 16, MG_MODE_REPLICATED, MG_ALGO_IVF_PQ, 40, 256, L2Expanded}, + {1000, 5000, 8, 16, MG_MODE_SHARDED, MG_ALGO_IVF_PQ, 40, 256, L2Expanded}, + + // CAGRA tests + {1000, 5000, 8, 16, MG_MODE_REPLICATED, MG_ALGO_CAGRA, 40, 256, L2Expanded}, + {1000, 5000, 8, 16, MG_MODE_SHARDED, MG_ALGO_CAGRA, 40, 256, L2Expanded}, +}; + +INSTANTIATE_TEST_SUITE_P(MgCTests, MgCTest, ::testing::ValuesIn(test_inputs)); diff --git a/cpp/tests/neighbors/run_mg_c.c b/cpp/tests/neighbors/run_mg_c.c new file mode 100644 index 0000000000..e5d35b82ca --- /dev/null +++ b/cpp/tests/neighbors/run_mg_c.c @@ -0,0 +1,576 @@ +/* + * Copyright (c) 2025, NVIDIA CORPORATION. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include +#include +#include +#include +#include +#include +#include + +typedef enum { MG_ALGO_IVF_FLAT, MG_ALGO_IVF_PQ, MG_ALGO_CAGRA } mg_algo_t; + +typedef enum { MG_MODE_REPLICATED, MG_MODE_SHARDED, MG_MODE_LOCAL_THEN_DISTRIBUTED } mg_mode_t; + +typedef struct { + int64_t num_queries; + int64_t num_db_vecs; + int64_t dim; + int64_t k; + mg_mode_t mode; + mg_algo_t algo; + int64_t nprobe; + int64_t nlist; + cuvsDistanceType metric; +} mg_test_params; + +// Test IVF-Flat multi-GPU functionality +int run_mg_ivf_flat_test(mg_test_params params, + float* index_data, + float* query_data, + float* distances_data, + int64_t* neighbors_data, + float* ref_distances_data, + int64_t* ref_neighbors_data) +{ + printf("Running MG IVF-Flat test (mode=%d, queries=%ld, db_vecs=%ld, dim=%ld, k=%ld)\n", + params.mode, + params.num_queries, + params.num_db_vecs, + params.dim, + params.k); + + // Create multi-GPU resources + cuvsResources_t res; + cuvsMultiGpuResourcesCreate(&res); + + // Create dataset tensor + DLManagedTensor dataset_tensor; + dataset_tensor.dl_tensor.data = index_data; + dataset_tensor.dl_tensor.device.device_type = kDLCPU; // Multi-GPU requires host memory + dataset_tensor.dl_tensor.ndim = 2; + dataset_tensor.dl_tensor.dtype.code = kDLFloat; + dataset_tensor.dl_tensor.dtype.bits = 32; + dataset_tensor.dl_tensor.dtype.lanes = 1; + int64_t dataset_shape[2] = {params.num_db_vecs, params.dim}; + dataset_tensor.dl_tensor.shape = dataset_shape; + dataset_tensor.dl_tensor.strides = NULL; + + // Create index + cuvsMultiGpuIvfFlatIndex_t index; + cuvsMultiGpuIvfFlatIndexCreate(&index); + + // Build index + cuvsMultiGpuIvfFlatIndexParams_t build_params; + cuvsMultiGpuIvfFlatIndexParamsCreate(&build_params); + build_params->base_params->metric = params.metric; + build_params->base_params->n_lists = params.nlist; + build_params->base_params->add_data_on_build = false; + build_params->base_params->kmeans_trainset_fraction = 1.0; + build_params->base_params->metric_arg = 0; + + if (params.mode == MG_MODE_REPLICATED) { + build_params->mode = CUVS_NEIGHBORS_MG_REPLICATED; + } else { + build_params->mode = CUVS_NEIGHBORS_MG_SHARDED; + } + + cuvsError_t build_result = cuvsMultiGpuIvfFlatBuild(res, build_params, &dataset_tensor, index); + if (build_result != CUVS_SUCCESS) { + const char* error_msg = cuvsGetLastErrorText(); + printf("MG IVF-Flat build failed: %s\n", error_msg ? error_msg : "Unknown error"); + goto cleanup; + } + + // Extend index with data + cuvsError_t extend_result = cuvsMultiGpuIvfFlatExtend(res, index, &dataset_tensor, NULL); + if (extend_result != CUVS_SUCCESS) { + printf("MG IVF-Flat extend failed\n"); + goto cleanup; + } + + // Serialize index to test serialization/deserialization + const char* serialize_filename = "/tmp/mg_ivf_flat_test_index.bin"; + cuvsError_t serialize_result = cuvsMultiGpuIvfFlatSerialize(res, index, serialize_filename); + if (serialize_result != CUVS_SUCCESS) { + printf("MG IVF-Flat serialize failed\n"); + goto cleanup; + } + + // Destroy current index and recreate it + cuvsMultiGpuIvfFlatIndexDestroy(index); + cuvsMultiGpuIvfFlatIndexCreate(&index); + + // Deserialize index from file + cuvsError_t deserialize_result = cuvsMultiGpuIvfFlatDeserialize(res, serialize_filename, index); + if (deserialize_result != CUVS_SUCCESS) { + printf("MG IVF-Flat deserialize failed\n"); + goto cleanup; + } + + // Create queries tensor + DLManagedTensor queries_tensor; + queries_tensor.dl_tensor.data = query_data; + queries_tensor.dl_tensor.device.device_type = kDLCPU; // Multi-GPU requires host memory + queries_tensor.dl_tensor.ndim = 2; + queries_tensor.dl_tensor.dtype.code = kDLFloat; + queries_tensor.dl_tensor.dtype.bits = 32; + queries_tensor.dl_tensor.dtype.lanes = 1; + int64_t queries_shape[2] = {params.num_queries, params.dim}; + queries_tensor.dl_tensor.shape = queries_shape; + queries_tensor.dl_tensor.strides = NULL; + + // Create neighbors tensor + DLManagedTensor neighbors_tensor; + neighbors_tensor.dl_tensor.data = neighbors_data; + neighbors_tensor.dl_tensor.device.device_type = kDLCPU; // Multi-GPU requires host memory + neighbors_tensor.dl_tensor.ndim = 2; + neighbors_tensor.dl_tensor.dtype.code = kDLInt; + neighbors_tensor.dl_tensor.dtype.bits = 64; + neighbors_tensor.dl_tensor.dtype.lanes = 1; + int64_t neighbors_shape[2] = {params.num_queries, params.k}; + neighbors_tensor.dl_tensor.shape = neighbors_shape; + neighbors_tensor.dl_tensor.strides = NULL; + + // Create distances tensor + DLManagedTensor distances_tensor; + distances_tensor.dl_tensor.data = distances_data; + distances_tensor.dl_tensor.device.device_type = kDLCPU; // Multi-GPU requires host memory + distances_tensor.dl_tensor.ndim = 2; + distances_tensor.dl_tensor.dtype.code = kDLFloat; + distances_tensor.dl_tensor.dtype.bits = 32; + distances_tensor.dl_tensor.dtype.lanes = 1; + int64_t distances_shape[2] = {params.num_queries, params.k}; + distances_tensor.dl_tensor.shape = distances_shape; + distances_tensor.dl_tensor.strides = NULL; + + // Search index + cuvsMultiGpuIvfFlatSearchParams_t search_params; + cuvsMultiGpuIvfFlatSearchParamsCreate(&search_params); + search_params->base_params->n_probes = params.nprobe; + search_params->search_mode = CUVS_NEIGHBORS_MG_LOAD_BALANCER; + search_params->merge_mode = CUVS_NEIGHBORS_MG_TREE_MERGE; + search_params->n_rows_per_batch = 3000; + + cuvsError_t search_result = cuvsMultiGpuIvfFlatSearch( + res, search_params, index, &queries_tensor, &neighbors_tensor, &distances_tensor); + if (search_result != CUVS_SUCCESS) { + printf("MG IVF-Flat search failed\n"); + goto cleanup; + } + + printf("MG IVF-Flat test completed successfully\n"); + +cleanup: + cuvsMultiGpuIvfFlatSearchParamsDestroy(search_params); + cuvsMultiGpuIvfFlatIndexParamsDestroy(build_params); + cuvsMultiGpuIvfFlatIndexDestroy(index); + cuvsMultiGpuResourcesDestroy(res); + + // Clean up temporary serialization file + remove("/tmp/mg_ivf_flat_test_index.bin"); + + return (build_result == CUVS_SUCCESS && extend_result == CUVS_SUCCESS && + serialize_result == CUVS_SUCCESS && deserialize_result == CUVS_SUCCESS && + search_result == CUVS_SUCCESS) + ? 0 + : 1; +} + +// Test IVF-PQ multi-GPU functionality +int run_mg_ivf_pq_test(mg_test_params params, + float* index_data, + float* query_data, + float* distances_data, + int64_t* neighbors_data, + float* ref_distances_data, + int64_t* ref_neighbors_data) +{ + printf("Running MG IVF-PQ test (mode=%d, queries=%ld, db_vecs=%ld, dim=%ld, k=%ld)\n", + params.mode, + params.num_queries, + params.num_db_vecs, + params.dim, + params.k); + + // Create multi-GPU resources + cuvsResources_t res; + cuvsMultiGpuResourcesCreate(&res); + + // Create dataset tensor + DLManagedTensor dataset_tensor; + dataset_tensor.dl_tensor.data = index_data; + dataset_tensor.dl_tensor.device.device_type = kDLCPU; // Multi-GPU requires host memory + dataset_tensor.dl_tensor.ndim = 2; + dataset_tensor.dl_tensor.dtype.code = kDLFloat; + dataset_tensor.dl_tensor.dtype.bits = 32; + dataset_tensor.dl_tensor.dtype.lanes = 1; + int64_t dataset_shape[2] = {params.num_db_vecs, params.dim}; + dataset_tensor.dl_tensor.shape = dataset_shape; + dataset_tensor.dl_tensor.strides = NULL; + + // Create index + cuvsMultiGpuIvfPqIndex_t index; + cuvsMultiGpuIvfPqIndexCreate(&index); + + // Build index + cuvsMultiGpuIvfPqIndexParams_t build_params; + cuvsMultiGpuIvfPqIndexParamsCreate(&build_params); + build_params->base_params->metric = params.metric; + build_params->base_params->n_lists = params.nlist; + build_params->base_params->add_data_on_build = false; + build_params->base_params->kmeans_trainset_fraction = 1.0; + build_params->base_params->metric_arg = 0; + + if (params.mode == MG_MODE_REPLICATED) { + build_params->mode = CUVS_NEIGHBORS_MG_REPLICATED; + } else { + build_params->mode = CUVS_NEIGHBORS_MG_SHARDED; + } + + cuvsError_t build_result = cuvsMultiGpuIvfPqBuild(res, build_params, &dataset_tensor, index); + if (build_result != CUVS_SUCCESS) { + const char* error_msg = cuvsGetLastErrorText(); + printf("MG IVF-PQ build failed: %s\n", error_msg ? error_msg : "Unknown error"); + goto cleanup; + } + + // Extend index with data + cuvsError_t extend_result = cuvsMultiGpuIvfPqExtend(res, index, &dataset_tensor, NULL); + if (extend_result != CUVS_SUCCESS) { + printf("MG IVF-PQ extend failed\n"); + goto cleanup; + } + + // Serialize index to test serialization/deserialization + const char* serialize_filename = "/tmp/mg_ivf_pq_test_index.bin"; + cuvsError_t serialize_result = cuvsMultiGpuIvfPqSerialize(res, index, serialize_filename); + if (serialize_result != CUVS_SUCCESS) { + printf("MG IVF-PQ serialize failed\n"); + goto cleanup; + } + + // Destroy current index and recreate it + cuvsMultiGpuIvfPqIndexDestroy(index); + cuvsMultiGpuIvfPqIndexCreate(&index); + + // Deserialize index from file + cuvsError_t deserialize_result = cuvsMultiGpuIvfPqDeserialize(res, serialize_filename, index); + if (deserialize_result != CUVS_SUCCESS) { + printf("MG IVF-PQ deserialize failed\n"); + goto cleanup; + } + + // Create queries tensor + DLManagedTensor queries_tensor; + queries_tensor.dl_tensor.data = query_data; + queries_tensor.dl_tensor.device.device_type = kDLCPU; // Multi-GPU requires host memory + queries_tensor.dl_tensor.ndim = 2; + queries_tensor.dl_tensor.dtype.code = kDLFloat; + queries_tensor.dl_tensor.dtype.bits = 32; + queries_tensor.dl_tensor.dtype.lanes = 1; + int64_t queries_shape[2] = {params.num_queries, params.dim}; + queries_tensor.dl_tensor.shape = queries_shape; + queries_tensor.dl_tensor.strides = NULL; + + // Create neighbors tensor + DLManagedTensor neighbors_tensor; + neighbors_tensor.dl_tensor.data = neighbors_data; + neighbors_tensor.dl_tensor.device.device_type = kDLCPU; // Multi-GPU requires host memory + neighbors_tensor.dl_tensor.ndim = 2; + neighbors_tensor.dl_tensor.dtype.code = kDLInt; + neighbors_tensor.dl_tensor.dtype.bits = 64; + neighbors_tensor.dl_tensor.dtype.lanes = 1; + int64_t neighbors_shape[2] = {params.num_queries, params.k}; + neighbors_tensor.dl_tensor.shape = neighbors_shape; + neighbors_tensor.dl_tensor.strides = NULL; + + // Create distances tensor + DLManagedTensor distances_tensor; + distances_tensor.dl_tensor.data = distances_data; + distances_tensor.dl_tensor.device.device_type = kDLCPU; // Multi-GPU requires host memory + distances_tensor.dl_tensor.ndim = 2; + distances_tensor.dl_tensor.dtype.code = kDLFloat; + distances_tensor.dl_tensor.dtype.bits = 32; + distances_tensor.dl_tensor.dtype.lanes = 1; + int64_t distances_shape[2] = {params.num_queries, params.k}; + distances_tensor.dl_tensor.shape = distances_shape; + distances_tensor.dl_tensor.strides = NULL; + + // Search index + cuvsMultiGpuIvfPqSearchParams_t search_params; + cuvsMultiGpuIvfPqSearchParamsCreate(&search_params); + search_params->base_params->n_probes = params.nprobe; + search_params->search_mode = CUVS_NEIGHBORS_MG_LOAD_BALANCER; + search_params->merge_mode = CUVS_NEIGHBORS_MG_TREE_MERGE; + search_params->n_rows_per_batch = 3000; + + cuvsError_t search_result = cuvsMultiGpuIvfPqSearch( + res, search_params, index, &queries_tensor, &neighbors_tensor, &distances_tensor); + if (search_result != CUVS_SUCCESS) { + printf("MG IVF-PQ search failed\n"); + goto cleanup; + } + + printf("MG IVF-PQ test completed successfully\n"); + +cleanup: + cuvsMultiGpuIvfPqSearchParamsDestroy(search_params); + cuvsMultiGpuIvfPqIndexParamsDestroy(build_params); + cuvsMultiGpuIvfPqIndexDestroy(index); + cuvsMultiGpuResourcesDestroy(res); + + // Clean up temporary serialization file + remove("/tmp/mg_ivf_pq_test_index.bin"); + + return (build_result == CUVS_SUCCESS && extend_result == CUVS_SUCCESS && + serialize_result == CUVS_SUCCESS && deserialize_result == CUVS_SUCCESS && + search_result == CUVS_SUCCESS) + ? 0 + : 1; +} + +// Test CAGRA multi-GPU functionality +int run_mg_cagra_test(mg_test_params params, + float* index_data, + float* query_data, + float* distances_data, + int64_t* neighbors_data, + float* ref_distances_data, + int64_t* ref_neighbors_data) +{ + printf("Running MG CAGRA test (mode=%d, queries=%ld, db_vecs=%ld, dim=%ld, k=%ld)\n", + params.mode, + params.num_queries, + params.num_db_vecs, + params.dim, + params.k); + + // Create multi-GPU resources + cuvsResources_t res; + cuvsMultiGpuResourcesCreate(&res); + + // Create dataset tensor + DLManagedTensor dataset_tensor; + dataset_tensor.dl_tensor.data = index_data; + dataset_tensor.dl_tensor.device.device_type = kDLCPU; // Multi-GPU requires host memory + dataset_tensor.dl_tensor.ndim = 2; + dataset_tensor.dl_tensor.dtype.code = kDLFloat; + dataset_tensor.dl_tensor.dtype.bits = 32; + dataset_tensor.dl_tensor.dtype.lanes = 1; + int64_t dataset_shape[2] = {params.num_db_vecs, params.dim}; + dataset_tensor.dl_tensor.shape = dataset_shape; + dataset_tensor.dl_tensor.strides = NULL; + + // Create index + cuvsMultiGpuCagraIndex_t index; + cuvsMultiGpuCagraIndexCreate(&index); + + // Build index + cuvsMultiGpuCagraIndexParams_t build_params; + cuvsMultiGpuCagraIndexParamsCreate(&build_params); + + if (params.mode == MG_MODE_REPLICATED) { + build_params->mode = CUVS_NEIGHBORS_MG_REPLICATED; + } else { + build_params->mode = CUVS_NEIGHBORS_MG_SHARDED; + } + + cuvsError_t build_result = cuvsMultiGpuCagraBuild(res, build_params, &dataset_tensor, index); + if (build_result != CUVS_SUCCESS) { + const char* error_msg = cuvsGetLastErrorText(); + printf("MG CAGRA build failed: %s\n", error_msg ? error_msg : "Unknown error"); + goto cleanup; + } + + // Serialize index to test serialization/deserialization + const char* serialize_filename = "/tmp/mg_cagra_test_index.bin"; + cuvsError_t serialize_result = cuvsMultiGpuCagraSerialize(res, index, serialize_filename); + if (serialize_result != CUVS_SUCCESS) { + printf("MG CAGRA serialize failed\n"); + goto cleanup; + } + + // Destroy current index and recreate it + cuvsMultiGpuCagraIndexDestroy(index); + cuvsMultiGpuCagraIndexCreate(&index); + + // Deserialize index from file + cuvsError_t deserialize_result = cuvsMultiGpuCagraDeserialize(res, serialize_filename, index); + if (deserialize_result != CUVS_SUCCESS) { + printf("MG CAGRA deserialize failed\n"); + goto cleanup; + } + + // Create queries tensor + DLManagedTensor queries_tensor; + queries_tensor.dl_tensor.data = query_data; + queries_tensor.dl_tensor.device.device_type = kDLCPU; // Multi-GPU requires host memory + queries_tensor.dl_tensor.ndim = 2; + queries_tensor.dl_tensor.dtype.code = kDLFloat; + queries_tensor.dl_tensor.dtype.bits = 32; + queries_tensor.dl_tensor.dtype.lanes = 1; + int64_t queries_shape[2] = {params.num_queries, params.dim}; + queries_tensor.dl_tensor.shape = queries_shape; + queries_tensor.dl_tensor.strides = NULL; + + // Create neighbors tensor + DLManagedTensor neighbors_tensor; + neighbors_tensor.dl_tensor.data = neighbors_data; + neighbors_tensor.dl_tensor.device.device_type = kDLCPU; // Multi-GPU requires host memory + neighbors_tensor.dl_tensor.ndim = 2; + neighbors_tensor.dl_tensor.dtype.code = kDLInt; + neighbors_tensor.dl_tensor.dtype.bits = 64; + neighbors_tensor.dl_tensor.dtype.lanes = 1; + int64_t neighbors_shape[2] = {params.num_queries, params.k}; + neighbors_tensor.dl_tensor.shape = neighbors_shape; + neighbors_tensor.dl_tensor.strides = NULL; + + // Create distances tensor + DLManagedTensor distances_tensor; + distances_tensor.dl_tensor.data = distances_data; + distances_tensor.dl_tensor.device.device_type = kDLCPU; // Multi-GPU requires host memory + distances_tensor.dl_tensor.ndim = 2; + distances_tensor.dl_tensor.dtype.code = kDLFloat; + distances_tensor.dl_tensor.dtype.bits = 32; + distances_tensor.dl_tensor.dtype.lanes = 1; + int64_t distances_shape[2] = {params.num_queries, params.k}; + distances_tensor.dl_tensor.shape = distances_shape; + distances_tensor.dl_tensor.strides = NULL; + + // Search index + cuvsMultiGpuCagraSearchParams_t search_params; + cuvsMultiGpuCagraSearchParamsCreate(&search_params); + search_params->search_mode = CUVS_NEIGHBORS_MG_LOAD_BALANCER; + search_params->merge_mode = CUVS_NEIGHBORS_MG_TREE_MERGE; + search_params->n_rows_per_batch = 3000; + + cuvsError_t search_result = cuvsMultiGpuCagraSearch( + res, search_params, index, &queries_tensor, &neighbors_tensor, &distances_tensor); + if (search_result != CUVS_SUCCESS) { + printf("MG CAGRA search failed\n"); + goto cleanup; + } + + printf("MG CAGRA test completed successfully\n"); + +cleanup: + cuvsMultiGpuCagraSearchParamsDestroy(search_params); + cuvsMultiGpuCagraIndexParamsDestroy(build_params); + cuvsMultiGpuCagraIndexDestroy(index); + cuvsMultiGpuResourcesDestroy(res); + + // Clean up temporary serialization file + remove("/tmp/mg_cagra_test_index.bin"); + + return (build_result == CUVS_SUCCESS && serialize_result == CUVS_SUCCESS && + deserialize_result == CUVS_SUCCESS && search_result == CUVS_SUCCESS) + ? 0 + : 1; +} + +// Generate reference results using brute force +int generate_reference_results(mg_test_params params, + float* index_data, + float* query_data, + float* ref_distances_data, + int64_t* ref_neighbors_data) +{ + printf("Generating reference results using brute force\n"); + + // Create resources + cuvsResources_t res; + cuvsResourcesCreate(&res); + + // Create dataset tensor + DLManagedTensor dataset_tensor; + dataset_tensor.dl_tensor.data = index_data; + dataset_tensor.dl_tensor.device.device_type = kDLCUDA; // Brute force can use device memory + dataset_tensor.dl_tensor.ndim = 2; + dataset_tensor.dl_tensor.dtype.code = kDLFloat; + dataset_tensor.dl_tensor.dtype.bits = 32; + dataset_tensor.dl_tensor.dtype.lanes = 1; + int64_t dataset_shape[2] = {params.num_db_vecs, params.dim}; + dataset_tensor.dl_tensor.shape = dataset_shape; + dataset_tensor.dl_tensor.strides = NULL; + + // Create queries tensor + DLManagedTensor queries_tensor; + queries_tensor.dl_tensor.data = query_data; + queries_tensor.dl_tensor.device.device_type = kDLCUDA; // Brute force can use device memory + queries_tensor.dl_tensor.ndim = 2; + queries_tensor.dl_tensor.dtype.code = kDLFloat; + queries_tensor.dl_tensor.dtype.bits = 32; + queries_tensor.dl_tensor.dtype.lanes = 1; + int64_t queries_shape[2] = {params.num_queries, params.dim}; + queries_tensor.dl_tensor.shape = queries_shape; + queries_tensor.dl_tensor.strides = NULL; + + // Create neighbors tensor + DLManagedTensor neighbors_tensor; + neighbors_tensor.dl_tensor.data = ref_neighbors_data; + neighbors_tensor.dl_tensor.device.device_type = kDLCUDA; + neighbors_tensor.dl_tensor.ndim = 2; + neighbors_tensor.dl_tensor.dtype.code = kDLInt; + neighbors_tensor.dl_tensor.dtype.bits = 64; + neighbors_tensor.dl_tensor.dtype.lanes = 1; + int64_t neighbors_shape[2] = {params.num_queries, params.k}; + neighbors_tensor.dl_tensor.shape = neighbors_shape; + neighbors_tensor.dl_tensor.strides = NULL; + + // Create distances tensor + DLManagedTensor distances_tensor; + distances_tensor.dl_tensor.data = ref_distances_data; + distances_tensor.dl_tensor.device.device_type = kDLCUDA; + distances_tensor.dl_tensor.ndim = 2; + distances_tensor.dl_tensor.dtype.code = kDLFloat; + distances_tensor.dl_tensor.dtype.bits = 32; + distances_tensor.dl_tensor.dtype.lanes = 1; + int64_t distances_shape[2] = {params.num_queries, params.k}; + distances_tensor.dl_tensor.shape = distances_shape; + distances_tensor.dl_tensor.strides = NULL; + + // Build brute force index + cuvsBruteForceIndex_t index; + cuvsBruteForceIndexCreate(&index); + + cuvsError_t build_result = cuvsBruteForceBuild(res, &dataset_tensor, params.metric, 0.0f, index); + if (build_result != CUVS_SUCCESS) { + printf("Brute force build failed\n"); + goto cleanup; + } + + // Search with brute force + cuvsFilter filter; + filter.type = NO_FILTER; + filter.addr = (uintptr_t)NULL; + + cuvsError_t search_result = + cuvsBruteForceSearch(res, index, &queries_tensor, &neighbors_tensor, &distances_tensor, filter); + if (search_result != CUVS_SUCCESS) { + printf("Brute force search failed\n"); + goto cleanup; + } + + printf("Reference results generated successfully\n"); + +cleanup: + cuvsBruteForceIndexDestroy(index); + cuvsResourcesDestroy(res); + + return (build_result == CUVS_SUCCESS && search_result == CUVS_SUCCESS) ? 0 : 1; +} diff --git a/docs/source/c_api/neighbors.rst b/docs/source/c_api/neighbors.rst index 9c3fce672a..9de9c40a96 100644 --- a/docs/source/c_api/neighbors.rst +++ b/docs/source/c_api/neighbors.rst @@ -14,3 +14,4 @@ Nearest Neighbors neighbors_ivf_pq_c.rst neighbors_cagra_c.rst neighbors_hnsw_c.rst + neighbors_mg.rst diff --git a/docs/source/c_api/neighbors_mg.rst b/docs/source/c_api/neighbors_mg.rst new file mode 100644 index 0000000000..bffe3fc4c5 --- /dev/null +++ b/docs/source/c_api/neighbors_mg.rst @@ -0,0 +1,257 @@ +Multi-GPU Nearest Neighbors +=========================== + +The Multi-GPU (SNMG - single-node multi-GPUs) C API provides a set of functions to deploy ANN indexes across multiple GPUs for improved performance and scalability. + +.. role:: py(code) + :language: c + :class: highlight + +Common Types and Enums +====================== + +Common types and enums used across multi-GPU ANN algorithms. + +``#include `` + +.. doxygengroup:: mg_c_common_types + :project: cuvs + :members: + :content-only: + +Multi-GPU IVF-Flat +================== + +The Multi-GPU IVF-Flat method extends the IVF-Flat ANN algorithm to work across multiple GPUs. It provides two distribution modes: replicated (for higher throughput) and sharded (for handling larger datasets). + +``#include `` + +IVF-Flat Index Build Parameters +------------------------------- + +.. doxygengroup:: mg_ivf_flat_c_index_params + :project: cuvs + :members: + :content-only: + +IVF-Flat Index Search Parameters +-------------------------------- + +.. doxygengroup:: mg_ivf_flat_c_search_params + :project: cuvs + :members: + :content-only: + +IVF-Flat Index +-------------- + +.. doxygengroup:: mg_ivf_flat_c_index + :project: cuvs + :members: + :content-only: + +IVF-Flat Index Build +-------------------- + +.. doxygengroup:: mg_ivf_flat_c_index_build + :project: cuvs + :members: + :content-only: + +IVF-Flat Index Search +--------------------- + +.. doxygengroup:: mg_ivf_flat_c_index_search + :project: cuvs + :members: + :content-only: + +IVF-Flat Index Extend +--------------------- + +.. doxygengroup:: mg_ivf_flat_c_index_extend + :project: cuvs + :members: + :content-only: + +IVF-Flat Index Serialize +------------------------ + +.. doxygengroup:: mg_ivf_flat_c_index_serialize + :project: cuvs + :members: + :content-only: + +IVF-Flat Index Deserialize +--------------------------- + +.. doxygengroup:: mg_ivf_flat_c_index_deserialize + :project: cuvs + :members: + :content-only: + +IVF-Flat Index Distribute +-------------------------- + +.. doxygengroup:: mg_ivf_flat_c_index_distribute + :project: cuvs + :members: + :content-only: + +Multi-GPU IVF-PQ +================= + +The Multi-GPU IVF-PQ method extends the IVF-PQ ANN algorithm to work across multiple GPUs. It provides two distribution modes: replicated (for higher throughput) and sharded (for handling larger datasets). + +``#include `` + +IVF-PQ Index Build Parameters +----------------------------- + +.. doxygengroup:: mg_ivf_pq_c_index_params + :project: cuvs + :members: + :content-only: + +IVF-PQ Index Search Parameters +------------------------------ + +.. doxygengroup:: mg_ivf_pq_c_search_params + :project: cuvs + :members: + :content-only: + +IVF-PQ Index +------------ + +.. doxygengroup:: mg_ivf_pq_c_index + :project: cuvs + :members: + :content-only: + +IVF-PQ Index Build +------------------ + +.. doxygengroup:: mg_ivf_pq_c_index_build + :project: cuvs + :members: + :content-only: + +IVF-PQ Index Search +------------------- + +.. doxygengroup:: mg_ivf_pq_c_index_search + :project: cuvs + :members: + :content-only: + +IVF-PQ Index Extend +------------------- + +.. doxygengroup:: mg_ivf_pq_c_index_extend + :project: cuvs + :members: + :content-only: + +IVF-PQ Index Serialize +---------------------- + +.. doxygengroup:: mg_ivf_pq_c_index_serialize + :project: cuvs + :members: + :content-only: + +IVF-PQ Index Deserialize +------------------------ + +.. doxygengroup:: mg_ivf_pq_c_index_deserialize + :project: cuvs + :members: + :content-only: + +IVF-PQ Index Distribute +----------------------- + +.. doxygengroup:: mg_ivf_pq_c_index_distribute + :project: cuvs + :members: + :content-only: + +Multi-GPU CAGRA +================ + +The Multi-GPU CAGRA method extends the CAGRA graph-based ANN algorithm to work across multiple GPUs. It provides two distribution modes: replicated (for higher throughput) and sharded (for handling larger datasets). + +``#include `` + +CAGRA Index Build Parameters +---------------------------- + +.. doxygengroup:: mg_cagra_c_index_params + :project: cuvs + :members: + :content-only: + +CAGRA Index Search Parameters +----------------------------- + +.. doxygengroup:: mg_cagra_c_search_params + :project: cuvs + :members: + :content-only: + +CAGRA Index +----------- + +.. doxygengroup:: mg_cagra_c_index + :project: cuvs + :members: + :content-only: + +CAGRA Index Build +----------------- + +.. doxygengroup:: mg_cagra_c_index_build + :project: cuvs + :members: + :content-only: + +CAGRA Index Search +------------------ + +.. doxygengroup:: mg_cagra_c_index_search + :project: cuvs + :members: + :content-only: + +CAGRA Index Extend +------------------ + +.. doxygengroup:: mg_cagra_c_index_extend + :project: cuvs + :members: + :content-only: + +CAGRA Index Serialize +--------------------- + +.. doxygengroup:: mg_cagra_c_index_serialize + :project: cuvs + :members: + :content-only: + +CAGRA Index Deserialize +----------------------- + +.. doxygengroup:: mg_cagra_c_index_deserialize + :project: cuvs + :members: + :content-only: + +CAGRA Index Distribute +---------------------- + +.. doxygengroup:: mg_cagra_c_index_distribute + :project: cuvs + :members: + :content-only: From 3ca1ed9096f32112d3ceeabd558fc137ee10e2e9 Mon Sep 17 00:00:00 2001 From: Vinay Deshpande Date: Tue, 2 Sep 2025 16:50:12 +0100 Subject: [PATCH 51/93] Fix redundant memset (#1305) The workspace is already being reset in the calling function `fusedDistanceNNImpl`, see https://github.com/rapidsai/cuvs/blob/branch-25.10/cpp/src/distance/detail/fused_distance_nn.cuh#L73 Therefore removing the memset in `fusedL2NNImpl` Authors: - Vinay Deshpande (https://github.com/vinaydes) Approvers: - Tamas Bela Feher (https://github.com/tfeher) URL: https://github.com/rapidsai/cuvs/pull/1305 --- cpp/src/distance/detail/fused_distance_nn/fused_l2_nn.cuh | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/cpp/src/distance/detail/fused_distance_nn/fused_l2_nn.cuh b/cpp/src/distance/detail/fused_distance_nn/fused_l2_nn.cuh index 59fe10ea00..0d194f3833 100644 --- a/cpp/src/distance/detail/fused_distance_nn/fused_l2_nn.cuh +++ b/cpp/src/distance/detail/fused_distance_nn/fused_l2_nn.cuh @@ -1,5 +1,5 @@ /* - * Copyright (c) 2021-2024, NVIDIA CORPORATION. + * Copyright (c) 2021-2025, NVIDIA CORPORATION. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -64,7 +64,6 @@ void fusedL2NNImpl(OutT* min, constexpr auto maxVal = std::numeric_limits::max(); typedef raft::KeyValuePair KVPair; - RAFT_CUDA_TRY(cudaMemsetAsync(workspace, 0, sizeof(int) * m, stream)); if (initOutBuffer) { initKernel <<>>(min, m, maxVal, redOp); From 1f482cac6ac3f934eba42a4929d1b34b1af77b61 Mon Sep 17 00:00:00 2001 From: Ben Frederickson Date: Wed, 3 Sep 2025 13:39:33 -0700 Subject: [PATCH 52/93] Expose pq-centers to C and Python (#1303) Authors: - Ben Frederickson (https://github.com/benfred) Approvers: - Tarang Jain (https://github.com/tarang-jain) - Corey J. Nolet (https://github.com/cjnolet) URL: https://github.com/rapidsai/cuvs/pull/1303 --- cpp/include/cuvs/neighbors/ivf_flat.h | 7 +-- cpp/include/cuvs/neighbors/ivf_pq.h | 22 +++++++--- cpp/src/neighbors/ivf_flat_c.cpp | 41 ++++-------------- cpp/src/neighbors/ivf_pq_c.cpp | 43 +++++++++++-------- .../cuvs/cuvs/neighbors/ivf_flat/ivf_flat.pxd | 5 +-- .../cuvs/cuvs/neighbors/ivf_flat/ivf_flat.pyx | 17 +++----- python/cuvs/cuvs/neighbors/ivf_pq/ivf_pq.pxd | 8 ++-- python/cuvs/cuvs/neighbors/ivf_pq/ivf_pq.pyx | 30 ++++++++----- python/cuvs/cuvs/tests/test_ivf_pq.py | 7 ++- 9 files changed, 91 insertions(+), 89 deletions(-) diff --git a/cpp/include/cuvs/neighbors/ivf_flat.h b/cpp/include/cuvs/neighbors/ivf_flat.h index e274af93b9..7d647829ad 100644 --- a/cpp/include/cuvs/neighbors/ivf_flat.h +++ b/cpp/include/cuvs/neighbors/ivf_flat.h @@ -1,5 +1,5 @@ /* - * Copyright (c) 2024, NVIDIA CORPORATION. + * Copyright (c) 2024-2025, NVIDIA CORPORATION. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -174,14 +174,11 @@ uint32_t cuvsIvfFlatIndexGetDim(cuvsIvfFlatIndex_t index); /** * @brief Get the cluster centers corresponding to the lists [n_lists, dim] * - * @param[in] res cuvsResources_t opaque C handle * @param[in] index cuvsIvfFlatIndex_t Built Ivf-Flat Index * @param[out] centers Preallocated array on host or device memory to store output, [n_lists, dim] * @return cuvsError_t */ -cuvsError_t cuvsIvfFlatIndexGetCenters(cuvsResources_t res, - cuvsIvfFlatIndex_t index, - DLManagedTensor* centers); +cuvsError_t cuvsIvfFlatIndexGetCenters(cuvsIvfFlatIndex_t index, DLManagedTensor* centers); /** * @} diff --git a/cpp/include/cuvs/neighbors/ivf_pq.h b/cpp/include/cuvs/neighbors/ivf_pq.h index 40b9e9b839..5c349be5cb 100644 --- a/cpp/include/cuvs/neighbors/ivf_pq.h +++ b/cpp/include/cuvs/neighbors/ivf_pq.h @@ -1,5 +1,5 @@ /* - * Copyright (c) 2024, NVIDIA CORPORATION. + * Copyright (c) 2024-2025, NVIDIA CORPORATION. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -270,15 +270,23 @@ uint32_t cuvsIvfPqIndexGetDim(cuvsIvfPqIndex_t index); /** * @brief Get the cluster centers corresponding to the lists in the original space * - * @param[in] res cuvsResources_t opaque C handle * @param[in] index cuvsIvfPqIndex_t Built Ivf-Pq index - * @param[out] centers Preallocated array on host or device memory to store output, - * dimensions [n_lists, dim] + * @param[out] centers Output tensor that will be populated with a non-owning view of the data + * @return cuvsError_t + */ +cuvsError_t cuvsIvfPqIndexGetCenters(cuvsIvfPqIndex_t index, DLManagedTensor* centers); + +/** + * @brief Get the PQ cluster centers + * + * - codebook_gen::PER_SUBSPACE: [pq_dim , pq_len, pq_book_size] + * - codebook_gen::PER_CLUSTER: [n_lists, pq_len, pq_book_size] + * + * @param[in] index cuvsIvfPqIndex_t Built Ivf-Pq index + * @param[out] pq_centers Output tensor that will be populated with a non-owning view of the data * @return cuvsError_t */ -cuvsError_t cuvsIvfPqIndexGetCenters(cuvsResources_t res, - cuvsIvfPqIndex_t index, - DLManagedTensor* centers); +cuvsError_t cuvsIvfPqIndexGetPqCenters(cuvsIvfPqIndex_t index, DLManagedTensor* pq_centers); /** * @} */ diff --git a/cpp/src/neighbors/ivf_flat_c.cpp b/cpp/src/neighbors/ivf_flat_c.cpp index 2e660f6377..91ad9abbed 100644 --- a/cpp/src/neighbors/ivf_flat_c.cpp +++ b/cpp/src/neighbors/ivf_flat_c.cpp @@ -1,6 +1,6 @@ /* - * Copyright (c) 2024, NVIDIA CORPORATION. + * Copyright (c) 2024-2025, NVIDIA CORPORATION. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -154,33 +154,11 @@ void _extend(cuvsResources_t res, cuvs::neighbors::ivf_flat::extend(*res_ptr, vectors_mds, indices_mds, index_ptr); } -template -void _get_centers(cuvsResources_t res, cuvsIvfFlatIndex index, DLManagedTensor* centers) -{ - auto res_ptr = reinterpret_cast(res); - auto index_ptr = reinterpret_cast*>(index.addr); - auto dst = cuvs::core::from_dlpack(centers); - auto src = index_ptr->centers(); - - RAFT_EXPECTS(src.extent(0) == dst.extent(0), "Output centers has incorrect number of rows"); - RAFT_EXPECTS(src.extent(1) == dst.extent(1), "Output centers has incorrect number of cols"); - - raft::copy(dst.data_handle(), - src.data_handle(), - dst.extent(0) * dst.extent(1), - raft::resource::get_cuda_stream(*res_ptr)); -} - template -void get_centers(cuvsResources_t res, cuvsIvfFlatIndex index, DLManagedTensor* centers) +void get_centers(cuvsIvfFlatIndex index, DLManagedTensor* centers) { - if (cuvs::core::is_dlpack_device_compatible(centers->dl_tensor)) { - using output_mdspan_type = raft::device_matrix_view; - _get_centers(res, index, centers); - } else { - using output_mdspan_type = raft::host_matrix_view; - _get_centers(res, index, centers); - } + auto index_ptr = reinterpret_cast*>(index.addr); + cuvs::core::to_dlpack(index_ptr->centers(), centers); } } // namespace @@ -441,19 +419,18 @@ extern "C" uint32_t cuvsIvfFlatIndexGetDim(cuvsIvfFlatIndex_t index) } } -extern "C" cuvsError_t cuvsIvfFlatIndexGetCenters(cuvsResources_t res, - cuvsIvfFlatIndex_t index, +extern "C" cuvsError_t cuvsIvfFlatIndexGetCenters(cuvsIvfFlatIndex_t index, DLManagedTensor* centers) { return cuvs::core::translate_exceptions([=] { if (index->dtype.code == kDLFloat && index->dtype.bits == 32) { - get_centers(res, *index, centers); + get_centers(*index, centers); } else if (index->dtype.code == kDLFloat && index->dtype.bits == 16) { - get_centers(res, *index, centers); + get_centers(*index, centers); } else if (index->dtype.code == kDLInt && index->dtype.bits == 8) { - get_centers(res, *index, centers); + get_centers(*index, centers); } else if (index->dtype.code == kDLUInt && index->dtype.bits == 8) { - get_centers(res, *index, centers); + get_centers(*index, centers); } else { RAFT_FAIL("Unsupported index dtype: %d and bits: %d", index->dtype.code, index->dtype.bits); } diff --git a/cpp/src/neighbors/ivf_pq_c.cpp b/cpp/src/neighbors/ivf_pq_c.cpp index 40871f7ef8..4abd080582 100644 --- a/cpp/src/neighbors/ivf_pq_c.cpp +++ b/cpp/src/neighbors/ivf_pq_c.cpp @@ -1,5 +1,5 @@ /* - * Copyright (c) 2024, NVIDIA CORPORATION. + * Copyright (c) 2024-2025, NVIDIA CORPORATION. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -157,14 +157,25 @@ void _extend(cuvsResources_t res, } } -template -void _get_centers(cuvsResources_t res, cuvsIvfPqIndex index, DLManagedTensor* centers) +template +void _get_centers(cuvsIvfPqIndex index, DLManagedTensor* output) { - auto res_ptr = reinterpret_cast(res); auto index_ptr = reinterpret_cast*>(index.addr); - auto dst = cuvs::core::from_dlpack(centers); - cuvs::neighbors::ivf_pq::helpers::extract_centers(*res_ptr, *index_ptr, dst); + auto centers = index_ptr->centers(); + + // mimic the behaviour of `extract_centers` without performing a copy of the data + auto strided_centers = raft::make_device_strided_matrix_view( + centers.data_handle(), centers.extent(0), index_ptr->dim(), index_ptr->dim_ext()); + + cuvs::core::to_dlpack(strided_centers, output); +} + +template +void _get_pq_centers(cuvsIvfPqIndex index, DLManagedTensor* centers) +{ + auto index_ptr = reinterpret_cast*>(index.addr); + cuvs::core::to_dlpack(index_ptr->pq_centers(), centers); } } // namespace @@ -348,17 +359,13 @@ extern "C" uint32_t cuvsIvfPqIndexGetDim(cuvsIvfPqIndex_t index) return index_ptr->dim(); } -extern "C" cuvsError_t cuvsIvfPqIndexGetCenters(cuvsResources_t res, - cuvsIvfPqIndex_t index, - DLManagedTensor* centers) +extern "C" cuvsError_t cuvsIvfPqIndexGetCenters(cuvsIvfPqIndex_t index, DLManagedTensor* centers) { - return cuvs::core::translate_exceptions([=] { - if (cuvs::core::is_dlpack_device_compatible(centers->dl_tensor)) { - using output_mdspan_type = raft::device_matrix_view; - _get_centers(res, *index, centers); - } else { - using output_mdspan_type = raft::host_matrix_view; - _get_centers(res, *index, centers); - } - }); + return cuvs::core::translate_exceptions([=] { _get_centers(*index, centers); }); +} + +extern "C" cuvsError_t cuvsIvfPqIndexGetPqCenters(cuvsIvfPqIndex_t index, + DLManagedTensor* pq_centers) +{ + return cuvs::core::translate_exceptions([=] { _get_pq_centers(*index, pq_centers); }); } diff --git a/python/cuvs/cuvs/neighbors/ivf_flat/ivf_flat.pxd b/python/cuvs/cuvs/neighbors/ivf_flat/ivf_flat.pxd index b97faa1ad9..6e30ac0a07 100644 --- a/python/cuvs/cuvs/neighbors/ivf_flat/ivf_flat.pxd +++ b/python/cuvs/cuvs/neighbors/ivf_flat/ivf_flat.pxd @@ -1,5 +1,5 @@ # -# Copyright (c) 2024, NVIDIA CORPORATION. +# Copyright (c) 2024-2025, NVIDIA CORPORATION. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -66,8 +66,7 @@ cdef extern from "cuvs/neighbors/ivf_flat.h" nogil: uint32_t cuvsIvfFlatIndexGetDim(cuvsIvfFlatIndex_t index) - cuvsError_t cuvsIvfFlatIndexGetCenters(cuvsResources_t res, - cuvsIvfFlatIndex_t index, + cuvsError_t cuvsIvfFlatIndexGetCenters(cuvsIvfFlatIndex_t index, DLManagedTensor * centers) cuvsError_t cuvsIvfFlatBuild(cuvsResources_t res, diff --git a/python/cuvs/cuvs/neighbors/ivf_flat/ivf_flat.pyx b/python/cuvs/cuvs/neighbors/ivf_flat/ivf_flat.pyx index 1ad8167cdc..6fb4279db1 100644 --- a/python/cuvs/cuvs/neighbors/ivf_flat/ivf_flat.pyx +++ b/python/cuvs/cuvs/neighbors/ivf_flat/ivf_flat.pyx @@ -1,5 +1,5 @@ # -# Copyright (c) 2024, NVIDIA CORPORATION. +# Copyright (c) 2024-2025, NVIDIA CORPORATION. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -32,6 +32,7 @@ from pylibraft.common import auto_convert_output, cai_wrapper, device_ndarray from pylibraft.common.cai_wrapper import wrap_array from pylibraft.common.interruptible import cuda_interruptible +from cuvs.common.device_tensor_view import DeviceTensorView from cuvs.distance import DISTANCE_NAMES, DISTANCE_TYPES from cuvs.neighbors.common import _check_input_array from cuvs.neighbors.filters import no_filter @@ -196,19 +197,15 @@ cdef class Index: def centers(self): """ Get the cluster centers corresponding to the lists in the original space """ - return self._get_centers() - - @auto_sync_resources - def _get_centers(self, resources=None): if not self.trained: raise ValueError("Index needs to be built before getting centers") - cdef cuvsResources_t res = resources.get_c_obj() + output = DeviceTensorView() + cdef cydlpack.DLManagedTensor * tensor = \ + output.get_handle() + check_cuvs(cuvsIvfFlatIndexGetCenters(self.index, tensor)) + output.parent = self - output = np.empty((self.n_lists, self.dim), dtype='float32') - ai = wrap_array(output) - cdef cydlpack.DLManagedTensor* output_dlpack = cydlpack.dlpack_c(ai) - check_cuvs(cuvsIvfFlatIndexGetCenters(res, self.index, output_dlpack)) return output diff --git a/python/cuvs/cuvs/neighbors/ivf_pq/ivf_pq.pxd b/python/cuvs/cuvs/neighbors/ivf_pq/ivf_pq.pxd index 06932ac601..897e4951a9 100644 --- a/python/cuvs/cuvs/neighbors/ivf_pq/ivf_pq.pxd +++ b/python/cuvs/cuvs/neighbors/ivf_pq/ivf_pq.pxd @@ -1,5 +1,5 @@ # -# Copyright (c) 2024, NVIDIA CORPORATION. +# Copyright (c) 2024-2025, NVIDIA CORPORATION. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -88,10 +88,12 @@ cdef extern from "cuvs/neighbors/ivf_pq.h" nogil: uint32_t cuvsIvfPqIndexGetDim(cuvsIvfPqIndex_t index) - cuvsError_t cuvsIvfPqIndexGetCenters(cuvsResources_t res, - cuvsIvfPqIndex_t index, + cuvsError_t cuvsIvfPqIndexGetCenters(cuvsIvfPqIndex_t index, DLManagedTensor * centers) + cuvsError_t cuvsIvfPqIndexGetPqCenters(cuvsIvfPqIndex_t index, + DLManagedTensor * centers) + cuvsError_t cuvsIvfPqBuild(cuvsResources_t res, cuvsIvfPqIndexParams* params, DLManagedTensor* dataset, diff --git a/python/cuvs/cuvs/neighbors/ivf_pq/ivf_pq.pyx b/python/cuvs/cuvs/neighbors/ivf_pq/ivf_pq.pyx index a2208b5148..40f22c236b 100644 --- a/python/cuvs/cuvs/neighbors/ivf_pq/ivf_pq.pyx +++ b/python/cuvs/cuvs/neighbors/ivf_pq/ivf_pq.pyx @@ -1,5 +1,5 @@ # -# Copyright (c) 2024, NVIDIA CORPORATION. +# Copyright (c) 2024-2025, NVIDIA CORPORATION. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -32,6 +32,7 @@ from pylibraft.common import auto_convert_output, cai_wrapper, device_ndarray from pylibraft.common.cai_wrapper import wrap_array from pylibraft.common.interruptible import cuda_interruptible +from cuvs.common.device_tensor_view import DeviceTensorView from cuvs.distance import DISTANCE_NAMES, DISTANCE_TYPES from cuvs.neighbors.common import _check_input_array @@ -257,19 +258,28 @@ cdef class Index: def centers(self): """ Get the cluster centers corresponding to the lists in the original space """ - return self._get_centers() - - @auto_sync_resources - def _get_centers(self, resources=None): if not self.trained: raise ValueError("Index needs to be built before getting centers") - cdef cuvsResources_t res = resources.get_c_obj() + output = DeviceTensorView() + cdef cydlpack.DLManagedTensor * tensor = \ + output.get_handle() + check_cuvs(cuvsIvfPqIndexGetCenters(self.index, tensor)) + output.parent = self + return output - output = np.empty((self.n_lists, self.dim), dtype='float32') - ai = wrap_array(output) - cdef cydlpack.DLManagedTensor* output_dlpack = cydlpack.dlpack_c(ai) - check_cuvs(cuvsIvfPqIndexGetCenters(res, self.index, output_dlpack)) + @property + def pq_centers(self): + """ Get the PQ cluster centers """ + if not self.trained: + raise ValueError("Index needs to be built before getting" + " pq centers") + + output = DeviceTensorView() + cdef cydlpack.DLManagedTensor * tensor = \ + output.get_handle() + check_cuvs(cuvsIvfPqIndexGetPqCenters(self.index, tensor)) + output.parent = self return output diff --git a/python/cuvs/cuvs/tests/test_ivf_pq.py b/python/cuvs/cuvs/tests/test_ivf_pq.py index b6365c5485..a7367b9736 100644 --- a/python/cuvs/cuvs/tests/test_ivf_pq.py +++ b/python/cuvs/cuvs/tests/test_ivf_pq.py @@ -1,4 +1,4 @@ -# Copyright (c) 2024, NVIDIA CORPORATION. +# Copyright (c) 2024-2025, NVIDIA CORPORATION. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -112,6 +112,11 @@ def run_ivf_pq_build_search_test( centers = index.centers assert centers.shape[0] == n_lists + assert centers.shape[1] == index.dim + + pq_centers = index.pq_centers + assert len(pq_centers.shape) == 3 + assert pq_centers.shape[2] == 1 << pq_bits if not compare: return From d24564a0a8e0dd6fbd4bcec46b04e19dfc8f2564 Mon Sep 17 00:00:00 2001 From: stic <1974622+stic@users.noreply.github.com> Date: Fri, 5 Sep 2025 16:28:31 +0200 Subject: [PATCH 53/93] Add note about Linux-only support for cuVS pre-compiled packages (#1306) As I've dived in into cuVS from outside RAPIDS project I wasn't aware of target platform. Hence lost coffee break trying to troubleshoot (along hallucinating GPT insisting it's issue with latest CUDA toolkit). Added one liner to reference to RAPIDS docs on WSL2 support Authors: - https://github.com/stic Approvers: - Corey J. Nolet (https://github.com/cjnolet) URL: https://github.com/rapidsai/cuvs/pull/1306 --- docs/source/build.rst | 2 ++ 1 file changed, 2 insertions(+) diff --git a/docs/source/build.rst b/docs/source/build.rst index 26c361dfbc..aee49bbf20 100644 --- a/docs/source/build.rst +++ b/docs/source/build.rst @@ -31,6 +31,8 @@ The cuVS software development kit provides APIs for C, C++, Python, and Rust lan Installing Pre-compiled Packages -------------------------------- +**Note:** The cuVS pre-compiled packages are available for **Linux** only (x86_64 and aarch64 architectures). Native Windows support is not available at this time. On Windows, use **WSL2** with GPU passthrough. See the `RAPIDS WSL2 guide `_. + C, C++, and Python through Conda ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ From 2c0e124e3d38575f7f8d920679166dcf7258474c Mon Sep 17 00:00:00 2001 From: James Lamb Date: Fri, 5 Sep 2025 14:58:27 -0500 Subject: [PATCH 54/93] Use branch-25.10 again (#1319) Contributes to https://github.com/rapidsai/build-planning/issues/208 Now that https://github.com/rapidsai/shared-workflows/pull/413 is merged, this converts all GitHub Actions references from `@cuda13.0` back to `branch-25.10`. ## Notes for Reviewers This is safe to admin-merge because the change is a no-op... configs on those 2 branches are identical. --- .github/workflows/build.yaml | 22 ++++++------- .github/workflows/pr.yaml | 32 +++++++++---------- .github/workflows/test.yaml | 10 +++--- .../trigger-breaking-change-alert.yaml | 2 +- 4 files changed, 33 insertions(+), 33 deletions(-) diff --git a/.github/workflows/build.yaml b/.github/workflows/build.yaml index ccb75dbdb7..892764844b 100644 --- a/.github/workflows/build.yaml +++ b/.github/workflows/build.yaml @@ -34,7 +34,7 @@ concurrency: jobs: cpp-build: secrets: inherit - uses: rapidsai/shared-workflows/.github/workflows/conda-cpp-build.yaml@cuda13.0 + uses: rapidsai/shared-workflows/.github/workflows/conda-cpp-build.yaml@branch-25.10 with: build_type: ${{ inputs.build_type || 'branch' }} branch: ${{ inputs.branch }} @@ -44,7 +44,7 @@ jobs: rust-build: needs: cpp-build secrets: inherit - uses: rapidsai/shared-workflows/.github/workflows/custom-job.yaml@cuda13.0 + uses: rapidsai/shared-workflows/.github/workflows/custom-job.yaml@branch-25.10 # Artifacts are not published from these jobs, so it's safe to run for multiple CUDA versions. # If these jobs start producing artifacts, the names will have to differentiate between CUDA versions. strategy: @@ -65,7 +65,7 @@ jobs: go-build: needs: cpp-build secrets: inherit - uses: rapidsai/shared-workflows/.github/workflows/custom-job.yaml@cuda13.0 + uses: rapidsai/shared-workflows/.github/workflows/custom-job.yaml@branch-25.10 # Artifacts are not published from these jobs, so it's safe to run for multiple CUDA versions. # If these jobs start producing artifacts, the names will have to differentiate between CUDA versions. strategy: @@ -86,7 +86,7 @@ jobs: java-build: needs: cpp-build secrets: inherit - uses: rapidsai/shared-workflows/.github/workflows/custom-job.yaml@cuda13.0 + uses: rapidsai/shared-workflows/.github/workflows/custom-job.yaml@branch-25.10 # Artifacts are not published from these jobs, so it's safe to run for multiple CUDA versions. # If these jobs start producing artifacts, the names will have to differentiate between CUDA versions. strategy: @@ -108,7 +108,7 @@ jobs: python-build: needs: [cpp-build] secrets: inherit - uses: rapidsai/shared-workflows/.github/workflows/conda-python-build.yaml@cuda13.0 + uses: rapidsai/shared-workflows/.github/workflows/conda-python-build.yaml@branch-25.10 with: build_type: ${{ inputs.build_type || 'branch' }} branch: ${{ inputs.branch }} @@ -118,7 +118,7 @@ jobs: upload-conda: needs: [cpp-build, python-build] secrets: inherit - uses: rapidsai/shared-workflows/.github/workflows/conda-upload-packages.yaml@cuda13.0 + uses: rapidsai/shared-workflows/.github/workflows/conda-upload-packages.yaml@branch-25.10 with: build_type: ${{ inputs.build_type || 'branch' }} branch: ${{ inputs.branch }} @@ -133,7 +133,7 @@ jobs: if: github.ref_type == 'branch' needs: python-build secrets: inherit - uses: rapidsai/shared-workflows/.github/workflows/custom-job.yaml@cuda13.0 + uses: rapidsai/shared-workflows/.github/workflows/custom-job.yaml@branch-25.10 with: arch: "amd64" branch: ${{ inputs.branch }} @@ -145,7 +145,7 @@ jobs: sha: ${{ inputs.sha }} wheel-build-libcuvs: secrets: inherit - uses: rapidsai/shared-workflows/.github/workflows/wheels-build.yaml@cuda13.0 + uses: rapidsai/shared-workflows/.github/workflows/wheels-build.yaml@branch-25.10 with: build_type: ${{ inputs.build_type || 'branch' }} branch: ${{ inputs.branch }} @@ -159,7 +159,7 @@ jobs: wheel-publish-libcuvs: needs: wheel-build-libcuvs secrets: inherit - uses: rapidsai/shared-workflows/.github/workflows/wheels-publish.yaml@cuda13.0 + uses: rapidsai/shared-workflows/.github/workflows/wheels-publish.yaml@branch-25.10 with: build_type: ${{ inputs.build_type || 'branch' }} branch: ${{ inputs.branch }} @@ -170,7 +170,7 @@ jobs: wheel-build-cuvs: needs: wheel-build-libcuvs secrets: inherit - uses: rapidsai/shared-workflows/.github/workflows/wheels-build.yaml@cuda13.0 + uses: rapidsai/shared-workflows/.github/workflows/wheels-build.yaml@branch-25.10 with: build_type: ${{ inputs.build_type || 'branch' }} branch: ${{ inputs.branch }} @@ -182,7 +182,7 @@ jobs: wheel-publish-cuvs: needs: wheel-build-cuvs secrets: inherit - uses: rapidsai/shared-workflows/.github/workflows/wheels-publish.yaml@cuda13.0 + uses: rapidsai/shared-workflows/.github/workflows/wheels-publish.yaml@branch-25.10 with: build_type: ${{ inputs.build_type || 'branch' }} branch: ${{ inputs.branch }} diff --git a/.github/workflows/pr.yaml b/.github/workflows/pr.yaml index 2c7a61fd39..15b3f35efb 100644 --- a/.github/workflows/pr.yaml +++ b/.github/workflows/pr.yaml @@ -27,7 +27,7 @@ jobs: - devcontainer - telemetry-setup secrets: inherit - uses: rapidsai/shared-workflows/.github/workflows/pr-builder.yaml@cuda13.0 + uses: rapidsai/shared-workflows/.github/workflows/pr-builder.yaml@branch-25.10 if: always() with: needs: ${{ toJSON(needs) }} @@ -54,7 +54,7 @@ jobs: changed-files: needs: telemetry-setup secrets: inherit - uses: rapidsai/shared-workflows/.github/workflows/changed-files.yaml@cuda13.0 + uses: rapidsai/shared-workflows/.github/workflows/changed-files.yaml@branch-25.10 with: files_yaml: | test_cpp: @@ -106,14 +106,14 @@ jobs: checks: needs: telemetry-setup secrets: inherit - uses: rapidsai/shared-workflows/.github/workflows/checks.yaml@cuda13.0 + uses: rapidsai/shared-workflows/.github/workflows/checks.yaml@branch-25.10 with: enable_check_generated_files: false ignored_pr_jobs: "telemetry-summarize" conda-cpp-build: needs: checks secrets: inherit - uses: rapidsai/shared-workflows/.github/workflows/conda-cpp-build.yaml@cuda13.0 + uses: rapidsai/shared-workflows/.github/workflows/conda-cpp-build.yaml@branch-25.10 with: build_type: pull-request node_type: cpu16 @@ -121,7 +121,7 @@ jobs: conda-cpp-tests: needs: [conda-cpp-build, changed-files] secrets: inherit - uses: rapidsai/shared-workflows/.github/workflows/conda-cpp-tests.yaml@cuda13.0 + uses: rapidsai/shared-workflows/.github/workflows/conda-cpp-tests.yaml@branch-25.10 if: fromJSON(needs.changed-files.outputs.changed_file_groups).test_cpp with: build_type: pull-request @@ -129,21 +129,21 @@ jobs: conda-cpp-checks: needs: conda-cpp-build secrets: inherit - uses: rapidsai/shared-workflows/.github/workflows/conda-cpp-post-build-checks.yaml@cuda13.0 + uses: rapidsai/shared-workflows/.github/workflows/conda-cpp-post-build-checks.yaml@branch-25.10 with: build_type: pull-request symbol_exclusions: (void (thrust::|cub::)) conda-python-build: needs: conda-cpp-build secrets: inherit - uses: rapidsai/shared-workflows/.github/workflows/conda-python-build.yaml@cuda13.0 + uses: rapidsai/shared-workflows/.github/workflows/conda-python-build.yaml@branch-25.10 with: build_type: pull-request script: ci/build_python.sh conda-python-tests: needs: [conda-python-build, changed-files] secrets: inherit - uses: rapidsai/shared-workflows/.github/workflows/conda-python-tests.yaml@cuda13.0 + uses: rapidsai/shared-workflows/.github/workflows/conda-python-tests.yaml@branch-25.10 if: fromJSON(needs.changed-files.outputs.changed_file_groups).test_python with: build_type: pull-request @@ -151,7 +151,7 @@ jobs: conda-java-tests: needs: [conda-cpp-build, changed-files] secrets: inherit - uses: rapidsai/shared-workflows/.github/workflows/custom-job.yaml@cuda13.0 + uses: rapidsai/shared-workflows/.github/workflows/custom-job.yaml@branch-25.10 if: fromJSON(needs.changed-files.outputs.changed_file_groups).test_java # Artifacts are not published from these jobs, so it's safe to run for multiple CUDA versions. # If these jobs start producing artifacts, the names will have to differentiate between CUDA versions. @@ -172,7 +172,7 @@ jobs: docs-build: needs: conda-python-build secrets: inherit - uses: rapidsai/shared-workflows/.github/workflows/custom-job.yaml@cuda13.0 + uses: rapidsai/shared-workflows/.github/workflows/custom-job.yaml@branch-25.10 with: build_type: pull-request node_type: "gpu-l4-latest-1" @@ -182,7 +182,7 @@ jobs: rust-build: needs: conda-cpp-build secrets: inherit - uses: rapidsai/shared-workflows/.github/workflows/custom-job.yaml@cuda13.0 + uses: rapidsai/shared-workflows/.github/workflows/custom-job.yaml@branch-25.10 # Artifacts are not published from these jobs, so it's safe to run for multiple CUDA versions. # If these jobs start producing artifacts, the names will have to differentiate between CUDA versions. strategy: @@ -200,7 +200,7 @@ jobs: go-build: needs: conda-cpp-build secrets: inherit - uses: rapidsai/shared-workflows/.github/workflows/custom-job.yaml@cuda13.0 + uses: rapidsai/shared-workflows/.github/workflows/custom-job.yaml@branch-25.10 # Artifacts are not published from these jobs, so it's safe to run for multiple CUDA versions. # If these jobs start producing artifacts, the names will have to differentiate between CUDA versions. strategy: @@ -218,7 +218,7 @@ jobs: wheel-build-libcuvs: needs: checks secrets: inherit - uses: rapidsai/shared-workflows/.github/workflows/wheels-build.yaml@cuda13.0 + uses: rapidsai/shared-workflows/.github/workflows/wheels-build.yaml@branch-25.10 with: build_type: pull-request script: ci/build_wheel_libcuvs.sh @@ -229,7 +229,7 @@ jobs: wheel-build-cuvs: needs: wheel-build-libcuvs secrets: inherit - uses: rapidsai/shared-workflows/.github/workflows/wheels-build.yaml@cuda13.0 + uses: rapidsai/shared-workflows/.github/workflows/wheels-build.yaml@branch-25.10 with: build_type: pull-request script: ci/build_wheel_cuvs.sh @@ -238,7 +238,7 @@ jobs: wheel-tests-cuvs: needs: [wheel-build-cuvs, changed-files] secrets: inherit - uses: rapidsai/shared-workflows/.github/workflows/wheels-test.yaml@cuda13.0 + uses: rapidsai/shared-workflows/.github/workflows/wheels-test.yaml@branch-25.10 if: fromJSON(needs.changed-files.outputs.changed_file_groups).test_python with: build_type: pull-request @@ -246,7 +246,7 @@ jobs: devcontainer: secrets: inherit needs: telemetry-setup - uses: rapidsai/shared-workflows/.github/workflows/build-in-devcontainer.yaml@cuda13.0 + uses: rapidsai/shared-workflows/.github/workflows/build-in-devcontainer.yaml@branch-25.10 with: arch: '["amd64", "arm64"]' cuda: '["13.0"]' diff --git a/.github/workflows/test.yaml b/.github/workflows/test.yaml index 51418cbd22..c5892ca7ea 100644 --- a/.github/workflows/test.yaml +++ b/.github/workflows/test.yaml @@ -25,7 +25,7 @@ on: jobs: conda-cpp-checks: secrets: inherit - uses: rapidsai/shared-workflows/.github/workflows/conda-cpp-post-build-checks.yaml@cuda13.0 + uses: rapidsai/shared-workflows/.github/workflows/conda-cpp-post-build-checks.yaml@branch-25.10 with: build_type: ${{ inputs.build_type }} branch: ${{ inputs.branch }} @@ -34,7 +34,7 @@ jobs: symbol_exclusions: (void (thrust::|cub::)) conda-cpp-tests: secrets: inherit - uses: rapidsai/shared-workflows/.github/workflows/conda-cpp-tests.yaml@cuda13.0 + uses: rapidsai/shared-workflows/.github/workflows/conda-cpp-tests.yaml@branch-25.10 with: build_type: ${{ inputs.build_type }} branch: ${{ inputs.branch }} @@ -43,7 +43,7 @@ jobs: sha: ${{ inputs.sha }} conda-python-tests: secrets: inherit - uses: rapidsai/shared-workflows/.github/workflows/conda-python-tests.yaml@cuda13.0 + uses: rapidsai/shared-workflows/.github/workflows/conda-python-tests.yaml@branch-25.10 with: build_type: ${{ inputs.build_type }} branch: ${{ inputs.branch }} @@ -52,7 +52,7 @@ jobs: sha: ${{ inputs.sha }} conda-java-tests: secrets: inherit - uses: rapidsai/shared-workflows/.github/workflows/custom-job.yaml@cuda13.0 + uses: rapidsai/shared-workflows/.github/workflows/custom-job.yaml@branch-25.10 # Artifacts are not published from these jobs, so it's safe to run for multiple CUDA versions. # If these jobs start producing artifacts, the names will have to differentiate between CUDA versions. strategy: @@ -72,7 +72,7 @@ jobs: script: "ci/test_java.sh" wheel-tests-cuvs: secrets: inherit - uses: rapidsai/shared-workflows/.github/workflows/wheels-test.yaml@cuda13.0 + uses: rapidsai/shared-workflows/.github/workflows/wheels-test.yaml@branch-25.10 with: build_type: ${{ inputs.build_type }} branch: ${{ inputs.branch }} diff --git a/.github/workflows/trigger-breaking-change-alert.yaml b/.github/workflows/trigger-breaking-change-alert.yaml index 72751d071b..48bf37afc4 100644 --- a/.github/workflows/trigger-breaking-change-alert.yaml +++ b/.github/workflows/trigger-breaking-change-alert.yaml @@ -12,7 +12,7 @@ jobs: trigger-notifier: if: contains(github.event.pull_request.labels.*.name, 'breaking') secrets: inherit - uses: rapidsai/shared-workflows/.github/workflows/breaking-change-alert.yaml@cuda13.0 + uses: rapidsai/shared-workflows/.github/workflows/breaking-change-alert.yaml@branch-25.10 with: sender_login: ${{ github.event.sender.login }} sender_avatar: ${{ github.event.sender.avatar_url }} From 03d62f663d8f9dbed85dacbb353bed8cd3d38dc9 Mon Sep 17 00:00:00 2001 From: rmaschal Date: Tue, 9 Sep 2025 22:44:59 -0700 Subject: [PATCH 55/93] ScaNN: Overlapped gather for AVQ (#1286) Adds a class cluster_loader for AVQ that enables overlapping the gather operation and HtoD copy with GPU computation. There are two scenarios: 1. dataset on device: This is identical to the previous code, using raft::matrix::gather to perform the gather on device. 2. dataset on host: cluster_loader allocates to pinned buffers in host for fast (and possibly async) copies of cluster vectors DtoH. The actual gather operation is performed on cpu, into the pinned buffer. Copies can be overlapped with GPU work (namely AVQ update of the previous cluster) if scheduled on a separate stream. Authors: - https://github.com/rmaschal Approvers: - Tamas Bela Feher (https://github.com/tfeher) - Ben Karsin (https://github.com/bkarsin) URL: https://github.com/rapidsai/cuvs/pull/1286 --- cpp/src/neighbors/scann/detail/scann_avq.cuh | 253 ++++++++++++++++-- .../neighbors/scann/detail/scann_build.cuh | 8 +- 2 files changed, 235 insertions(+), 26 deletions(-) diff --git a/cpp/src/neighbors/scann/detail/scann_avq.cuh b/cpp/src/neighbors/scann/detail/scann_avq.cuh index cdfd4fbf28..1c42f6d2b9 100644 --- a/cpp/src/neighbors/scann/detail/scann_avq.cuh +++ b/cpp/src/neighbors/scann/detail/scann_avq.cuh @@ -68,7 +68,8 @@ __global__ void build_clusters( template void compute_cluster_offsets(raft::resources const& dev_resources, raft::device_vector_view clusters, - raft::device_vector_view cluster_sizes) + raft::device_vector_view cluster_sizes, + int64_t& max_cluster_size) { cudaStream_t stream = raft::resource::get_cuda_stream(dev_resources); rmm::device_async_resource_ref device_memory = @@ -103,10 +104,30 @@ void compute_cluster_offsets(raft::resources const& dev_resources, clusters.extent(0), stream); + int num_items = cluster_sizes.extent(0); + + // Compute max cluster size + // auto d_max_cluster_size = rmmo::make_device_scalar(dev_resources, 0); + rmm::device_scalar d_max_cluster_size(stream); + temp_storage_bytes = 0; + + cub::DeviceReduce::Max( + nullptr, temp_storage_bytes, cluster_sizes.data_handle(), d_max_cluster_size.data(), num_items); + + rmm::device_uvector temp_storage_max(temp_storage_bytes, stream, device_memory); + + cub::DeviceReduce::Max(temp_storage_max.data(), + temp_storage_bytes, + cluster_sizes.data_handle(), + d_max_cluster_size.data(), + num_items); + + max_cluster_size = d_max_cluster_size.value(stream); + // Scan to sum cluster sizes and get cluster start ptrs in flat array // Done in place - int num_items = cluster_sizes.extent(0); + temp_storage_bytes = 0; cub::DeviceScan::ExclusiveSum(nullptr, temp_storage_bytes, @@ -368,6 +389,190 @@ void rescale_avq_centroids(raft::resources const& dev_resources, }); } +/** + * A class for loading clusters into a compact matrix (sparse gather) + * for use in AVQ. + * + * There are two possible scenarios: + * 1. Dataset is stored in device memory: No host buffers are allocated, + * and the gather is performed on device + * 2. Dataset is stored in host memory: Two pinned buffers are allocated + * in host for fast DtoH copies of cluster ids, and fast HtoD copy of the + * cluster matrix, while amortizing the cost of allocating pinned memory. + * The gather is performed on cpu, overlapping with GPU compute. Copies are + * allocated on the provided stream, allowing for overlapping with + * other work on other streams. + */ +template +class cluster_loader { + private: + raft::pinned_matrix cluster_buf_; + raft::pinned_vector cluster_ids_buf_; + raft::device_matrix d_cluster_buf_; + raft::device_matrix d_cluster_copy_buf_; + const T* dataset_ptr_; + raft::host_vector_view h_cluster_offsets_; + raft::device_vector_view cluster_ids_; + cudaStream_t stream_; + int64_t dim_; + int64_t n_rows_; + bool needs_copy_; + + int64_t cur_idx_ = -1; + int64_t copy_idx_ = -1; + + size_t cluster_size(LabelT idx) + { + if (idx + 1 < h_cluster_offsets_.extent(0)) { + return h_cluster_offsets_(idx + 1) - h_cluster_offsets_(idx); + } + return n_rows_ - h_cluster_offsets_(idx); + } + + cluster_loader(raft::resources const& res, + const T* dataset_ptr, + int64_t dim, + int64_t n_rows, + int64_t max_cluster_size, + int64_t h_buf_size, + raft::host_vector_view h_cluster_offsets, + raft::device_vector_view cluster_ids, + bool needs_copy, + cudaStream_t stream) + : dim_(dim), + n_rows_(n_rows), + dataset_ptr_(dataset_ptr), + cluster_buf_(raft::make_pinned_matrix(res, h_buf_size, dim)), + cluster_ids_buf_(raft::make_pinned_vector(res, h_buf_size)), + d_cluster_buf_(raft::make_device_matrix(res, max_cluster_size, dim)), + d_cluster_copy_buf_(raft::make_device_matrix(res, max_cluster_size, dim)), + h_cluster_offsets_(h_cluster_offsets), + cluster_ids_(cluster_ids), + needs_copy_(needs_copy), + stream_(stream) + { + } + + public: + cluster_loader(raft::resources const& res, + raft::device_matrix_view dataset_view, + raft::host_vector_view h_cluster_offsets, + raft::device_vector_view cluster_ids, + int64_t max_cluster_size, + cudaStream_t stream) + : cluster_loader(res, + dataset_view.data_handle(), + dataset_view.extent(1), + dataset_view.extent(0), + max_cluster_size, + 0, + h_cluster_offsets, + cluster_ids, + false, + stream) + + { + } + + cluster_loader(raft::resources const& res, + raft::host_matrix_view dataset_view, + raft::host_vector_view h_cluster_offsets, + raft::device_vector_view cluster_ids, + int64_t max_cluster_size, + cudaStream_t stream) + : cluster_loader(res, + dataset_view.data_handle(), + dataset_view.extent(1), + dataset_view.extent(0), + max_cluster_size, + max_cluster_size, + h_cluster_offsets, + cluster_ids, + true, + stream) + + { + } + + /** + * @brief load and return a view of the provided cluster + * + * @param res: the raft resources; + * @param cluster_idx: the index of the cluster to be loaded + * @return device_matrix_view of the cluster vectors + */ + raft::device_matrix_view load_cluster(raft::resources const& res, LabelT cluster_idx) + { + size_t size = cluster_size(cluster_idx); + + // Check if cluster is already loaded + if (cur_idx_ != cluster_idx) { + // If not, load the cluster + if (copy_idx_ != cluster_idx) { prefetch_cluster(res, cluster_idx); } + + // swap buffers + std::swap(d_cluster_buf_, d_cluster_copy_buf_); + std::swap(cur_idx_, copy_idx_); + } + + return raft::make_device_matrix_view(d_cluster_buf_.data_handle(), size, dim_); + } + + /** @brief Perform gather operation on stream_ + * + * @param res: the raft resources + * @param cluster_idx: the index of the cluster + */ + void prefetch_cluster(raft::resources const& res, LabelT cluster_idx) + { + if (cluster_idx >= h_cluster_offsets_.extent(0)) { return; } + + size_t size = cluster_size(cluster_idx); + + auto cluster_ids = raft::make_device_vector_view( + cluster_ids_.data_handle() + h_cluster_offsets_(cluster_idx), size); + + auto cluster_vectors = + raft::make_device_matrix_view(d_cluster_copy_buf_.data_handle(), size, dim_); + + if (needs_copy_) { + // htod + auto h_cluster_ids = + raft::make_pinned_vector_view(cluster_ids_buf_.data_handle(), size); + + raft::copy( + h_cluster_ids.data_handle(), cluster_ids.data_handle(), cluster_ids.size(), stream_); + raft::resource::sync_stream(res, stream_); + + auto pinned_cluster = raft::make_pinned_matrix_view( + cluster_buf_.data_handle(), cluster_vectors.extent(0), cluster_vectors.extent(1)); + + int n_threads = std::min(omp_get_max_threads(), 32); +#pragma omp parallel for num_threads(n_threads) + for (int i = 0; i < h_cluster_ids.extent(0); i++) { + memcpy(pinned_cluster.data_handle() + i * pinned_cluster.extent(1), + dataset_ptr_ + h_cluster_ids(i) * dim_, + sizeof(T) * dim_); + } + + raft::copy(cluster_vectors.data_handle(), + pinned_cluster.data_handle(), + pinned_cluster.size(), + stream_); + raft::resource::sync_stream(res, stream_); + + } else { + // dtod + auto dataset_view = + raft::make_device_matrix_view(dataset_ptr_, n_rows_, dim_); + + raft::matrix::gather(res, dataset_view, cluster_ids, cluster_vectors); + } + + copy_idx_ = cluster_idx; + } +}; + /** * @brief Perform AVQ adjustment on cluster centers * @@ -396,26 +601,28 @@ void apply_avq(raft::resources const& res, raft::mdspan, raft::row_major, Accessor> dataset, raft::device_matrix_view centroids_view, raft::device_vector_view labels_view, - float eta) + float eta, + cudaStream_t copy_stream) { // Compute clusters - cudaStream_t stream = raft::resource::get_cuda_stream(res); - auto cluster_ptrs = raft::make_device_vector(res, centroids_view.extent(0)); - auto clusters = raft::make_device_vector(res, dataset.extent(0)); - - compute_cluster_offsets(res, labels_view, cluster_ptrs.view()); + cudaStream_t stream = raft::resource::get_cuda_stream(res); + 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; - auto h_cluster_ptrs = raft::make_host_vector(cluster_ptrs.extent(0)); + compute_cluster_offsets(res, labels_view, cluster_offsets.view(), max_cluster_size); + auto h_cluster_offsets = raft::make_host_vector(cluster_offsets.extent(0)); - raft::copy(h_cluster_ptrs.data_handle(), cluster_ptrs.data_handle(), cluster_ptrs.size(), stream); + raft::copy( + h_cluster_offsets.data_handle(), cluster_offsets.data_handle(), cluster_offsets.size(), stream); dim3 block(32, 1, 1); dim3 grid((dataset.extent(0) + block.x - 1) / block.x, 1, 1); build_clusters<<>>(labels_view.data_handle(), clusters.view().data_handle(), - cluster_ptrs.view().data_handle(), + cluster_offsets.view().data_handle(), dataset.extent(0), labels_view.extent(0)); RAFT_CUDA_TRY(cudaPeekAtLastError()); @@ -423,36 +630,34 @@ void apply_avq(raft::resources const& res, auto rescale_num = raft::make_device_vector(res, centroids_view.extent(0)); auto rescale_denom = raft::make_device_vector(res, centroids_view.extent(0)); + cluster_loader loader( + res, dataset, h_cluster_offsets.view(), clusters.view(), max_cluster_size, copy_stream); raft::resource::sync_stream(res); RAFT_LOG_DEBUG("Compute AVQ centroids\n"); - for (int i = 0; i < h_cluster_ptrs.extent(0); i++) { - int cluster_size = i + 1 < h_cluster_ptrs.extent(0) ? h_cluster_ptrs(i + 1) - h_cluster_ptrs(i) - : dataset.extent(0) - h_cluster_ptrs(i); + for (int i = 0; i < h_cluster_offsets.extent(0); i++) { + auto cluster_vectors = loader.load_cluster(res, i); - if (cluster_size == 0) { continue; } - auto cluster_ids = raft::make_device_vector_view( - clusters.data_handle() + h_cluster_ptrs(i), cluster_size); - auto cluster_vectors = - raft::make_device_matrix(res, cluster_size, dataset.extent(1)); auto avq_centroid = raft::make_device_vector_view( centroids_view.data_handle() + i * dataset.extent(1), dataset.extent(1)); auto rescale_num_view = raft::make_device_scalar_view(rescale_num.data_handle() + i); auto rescale_denom_view = raft::make_device_scalar_view(rescale_denom.data_handle() + i); - gather_functor{}( - res, dataset, cluster_ids, cluster_vectors.view(), raft::resource::get_cuda_stream(res)); - compute_avq_centroid( - res, cluster_vectors.view(), avq_centroid, rescale_num_view, rescale_denom_view, eta); + res, cluster_vectors, avq_centroid, rescale_num_view, rescale_denom_view, eta); + + loader.prefetch_cluster(res, i + 1); + + // make sure work is done before swapping buffers in cluster_loader + raft::resource::sync_stream(res); } rescale_avq_centroids(res, centroids_view, rescale_num.view(), rescale_denom.view(), - cluster_ptrs.view(), + cluster_offsets.view(), dataset.extent(0)); raft::resource::sync_stream(res); diff --git a/cpp/src/neighbors/scann/detail/scann_build.cuh b/cpp/src/neighbors/scann/detail/scann_build.cuh index da723d56e9..62ab8a317f 100644 --- a/cpp/src/neighbors/scann/detail/scann_build.cuh +++ b/cpp/src/neighbors/scann/detail/scann_build.cuh @@ -179,8 +179,12 @@ index build( } // AVQ update of KMeans centroids - apply_avq( - res, dataset, centroids_view, raft::make_const_mdspan(labels_view), params.partitioning_eta); + apply_avq(res, + dataset, + centroids_view, + raft::make_const_mdspan(labels_view), + params.partitioning_eta, + copy_stream); raft::device_vector_view soar_labels_view = idx.soar_labels(); From 21357feac3742c246af187ac30ae2b35178789b0 Mon Sep 17 00:00:00 2001 From: James Lamb Date: Wed, 10 Sep 2025 13:40:55 -0500 Subject: [PATCH 56/93] update dependencies: use cuda-toolkit wheels (#1326) Contributes to https://github.com/rapidsai/build-planning/issues/213 * switches uses of `nvidia-cuda-*-cu{12,13}` wheels to https://pypi.org/project/cuda-toolkit/ ## Notes for Reviewers ### Benefits of these changes Using the `cuda-toolkit` metapackage means we don't need to maintain our own pins for mutually compatible versions of CTK sub-components (which each follow their own different versioning conventions). Authors: - James Lamb (https://github.com/jameslamb) Approvers: - Kyle Edwards (https://github.com/KyleFromNVIDIA) - Mike Sarahan (https://github.com/msarahan) URL: https://github.com/rapidsai/cuvs/pull/1326 --- dependencies.yaml | 17 ++++------------- python/libcuvs/CMakeLists.txt | 19 ++++++++++++------- python/libcuvs/pyproject.toml | 5 +---- 3 files changed, 17 insertions(+), 24 deletions(-) diff --git a/dependencies.yaml b/dependencies.yaml index c370e90f88..adb923afab 100644 --- a/dependencies.yaml +++ b/dependencies.yaml @@ -329,29 +329,20 @@ dependencies: cuda: "12.*" use_cuda_wheels: "true" packages: - - nvidia-cublas-cu12 - - nvidia-curand-cu12 - - nvidia-cusolver-cu12 - - nvidia-cusparse-cu12 + - cuda-toolkit[cublas,curand,cusolver,cusparse]==12.* - matrix: cuda: "13.*" use_cuda_wheels: "true" packages: - - nvidia-cublas-cu13 - - nvidia-curand-cu13 - - nvidia-cusolver-cu13 - - nvidia-cusparse-cu13 + - cuda-toolkit[cublas,curand,cusolver,cusparse]==13.* - matrix: use_cuda_wheels: "false" packages: - # if no matching matrix selectors passed, list the unsuffixed packages + # if no matching matrix selectors passed, list a range # (just as a source of documentation, as this populates pyproject.toml in source control) - matrix: packages: - - nvidia-cublas - - nvidia-curand - - nvidia-cusolver - - nvidia-cusparse + - cuda-toolkit[cublas,curand,cusolver,cusparse]>=12,<14 depends_on_cupy: common: - output_types: conda diff --git a/python/libcuvs/CMakeLists.txt b/python/libcuvs/CMakeLists.txt index da43cb48a8..02574ac104 100644 --- a/python/libcuvs/CMakeLists.txt +++ b/python/libcuvs/CMakeLists.txt @@ -52,13 +52,18 @@ set(CUVS_COMPILE_DYNAMIC_ONLY ON) add_subdirectory(../../cpp cuvs-cpp) # assumes libcuvs.so is installed 2 levels deep, e.g. site-packages/libcuvs/lib64/libcuvs.so -set(rpaths - "$ORIGIN/../../nvidia/cublas/lib" - "$ORIGIN/../../nvidia/curand/lib" - "$ORIGIN/../../nvidia/cusolver/lib" - "$ORIGIN/../../nvidia/cusparse/lib" - "$ORIGIN/../../nvidia/nvjitlink/lib" -) +if(CMAKE_CUDA_COMPILER_VERSION VERSION_GREATER_EQUAL 13.0) + # starting with CTK 13 wheels, all libraries are grouped in this 'nvidia/cu13/lib' directory + set(rpaths "$ORIGIN/../../nvidia/cu13/lib") +else() + set(rpaths + "$ORIGIN/../../nvidia/cublas/lib" + "$ORIGIN/../../nvidia/curand/lib" + "$ORIGIN/../../nvidia/cusolver/lib" + "$ORIGIN/../../nvidia/cusparse/lib" + "$ORIGIN/../../nvidia/nvjitlink/lib" + ) +endif() if(USE_NCCL_RUNTIME_WHEEL) list(APPEND rpaths "$ORIGIN/../../nvidia/nccl/lib") diff --git a/python/libcuvs/pyproject.toml b/python/libcuvs/pyproject.toml index 51bb46316b..f19f2d33df 100644 --- a/python/libcuvs/pyproject.toml +++ b/python/libcuvs/pyproject.toml @@ -30,12 +30,9 @@ authors = [ license = { text = "Apache-2.0" } requires-python = ">=3.10" dependencies = [ + "cuda-toolkit[cublas,curand,cusolver,cusparse]>=12,<14", "libraft==25.10.*,>=0.0.0a0", "librmm==25.10.*,>=0.0.0a0", - "nvidia-cublas", - "nvidia-curand", - "nvidia-cusolver", - "nvidia-cusparse", ] # This list was generated by `rapids-dependency-file-generator`. To make changes, edit ../../dependencies.yaml and run `rapids-dependency-file-generator`. classifiers = [ "Intended Audience :: Developers", From 60244ca8201119ea6de0e947ed9bf4f5862f3712 Mon Sep 17 00:00:00 2001 From: MithunR Date: Wed, 10 Sep 2025 13:56:52 -0700 Subject: [PATCH 57/93] [REVIEW] [Java] Option to build fat-jars with native dependencies included (#1296) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit This commit introduces an _option_ to include the native libraries as part of a new Java JAR artifact. In addition, this commit also adds scripts to build the libraries included in the fat-jars using `gcc-toolset`, to allow the libraries to be portable across several Linux / `libstdc++` versions. (This was earlier attempted in #1264, but will now reside in this commit.) Note that for the initial cut, the "fat" jars will include only the following libraries: 1. `libcuvs.so` 2. `libcuvs_c.so` 3. `librmm.so` 4. `librapids_logger.so` The resultant JARs will still be dependent on `LD_LIBRARY_PATH` for other dependencies (`cublas`, `cusparse`, `cusolver`, `nccl`, etc.). Two new profiles have been introduced in the `pom.xml`: 1. `x86_64-cuda12` 5. `x86_64-cuda13` (Although this is more of an example than anything.) The main JAR artifact (`cuvs-java-25.x.x.jar`) remains unmodified. But when `-P x86_64-cuda12` is employed, an additional `cuvs-java-25.x.x-x86_64-cuda12.jar` is produced, containing `libcuvs.so`, `libcuvs_c.so`, and some (minimal) additional dependencies. The idea is that the JAR artifact build for `x86_64` + `cuda12` would look something like: ```bash # On an x86 build box, using the cuda-12 conda env, from the project root directory: # Build `libcuvs.so` first. LIBCUVS_BUILD_DIR=`pwd`/cpp/build/cuda12 ./build.sh libcuvs # Now the Java build. cd java/ CMAKE_PREFIX_PATH=`pwd`/../cpp/build/cuda12 ./build.sh ``` The `java/build.sh` detects the CPU platform and the CUDA version to automatically choose the profile (`x86-cuda12`). Note that there are tangential changes to the `pom.xml`: 1. Fixes #1293: The Java 22 portion of the build will now follow the Java 21, to prevent races between the builds. 2. The `maven-compiler-plugin` version has been dropped to `3.11.0`, to prevent build errors regarding a "0-byte module-info.class". This is apparently a known issue in `3.13`, to be fixed in `3.14`. Authors: - MithunR (https://github.com/mythrocks) Approvers: - Lorenzo Dematté (https://github.com/ldematte) - Mike Sarahan (https://github.com/msarahan) - Ben Frederickson (https://github.com/benfred) URL: https://github.com/rapidsai/cuvs/pull/1296 --- java/build.sh | 18 +- java/cuvs-java/pom.xml | 198 +++++++++++++++--- .../src/assembly/native-with-deps.xml | 52 +++++ .../com/nvidia/cuvs/spi/JDKProvider.java | 4 + .../spi/OptionalNativeDependencyLoader.java | 81 +++++++ java/docker-build/Dockerfile | 104 +++++++++ java/docker-build/README.md | 152 ++++++++++++++ java/docker-build/build-in-docker | 44 ++++ java/docker-build/run-in-docker | 161 ++++++++++++++ 9 files changed, 785 insertions(+), 29 deletions(-) create mode 100644 java/cuvs-java/src/assembly/native-with-deps.xml create mode 100644 java/cuvs-java/src/main/java22/com/nvidia/cuvs/spi/OptionalNativeDependencyLoader.java create mode 100644 java/docker-build/Dockerfile create mode 100644 java/docker-build/README.md create mode 100755 java/docker-build/build-in-docker create mode 100755 java/docker-build/run-in-docker diff --git a/java/build.sh b/java/build.sh index 7d2628d7e5..cbb1b9ed48 100755 --- a/java/build.sh +++ b/java/build.sh @@ -7,10 +7,18 @@ set -e -u -o pipefail ARGS="$*" NUMARGS=$# -CURDIR=$(cd "$(dirname "$0")"; pwd) VERSION="25.10.0" # Note: The version is updated automatically when ci/release/update-version.sh is invoked GROUP_ID="com.nvidia.cuvs" +# Identify CUDA major version. +CUDA_VERSION_FROM_NVCC=$(nvcc --version | grep -oP 'release [0-9]+' | awk '{print $2}') +CUDA_MAJOR_VERSION=${CUDA_VERSION_FROM_NVCC:-12} + +# Identify architecture. +ARCH=$(uname -m) + +BUILD_PROFILE="$ARCH-cuda$CUDA_MAJOR_VERSION" + if [ -z "${CMAKE_PREFIX_PATH:=}" ]; then CMAKE_PREFIX_PATH="$(pwd)/../cpp/build" export CMAKE_PREFIX_PATH @@ -33,12 +41,12 @@ fi # Build the java layer if [ -z ${LD_LIBRARY_PATH+x} ] -then export LD_LIBRARY_PATH=${CURDIR}/../cpp/build -else export LD_LIBRARY_PATH=${CURDIR}/../cpp/build:${LD_LIBRARY_PATH} +then export LD_LIBRARY_PATH=$CMAKE_PREFIX_PATH +else export LD_LIBRARY_PATH=$CMAKE_PREFIX_PATH:${LD_LIBRARY_PATH} fi -export LD_LIBRARY_PATH=${CURDIR}/../cpp/build:${LD_LIBRARY_PATH} cd cuvs-java -mvn verify "${MAVEN_VERIFY_ARGS[@]}" \ +mvn clean verify "${MAVEN_VERIFY_ARGS[@]}" -P "$BUILD_PROFILE" \ && mvn install:install-file -Dfile=./target/cuvs-java-$VERSION.jar -DgroupId=$GROUP_ID -DartifactId=cuvs-java -Dversion=$VERSION -Dpackaging=jar \ + && mvn install:install-file -Dfile=./target/cuvs-java-$VERSION-"$BUILD_PROFILE".jar -DgroupId=$GROUP_ID -DartifactId=cuvs-java -Dversion=$VERSION -Dclassifier="$BUILD_PROFILE" -Dpackaging=jar \ && cp pom.xml ./target/ diff --git a/java/cuvs-java/pom.xml b/java/cuvs-java/pom.xml index b3b862d291..ca27c67635 100644 --- a/java/cuvs-java/pom.xml +++ b/java/cuvs-java/pom.xml @@ -62,6 +62,7 @@ 22 UTF-8 UTF-8 + ${project.build.directory}/../../../cpp/build @@ -108,10 +109,11 @@ org.apache.maven.plugins maven-compiler-plugin - 3.13.0 + 3.11.0 default-compile + compile compile @@ -134,6 +136,9 @@ ${project.basedir}/src/main/java22 true + + module-info.java + @@ -155,29 +160,6 @@ - - org.apache.maven.plugins - maven-assembly-plugin - 3.4.2 - - - jar-with-dependencies - - merge - - add - - - - - assemble-all - package - - single - - - - org.apache.maven.plugins maven-jar-plugin @@ -277,4 +259,172 @@ + + + + x86_64-cuda12 + + + cuda.version + 12 + + + + x86_64-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 + com.nvidia.cuvs.examples.CagraExample + + + + + + assemble-native + package + + single + + + + + + + + + + x86_64-cuda13 + + + cuda.version + 13 + + + + x86_64-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 + com.nvidia.cuvs.examples.CagraExample + + + + + + assemble-native + package + + single + + + + + + + + diff --git a/java/cuvs-java/src/assembly/native-with-deps.xml b/java/cuvs-java/src/assembly/native-with-deps.xml new file mode 100644 index 0000000000..91e1ad85de --- /dev/null +++ b/java/cuvs-java/src/assembly/native-with-deps.xml @@ -0,0 +1,52 @@ + + + + + ${native.classifier} + + jar + + + false + + + + + / + true + true + runtime + + + + + + + ${project.build.directory}/native-libs + / + + **/* + + + + + diff --git a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/spi/JDKProvider.java b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/spi/JDKProvider.java index 0ee0cf1d10..cfbc8acefa 100644 --- a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/spi/JDKProvider.java +++ b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/spi/JDKProvider.java @@ -31,6 +31,10 @@ final class JDKProvider implements CuVSProvider { + static { + OptionalNativeDependencyLoader.loadLibraries(); + } + private static final MethodHandle createNativeDataset$mh = createNativeDatasetBuilder(); static MethodHandle createNativeDatasetBuilder() { diff --git a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/spi/OptionalNativeDependencyLoader.java b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/spi/OptionalNativeDependencyLoader.java new file mode 100644 index 0000000000..facbb670a3 --- /dev/null +++ b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/spi/OptionalNativeDependencyLoader.java @@ -0,0 +1,81 @@ +/* + * Copyright (c) 2025, NVIDIA CORPORATION. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.nvidia.cuvs.spi; + +import java.io.*; +import java.net.URL; +import java.util.stream.*; + +/** + * A class that loads native dependencies if they are available in the jar. + */ +public class OptionalNativeDependencyLoader { + + private static final ClassLoader loader = JDKProvider.class.getClassLoader(); + + private static boolean loaded = false; + + private static final String[] FILES_TO_LOAD = { + "rapids_logger", "rmm", "cuvs", "cuvs_c", + }; + + public static void loadLibraries() { + if (!loaded) { + String os = System.getProperty("os.name"); + String arch = System.getProperty("os.arch"); + + Stream.of(FILES_TO_LOAD) + .forEach( + file -> { + // Uncomment the following line to trace the loading of native dependencies. + // System.out.println("Loading native dependency: " + file); + try { + System.load(createFile(os, arch, file).getAbsolutePath()); + } catch (Throwable t) { + System.err.println( + "Continuing despite failure to load native dependency: " + + System.mapLibraryName(file) + + ".so: " + + t.getMessage()); + } + }); + + loaded = true; + } + } + + /** Extract the contents of a library resource into a temporary file */ + private static File createFile(String os, String arch, String baseName) throws IOException { + String path = arch + "/" + os + "/" + System.mapLibraryName(baseName); + File loc; + URL resource = loader.getResource(path); + if (resource == null) { + throw new FileNotFoundException("Could not locate native dependency " + path); + } + try (InputStream in = resource.openStream()) { + loc = File.createTempFile(baseName, ".so"); + loc.deleteOnExit(); + try (OutputStream out = new FileOutputStream(loc)) { + byte[] buffer = new byte[1024 * 16]; + int read = 0; + while ((read = in.read(buffer)) >= 0) { + out.write(buffer, 0, read); + } + } + } + return loc; + } +} diff --git a/java/docker-build/Dockerfile b/java/docker-build/Dockerfile new file mode 100644 index 0000000000..84e3d1f712 --- /dev/null +++ b/java/docker-build/Dockerfile @@ -0,0 +1,104 @@ +# +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +### +# Build the image for cuVS Java API development environment. +# +# Arguments: CUDA_VERSION=[12.X.Y], OS_RELEASE=[8, 9], TARGETPLATFORM=[linux/amd64, linux/arm64] +# +### +ARG CUDA_VERSION=12.9.1 +ARG OS_RELEASE=9 +ARG TARGETPLATFORM=linux/amd64 + +# multi-platform build with: docker buildx build --platform linux/arm64,linux/amd64 on either amd64 or arm64 host +# check available official arm-based docker images at https://hub.docker.com/r/nvidia/cuda/tags (OS/ARCH) +FROM --platform=$TARGETPLATFORM nvidia/cuda:$CUDA_VERSION-devel-rockylinux$OS_RELEASE + +# Re-declare ARGs after FROM since they don't persist across FROM boundary +ARG CUDA_VERSION=12.9.1 +ARG OS_RELEASE=9 +ARG TARGETPLATFORM=linux/amd64 + +# If DEV_BUILD is ON, the gcc-toolset will be enabled by default for bash shell +ARG DEV_BUILD=OFF + +# Dependency versions +# Act as default GCC toolset in the image +ARG TOOLSET_VERSION=14 +ARG CMAKE_VERSION=3.30.4 +ARG CCACHE_VERSION=4.11.2 + +# Default x86_64 from x86 build, aarch64 cmake for arm build +ARG CMAKE_ARCH=x86_64 + +### Install basic requirements +RUN if [ "$OS_RELEASE" = "9" ]; then \ + dnf --enablerepo=crb install -y scl-utils gcc-toolset-${TOOLSET_VERSION} python39 \ + zlib-devel maven tar wget patch ninja-build git zip; \ + elif [ "$OS_RELEASE" = "8" ]; then \ + dnf --enablerepo=devel install -y scl-utils gcc-toolset-${TOOLSET_VERSION} python39 \ + zlib-devel maven tar wget patch ninja-build git zip; \ + else \ + echo "Unsupported OS version: $OS_RELEASE"; \ + fi + +# Enable the gcc-toolset by default for bash shell if DEV_BUILD is ON +RUN if [ "$DEV_BUILD" = "ON" ]; then \ + echo "source scl_source enable gcc-toolset-${TOOLSET_VERSION}" >> /etc/bashrc; \ + fi + +# Execute every time a new non-interactive bash shell is started +ENV BASH_ENV=/etc/bashrc + +## pre-create the CMAKE_INSTALL_PREFIX folder, set writable by any user for Jenkins +RUN mkdir -m 777 /usr/local/rapids /rapids + +# Fetch and install CMake. +RUN cd /usr/local && wget --quiet https://github.com/Kitware/CMake/releases/download/v${CMAKE_VERSION}/cmake-${CMAKE_VERSION}-linux-${CMAKE_ARCH}.tar.gz && \ + tar zxf cmake-${CMAKE_VERSION}-linux-${CMAKE_ARCH}.tar.gz && \ + rm cmake-${CMAKE_VERSION}-linux-${CMAKE_ARCH}.tar.gz + +# Make version-less alias for external reference such as when cmake is called by an IDE outside of the container +RUN ln -s /usr/local/cmake-${CMAKE_VERSION}-linux-${CMAKE_ARCH}/bin/cmake /usr/local/bin/cmake + +# ccache for interactive builds +RUN cd /tmp && wget --quiet https://github.com/ccache/ccache/releases/download/v${CCACHE_VERSION}/ccache-${CCACHE_VERSION}.tar.gz && \ + tar zxf ccache-${CCACHE_VERSION}.tar.gz && \ + rm ccache-${CCACHE_VERSION}.tar.gz && \ + cd ccache-${CCACHE_VERSION} && \ + mkdir build && \ + cd build && \ + scl enable gcc-toolset-${TOOLSET_VERSION} \ + "cmake .. \ + -DCMAKE_BUILD_TYPE=Release \ + -DZSTD_FROM_INTERNET=ON \ + -DREDIS_STORAGE_BACKEND=OFF && \ + cmake --build . --parallel ${PARALLEL_LEVEL} --target install" && \ + cd ../.. && \ + rm -rf ccache-${CCACHE_VERSION} + +# Fetch and install JDK 22. +RUN cd /usr/local && wget --quiet https://download.oracle.com/java/22/archive/jdk-22.0.2_linux-x64_bin.tar.gz && \ + tar zxf jdk-22.0.2_linux-x64_bin.tar.gz && \ + rm jdk-22.0.2_linux-x64_bin.tar.gz + +# Make Java22 available in path. +RUN ln -sf /usr/local/jdk-22.0.2/bin/java /usr/local/bin/java +ENV JAVA_HOME="/usr/local/jdk-22.0.2" + +# disable cuda container constraints to allow running w/ older drivers on datacenter GPUs +ENV NVIDIA_DISABLE_REQUIRE="true" diff --git a/java/docker-build/README.md b/java/docker-build/README.md new file mode 100644 index 0000000000..92cb616bf1 --- /dev/null +++ b/java/docker-build/README.md @@ -0,0 +1,152 @@ +# Docker Build Environment for cuVS Java API + +This directory contains the Docker-based build system for the cuVS Java API, providing a containerized environment with all necessary dependencies for building the project across different CUDA versions and architectures. + +## Overview + +The Docker build system provides: +- Consistent build environment across different host systems +- Support for multiple CUDA versions (12.x and 13.x) +- Cross-platform builds (x86_64 and aarch64) +- Pre-configured development tools (GCC toolset, CMake, Maven, JDK 22) + +## Quick Start + +### Default Build (CUDA 12, Rocky Linux 9, All GPU Architectures) + +```bash +./build-in-docker +``` + +This builds using the defaults: +- CUDA version: 12.9.1 +- OS version: Rocky Linux 9 +- Target: libcuvs java with all GPU architectures (`--allgpuarch`) + +### Building for CUDA 13, All GPU Architectures + +```bash +CUDA_VERSION=13.0.0 ./build-in-docker +``` + +### Building for Local GPU Architecture Only + +```bash +./build-in-docker libcuvs java +``` + +## Environment Variables + +### Core Configuration + +- **`CUDA_VERSION`**: CUDA toolkit version to use (default: `12.9.1`) + - Examples: `12.9.1`, `13.0.0`, `13.1.0` +- **`OS_VERSION`**: Rocky Linux OS version to use (default: `9`) + - Supported values: `8`, `9` +- **`CMAKE_GENERATOR`**: CMake generator to use (default: `Ninja`) +- **`LOCAL_MAVEN_REPO`**: Local Maven repository path (default: `$HOME/.m2/repository`) + +### Docker Configuration + +- **`JNI_DOCKER_DEV_BUILD`**: Set to `ON` for development builds with gcc-toolset enabled by default (default: `OFF`) +- **`DOCKER_CMD`**: Docker command to use (default: `docker`) +- **`DOCKER_BUILD_EXTRA_ARGS`**: Additional arguments for `docker build` +- **`DOCKER_RUN_EXTRA_ARGS`**: Additional arguments for `docker run` +- **`DOCKER_GPU_OPTS`**: GPU options for Docker (default: `--gpus all`) + +### Build Optimization + +- **`LOCAL_CCACHE_DIR`**: ccache directory for build acceleration (default: `$HOME/.ccache`) +- **`PARALLEL_LEVEL`**: Number of parallel build jobs +- **`VERBOSE`**: Enable verbose build output + +## Architecture Support + +The build system automatically detects the host architecture: +- **x86_64**: Uses `linux/amd64` platform and x86_64 CMake binaries +- **aarch64**: Uses `linux/arm64` platform and aarch64 CMake binaries + +## Files in This Directory + +- **`build-in-docker`**: Main entry point script for Docker-based builds +- **`run-in-docker`**: Lower-level script that handles Docker container execution +- **`Dockerfile`**: Multi-stage Docker image definition with CUDA, development tools, and dependencies +- **`env.sh`**: Environment configuration script + +## Examples + +### Development Build with Custom CUDA Version (Default: All GPU Architectures) + +```bash +CUDA_VERSION=13.1.0 JNI_DOCKER_DEV_BUILD=ON ./build-in-docker +``` + +### Production Build with Rocky Linux 8 + +```bash +OS_VERSION=8 ./build-in-docker +``` + +### Build with Custom Maven Repository and ccache + +```bash +LOCAL_MAVEN_REPO=/custom/maven/repo LOCAL_CCACHE_DIR=/custom/ccache CUDA_VERSION=13.0.0 ./build-in-docker +``` + +### Build with Additional Docker Arguments + +```bash +DOCKER_BUILD_EXTRA_ARGS="--no-cache" ./build-in-docker +``` + +### Interactive Development Session + +```bash +CUDA_VERSION=12.9.1 JNI_DOCKER_DEV_BUILD=ON ./docker-build/run-in-docker +``` + +This starts an interactive bash shell in the container for development work. + +## Docker Image Details + +The Docker image is based on `nvidia/cuda:{CUDA_VERSION}-devel-rockylinux9` and includes: + +- **CUDA Development Tools**: Complete CUDA toolkit for the specified version +- **GCC Toolset 14**: Modern C++ compiler with C++20 support +- **CMake 3.30.4**: Build system generator +- **JDK 22**: Java Development Kit with Panama FFM support +- **Maven**: Java project management and build tool +- **ccache**: Compiler cache for faster incremental builds +- **Additional Tools**: git, ninja-build, wget, tar, zip, patch + +## Troubleshooting + +### CUDA Version Compatibility + +Ensure your target CUDA version is supported by checking available tags at: +- [NVIDIA CUDA Docker Hub](https://hub.docker.com/r/nvidia/cuda/tags) + +### GPU Access Issues + +If you encounter GPU access problems, verify: +1. NVIDIA Docker runtime is installed (`nvidia-docker2` or `nvidia-container-toolkit`) +2. Your user has permissions to access Docker +3. GPU drivers are compatible with the CUDA version + +### Build Performance + +For tweaking build performance: +1. Use ccache: Ensure `LOCAL_CCACHE_DIR` is set to a persistent directory. (By default, the build uses ccache.) +2. Change parallelism: Set `PARALLEL_LEVEL` to match your CPU cores. (This happens by default.) +3. Use development builds: Set `JNI_DOCKER_DEV_BUILD=ON` for development + +### Memory Requirements + +The Docker build process requires significant memory. For large projects, ensure: +- At least 8GB of available RAM +- Sufficient disk space for Docker images and build artifacts + +## Related Documentation + +- [cuVS Java API README](../README.md): Main Java API documentation +- [cuVS Build Instructions](https://docs.rapids.ai/api/cuvs/stable/build/): Native build documentation diff --git a/java/docker-build/build-in-docker b/java/docker-build/build-in-docker new file mode 100755 index 0000000000..1c61429fa8 --- /dev/null +++ b/java/docker-build/build-in-docker @@ -0,0 +1,44 @@ +#!/bin/bash + +# +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +# Build the cuVS Java API artifact in a Docker container with devtoolset + +set -e + +# Base paths relative to this script's location +SCRIPTDIR=$(cd "$(dirname "$0")"; pwd) + +LOCAL_MAVEN_REPO=${LOCAL_MAVEN_REPO:-"$HOME/.m2/repository"} +export CMAKE_GENERATOR=${CMAKE_GENERATOR:-"Ninja"} +export CUDA_VERSION=${CUDA_VERSION:-12.9.1} +# CUDA_CLASSIFIER=cuda${CUDA_VERSION%%.*} # Currently unused, but may be needed for future CUDA version handling + +# Set env for arm64 build, The possible values of 'uname -m' : [x86_64/i386/aarch64/mips/...] +if [ "$(uname -m)" == "aarch64" ]; then + # ARM-specific settings. + echo "Set ARM Specific settings" +fi + +# TODO: Settings to build for different CUDA versions, different GPU archs, CPU archs, etc. + +# Default to building libcuvs java with all GPU architectures if no arguments provided +if [ $# -eq 0 ]; then + "$SCRIPTDIR/run-in-docker" "./build.sh" "libcuvs" "java" "--allgpuarch" +else + "$SCRIPTDIR/run-in-docker" "./build.sh" "$@" +fi diff --git a/java/docker-build/run-in-docker b/java/docker-build/run-in-docker new file mode 100755 index 0000000000..9c102b64a3 --- /dev/null +++ b/java/docker-build/run-in-docker @@ -0,0 +1,161 @@ +#!/bin/bash + +# +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +# Run a command in a Docker container with devtoolset + +set -e + +# Convert sclCMD to array for proper argument handling +if [[ -n "${sclCMD:-}" ]]; then + # Convert existing string to array by splitting on whitespace + read -ra sclCMD <<< "$sclCMD" +else + sclCMD=(scl enable gcc-toolset-14) +fi + +# Set this environment variable to ON to build the docker image for local development, +# with gcc-toolset enabled by default. +# This is necessary especially if the docker container is used for local development with an IDE. +JNI_DOCKER_DEV_BUILD=${JNI_DOCKER_DEV_BUILD:-OFF} + +REPODIR_REL=$(git rev-parse --show-toplevel) +REPODIR=$(realpath "$REPODIR_REL") +GIT_COMMON_DIR_REL=$(git rev-parse --git-common-dir) +GIT_COMMON_DIR=$(realpath "$GIT_COMMON_DIR_REL") +WORKDIR=${WORKDIR:-$REPODIR} +TZ=${TZ:-UTC} + +CUDA_VERSION=${CUDA_VERSION:-12.9.1} +CUDA_CLASSIFIER=cuda${CUDA_VERSION%%.*} +OS_VERSION=${OS_VERSION:-9} # Default to 9 + +DOCKER_CMD=${DOCKER_CMD:-docker} + +# Initialize DOCKER_BUILD_EXTRA_ARGS as an array +if [[ -n "${DOCKER_BUILD_EXTRA_ARGS:-}" ]]; then + # Convert existing string to array by splitting on whitespace + read -ra DOCKER_BUILD_EXTRA_ARGS <<< "$DOCKER_BUILD_EXTRA_ARGS" +else + DOCKER_BUILD_EXTRA_ARGS=() +fi + +if [ "$(uname -m)" == "aarch64" ]; then + DOCKER_BUILD_EXTRA_ARGS=(--build-arg TARGETPLATFORM=linux/arm64 --build-arg CMAKE_ARCH=aarch64 "${DOCKER_BUILD_EXTRA_ARGS[@]}") +else + DOCKER_BUILD_EXTRA_ARGS=(--build-arg TARGETPLATFORM=linux/amd64 --build-arg CMAKE_ARCH=x86_64 "${DOCKER_BUILD_EXTRA_ARGS[@]}") +fi +# Initialize DOCKER_RUN_EXTRA_ARGS as an array +if [[ -n "${DOCKER_RUN_EXTRA_ARGS:-}" ]]; then + # Convert existing string to array by splitting on whitespace + read -ra DOCKER_RUN_EXTRA_ARGS <<< "$DOCKER_RUN_EXTRA_ARGS" +else + DOCKER_RUN_EXTRA_ARGS=() +fi +LOCAL_CCACHE_DIR=${LOCAL_CCACHE_DIR:-"$HOME/.ccache"} +LOCAL_MAVEN_REPO=${LOCAL_MAVEN_REPO:-"$HOME/.m2/repository"} + +if [ "$JNI_DOCKER_DEV_BUILD" == "ON" ]; then + echo "Building docker image for local development, gcc-toolset is enabled by default..." + JNI_DOCKER_IMAGE="cuvs-jni-build:${CUDA_VERSION}-devel-rockylinux${OS_VERSION}" +else + echo "Building docker image for production, gcc-toolset is NOT enabled by default..." + JNI_DOCKER_IMAGE="cuvs-jni-build:${CUDA_VERSION}-rockylinux${OS_VERSION}" +fi + +# ensure directories exist +mkdir -p "$LOCAL_CCACHE_DIR" "$LOCAL_MAVEN_REPO" + +$DOCKER_CMD build "${DOCKER_BUILD_EXTRA_ARGS[@]}" -f "$REPODIR/java/docker-build/Dockerfile" \ + --build-arg CUDA_VERSION="$CUDA_VERSION" \ + --build-arg OS_RELEASE="$OS_VERSION" \ + --build-arg DEV_BUILD="$JNI_DOCKER_DEV_BUILD" \ + -t "$JNI_DOCKER_IMAGE" \ + "$REPODIR/java/docker-build" + +if [[ "$DOCKER_CMD" == "docker" ]]; then + # Initialize DOCKER_GPU_OPTS as an array + if [[ -n "${DOCKER_GPU_OPTS:-}" ]]; then + # Convert existing string to array by splitting on whitespace + read -ra DOCKER_GPU_OPTS <<< "$DOCKER_GPU_OPTS" + else + DOCKER_GPU_OPTS=(--gpus all) + fi +else + DOCKER_GPU_OPTS=() +fi + +# Initialize DOCKER_OPTS as an array +if [[ -n "${DOCKER_OPTS:-}" ]]; then + # Convert existing string to array by splitting on whitespace + read -ra DOCKER_OPTS <<< "$DOCKER_OPTS" +else + DOCKER_OPTS=() +fi + +if (( $# == 0 )); then + # no arguments gets an interactive shell + DOCKER_OPTS+=(-it) + RUN_CMD=("/bin/bash") +else + RUN_CMD=("$@") +fi + +MNT_ARGS=() + +RO_SRC=( + "/etc/group" + "/etc/passwd" + "/etc/shadow" + "/etc/sudoers.d" +) +RO_DST=("${RO_SRC[@]}") +if [[ "$HOST_CUDA_PATH" != "" ]]; then + RO_SRC+=("$HOST_CUDA_PATH") + RO_DST+=("/usr/local/cuda") +fi +for (( i=0; i<${#RO_SRC[@]}; i++)); do + MNT_ARGS+=(--mount "type=bind,src=${RO_SRC[$i]},dst=${RO_DST[$i]},ro") +done + +RW_SRC=( + "$GIT_COMMON_DIR" + "$WORKDIR" + "$LOCAL_CCACHE_DIR" + "$LOCAL_MAVEN_REPO" +) +for (( i=0; i<${#RW_SRC[@]}; i++)); do + MNT_ARGS+=(--mount "type=bind,src=${RW_SRC[$i]},dst=${RW_SRC[$i]}") +done + +$DOCKER_CMD run "${DOCKER_GPU_OPTS[@]}" "${DOCKER_RUN_EXTRA_ARGS[@]}" -u "$(id -u):$(id -g)" --rm \ + "${MNT_ARGS[@]}" \ + --workdir "$WORKDIR" \ + -e CCACHE_DIR="$LOCAL_CCACHE_DIR" \ + -e CMAKE_C_COMPILER_LAUNCHER="ccache" \ + -e CMAKE_CXX_COMPILER_LAUNCHER="ccache" \ + -e CMAKE_CUDA_COMPILER_LAUNCHER="ccache" \ + -e CMAKE_CXX_LINKER_LAUNCHER="ccache" \ + -e CMAKE_PREFIX_PATH="${WORKDIR}/cpp/build/$CUDA_CLASSIFIER" \ + -e LIBCUVS_BUILD_DIR="${WORKDIR}/cpp/build/$CUDA_CLASSIFIER" \ + -e CMAKE_GENERATOR \ + -e CUDA_VISIBLE_DEVICES \ + -e PARALLEL_LEVEL \ + -e VERBOSE \ + "${DOCKER_OPTS[@]}" \ + "$JNI_DOCKER_IMAGE" \ + "${sclCMD[@]}" -- "${RUN_CMD[@]}" From 792730c89e0c955442faab4d2755e7c0b45e3589 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Lorenzo=20Dematt=C3=A9?= Date: Thu, 11 Sep 2025 16:46:03 +0200 Subject: [PATCH 58/93] [Java] Add reason(s) to UnsupportedProvider/UnsupportedOperationExceptions (#1314) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit This PR adds some precise and readable messages specifying why a CuVSProvider is unsupported (it throws UnsupportedOperationException), to help with diagnosing cuvs-java startup issues. Authors: - Lorenzo Dematté (https://github.com/ldematte) - MithunR (https://github.com/mythrocks) Approvers: - Chris Hegarty (https://github.com/ChrisHegarty) - MithunR (https://github.com/mythrocks) URL: https://github.com/rapidsai/cuvs/pull/1314 --- .../nvidia/cuvs/spi/CuVSServiceProvider.java | 27 +++++++++------ .../nvidia/cuvs/spi/UnsupportedProvider.java | 34 +++++++++++-------- 2 files changed, 37 insertions(+), 24 deletions(-) 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 c8c886dc13..a4a5843adf 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 @@ -17,6 +17,8 @@ import java.lang.invoke.MethodHandles; import java.lang.invoke.MethodType; +import java.util.ArrayList; +import java.util.List; import java.util.ServiceLoader; /** @@ -43,7 +45,10 @@ private static CuVSProvider loadProvider() { } static CuVSProvider builtinProvider() { - if (Runtime.version().feature() > 21 && isLinuxAmd64()) { + var supportedJavaRuntime = Runtime.version().feature() > 21; + var supportedOs = System.getProperty("os.name").startsWith("Linux"); + var supportedArchitecture = System.getProperty("os.arch").equals("amd64"); + if (supportedJavaRuntime && supportedOs && supportedArchitecture) { try { var cls = Class.forName("com.nvidia.cuvs.spi.JDKProvider"); var ctr = MethodHandles.lookup().findConstructor(cls, MethodType.methodType(void.class)); @@ -52,16 +57,18 @@ static CuVSProvider builtinProvider() { throw new AssertionError(e); } } - return new UnsupportedProvider(); - } + List unsupportedReasons = new ArrayList<>(); + if (!supportedJavaRuntime) { + unsupportedReasons.add("cuvs-java requires Java Runtime version 22 or greater"); + } + if (!supportedOs) { + unsupportedReasons.add("cuvs-java supports only Linux"); + } + if (!supportedArchitecture) { + unsupportedReasons.add("cuvs-java supports only x86"); + } - /** - * Returns true iff the architecture is x64 (amd64) and the OS Linux - * (the * OS we currently support for the native lib). - */ - static boolean isLinuxAmd64() { - String name = System.getProperty("os.name"); - return (name.startsWith("Linux")) && System.getProperty("os.arch").equals("amd64"); + return new UnsupportedProvider(String.join("; ", unsupportedReasons)); } } } diff --git a/java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/UnsupportedProvider.java b/java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/UnsupportedProvider.java index 60509004c4..0f1edda99c 100644 --- a/java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/UnsupportedProvider.java +++ b/java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/UnsupportedProvider.java @@ -24,51 +24,57 @@ */ final class UnsupportedProvider implements CuVSProvider { + private final String reasons; + + public UnsupportedProvider(String reasons) { + this.reasons = reasons; + } + @Override public CuVSResources newCuVSResources(Path tempDirectory) { - throw new UnsupportedOperationException(); + throw new UnsupportedOperationException(reasons); } @Override public BruteForceIndex.Builder newBruteForceIndexBuilder(CuVSResources cuVSResources) { - throw new UnsupportedOperationException(); + throw new UnsupportedOperationException(reasons); } @Override public CagraIndex.Builder newCagraIndexBuilder(CuVSResources cuVSResources) { - throw new UnsupportedOperationException(); + throw new UnsupportedOperationException(reasons); } @Override public HnswIndex.Builder newHnswIndexBuilder(CuVSResources cuVSResources) { - throw new UnsupportedOperationException(); + throw new UnsupportedOperationException(reasons); } @Override public TieredIndex.Builder newTieredIndexBuilder(CuVSResources cuVSResources) { - throw new UnsupportedOperationException(); + throw new UnsupportedOperationException(reasons); } @Override public CagraIndex mergeCagraIndexes(CagraIndex[] indexes) { - throw new UnsupportedOperationException(); + throw new UnsupportedOperationException(reasons); } @Override public CuVSMatrix.Builder newHostMatrixBuilder( long size, long dimensions, CuVSMatrix.DataType dataType) { - throw new UnsupportedOperationException(); + throw new UnsupportedOperationException(reasons); } @Override public CuVSMatrix.Builder newDeviceMatrixBuilder( CuVSResources cuVSResources, long size, long dimensions, CuVSMatrix.DataType dataType) { - throw new UnsupportedOperationException(); + throw new UnsupportedOperationException(reasons); } @Override public GPUInfoProvider gpuInfoProvider() { - throw new UnsupportedOperationException(); + throw new UnsupportedOperationException(reasons); } @Override @@ -79,26 +85,26 @@ public CuVSMatrix.Builder newDeviceMatrixBuilder( int rowStride, int columnStride, CuVSMatrix.DataType dataType) { - throw new UnsupportedOperationException(); + throw new UnsupportedOperationException(reasons); } @Override public MethodHandle newNativeMatrixBuilder() { - throw new UnsupportedOperationException(); + throw new UnsupportedOperationException(reasons); } @Override public CuVSMatrix newMatrixFromArray(float[][] vectors) { - throw new UnsupportedOperationException(); + throw new UnsupportedOperationException(reasons); } @Override public CuVSMatrix newMatrixFromArray(int[][] vectors) { - throw new UnsupportedOperationException(); + throw new UnsupportedOperationException(reasons); } @Override public CuVSMatrix newMatrixFromArray(byte[][] vectors) { - throw new UnsupportedOperationException(); + throw new UnsupportedOperationException(reasons); } } From 607e9ebc4b32ae703bf7fb7d62609a27c5fc577d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Lorenzo=20Dematt=C3=A9?= Date: Fri, 12 Sep 2025 00:27:34 +0200 Subject: [PATCH 59/93] [Java] Make `cudaGetDeviceProperties` compatible with CUDA 12 and 13 based on symbol presence (#1323) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit In #1273 we addressed a signature change in CUDA 13 by binding different symbols based on a environment variable, `RAPIDS_CUDA_MAJOR`. That works, but forces users of cuvs-java with CUDA 12 to define this environment variable. This PR improves on it by making the symbol lookup dynamic, looking for the CUDA 13 symbol name, and falling back to the CUDA 12 exported name if we fail to locate the first. Authors: - Lorenzo Dematté (https://github.com/ldematte) Approvers: - MithunR (https://github.com/mythrocks) URL: https://github.com/rapidsai/cuvs/pull/1323 --- .../main/java22/com/nvidia/cuvs/internal/common/Util.java | 8 ++------ 1 file changed, 2 insertions(+), 6 deletions(-) diff --git a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/common/Util.java b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/common/Util.java index 41f6b3736d..28dfc5c020 100644 --- a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/common/Util.java +++ b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/common/Util.java @@ -65,15 +65,11 @@ private Util() {} LINKER.downcallHandle( cudaMemcpyAsync$address(), cudaMemcpyAsync$descriptor(), Linker.Option.critical(true)); - private static final String cudaGetDevicePropertiesSymbolName = - "12".equals(System.getenv("RAPIDS_CUDA_MAJOR")) - ? "cudaGetDeviceProperties_v2" - : "cudaGetDeviceProperties"; - private static final MethodHandle cudaGetDeviceProperties$mh = LINKER.downcallHandle( SYMBOL_LOOKUP - .find(cudaGetDevicePropertiesSymbolName) + .find("cudaGetDeviceProperties") // CUDA 13+ symbol name + .or(() -> SYMBOL_LOOKUP.find("cudaGetDeviceProperties_v2")) // CUDA 12 symbol name .orElseThrow(UnsatisfiedLinkError::new), FunctionDescriptor.of(headers_h.C_INT, headers_h.C_POINTER, headers_h.C_INT)); From e292a09efe2341e91da7d74e5f807563d52b2317 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Lorenzo=20Dematt=C3=A9?= Date: Sun, 14 Sep 2025 02:45:13 +0200 Subject: [PATCH 60/93] [Java] Add `libcuvs` <-> cuvs-java version check (#1327) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit This PR extends https://github.com/rapidsai/cuvs/pull/1314 to add version check (version reported by libcuvs vs. cuvs-java build version) It replaces #1315 which was already reviewed and approved, but got closed before I could rebase it. Authors: - Lorenzo Dematté (https://github.com/ldematte) Approvers: - MithunR (https://github.com/mythrocks) URL: https://github.com/rapidsai/cuvs/pull/1327 --- java/cuvs-java/pom.xml | 5 ++- .../nvidia/cuvs/spi/CuVSServiceProvider.java | 6 ++- .../spi/ProviderInitializationException.java | 26 +++++++++++ .../com/nvidia/cuvs/spi/JDKProvider.java | 45 +++++++++++++++++++ 4 files changed, 79 insertions(+), 3 deletions(-) create mode 100644 java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/ProviderInitializationException.java diff --git a/java/cuvs-java/pom.xml b/java/cuvs-java/pom.xml index ca27c67635..57da752660 100644 --- a/java/cuvs-java/pom.xml +++ b/java/cuvs-java/pom.xml @@ -166,11 +166,12 @@ 3.4.2 + + true + true true - - com.nvidia.cuvs.examples.CagraExample 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 a4a5843adf..ae4c9b083a 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 @@ -51,8 +51,12 @@ static CuVSProvider builtinProvider() { if (supportedJavaRuntime && supportedOs && supportedArchitecture) { try { var cls = Class.forName("com.nvidia.cuvs.spi.JDKProvider"); - var ctr = MethodHandles.lookup().findConstructor(cls, MethodType.methodType(void.class)); + var ctr = + MethodHandles.lookup() + .findStatic(cls, "create", MethodType.methodType(CuVSProvider.class)); return (CuVSProvider) ctr.invoke(); + } catch (ProviderInitializationException e) { + return new UnsupportedProvider("cannot create JDKProvider: " + e.getMessage()); } catch (Throwable e) { throw new AssertionError(e); } diff --git a/java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/ProviderInitializationException.java b/java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/ProviderInitializationException.java new file mode 100644 index 0000000000..1a6c34a2fd --- /dev/null +++ b/java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/ProviderInitializationException.java @@ -0,0 +1,26 @@ +/* + * Copyright (c) 2025, NVIDIA CORPORATION. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.nvidia.cuvs.spi; + +class ProviderInitializationException extends Exception { + ProviderInitializationException(String message, Throwable cause) { + super(message, cause); + } + + public ProviderInitializationException(String message) { + super(message); + } +} diff --git a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/spi/JDKProvider.java b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/spi/JDKProvider.java index cfbc8acefa..77eb246e80 100644 --- a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/spi/JDKProvider.java +++ b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/spi/JDKProvider.java @@ -16,10 +16,14 @@ package com.nvidia.cuvs.spi; import static com.nvidia.cuvs.internal.common.Util.*; +import static com.nvidia.cuvs.internal.panama.headers_h.cuvsVersionGet; +import static com.nvidia.cuvs.internal.panama.headers_h.uint16_t; import com.nvidia.cuvs.*; import com.nvidia.cuvs.internal.*; import com.nvidia.cuvs.internal.common.Util; +import java.io.IOException; +import java.lang.foreign.Arena; import java.lang.foreign.MemorySegment; import java.lang.invoke.MethodHandle; import java.lang.invoke.MethodHandles; @@ -28,6 +32,8 @@ import java.nio.file.Path; import java.util.Locale; import java.util.Objects; +import java.util.jar.JarFile; +import java.util.jar.Manifest; final class JDKProvider implements CuVSProvider { @@ -37,6 +43,45 @@ final class JDKProvider implements CuVSProvider { private static final MethodHandle createNativeDataset$mh = createNativeDatasetBuilder(); + static CuVSProvider create() throws Throwable { + var mavenVersion = readCuVSVersionFromManifest(); + + try (var localArena = Arena.ofConfined()) { + var majorPtr = localArena.allocate(uint16_t); + var minorPtr = localArena.allocate(uint16_t); + var patchPtr = localArena.allocate(uint16_t); + checkCuVSError(cuvsVersionGet(majorPtr, minorPtr, patchPtr), "cuvsVersionGet"); + var major = majorPtr.get(uint16_t, 0); + var minor = minorPtr.get(uint16_t, 0); + var patch = patchPtr.get(uint16_t, 0); + + var cuvsVersionString = String.format(Locale.ROOT, "%02d.%02d.%d", major, minor, patch); + if (mavenVersion != null && !cuvsVersionString.equals(mavenVersion)) { + throw new ProviderInitializationException( + String.format( + Locale.ROOT, + "libcuvs_c version mismatch: expected [%s], found [%s]", + mavenVersion, + cuvsVersionString)); + } + } + return new JDKProvider(); + } + + /** + * Read cuvs-java version from this Jar Manifest, or null if these are not available + */ + private static String readCuVSVersionFromManifest() { + try (var jarFile = + new JarFile( + JDKProvider.class.getProtectionDomain().getCodeSource().getLocation().getPath())) { + Manifest manifest = jarFile.getManifest(); + return manifest.getMainAttributes().getValue("Implementation-Version"); + } catch (IOException e) { + return null; + } + } + static MethodHandle createNativeDatasetBuilder() { try { var lookup = MethodHandles.lookup(); From 757b02a4293f25ade7f0cf3194f1f374007cee51 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Lorenzo=20Dematt=C3=A9?= Date: Mon, 15 Sep 2025 17:38:19 +0200 Subject: [PATCH 61/93] [Review][Java] Add detailed error message for `libcuvs` load failure to UnsupportedProvider/UnsupportedOperationExceptions (#1316) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit This PR further extends #1296 and #1314 to give meaningful error messages in case libcuvs fails to load. The `jextract` generated bindings we use in cuvs-java use `SymbolLookup#libraryLookup` to load the `cuvs_c` dynamic library; this uses `RawNativeLibraries#load` (see https://github.com/openjdk/jdk/blob/master/src/java.base/share/native/libjava/RawNativeLibraries.c#L58); `RawNativeLibraries#load` in turn calls `JVM_LoadLibrary`. `JVM_LoadLibrary` does a good job to put together a good error message (e.g. calling `dlerror`, trying to locate and inspect the file for platform mismatch, etc. Unfortunately, `RawNativeLibraries#load` calls it passing false to the `throwException` parameter, which means that the detailed error messages are not surfaced. This PR follows the pattern introduced in #1296 and preloads libcuvs (and dependencies) using `JVM_LoadLibrary` directly with `throwException` true; preloading it will also cause the OS to look for and load all dependencies. In case of error we can see what's broken in better detail; e.g. if `libcuvs_c.so` is present, but `librmm.so` is missing: ``` java.lang.UnsupportedOperationException: cannot create JDKProvider: libcuvs_c.so: librmm.so: cannot open shared object file: No such file or directory at com.nvidia.cuvs@25.10.0/com.nvidia.cuvs.spi.UnsupportedProvider.newCuVSResources(UnsupportedProvider.java:35) at com.nvidia.cuvs@25.10.0/com.nvidia.cuvs.CuVSResources.create(CuVSResources.java:90) at com.nvidia.cuvs@25.10.0/com.nvidia.cuvs.CuVSResources.create(CuVSResources.java:79) ``` Fixes https://github.com/rapidsai/cuvs/issues/1321 Authors: - Lorenzo Dematté (https://github.com/ldematte) Approvers: - Ishan Chattopadhyaya (https://github.com/chatman) - Chris Hegarty (https://github.com/ChrisHegarty) - MithunR (https://github.com/mythrocks) URL: https://github.com/rapidsai/cuvs/pull/1316 --- java/cuvs-java/pom.xml | 4 +- .../nvidia/cuvs/spi/CuVSServiceProvider.java | 2 +- .../internal/common/NativeLibraryUtils.java | 43 +++++ .../com/nvidia/cuvs/spi/JDKProvider.java | 10 +- .../cuvs/spi/NativeDependencyLoader.java | 163 ++++++++++++++++++ .../spi/OptionalNativeDependencyLoader.java | 81 --------- 6 files changed, 214 insertions(+), 89 deletions(-) create mode 100644 java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/common/NativeLibraryUtils.java create mode 100644 java/cuvs-java/src/main/java22/com/nvidia/cuvs/spi/NativeDependencyLoader.java delete mode 100644 java/cuvs-java/src/main/java22/com/nvidia/cuvs/spi/OptionalNativeDependencyLoader.java diff --git a/java/cuvs-java/pom.xml b/java/cuvs-java/pom.xml index 57da752660..5e79d0c595 100644 --- a/java/cuvs-java/pom.xml +++ b/java/cuvs-java/pom.xml @@ -327,7 +327,7 @@ true true - com.nvidia.cuvs.examples.CagraExample + 12 @@ -410,7 +410,7 @@ true true - com.nvidia.cuvs.examples.CagraExample + 13 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 ae4c9b083a..56db88098d 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 @@ -56,7 +56,7 @@ static CuVSProvider builtinProvider() { .findStatic(cls, "create", MethodType.methodType(CuVSProvider.class)); return (CuVSProvider) ctr.invoke(); } catch (ProviderInitializationException e) { - return new UnsupportedProvider("cannot create JDKProvider: " + e.getMessage()); + return new UnsupportedProvider("Cannot create JDKProvider: " + e.getMessage()); } catch (Throwable e) { throw new AssertionError(e); } diff --git a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/common/NativeLibraryUtils.java b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/common/NativeLibraryUtils.java new file mode 100644 index 0000000000..b09d528224 --- /dev/null +++ b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/common/NativeLibraryUtils.java @@ -0,0 +1,43 @@ +/* + * Copyright (c) 2025, NVIDIA CORPORATION. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.nvidia.cuvs.internal.common; + +import java.lang.foreign.*; +import java.lang.invoke.MethodHandle; + +public class NativeLibraryUtils { + + private NativeLibraryUtils() {} + + private static final SymbolLookup LOOKUP = + SymbolLookup.libraryLookup(System.mapLibraryName("jvm"), Arena.ofAuto()) + .or(SymbolLookup.loaderLookup()) + .or(Linker.nativeLinker().defaultLookup()); + + // void * JVM_LoadLibrary(const char *name, jboolean throwException); + public static MethodHandle JVM_LoadLibrary$mh = + Linker.nativeLinker() + .downcallHandle( + LOOKUP.find("JVM_LoadLibrary").orElseThrow(), + FunctionDescriptor.of( + ValueLayout.ADDRESS, ValueLayout.ADDRESS, ValueLayout.JAVA_BOOLEAN)); + // void JVM_UnloadLibrary(void * handle); + public static MethodHandle JVM_UnloadLibrary$mh = + Linker.nativeLinker() + .downcallHandle( + LOOKUP.find("JVM_UnloadLibrary").orElseThrow(), + FunctionDescriptor.ofVoid(ValueLayout.ADDRESS)); +} diff --git a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/spi/JDKProvider.java b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/spi/JDKProvider.java index 77eb246e80..6f3d3fa6e3 100644 --- a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/spi/JDKProvider.java +++ b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/spi/JDKProvider.java @@ -37,13 +37,13 @@ final class JDKProvider implements CuVSProvider { - static { - OptionalNativeDependencyLoader.loadLibraries(); - } - private static final MethodHandle createNativeDataset$mh = createNativeDatasetBuilder(); - static CuVSProvider create() throws Throwable { + private JDKProvider() {} + + static CuVSProvider create() throws ProviderInitializationException { + NativeDependencyLoader.loadLibraries(); + var mavenVersion = readCuVSVersionFromManifest(); try (var localArena = Arena.ofConfined()) { diff --git a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/spi/NativeDependencyLoader.java b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/spi/NativeDependencyLoader.java new file mode 100644 index 0000000000..436008ed52 --- /dev/null +++ b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/spi/NativeDependencyLoader.java @@ -0,0 +1,163 @@ +/* + * Copyright (c) 2025, NVIDIA CORPORATION. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.nvidia.cuvs.spi; + +import static com.nvidia.cuvs.internal.common.NativeLibraryUtils.JVM_LoadLibrary$mh; + +import java.io.*; +import java.lang.foreign.Arena; +import java.net.URL; +import java.nio.file.Files; +import java.nio.file.StandardCopyOption; +import java.util.jar.JarFile; +import java.util.jar.Manifest; + +/** + * A class that loads native dependencies if they are available in the jar. + */ +class NativeDependencyLoader { + + interface NativeDependencyLoaderStrategy { + void loadLibraries() throws ProviderInitializationException; + } + + private static final NativeDependencyLoaderStrategy LOADER_STRATEGY = createLoaderStrategy(); + + private static NativeDependencyLoaderStrategy createLoaderStrategy() { + if (jarHasNativeDependencies()) { + return new EmbeddedNativeDependencyLoaderStrategy(); + } else { + return new SystemNativeDependencyLoaderStrategy(); + } + } + + private static boolean jarHasNativeDependencies() { + try (var jarFile = + new JarFile( + JDKProvider.class.getProtectionDomain().getCodeSource().getLocation().getPath())) { + Manifest manifest = jarFile.getManifest(); + // TODO: use this variable to add a check on the installed CUDA version + // (which will be system-loaded in any case, even with the fat-jar) + var embeddedLibrariesCudaVersion = + manifest.getMainAttributes().getValue("Embedded-Libraries-Cuda-Version"); + return embeddedLibrariesCudaVersion != null; + } catch (IOException e) { + return false; + } + } + + private static boolean loaded = false; + + static void loadLibraries() throws ProviderInitializationException { + if (!loaded) { + try { + LOADER_STRATEGY.loadLibraries(); + } finally { + loaded = true; + } + } + } + + private static class EmbeddedNativeDependencyLoaderStrategy + implements NativeDependencyLoaderStrategy { + + private static final String OS = System.getProperty("os.name"); + private static final String ARCH = System.getProperty("os.arch"); + private static final ClassLoader CLASS_LOADER = JDKProvider.class.getClassLoader(); + + private static final String[] FILES_TO_LOAD = { + "rapids_logger", "rmm", "cuvs", "cuvs_c", + }; + + @Override + public void loadLibraries() throws ProviderInitializationException { + for (String file : FILES_TO_LOAD) { + // Uncomment the following line to trace the loading of native dependencies. + // System.out.println("Loading native dependency: " + file); + try { + System.load(createFile(file).getAbsolutePath()); + } catch (Throwable t) { + throw new ProviderInitializationException( + "Failed to load native dependency: " + + System.mapLibraryName(file) + + ".so: " + + t.getMessage(), + t); + } + } + } + + /** + * Extract the contents of a library resource into a temporary file + */ + private static File createFile(String baseName) throws IOException { + String path = + EmbeddedNativeDependencyLoaderStrategy.ARCH + + "/" + + EmbeddedNativeDependencyLoaderStrategy.OS + + "/" + + System.mapLibraryName(baseName); + File loc; + URL resource = CLASS_LOADER.getResource(path); + if (resource == null) { + throw new FileNotFoundException("Could not locate native dependency " + path); + } + try (InputStream in = resource.openStream()) { + loc = File.createTempFile(baseName, ".so"); + loc.deleteOnExit(); + + Files.copy(in, loc.toPath(), StandardCopyOption.REPLACE_EXISTING); + } + return loc; + } + } + + private static class SystemNativeDependencyLoaderStrategy + implements NativeDependencyLoaderStrategy { + + @Override + public void loadLibraries() throws ProviderInitializationException { + // Try load libcuvs using directly JVM_LoadLibrary with the correct flags for in-depth failure + // diagnosis. + // + // jextract loads the dynamic libraries it references with SymbolLookup.libraryLookup; this + // uses + // RawNativeLibraries::load + // https://github.com/openjdk/jdk/blob/master/src/java.base/share/native/libjava/RawNativeLibraries.c#L58 + // RawNativeLibraries::load in turn calls JVM_LoadLibrary. Unfortunately, it calls it with a + // JNI_FALSE parameter for throwException, which means that the detailed error messages are + // not surfaced. + // + // Here we invoke it with throwException true, so in case of error we can see what's broken + String cuvsLibraryName = System.mapLibraryName("cuvs_c"); + + final Object lib; + try (var localArena = Arena.ofConfined()) { + var name = localArena.allocateFrom(cuvsLibraryName); + lib = JVM_LoadLibrary$mh.invoke(name, true); + } catch (Throwable ex) { + if (ex instanceof UnsatisfiedLinkError ulex) { + throw new ProviderInitializationException(ulex.getMessage(), ulex); + } else { + throw new ProviderInitializationException("Error while loading " + cuvsLibraryName, ex); + } + } + if (lib == null) { + throw new ProviderInitializationException("Unspecified failure loading " + cuvsLibraryName); + } + } + } +} diff --git a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/spi/OptionalNativeDependencyLoader.java b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/spi/OptionalNativeDependencyLoader.java deleted file mode 100644 index facbb670a3..0000000000 --- a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/spi/OptionalNativeDependencyLoader.java +++ /dev/null @@ -1,81 +0,0 @@ -/* - * Copyright (c) 2025, NVIDIA CORPORATION. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ -package com.nvidia.cuvs.spi; - -import java.io.*; -import java.net.URL; -import java.util.stream.*; - -/** - * A class that loads native dependencies if they are available in the jar. - */ -public class OptionalNativeDependencyLoader { - - private static final ClassLoader loader = JDKProvider.class.getClassLoader(); - - private static boolean loaded = false; - - private static final String[] FILES_TO_LOAD = { - "rapids_logger", "rmm", "cuvs", "cuvs_c", - }; - - public static void loadLibraries() { - if (!loaded) { - String os = System.getProperty("os.name"); - String arch = System.getProperty("os.arch"); - - Stream.of(FILES_TO_LOAD) - .forEach( - file -> { - // Uncomment the following line to trace the loading of native dependencies. - // System.out.println("Loading native dependency: " + file); - try { - System.load(createFile(os, arch, file).getAbsolutePath()); - } catch (Throwable t) { - System.err.println( - "Continuing despite failure to load native dependency: " - + System.mapLibraryName(file) - + ".so: " - + t.getMessage()); - } - }); - - loaded = true; - } - } - - /** Extract the contents of a library resource into a temporary file */ - private static File createFile(String os, String arch, String baseName) throws IOException { - String path = arch + "/" + os + "/" + System.mapLibraryName(baseName); - File loc; - URL resource = loader.getResource(path); - if (resource == null) { - throw new FileNotFoundException("Could not locate native dependency " + path); - } - try (InputStream in = resource.openStream()) { - loc = File.createTempFile(baseName, ".so"); - loc.deleteOnExit(); - try (OutputStream out = new FileOutputStream(loc)) { - byte[] buffer = new byte[1024 * 16]; - int read = 0; - while ((read = in.read(buffer)) >= 0) { - out.write(buffer, 0, read); - } - } - } - return loc; - } -} From 703a1028e2b5c1eec8ef805a74ac718ac9a26086 Mon Sep 17 00:00:00 2001 From: Ben Karsin Date: Mon, 15 Sep 2025 18:10:18 -1000 Subject: [PATCH 62/93] Fix pq_compressed file created by SSD Vamana build (#1313) Fix issue with pq_compressed files being incorrect when GPU vamana build was being provided opq codebooks as an input to the build. Bug was caused by vq_centers not being set to 0s, despite OPQ compression not using vq centers, and only using pq_centers during quantized vector encoding. Also cleaned up vamana_example file to support other datatypes and function correctly when building the DiskANN index without quantization. Authors: - Ben Karsin (https://github.com/bkarsin) Approvers: - Corey J. Nolet (https://github.com/cjnolet) - Tarang Jain (https://github.com/tarang-jain) URL: https://github.com/rapidsai/cuvs/pull/1313 --- cpp/src/neighbors/detail/vamana/vamana_build.cuh | 1 + examples/cpp/src/vamana_example.cu | 12 ++++++------ 2 files changed, 7 insertions(+), 6 deletions(-) diff --git a/cpp/src/neighbors/detail/vamana/vamana_build.cuh b/cpp/src/neighbors/detail/vamana/vamana_build.cuh index 25c842947f..629319e0e3 100644 --- a/cpp/src/neighbors/detail/vamana/vamana_build.cuh +++ b/cpp/src/neighbors/detail/vamana/vamana_build.cuh @@ -570,6 +570,7 @@ auto quantize_all_vectors(raft::resources const& res, { auto dim = residuals.extent(1); auto vq_codebook = raft::make_device_matrix(res, 1, dim); + raft::matrix::fill(res, vq_codebook.view(), 0.0); auto codes = cuvs::neighbors::detail::process_and_fill_codes_subspaces( res, ps, residuals, raft::make_const_mdspan(vq_codebook.view()), pq_codebook); diff --git a/examples/cpp/src/vamana_example.cu b/examples/cpp/src/vamana_example.cu index 9715d0dc38..c708020831 100644 --- a/examples/cpp/src/vamana_example.cu +++ b/examples/cpp/src/vamana_example.cu @@ -63,12 +63,12 @@ void vamana_build_and_write(raft::device_resources const& dev_resources, std::cout << "Time to build index: " << elapsed_seconds.count() << "s\n"; - // Output index to file (in-memory format) - serialize(dev_resources, out_fname, index); - - // Output index to file (disk sector-aligned format) if (codebook_prefix != "") { - serialize(dev_resources, out_fname + ".sector_aligned", index, false, true); + // Output index to file (disk sector-aligned format) + serialize(dev_resources, out_fname, index, false, true); + } else { + // Output index to file (in-memory format) + serialize(dev_resources, out_fname, index); } } @@ -80,8 +80,8 @@ void usage() "codebook prefix>\n"); printf("Input file expected to be binary file of fp32 vectors.\n"); - printf("Graph degree sizes supported: 32, 64, 128, 256\n"); printf("Datatype of input dataset (int8 or float)\n"); + printf("Graph degree sizes supported: 32, 64, 128, 256\n"); printf("Visited_size must be > degree and a power of 2.\n"); printf("max_fraction > 0 and <= 1. Typical values are 0.06 or 0.1.\n"); printf("Default iterations = 1.0, increase for better quality graph.\n"); From 07e422b54d1f2ed43172949d0ef40d989229dbca Mon Sep 17 00:00:00 2001 From: Ben Frederickson Date: Tue, 16 Sep 2025 06:28:36 -0700 Subject: [PATCH 63/93] Documentation Fixes (#1301) This fixes all the warnings in the docs build, and enables warnings as errors to prevent future regressions Authors: - Ben Frederickson (https://github.com/benfred) - Corey J. Nolet (https://github.com/cjnolet) Approvers: - Corey J. Nolet (https://github.com/cjnolet) - Kyle Edwards (https://github.com/KyleFromNVIDIA) URL: https://github.com/rapidsai/cuvs/pull/1301 --- build.sh | 2 +- cpp/doxygen/Doxyfile | 1209 +++++++++++------ cpp/include/cuvs/neighbors/brute_force.h | 4 +- cpp/include/cuvs/neighbors/cagra.h | 2 +- cpp/include/cuvs/neighbors/cagra.hpp | 2 + cpp/include/cuvs/neighbors/common.hpp | 9 +- cpp/include/cuvs/neighbors/hnsw.h | 4 +- cpp/include/cuvs/neighbors/hnsw.hpp | 34 +- cpp/include/cuvs/neighbors/ivf_flat.h | 2 +- cpp/include/cuvs/neighbors/ivf_pq.h | 2 +- cpp/include/cuvs/neighbors/vamana.hpp | 27 +- .../cuvs/preprocessing/quantize/binary.hpp | 1 - .../detail/vamana/vamana_codebooks.cuh | 5 +- cpp/src/neighbors/vamana.cuh | 3 +- cpp/src/neighbors/vamana_codebooks_float.cu | 2 +- docs/source/api_basics.rst | 4 +- docs/source/api_interoperability.rst | 4 +- docs/source/c_api/neighbors_cagra_c.rst | 2 +- docs/source/conf.py | 4 +- docs/source/cpp_api.rst | 1 + docs/source/cpp_api/distance.rst | 4 +- .../cpp_api/neighbors_all_neighbors.rst | 6 +- docs/source/cpp_api/neighbors_cagra.rst | 8 - docs/source/cuvs_bench/index.rst | 17 +- docs/source/cuvs_bench/param_tuning.rst | 56 +- docs/source/cuvs_bench/wiki_all_dataset.rst | 1 + docs/source/filtering.rst | 2 +- docs/source/getting_started.rst | 2 +- docs/source/index.rst | 1 + docs/source/indexes/ivfpq.rst | 2 +- docs/source/indexes/vamana.rst | 2 +- docs/source/python_api/neighbors_cagra.rst | 5 - docs/source/python_api/neighbors_hnsw.rst | 2 +- .../source/python_api/neighbors_nn_decent.rst | 2 +- docs/source/python_api/preprocessing.rst | 13 +- docs/source/sphinxext/github_link.py | 68 +- docs/source/working_with_ann_indexes_c.rst | 2 - docs/source/working_with_ann_indexes_cpp.rst | 2 - .../working_with_ann_indexes_python.rst | 2 - docs/source/working_with_ann_indexes_rust.rst | 39 +- .../neighbors/brute_force/brute_force.pyx | 8 +- python/cuvs/cuvs/neighbors/cagra/cagra.pyx | 46 +- python/cuvs/cuvs/neighbors/hnsw/hnsw.pyx | 23 +- .../cuvs/cuvs/neighbors/ivf_flat/ivf_flat.pyx | 2 + python/cuvs/cuvs/neighbors/ivf_pq/ivf_pq.pyx | 2 + 45 files changed, 1047 insertions(+), 593 deletions(-) diff --git a/build.sh b/build.sh index 2caec74ca7..328c64b43b 100755 --- a/build.sh +++ b/build.sh @@ -517,7 +517,7 @@ if hasArg docs; then cd "${DOXYGEN_BUILD_DIR}" doxygen Doxyfile cd "${SPHINX_BUILD_DIR}" - sphinx-build -b html source _html + sphinx-build -W -b html source _html cd "${REPODIR}"/rust cargo doc -p cuvs --no-deps rsync -av "${RUST_BUILD_DIR}"/doc/ "${SPHINX_BUILD_DIR}"/_html/_static/rust diff --git a/cpp/doxygen/Doxyfile b/cpp/doxygen/Doxyfile index ce6741c1f2..f9ffb39a20 100644 --- a/cpp/doxygen/Doxyfile +++ b/cpp/doxygen/Doxyfile @@ -1,7 +1,7 @@ -# Doxyfile 1.8.20 +# Doxyfile 1.13.2 # This file describes the settings to be used by the documentation system -# doxygen (www.doxygen.org) for a project. +# Doxygen (www.doxygen.org) for a project. # # All text after a double hash (##) is considered a comment and is placed in # front of the TAG it is preceding. @@ -12,6 +12,16 @@ # For lists, items can also be appended using: # TAG += value [value, ...] # Values that contain spaces should be placed between quotes (\" \"). +# +# Note: +# +# Use Doxygen to compare the used configuration file with the template +# configuration file: +# doxygen -x [configFile] +# Use Doxygen to compare the used configuration file with the template +# configuration file without replacing the environment variables or CMake type +# replacement variables: +# doxygen -x_noenv [configFile] #--------------------------------------------------------------------------- # Project related configuration options @@ -38,10 +48,10 @@ PROJECT_NAME = "cuVS C++ API" # could be handy for archiving the generated documentation or if some version # control system is used. -PROJECT_NUMBER = "$(RAPIDS_VERSION_MAJOR_MINOR)" +PROJECT_NUMBER = $(RAPIDS_VERSION_MAJOR_MINOR) # Using the PROJECT_BRIEF tag one can provide an optional one line description -# for a project that appears at the top of each page and should give viewer a +# for a project that appears at the top of each page and should give viewers a # quick idea about the purpose of the project. Keep the description short. PROJECT_BRIEF = @@ -53,24 +63,42 @@ PROJECT_BRIEF = PROJECT_LOGO = +# With the PROJECT_ICON tag one can specify an icon that is included in the tabs +# when the HTML document is shown. Doxygen will copy the logo to the output +# directory. + +PROJECT_ICON = + # The OUTPUT_DIRECTORY tag is used to specify the (relative or absolute) path # into which the generated documentation will be written. If a relative path is -# entered, it will be relative to the location where doxygen was started. If +# entered, it will be relative to the location where Doxygen was started. If # left blank the current directory will be used. OUTPUT_DIRECTORY = -# If the CREATE_SUBDIRS tag is set to YES then doxygen will create 4096 sub- -# directories (in 2 levels) under the output directory of each output format and -# will distribute the generated files over these directories. Enabling this -# option can be useful when feeding doxygen a huge amount of source files, where -# putting all generated files in the same directory would otherwise causes -# performance problems for the file system. +# If the CREATE_SUBDIRS tag is set to YES then Doxygen will create up to 4096 +# sub-directories (in 2 levels) under the output directory of each output format +# and will distribute the generated files over these directories. Enabling this +# option can be useful when feeding Doxygen a huge amount of source files, where +# putting all generated files in the same directory would otherwise cause +# performance problems for the file system. Adapt CREATE_SUBDIRS_LEVEL to +# control the number of sub-directories. # The default value is: NO. CREATE_SUBDIRS = NO -# If the ALLOW_UNICODE_NAMES tag is set to YES, doxygen will allow non-ASCII +# Controls the number of sub-directories that will be created when +# CREATE_SUBDIRS tag is set to YES. Level 0 represents 16 directories, and every +# level increment doubles the number of directories, resulting in 4096 +# directories at level 8 which is the default and also the maximum value. The +# sub-directories are organized in 2 levels, the first level always has a fixed +# number of 16 directories. +# Minimum value: 0, maximum value: 8, default value: 8. +# This tag requires that the tag CREATE_SUBDIRS is set to YES. + +CREATE_SUBDIRS_LEVEL = 8 + +# If the ALLOW_UNICODE_NAMES tag is set to YES, Doxygen will allow non-ASCII # characters to appear in the names of generated files. If set to NO, non-ASCII # characters will be escaped, for example _xE3_x81_x84 will be used for Unicode # U+3044. @@ -79,36 +107,28 @@ CREATE_SUBDIRS = NO ALLOW_UNICODE_NAMES = NO # The OUTPUT_LANGUAGE tag is used to specify the language in which all -# documentation generated by doxygen is written. Doxygen will use this +# documentation generated by Doxygen is written. Doxygen will use this # information to generate all constant output in the proper language. -# Possible values are: Afrikaans, Arabic, Armenian, Brazilian, Catalan, Chinese, -# Chinese-Traditional, Croatian, Czech, Danish, Dutch, English (United States), -# Esperanto, Farsi (Persian), Finnish, French, German, Greek, Hungarian, -# Indonesian, Italian, Japanese, Japanese-en (Japanese with English messages), -# Korean, Korean-en (Korean with English messages), Latvian, Lithuanian, -# Macedonian, Norwegian, Persian (Farsi), Polish, Portuguese, Romanian, Russian, -# Serbian, Serbian-Cyrillic, Slovak, Slovene, Spanish, Swedish, Turkish, -# Ukrainian and Vietnamese. +# Possible values are: Afrikaans, Arabic, Armenian, Brazilian, Bulgarian, +# Catalan, Chinese, Chinese-Traditional, Croatian, Czech, Danish, Dutch, English +# (United States), Esperanto, Farsi (Persian), Finnish, French, German, Greek, +# Hindi, Hungarian, Indonesian, Italian, Japanese, Japanese-en (Japanese with +# English messages), Korean, Korean-en (Korean with English messages), Latvian, +# Lithuanian, Macedonian, Norwegian, Persian (Farsi), Polish, Portuguese, +# Romanian, Russian, Serbian, Serbian-Cyrillic, Slovak, Slovene, Spanish, +# Swedish, Turkish, Ukrainian and Vietnamese. # The default value is: English. OUTPUT_LANGUAGE = English -# The OUTPUT_TEXT_DIRECTION tag is used to specify the direction in which all -# documentation generated by doxygen is written. Doxygen will use this -# information to generate all generated output in the proper direction. -# Possible values are: None, LTR, RTL and Context. -# The default value is: None. - -OUTPUT_TEXT_DIRECTION = None - -# If the BRIEF_MEMBER_DESC tag is set to YES, doxygen will include brief member +# If the BRIEF_MEMBER_DESC tag is set to YES, Doxygen will include brief member # descriptions after the members that are listed in the file and class # documentation (similar to Javadoc). Set to NO to disable this. # The default value is: YES. BRIEF_MEMBER_DESC = YES -# If the REPEAT_BRIEF tag is set to YES, doxygen will prepend the brief +# If the REPEAT_BRIEF tag is set to YES, Doxygen will prepend the brief # description of a member or function before the detailed description # # Note: If both HIDE_UNDOC_MEMBERS and BRIEF_MEMBER_DESC are set to NO, the @@ -129,13 +149,13 @@ REPEAT_BRIEF = YES ABBREVIATE_BRIEF = # If the ALWAYS_DETAILED_SEC and REPEAT_BRIEF tags are both set to YES then -# doxygen will generate a detailed section even if there is only a brief +# Doxygen will generate a detailed section even if there is only a brief # description. # The default value is: NO. ALWAYS_DETAILED_SEC = NO -# If the INLINE_INHERITED_MEMB tag is set to YES, doxygen will show all +# If the INLINE_INHERITED_MEMB tag is set to YES, Doxygen will show all # inherited members of a class in the documentation of that class as if those # members were ordinary class members. Constructors, destructors and assignment # operators of the base classes will not be shown. @@ -143,7 +163,7 @@ ALWAYS_DETAILED_SEC = NO INLINE_INHERITED_MEMB = NO -# If the FULL_PATH_NAMES tag is set to YES, doxygen will prepend the full path +# If the FULL_PATH_NAMES tag is set to YES, Doxygen will prepend the full path # before files name in the file list and in the header files. If set to NO the # shortest path that makes the file name unique will be used # The default value is: YES. @@ -153,11 +173,11 @@ FULL_PATH_NAMES = YES # The STRIP_FROM_PATH tag can be used to strip a user-defined part of the path. # Stripping is only done if one of the specified strings matches the left-hand # part of the path. The tag can be used to show relative paths in the file list. -# If left blank the directory from which doxygen is run is used as the path to +# If left blank the directory from which Doxygen is run is used as the path to # strip. # # Note that you can specify absolute paths here, but also relative paths, which -# will be relative from the directory where doxygen is started. +# will be relative from the directory where Doxygen is started. # This tag requires that the tag FULL_PATH_NAMES is set to YES. STRIP_FROM_PATH = @@ -171,41 +191,42 @@ STRIP_FROM_PATH = STRIP_FROM_INC_PATH = -# If the SHORT_NAMES tag is set to YES, doxygen will generate much shorter (but -# less readable) file names. This can be useful is your file systems doesn't +# If the SHORT_NAMES tag is set to YES, Doxygen will generate much shorter (but +# less readable) file names. This can be useful if your file system doesn't # support long names like on DOS, Mac, or CD-ROM. # The default value is: NO. SHORT_NAMES = NO -# If the JAVADOC_AUTOBRIEF tag is set to YES then doxygen will interpret the -# first line (until the first dot) of a Javadoc-style comment as the brief -# description. If set to NO, the Javadoc-style will behave just like regular Qt- -# style comments (thus requiring an explicit @brief command for a brief -# description.) +# If the JAVADOC_AUTOBRIEF tag is set to YES then Doxygen will interpret the +# first line (until the first dot, question mark or exclamation mark) of a +# Javadoc-style comment as the brief description. If set to NO, the Javadoc- +# style will behave just like regular Qt-style comments (thus requiring an +# explicit @brief command for a brief description.) # The default value is: NO. JAVADOC_AUTOBRIEF = NO -# If the JAVADOC_BANNER tag is set to YES then doxygen will interpret a line +# If the JAVADOC_BANNER tag is set to YES then Doxygen will interpret a line # such as # /*************** # as being the beginning of a Javadoc-style comment "banner". If set to NO, the # Javadoc-style will behave just like regular comments and it will not be -# interpreted by doxygen. +# interpreted by Doxygen. # The default value is: NO. JAVADOC_BANNER = NO -# If the QT_AUTOBRIEF tag is set to YES then doxygen will interpret the first -# line (until the first dot) of a Qt-style comment as the brief description. If -# set to NO, the Qt-style will behave just like regular Qt-style comments (thus -# requiring an explicit \brief command for a brief description.) +# If the QT_AUTOBRIEF tag is set to YES then Doxygen will interpret the first +# line (until the first dot, question mark or exclamation mark) of a Qt-style +# comment as the brief description. If set to NO, the Qt-style will behave just +# like regular Qt-style comments (thus requiring an explicit \brief command for +# a brief description.) # The default value is: NO. QT_AUTOBRIEF = NO -# The MULTILINE_CPP_IS_BRIEF tag can be set to YES to make doxygen treat a +# The MULTILINE_CPP_IS_BRIEF tag can be set to YES to make Doxygen treat a # multi-line C++ special comment block (i.e. a block of //! or /// comments) as # a brief description. This used to be the default behavior. The new default is # to treat a multi-line C++ comment block as a detailed description. Set this @@ -217,10 +238,10 @@ QT_AUTOBRIEF = NO MULTILINE_CPP_IS_BRIEF = NO -# By default Python docstrings are displayed as preformatted text and doxygen's +# By default Python docstrings are displayed as preformatted text and Doxygen's # special commands cannot be used. By setting PYTHON_DOCSTRING to NO the -# doxygen's special commands can be used and the contents of the docstring -# documentation blocks is shown as doxygen documentation. +# Doxygen's special commands can be used and the contents of the docstring +# documentation blocks is shown as Doxygen documentation. # The default value is: YES. PYTHON_DOCSTRING = YES @@ -231,7 +252,7 @@ PYTHON_DOCSTRING = YES INHERIT_DOCS = YES -# If the SEPARATE_MEMBER_PAGES tag is set to YES then doxygen will produce a new +# If the SEPARATE_MEMBER_PAGES tag is set to YES then Doxygen will produce a new # page for each member. If set to NO, the documentation of a member will be part # of the file/class/namespace that contains it. # The default value is: NO. @@ -248,16 +269,16 @@ TAB_SIZE = 4 # the documentation. An alias has the form: # name=value # For example adding -# "sideeffect=@par Side Effects:\n" +# "sideeffect=@par Side Effects:^^" # will allow you to put the command \sideeffect (or @sideeffect) in the # documentation, which will result in a user-defined paragraph with heading -# "Side Effects:". You can put \n's in the value part of an alias to insert -# newlines (in the resulting output). You can put ^^ in the value part of an -# alias to insert a newline as if a physical newline was in the original file. -# When you need a literal { or } or , in the value part of an alias you have to -# escape them by means of a backslash (\), this can lead to conflicts with the -# commands \{ and \} for these it is advised to use the version @{ and @} or use -# a double escape (\\{ and \\}) +# "Side Effects:". Note that you cannot put \n's in the value part of an alias +# to insert newlines (in the resulting output). You can put ^^ in the value part +# of an alias to insert a newline as if a physical newline was in the original +# file. When you need a literal { or } or , in the value part of an alias you +# have to escape them by means of a backslash (\), this can lead to conflicts +# with the commands \{ and \} for these it is advised to use the version @{ and +# @} or use a double escape (\\{ and \\}) ALIASES = @@ -301,28 +322,31 @@ OPTIMIZE_OUTPUT_SLICE = NO # parses. With this tag you can assign which parser to use for a given # extension. Doxygen has a built-in mapping, but you can override or extend it # using this tag. The format is ext=language, where ext is a file extension, and -# language is one of the parsers supported by doxygen: IDL, Java, JavaScript, -# Csharp (C#), C, C++, D, PHP, md (Markdown), Objective-C, Python, Slice, VHDL, -# Fortran (fixed format Fortran: FortranFixed, free formatted Fortran: +# language is one of the parsers supported by Doxygen: IDL, Java, JavaScript, +# Csharp (C#), C, C++, Lex, D, PHP, md (Markdown), Objective-C, Python, Slice, +# VHDL, Fortran (fixed format Fortran: FortranFixed, free formatted Fortran: # FortranFree, unknown formatted Fortran: Fortran. In the later case the parser # tries to guess whether the code is fixed or free formatted code, this is the -# default for Fortran type files). For instance to make doxygen treat .inc files +# default for Fortran type files). For instance to make Doxygen treat .inc files # as Fortran files (default is PHP), and .f files as C (default is Fortran), # use: inc=Fortran f=C. # # Note: For files without extension you can use no_extension as a placeholder. # # Note that for custom extensions you also need to set FILE_PATTERNS otherwise -# the files are not read by doxygen. +# the files are not read by Doxygen. When specifying no_extension you should add +# * to the FILE_PATTERNS. +# +# Note see also the list of default file extension mappings. EXTENSION_MAPPING = cu=C++ \ cuh=C++ -# If the MARKDOWN_SUPPORT tag is enabled then doxygen pre-processes all comments +# If the MARKDOWN_SUPPORT tag is enabled then Doxygen pre-processes all comments # according to the Markdown format, which allows for more readable # documentation. See https://daringfireball.net/projects/markdown/ for details. -# The output of markdown processing is further processed by doxygen, so you can -# mix doxygen, HTML, and XML commands with Markdown formatting. Disable only in +# The output of markdown processing is further processed by Doxygen, so you can +# mix Doxygen, HTML, and XML commands with Markdown formatting. Disable only in # case of backward compatibilities issues. # The default value is: YES. @@ -332,25 +356,45 @@ MARKDOWN_SUPPORT = YES # to that level are automatically included in the table of contents, even if # they do not have an id attribute. # Note: This feature currently applies only to Markdown headings. -# Minimum value: 0, maximum value: 99, default value: 5. +# Minimum value: 0, maximum value: 99, default value: 6. # This tag requires that the tag MARKDOWN_SUPPORT is set to YES. TOC_INCLUDE_HEADINGS = 5 -# When enabled doxygen tries to link words that correspond to documented +# The MARKDOWN_ID_STYLE tag can be used to specify the algorithm used to +# generate identifiers for the Markdown headings. Note: Every identifier is +# unique. +# Possible values are: DOXYGEN use a fixed 'autotoc_md' string followed by a +# sequence number starting at 0 and GITHUB use the lower case version of title +# with any whitespace replaced by '-' and punctuation characters removed. +# The default value is: DOXYGEN. +# This tag requires that the tag MARKDOWN_SUPPORT is set to YES. + +MARKDOWN_ID_STYLE = DOXYGEN + +# When enabled Doxygen tries to link words that correspond to documented # classes, or namespaces to their corresponding documentation. Such a link can # be prevented in individual cases by putting a % sign in front of the word or -# globally by setting AUTOLINK_SUPPORT to NO. +# globally by setting AUTOLINK_SUPPORT to NO. Words listed in the +# AUTOLINK_IGNORE_WORDS tag are excluded from automatic linking. # The default value is: YES. AUTOLINK_SUPPORT = YES +# This tag specifies a list of words that, when matching the start of a word in +# the documentation, will suppress auto links generation, if it is enabled via +# AUTOLINK_SUPPORT. This list does not affect affect links explicitly created +# using \# or the \link or commands. +# This tag requires that the tag AUTOLINK_SUPPORT is set to YES. + +AUTOLINK_IGNORE_WORDS = + # If you use STL classes (i.e. std::string, std::vector, etc.) but do not want # to include (a tag file for) the STL sources as input, then you should set this -# tag to YES in order to let doxygen match functions declarations and +# tag to YES in order to let Doxygen match functions declarations and # definitions whose arguments contain STL classes (e.g. func(std::string); -# versus func(std::string) {}). This also make the inheritance and collaboration -# diagrams that involve STL classes more complete and accurate. +# versus func(std::string) {}). This also makes the inheritance and +# collaboration diagrams that involve STL classes more complete and accurate. # The default value is: NO. BUILTIN_STL_SUPPORT = NO @@ -362,16 +406,16 @@ BUILTIN_STL_SUPPORT = NO CPP_CLI_SUPPORT = NO # Set the SIP_SUPPORT tag to YES if your project consists of sip (see: -# https://www.riverbankcomputing.com/software/sip/intro) sources only. Doxygen -# will parse them like normal C++ but will assume all classes use public instead -# of private inheritance when no explicit protection keyword is present. +# https://www.riverbankcomputing.com/software) sources only. Doxygen will parse +# them like normal C++ but will assume all classes use public instead of private +# inheritance when no explicit protection keyword is present. # The default value is: NO. SIP_SUPPORT = NO # For Microsoft's IDL there are propget and propput attributes to indicate # getter and setter methods for a property. Setting this option to YES will make -# doxygen to replace the get and set methods by a property in the documentation. +# Doxygen to replace the get and set methods by a property in the documentation. # This will only work if the methods are indeed getting or setting a simple # type. If this is not the case, or you want to show the methods anyway, you # should set this option to NO. @@ -380,7 +424,7 @@ SIP_SUPPORT = NO IDL_PROPERTY_SUPPORT = YES # If member grouping is used in the documentation and the DISTRIBUTE_GROUP_DOC -# tag is set to YES then doxygen will reuse the documentation of the first +# tag is set to YES then Doxygen will reuse the documentation of the first # member in the group (if any) for the other members of the group. By default # all members of a group must be documented explicitly. # The default value is: NO. @@ -438,18 +482,18 @@ TYPEDEF_HIDES_STRUCT = NO # The size of the symbol lookup cache can be set using LOOKUP_CACHE_SIZE. This # cache is used to resolve symbols given their name and scope. Since this can be # an expensive process and often the same symbol appears multiple times in the -# code, doxygen keeps a cache of pre-resolved symbols. If the cache is too small -# doxygen will become slower. If the cache is too large, memory is wasted. The +# code, Doxygen keeps a cache of pre-resolved symbols. If the cache is too small +# Doxygen will become slower. If the cache is too large, memory is wasted. The # cache size is given by this formula: 2^(16+LOOKUP_CACHE_SIZE). The valid range # is 0..9, the default is 0, corresponding to a cache size of 2^16=65536 -# symbols. At the end of a run doxygen will report the cache usage and suggest +# symbols. At the end of a run Doxygen will report the cache usage and suggest # the optimal cache size from a speed point of view. # Minimum value: 0, maximum value: 9, default value: 0. LOOKUP_CACHE_SIZE = 0 -# The NUM_PROC_THREADS specifies the number threads doxygen is allowed to use -# during processing. When set to 0 doxygen will based this on the number of +# The NUM_PROC_THREADS specifies the number of threads Doxygen is allowed to use +# during processing. When set to 0 Doxygen will based this on the number of # cores available in the system. You can set it explicitly to a value larger # than 0 to get more control over the balance between CPU load and processing # speed. At this moment only the input processing can be done using multiple @@ -461,11 +505,19 @@ LOOKUP_CACHE_SIZE = 0 NUM_PROC_THREADS = 0 +# If the TIMESTAMP tag is set different from NO then each generated page will +# contain the date or date and time when the page was generated. Setting this to +# NO can help when comparing the output of multiple runs. +# Possible values are: YES, NO, DATETIME and DATE. +# The default value is: NO. + +TIMESTAMP = NO + #--------------------------------------------------------------------------- # Build related configuration options #--------------------------------------------------------------------------- -# If the EXTRACT_ALL tag is set to YES, doxygen will assume all entities in +# If the EXTRACT_ALL tag is set to YES, Doxygen will assume all entities in # documentation are documented, even if no documentation was available. Private # class members and static file members will be hidden unless the # EXTRACT_PRIVATE respectively EXTRACT_STATIC tags are set to YES. @@ -524,7 +576,14 @@ EXTRACT_LOCAL_METHODS = NO EXTRACT_ANON_NSPACES = NO -# If the HIDE_UNDOC_MEMBERS tag is set to YES, doxygen will hide all +# If this flag is set to YES, the name of an unnamed parameter in a declaration +# will be determined by the corresponding definition. By default unnamed +# parameters remain unnamed in the output. +# The default value is: YES. + +RESOLVE_UNNAMED_PARAMS = YES + +# If the HIDE_UNDOC_MEMBERS tag is set to YES, Doxygen will hide all # undocumented members inside documented classes or files. If set to NO these # members will be included in the various overviews, but no documentation # section is generated. This option has no effect if EXTRACT_ALL is enabled. @@ -532,22 +591,31 @@ EXTRACT_ANON_NSPACES = NO HIDE_UNDOC_MEMBERS = NO -# If the HIDE_UNDOC_CLASSES tag is set to YES, doxygen will hide all +# If the HIDE_UNDOC_CLASSES tag is set to YES, Doxygen will hide all # undocumented classes that are normally visible in the class hierarchy. If set # to NO, these classes will be included in the various overviews. This option -# has no effect if EXTRACT_ALL is enabled. +# will also hide undocumented C++ concepts if enabled. This option has no effect +# if EXTRACT_ALL is enabled. # The default value is: NO. HIDE_UNDOC_CLASSES = NO -# If the HIDE_FRIEND_COMPOUNDS tag is set to YES, doxygen will hide all friend +# If the HIDE_UNDOC_NAMESPACES tag is set to YES, Doxygen will hide all +# undocumented namespaces that are normally visible in the namespace hierarchy. +# If set to NO, these namespaces will be included in the various overviews. This +# option has no effect if EXTRACT_ALL is enabled. +# The default value is: YES. + +HIDE_UNDOC_NAMESPACES = YES + +# If the HIDE_FRIEND_COMPOUNDS tag is set to YES, Doxygen will hide all friend # declarations. If set to NO, these declarations will be included in the # documentation. # The default value is: NO. HIDE_FRIEND_COMPOUNDS = NO -# If the HIDE_IN_BODY_DOCS tag is set to YES, doxygen will hide any +# If the HIDE_IN_BODY_DOCS tag is set to YES, Doxygen will hide any # documentation blocks found inside the body of a function. If set to NO, these # blocks will be appended to the function's detailed documentation block. # The default value is: NO. @@ -561,30 +629,44 @@ HIDE_IN_BODY_DOCS = NO INTERNAL_DOCS = NO -# If the CASE_SENSE_NAMES tag is set to NO then doxygen will only generate file -# names in lower-case letters. If set to YES, upper-case letters are also -# allowed. This is useful if you have classes or files whose names only differ -# in case and if your file system supports case sensitive file names. Windows -# (including Cygwin) and Mac users are advised to set this option to NO. -# The default value is: system dependent. +# With the correct setting of option CASE_SENSE_NAMES Doxygen will better be +# able to match the capabilities of the underlying filesystem. In case the +# filesystem is case sensitive (i.e. it supports files in the same directory +# whose names only differ in casing), the option must be set to YES to properly +# deal with such files in case they appear in the input. For filesystems that +# are not case sensitive the option should be set to NO to properly deal with +# output files written for symbols that only differ in casing, such as for two +# classes, one named CLASS and the other named Class, and to also support +# references to files without having to specify the exact matching casing. On +# Windows (including Cygwin) and macOS, users should typically set this option +# to NO, whereas on Linux or other Unix flavors it should typically be set to +# YES. +# Possible values are: SYSTEM, NO and YES. +# The default value is: SYSTEM. CASE_SENSE_NAMES = YES -# If the HIDE_SCOPE_NAMES tag is set to NO then doxygen will show members with +# If the HIDE_SCOPE_NAMES tag is set to NO then Doxygen will show members with # their full class and namespace scopes in the documentation. If set to YES, the # scope will be hidden. # The default value is: NO. HIDE_SCOPE_NAMES = NO -# If the HIDE_COMPOUND_REFERENCE tag is set to NO (default) then doxygen will +# If the HIDE_COMPOUND_REFERENCE tag is set to NO (default) then Doxygen will # append additional text to a page's title, such as Class Reference. If set to # YES the compound reference will be hidden. # The default value is: NO. HIDE_COMPOUND_REFERENCE= NO -# If the SHOW_INCLUDE_FILES tag is set to YES then doxygen will put a list of +# If the SHOW_HEADERFILE tag is set to YES then the documentation for a class +# will show which file needs to be included to use the class. +# The default value is: YES. + +SHOW_HEADERFILE = YES + +# If the SHOW_INCLUDE_FILES tag is set to YES then Doxygen will put a list of # the files that are included by a file in the documentation of that file. # The default value is: YES. @@ -597,7 +679,7 @@ SHOW_INCLUDE_FILES = YES SHOW_GROUPED_MEMB_INC = NO -# If the FORCE_LOCAL_INCLUDES tag is set to YES then doxygen will list include +# If the FORCE_LOCAL_INCLUDES tag is set to YES then Doxygen will list include # files with double quotes in the documentation rather than with sharp brackets. # The default value is: NO. @@ -609,14 +691,14 @@ FORCE_LOCAL_INCLUDES = NO INLINE_INFO = YES -# If the SORT_MEMBER_DOCS tag is set to YES then doxygen will sort the +# If the SORT_MEMBER_DOCS tag is set to YES then Doxygen will sort the # (detailed) documentation of file and class members alphabetically by member # name. If set to NO, the members will appear in declaration order. # The default value is: YES. SORT_MEMBER_DOCS = YES -# If the SORT_BRIEF_DOCS tag is set to YES then doxygen will sort the brief +# If the SORT_BRIEF_DOCS tag is set to YES then Doxygen will sort the brief # descriptions of file, namespace and class members alphabetically by member # name. If set to NO, the members will appear in declaration order. Note that # this will also influence the order of the classes in the class list. @@ -624,7 +706,7 @@ SORT_MEMBER_DOCS = YES SORT_BRIEF_DOCS = NO -# If the SORT_MEMBERS_CTORS_1ST tag is set to YES then doxygen will sort the +# If the SORT_MEMBERS_CTORS_1ST tag is set to YES then Doxygen will sort the # (brief and detailed) documentation of class members so that constructors and # destructors are listed first. If set to NO the constructors will appear in the # respective orders defined by SORT_BRIEF_DOCS and SORT_MEMBER_DOCS. @@ -636,7 +718,7 @@ SORT_BRIEF_DOCS = NO SORT_MEMBERS_CTORS_1ST = NO -# If the SORT_GROUP_NAMES tag is set to YES then doxygen will sort the hierarchy +# If the SORT_GROUP_NAMES tag is set to YES then Doxygen will sort the hierarchy # of group names into alphabetical order. If set to NO the group names will # appear in their defined order. # The default value is: NO. @@ -653,11 +735,11 @@ SORT_GROUP_NAMES = NO SORT_BY_SCOPE_NAME = NO -# If the STRICT_PROTO_MATCHING option is enabled and doxygen fails to do proper +# If the STRICT_PROTO_MATCHING option is enabled and Doxygen fails to do proper # type resolution of all parameters of a function it will reject a match between # the prototype and the implementation of a member function even if there is # only one candidate or it is obvious which candidate to choose by doing a -# simple string match. By disabling STRICT_PROTO_MATCHING doxygen will still +# simple string match. By disabling STRICT_PROTO_MATCHING Doxygen will still # accept a match between prototype and implementation in such cases. # The default value is: NO. @@ -727,24 +809,25 @@ SHOW_FILES = YES SHOW_NAMESPACES = YES # The FILE_VERSION_FILTER tag can be used to specify a program or script that -# doxygen should invoke to get the current version for each file (typically from +# Doxygen should invoke to get the current version for each file (typically from # the version control system). Doxygen will invoke the program by executing (via # popen()) the command command input-file, where command is the value of the # FILE_VERSION_FILTER tag, and input-file is the name of an input file provided -# by doxygen. Whatever the program writes to standard output is used as the file +# by Doxygen. Whatever the program writes to standard output is used as the file # version. For an example see the documentation. FILE_VERSION_FILTER = # The LAYOUT_FILE tag can be used to specify a layout file which will be parsed -# by doxygen. The layout file controls the global structure of the generated +# by Doxygen. The layout file controls the global structure of the generated # output files in an output format independent way. To create the layout file -# that represents doxygen's defaults, run doxygen with the -l option. You can +# that represents Doxygen's defaults, run Doxygen with the -l option. You can # optionally specify a file name after the option, if omitted DoxygenLayout.xml -# will be used as the name of the layout file. +# will be used as the name of the layout file. See also section "Changing the +# layout of pages" for information. # -# Note that if you run doxygen from a directory containing a file called -# DoxygenLayout.xml, doxygen will parse it automatically even if the LAYOUT_FILE +# Note that if you run Doxygen from a directory containing a file called +# DoxygenLayout.xml, Doxygen will parse it automatically even if the LAYOUT_FILE # tag is left empty. LAYOUT_FILE = @@ -759,19 +842,35 @@ LAYOUT_FILE = CITE_BIB_FILES = +# The EXTERNAL_TOOL_PATH tag can be used to extend the search path (PATH +# environment variable) so that external tools such as latex and gs can be +# found. +# Note: Directories specified with EXTERNAL_TOOL_PATH are added in front of the +# path already specified by the PATH variable, and are added in the order +# specified. +# Note: This option is particularly useful for macOS version 14 (Sonoma) and +# higher, when running Doxygen from Doxywizard, because in this case any user- +# defined changes to the PATH are ignored. A typical example on macOS is to set +# EXTERNAL_TOOL_PATH = /Library/TeX/texbin /usr/local/bin +# together with the standard path, the full search path used by doxygen when +# launching external tools will then become +# PATH=/Library/TeX/texbin:/usr/local/bin:/usr/bin:/bin:/usr/sbin:/sbin + +EXTERNAL_TOOL_PATH = + #--------------------------------------------------------------------------- # Configuration options related to warning and progress messages #--------------------------------------------------------------------------- # The QUIET tag can be used to turn on/off the messages that are generated to -# standard output by doxygen. If QUIET is set to YES this implies that the +# standard output by Doxygen. If QUIET is set to YES this implies that the # messages are off. # The default value is: NO. QUIET = NO # The WARNINGS tag can be used to turn on/off the warning messages that are -# generated to standard error (stderr) by doxygen. If WARNINGS is set to YES +# generated to standard error (stderr) by Doxygen. If WARNINGS is set to YES # this implies that the warnings are on. # # Tip: Turn warnings on while writing the documentation. @@ -779,49 +878,97 @@ QUIET = NO WARNINGS = YES -# If the WARN_IF_UNDOCUMENTED tag is set to YES then doxygen will generate +# If the WARN_IF_UNDOCUMENTED tag is set to YES then Doxygen will generate # warnings for undocumented members. If EXTRACT_ALL is set to YES then this flag # will automatically be disabled. # The default value is: YES. WARN_IF_UNDOCUMENTED = YES -# If the WARN_IF_DOC_ERROR tag is set to YES, doxygen will generate warnings for -# potential errors in the documentation, such as not documenting some parameters -# in a documented function, or documenting parameters that don't exist or using -# markup commands wrongly. +# If the WARN_IF_DOC_ERROR tag is set to YES, Doxygen will generate warnings for +# potential errors in the documentation, such as documenting some parameters in +# a documented function twice, or documenting parameters that don't exist or +# using markup commands wrongly. # The default value is: YES. WARN_IF_DOC_ERROR = YES +# If WARN_IF_INCOMPLETE_DOC is set to YES, Doxygen will warn about incomplete +# function parameter documentation. If set to NO, Doxygen will accept that some +# parameters have no documentation without warning. +# The default value is: YES. + +WARN_IF_INCOMPLETE_DOC = YES + # This WARN_NO_PARAMDOC option can be enabled to get warnings for functions that # are documented, but have no documentation for their parameters or return -# value. If set to NO, doxygen will only warn about wrong or incomplete -# parameter documentation, but not about the absence of documentation. If -# EXTRACT_ALL is set to YES then this flag will automatically be disabled. +# value. If set to NO, Doxygen will only warn about wrong parameter +# documentation, but not about the absence of documentation. If EXTRACT_ALL is +# set to YES then this flag will automatically be disabled. See also +# WARN_IF_INCOMPLETE_DOC # The default value is: NO. WARN_NO_PARAMDOC = YES -# If the WARN_AS_ERROR tag is set to YES then doxygen will immediately stop when -# a warning is encountered. +# If WARN_IF_UNDOC_ENUM_VAL option is set to YES, Doxygen will warn about +# undocumented enumeration values. If set to NO, Doxygen will accept +# undocumented enumeration values. If EXTRACT_ALL is set to YES then this flag +# will automatically be disabled. +# The default value is: NO. + +WARN_IF_UNDOC_ENUM_VAL = NO + +# If WARN_LAYOUT_FILE option is set to YES, Doxygen will warn about issues found +# while parsing the user defined layout file, such as missing or wrong elements. +# See also LAYOUT_FILE for details. If set to NO, problems with the layout file +# will be suppressed. +# The default value is: YES. + +WARN_LAYOUT_FILE = YES + +# If the WARN_AS_ERROR tag is set to YES then Doxygen will immediately stop when +# a warning is encountered. If the WARN_AS_ERROR tag is set to FAIL_ON_WARNINGS +# then Doxygen will continue running as if WARN_AS_ERROR tag is set to NO, but +# at the end of the Doxygen process Doxygen will return with a non-zero status. +# If the WARN_AS_ERROR tag is set to FAIL_ON_WARNINGS_PRINT then Doxygen behaves +# like FAIL_ON_WARNINGS but in case no WARN_LOGFILE is defined Doxygen will not +# write the warning messages in between other messages but write them at the end +# of a run, in case a WARN_LOGFILE is defined the warning messages will be +# besides being in the defined file also be shown at the end of a run, unless +# the WARN_LOGFILE is defined as - i.e. standard output (stdout) in that case +# the behavior will remain as with the setting FAIL_ON_WARNINGS. +# Possible values are: NO, YES, FAIL_ON_WARNINGS and FAIL_ON_WARNINGS_PRINT. # The default value is: NO. WARN_AS_ERROR = YES -# The WARN_FORMAT tag determines the format of the warning messages that doxygen +# The WARN_FORMAT tag determines the format of the warning messages that Doxygen # can produce. The string should contain the $file, $line, and $text tags, which # will be replaced by the file and line number from which the warning originated # and the warning text. Optionally the format may contain $version, which will # be replaced by the version of the file (if it could be obtained via # FILE_VERSION_FILTER) +# See also: WARN_LINE_FORMAT # The default value is: $file:$line: $text. WARN_FORMAT = "$file:$line: $text" +# In the $text part of the WARN_FORMAT command it is possible that a reference +# to a more specific place is given. To make it easier to jump to this place +# (outside of Doxygen) the user can define a custom "cut" / "paste" string. +# Example: +# WARN_LINE_FORMAT = "'vi $file +$line'" +# See also: WARN_FORMAT +# The default value is: at line $line of file $file. + +WARN_LINE_FORMAT = "at line $line of file $file" + # The WARN_LOGFILE tag can be used to specify a file to which warning and error # messages should be written. If left blank the output is written to standard -# error (stderr). +# error (stderr). In case the file specified cannot be opened for writing the +# warning and error messages are written to standard error. When as file - is +# specified the warning and error messages are written to standard output +# (stdout). WARN_LOGFILE = @@ -839,29 +986,42 @@ INPUT = main_page.md \ ../include # This tag can be used to specify the character encoding of the source files -# that doxygen parses. Internally doxygen uses the UTF-8 encoding. Doxygen uses +# that Doxygen parses. Internally Doxygen uses the UTF-8 encoding. Doxygen uses # libiconv (or the iconv built into libc) for the transcoding. See the libiconv -# documentation (see: https://www.gnu.org/software/libiconv/) for the list of -# possible encodings. +# documentation (see: +# https://www.gnu.org/software/libiconv/) for the list of possible encodings. +# See also: INPUT_FILE_ENCODING # The default value is: UTF-8. INPUT_ENCODING = UTF-8 +# This tag can be used to specify the character encoding of the source files +# that Doxygen parses. The INPUT_FILE_ENCODING tag can be used to specify +# character encoding on a per file pattern basis. Doxygen will compare the file +# name with each pattern and apply the encoding instead of the default +# INPUT_ENCODING if there is a match. The character encodings are a list of the +# form: pattern=encoding (like *.php=ISO-8859-1). +# See also: INPUT_ENCODING for further information on supported encodings. + +INPUT_FILE_ENCODING = + # If the value of the INPUT tag contains directories, you can use the # FILE_PATTERNS tag to specify one or more wildcard patterns (like *.cpp and # *.h) to filter out the source-files in the directories. # # Note that for custom extensions or not directly supported extensions you also # need to set EXTENSION_MAPPING for the extension otherwise the files are not -# read by doxygen. +# read by Doxygen. # -# If left blank the following patterns are tested:*.c, *.cc, *.cxx, *.cpp, -# *.c++, *.java, *.ii, *.ixx, *.ipp, *.i++, *.inl, *.idl, *.ddl, *.odl, *.h, -# *.hh, *.hxx, *.hpp, *.h++, *.cs, *.d, *.php, *.php4, *.php5, *.phtml, *.inc, -# *.m, *.markdown, *.md, *.mm, *.dox (to be provided as doxygen C comment), -# *.doc (to be provided as doxygen C comment), *.txt (to be provided as doxygen -# C comment), *.py, *.pyw, *.f90, *.f95, *.f03, *.f08, *.f18, *.f, *.for, *.vhd, -# *.vhdl, *.ucf, *.qsf and *.ice. +# Note the list of default checked file patterns might differ from the list of +# default file extension mappings. +# +# If left blank the following patterns are tested:*.c, *.cc, *.cxx, *.cxxm, +# *.cpp, *.cppm, *.ccm, *.c++, *.c++m, *.java, *.ii, *.ixx, *.ipp, *.i++, *.inl, +# *.idl, *.ddl, *.odl, *.h, *.hh, *.hxx, *.hpp, *.h++, *.ixx, *.l, *.cs, *.d, +# *.php, *.php4, *.php5, *.phtml, *.inc, *.m, *.markdown, *.md, *.mm, *.dox (to +# be provided as Doxygen C comment), *.py, *.pyw, *.f90, *.f95, *.f03, *.f08, +# *.f18, *.f, *.for, *.vhd, *.vhdl, *.ucf, *.qsf and *.ice. FILE_PATTERNS = *.hpp \ *.h \ @@ -877,10 +1037,13 @@ RECURSIVE = YES # excluded from the INPUT source files. This way you can easily exclude a # subdirectory from a directory tree whose root is specified with the INPUT tag. # -# Note that relative paths are relative to the directory from which doxygen is +# Note that relative paths are relative to the directory from which Doxygen is # run. -EXCLUDE = ../include/cuvs/sparse/selection +EXCLUDE = ../include/cuvs/sparse/selection \ + ../build \ + ../src + # The EXCLUDE_SYMLINKS tag can be used to select whether or not files or # directories that are symbolic links (a Unix file system feature) are excluded @@ -896,7 +1059,6 @@ EXCLUDE_SYMLINKS = NO # Note that the wildcards are matched against the file with absolute path, so to # exclude all test directories for example use the pattern */test/* -# TODO: remove specializations from exclude patterns when headers have been removed. EXCLUDE_PATTERNS = */detail/* \ */specializations/* \ */thirdparty/* @@ -905,10 +1067,7 @@ EXCLUDE_PATTERNS = */detail/* \ # (namespaces, classes, functions, etc.) that should be excluded from the # output. The symbol name can be a fully qualified name, a word, or if the # wildcard * is used, a substring. Examples: ANamespace, AClass, -# AClass::ANamespace, ANamespace::*Test -# -# Note that the wildcards are matched against the file with absolute path, so to -# exclude all test directories use the pattern */test/* +# ANamespace::AClass, ANamespace::*Test EXCLUDE_SYMBOLS = detail @@ -938,7 +1097,7 @@ EXAMPLE_RECURSIVE = NO IMAGE_PATH = -# The INPUT_FILTER tag can be used to specify a program that doxygen should +# The INPUT_FILTER tag can be used to specify a program that Doxygen should # invoke to filter for each input file. Doxygen will invoke the filter program # by executing (via popen()) the command: # @@ -953,9 +1112,14 @@ IMAGE_PATH = # code is scanned, but not when the output code is generated. If lines are added # or removed, the anchors will not be placed correctly. # +# Note that Doxygen will use the data processed and written to standard output +# for further processing, therefore nothing else, like debug statements or used +# commands (so in case of a Windows batch file always use @echo OFF), should be +# written to standard output. +# # Note that for custom extensions or not directly supported extensions you also # need to set EXTENSION_MAPPING for the extension otherwise the files are not -# properly processed by doxygen. +# properly processed by Doxygen. INPUT_FILTER = @@ -968,7 +1132,7 @@ INPUT_FILTER = # # Note that for custom extensions or not directly supported extensions you also # need to set EXTENSION_MAPPING for the extension otherwise the files are not -# properly processed by doxygen. +# properly processed by Doxygen. FILTER_PATTERNS = @@ -990,10 +1154,28 @@ FILTER_SOURCE_PATTERNS = # If the USE_MDFILE_AS_MAINPAGE tag refers to the name of a markdown file that # is part of the input, its contents will be placed on the main page # (index.html). This can be useful if you have a project on for instance GitHub -# and want to reuse the introduction page also for the doxygen output. +# and want to reuse the introduction page also for the Doxygen output. USE_MDFILE_AS_MAINPAGE = main_page.md +# If the IMPLICIT_DIR_DOCS tag is set to YES, any README.md file found in sub- +# directories of the project's root, is used as the documentation for that sub- +# directory, except when the README.md starts with a \dir, \page or \mainpage +# command. If set to NO, the README.md file needs to start with an explicit \dir +# command in order to be used as directory documentation. +# The default value is: YES. + +IMPLICIT_DIR_DOCS = YES + +# The Fortran standard specifies that for fixed formatted Fortran code all +# characters from position 72 are to be considered as comment. A common +# extension is to allow longer lines before the automatic comment starts. The +# setting FORTRAN_COMMENT_AFTER will also make it possible that longer lines can +# be processed before the automatic comment starts. +# Minimum value: 7, maximum value: 10000, default value: 72. + +FORTRAN_COMMENT_AFTER = 72 + #--------------------------------------------------------------------------- # Configuration options related to source browsing #--------------------------------------------------------------------------- @@ -1008,12 +1190,13 @@ USE_MDFILE_AS_MAINPAGE = main_page.md SOURCE_BROWSER = NO # Setting the INLINE_SOURCES tag to YES will include the body of functions, -# classes and enums directly into the documentation. +# multi-line macros, enums or list initialized variables directly into the +# documentation. # The default value is: NO. INLINE_SOURCES = NO -# Setting the STRIP_CODE_COMMENTS tag to YES will instruct doxygen to hide any +# Setting the STRIP_CODE_COMMENTS tag to YES will instruct Doxygen to hide any # special comment blocks from generated source code fragments. Normal C, C++ and # Fortran comments will always remain visible. # The default value is: YES. @@ -1051,7 +1234,7 @@ REFERENCES_LINK_SOURCE = YES SOURCE_TOOLTIPS = YES # If the USE_HTAGS tag is set to YES then the references to source code will -# point to the HTML generated by the htags(1) tool instead of doxygen built-in +# point to the HTML generated by the htags(1) tool instead of Doxygen built-in # source browser. The htags tool is part of GNU's global source tagging system # (see https://www.gnu.org/software/global/global.html). You will need version # 4.8.6 or higher. @@ -1065,14 +1248,14 @@ SOURCE_TOOLTIPS = YES # Doxygen will invoke htags (and that will in turn invoke gtags), so these # tools must be available from the command line (i.e. in the search path). # -# The result: instead of the source browser generated by doxygen, the links to +# The result: instead of the source browser generated by Doxygen, the links to # source code will now point to the output of htags. # The default value is: NO. # This tag requires that the tag SOURCE_BROWSER is set to YES. USE_HTAGS = NO -# If the VERBATIM_HEADERS tag is set the YES then doxygen will generate a +# If the VERBATIM_HEADERS tag is set the YES then Doxygen will generate a # verbatim copy of the header file for each class for which an include is # specified. Set to NO to disable this. # See also: Section \class. @@ -1091,17 +1274,11 @@ VERBATIM_HEADERS = YES ALPHABETICAL_INDEX = YES -# The COLS_IN_ALPHA_INDEX tag can be used to specify the number of columns in -# which the alphabetical index list will be split. -# Minimum value: 1, maximum value: 20, default value: 5. -# This tag requires that the tag ALPHABETICAL_INDEX is set to YES. - -COLS_IN_ALPHA_INDEX = 5 - -# In case all classes in a project start with a common prefix, all classes will -# be put under the same header in the alphabetical index. The IGNORE_PREFIX tag -# can be used to specify a prefix (or a list of prefixes) that should be ignored -# while generating the index headers. +# The IGNORE_PREFIX tag can be used to specify a prefix (or a list of prefixes) +# that should be ignored while generating the index headers. The IGNORE_PREFIX +# tag works for classes, function and member names. The entity will be placed in +# the alphabetical list under the first letter of the entity name that remains +# after removing the prefix. # This tag requires that the tag ALPHABETICAL_INDEX is set to YES. IGNORE_PREFIX = @@ -1110,7 +1287,7 @@ IGNORE_PREFIX = # Configuration options related to the HTML output #--------------------------------------------------------------------------- -# If the GENERATE_HTML tag is set to YES, doxygen will generate HTML output +# If the GENERATE_HTML tag is set to YES, Doxygen will generate HTML output # The default value is: YES. GENERATE_HTML = NO @@ -1131,40 +1308,40 @@ HTML_OUTPUT = html HTML_FILE_EXTENSION = .html # The HTML_HEADER tag can be used to specify a user-defined HTML header file for -# each generated HTML page. If the tag is left blank doxygen will generate a +# each generated HTML page. If the tag is left blank Doxygen will generate a # standard header. # # To get valid HTML the header file that includes any scripts and style sheets -# that doxygen needs, which is dependent on the configuration options used (e.g. +# that Doxygen needs, which is dependent on the configuration options used (e.g. # the setting GENERATE_TREEVIEW). It is highly recommended to start with a # default header using # doxygen -w html new_header.html new_footer.html new_stylesheet.css # YourConfigFile # and then modify the file new_header.html. See also section "Doxygen usage" -# for information on how to generate the default header that doxygen normally +# for information on how to generate the default header that Doxygen normally # uses. # Note: The header is subject to change so you typically have to regenerate the -# default header when upgrading to a newer version of doxygen. For a description +# default header when upgrading to a newer version of Doxygen. For a description # of the possible markers and block names see the documentation. # This tag requires that the tag GENERATE_HTML is set to YES. HTML_HEADER = header.html # The HTML_FOOTER tag can be used to specify a user-defined HTML footer for each -# generated HTML page. If the tag is left blank doxygen will generate a standard +# generated HTML page. If the tag is left blank Doxygen will generate a standard # footer. See HTML_HEADER for more information on how to generate a default # footer and what special commands can be used inside the footer. See also # section "Doxygen usage" for information on how to generate the default footer -# that doxygen normally uses. +# that Doxygen normally uses. # This tag requires that the tag GENERATE_HTML is set to YES. HTML_FOOTER = # The HTML_STYLESHEET tag can be used to specify a user-defined cascading style # sheet that is used by each HTML page. It can be used to fine-tune the look of -# the HTML output. If left blank doxygen will generate a default style sheet. +# the HTML output. If left blank Doxygen will generate a default style sheet. # See also section "Doxygen usage" for information on how to generate the style -# sheet that doxygen normally uses. +# sheet that Doxygen normally uses. # Note: It is recommended to use HTML_EXTRA_STYLESHEET instead of this tag, as # it is more robust and this tag (HTML_STYLESHEET) will in the future become # obsolete. @@ -1174,13 +1351,18 @@ HTML_STYLESHEET = # The HTML_EXTRA_STYLESHEET tag can be used to specify additional user-defined # cascading style sheets that are included after the standard style sheets -# created by doxygen. Using this option one can overrule certain style aspects. +# created by Doxygen. Using this option one can overrule certain style aspects. # This is preferred over using HTML_STYLESHEET since it does not replace the # standard style sheet and is therefore more robust against future updates. # Doxygen will copy the style sheet files to the output directory. # Note: The order of the extra style sheet files is of importance (e.g. the last # style sheet in the list overrules the setting of the previous ones in the -# list). For an example see the documentation. +# list). +# Note: Since the styling of scrollbars can currently not be overruled in +# Webkit/Chromium, the styling will be left out of the default doxygen.css if +# one or more extra stylesheets have been specified. So if scrollbar +# customization is desired it has to be added explicitly. For an example see the +# documentation. # This tag requires that the tag GENERATE_HTML is set to YES. HTML_EXTRA_STYLESHEET = @@ -1195,9 +1377,22 @@ HTML_EXTRA_STYLESHEET = HTML_EXTRA_FILES = +# The HTML_COLORSTYLE tag can be used to specify if the generated HTML output +# should be rendered with a dark or light theme. +# Possible values are: LIGHT always generates light mode output, DARK always +# generates dark mode output, AUTO_LIGHT automatically sets the mode according +# to the user preference, uses light mode if no preference is set (the default), +# AUTO_DARK automatically sets the mode according to the user preference, uses +# dark mode if no preference is set and TOGGLE allows a user to switch between +# light and dark mode via a button. +# The default value is: AUTO_LIGHT. +# This tag requires that the tag GENERATE_HTML is set to YES. + +HTML_COLORSTYLE = AUTO_LIGHT + # The HTML_COLORSTYLE_HUE tag controls the color of the HTML output. Doxygen # will adjust the colors in the style sheet and background images according to -# this color. Hue is specified as an angle on a colorwheel, see +# this color. Hue is specified as an angle on a color-wheel, see # https://en.wikipedia.org/wiki/Hue for more information. For instance the value # 0 represents red, 60 is yellow, 120 is green, 180 is cyan, 240 is blue, 300 # purple, and 360 is red again. @@ -1207,7 +1402,7 @@ HTML_EXTRA_FILES = HTML_COLORSTYLE_HUE = 266 # The HTML_COLORSTYLE_SAT tag controls the purity (or saturation) of the colors -# in the HTML output. For a value of 0 the output will use grayscales only. A +# in the HTML output. For a value of 0 the output will use gray-scales only. A # value of 255 will produce the most vivid colors. # Minimum value: 0, maximum value: 255, default value: 100. # This tag requires that the tag GENERATE_HTML is set to YES. @@ -1225,15 +1420,6 @@ HTML_COLORSTYLE_SAT = 255 HTML_COLORSTYLE_GAMMA = 52 -# If the HTML_TIMESTAMP tag is set to YES then the footer of each generated HTML -# page will contain the date and time when the page was generated. Setting this -# to YES can help to show when doxygen was last run and thus if the -# documentation is up to date. -# The default value is: NO. -# This tag requires that the tag GENERATE_HTML is set to YES. - -HTML_TIMESTAMP = NO - # If the HTML_DYNAMIC_MENUS tag is set to YES then the generated HTML # documentation will contain a main index with vertical navigation menus that # are dynamically created via JavaScript. If disabled, the navigation index will @@ -1253,6 +1439,33 @@ HTML_DYNAMIC_MENUS = YES HTML_DYNAMIC_SECTIONS = NO +# If the HTML_CODE_FOLDING tag is set to YES then classes and functions can be +# dynamically folded and expanded in the generated HTML source code. +# The default value is: YES. +# This tag requires that the tag GENERATE_HTML is set to YES. + +HTML_CODE_FOLDING = YES + +# If the HTML_COPY_CLIPBOARD tag is set to YES then Doxygen will show an icon in +# the top right corner of code and text fragments that allows the user to copy +# its content to the clipboard. Note this only works if supported by the browser +# and the web page is served via a secure context (see: +# https://www.w3.org/TR/secure-contexts/), i.e. using the https: or file: +# protocol. +# The default value is: YES. +# This tag requires that the tag GENERATE_HTML is set to YES. + +HTML_COPY_CLIPBOARD = YES + +# Doxygen stores a couple of settings persistently in the browser (via e.g. +# cookies). By default these settings apply to all HTML pages generated by +# Doxygen across all projects. The HTML_PROJECT_COOKIE tag can be used to store +# the settings under a project specific key, such that the user preferences will +# be stored separately. +# This tag requires that the tag GENERATE_HTML is set to YES. + +HTML_PROJECT_COOKIE = + # With HTML_INDEX_NUM_ENTRIES one can control the preferred number of entries # shown in the various tree structured indices initially; the user can expand # and collapse entries dynamically later on. Doxygen will expand the tree to @@ -1268,10 +1481,11 @@ HTML_INDEX_NUM_ENTRIES = 100 # If the GENERATE_DOCSET tag is set to YES, additional index files will be # generated that can be used as input for Apple's Xcode 3 integrated development -# environment (see: https://developer.apple.com/xcode/), introduced with OSX -# 10.5 (Leopard). To create a documentation set, doxygen will generate a -# Makefile in the HTML output directory. Running make will produce the docset in -# that directory and running make install will install the docset in +# environment (see: +# https://developer.apple.com/xcode/), introduced with OSX 10.5 (Leopard). To +# create a documentation set, Doxygen will generate a Makefile in the HTML +# output directory. Running make will produce the docset in that directory and +# running make install will install the docset in # ~/Library/Developer/Shared/Documentation/DocSets so that Xcode will find it at # startup. See https://developer.apple.com/library/archive/featuredarticles/Doxy # genXcode/_index.html for more information. @@ -1288,6 +1502,13 @@ GENERATE_DOCSET = NO DOCSET_FEEDNAME = "Doxygen generated docs" +# This tag determines the URL of the docset feed. A documentation feed provides +# an umbrella under which multiple documentation sets from a single provider +# (such as a company or product suite) can be grouped. +# This tag requires that the tag GENERATE_DOCSET is set to YES. + +DOCSET_FEEDURL = + # This tag specifies a string that should uniquely identify the documentation # set bundle. This should be a reverse domain-name style string, e.g. # com.mycompany.MyDocSet. Doxygen will append .docset to the name. @@ -1310,14 +1531,18 @@ DOCSET_PUBLISHER_ID = org.doxygen.Publisher DOCSET_PUBLISHER_NAME = Publisher -# If the GENERATE_HTMLHELP tag is set to YES then doxygen generates three +# If the GENERATE_HTMLHELP tag is set to YES then Doxygen generates three # additional HTML index files: index.hhp, index.hhc, and index.hhk. The # index.hhp is a project file that can be read by Microsoft's HTML Help Workshop -# (see: https://www.microsoft.com/en-us/download/details.aspx?id=21138) on -# Windows. +# on Windows. In the beginning of 2021 Microsoft took the original page, with +# a.o. the download links, offline (the HTML help workshop was already many +# years in maintenance mode). You can download the HTML help workshop from the +# web archives at Installation executable (see: +# http://web.archive.org/web/20160201063255/http://download.microsoft.com/downlo +# ad/0/A/9/0A939EF6-E31C-430F-A3DF-DFAE7960D564/htmlhelp.exe). # # The HTML Help Workshop contains a compiler that can convert all HTML output -# generated by doxygen into a single compiled HTML file (.chm). Compiled HTML +# generated by Doxygen into a single compiled HTML file (.chm). Compiled HTML # files are now used as the Windows 98 help format, and will replace the old # Windows help format (.hlp) on all Windows platforms in the future. Compressed # HTML files also contain an index, a table of contents, and you can search for @@ -1337,7 +1562,7 @@ CHM_FILE = # The HHC_LOCATION tag can be used to specify the location (absolute path # including file name) of the HTML help compiler (hhc.exe). If non-empty, -# doxygen will try to run the HTML help compiler on the generated index.hhp. +# Doxygen will try to run the HTML help compiler on the generated index.hhp. # The file has to be specified with full path. # This tag requires that the tag GENERATE_HTMLHELP is set to YES. @@ -1371,6 +1596,16 @@ BINARY_TOC = NO TOC_EXPAND = NO +# The SITEMAP_URL tag is used to specify the full URL of the place where the +# generated documentation will be placed on the server by the user during the +# deployment of the documentation. The generated sitemap is called sitemap.xml +# and placed on the directory specified by HTML_OUTPUT. In case no SITEMAP_URL +# is specified no sitemap is generated. For information about the sitemap +# protocol see https://www.sitemaps.org +# This tag requires that the tag GENERATE_HTML is set to YES. + +SITEMAP_URL = + # If the GENERATE_QHP tag is set to YES and both QHP_NAMESPACE and # QHP_VIRTUAL_FOLDER are set, an additional index file will be generated that # can be used as input for Qt's qhelpgenerator to generate a Qt Compressed Help @@ -1389,7 +1624,8 @@ QCH_FILE = # The QHP_NAMESPACE tag specifies the namespace to use when generating Qt Help # Project output. For more information please see Qt Help Project / Namespace -# (see: https://doc.qt.io/archives/qt-4.8/qthelpproject.html#namespace). +# (see: +# https://doc.qt.io/archives/qt-4.8/qthelpproject.html#namespace). # The default value is: org.doxygen.Project. # This tag requires that the tag GENERATE_QHP is set to YES. @@ -1397,8 +1633,8 @@ QHP_NAMESPACE = org.doxygen.Project # The QHP_VIRTUAL_FOLDER tag specifies the namespace to use when generating Qt # Help Project output. For more information please see Qt Help Project / Virtual -# Folders (see: https://doc.qt.io/archives/qt-4.8/qthelpproject.html#virtual- -# folders). +# Folders (see: +# https://doc.qt.io/archives/qt-4.8/qthelpproject.html#virtual-folders). # The default value is: doc. # This tag requires that the tag GENERATE_QHP is set to YES. @@ -1406,16 +1642,16 @@ QHP_VIRTUAL_FOLDER = doc # If the QHP_CUST_FILTER_NAME tag is set, it specifies the name of a custom # filter to add. For more information please see Qt Help Project / Custom -# Filters (see: https://doc.qt.io/archives/qt-4.8/qthelpproject.html#custom- -# filters). +# Filters (see: +# https://doc.qt.io/archives/qt-4.8/qthelpproject.html#custom-filters). # This tag requires that the tag GENERATE_QHP is set to YES. QHP_CUST_FILTER_NAME = # The QHP_CUST_FILTER_ATTRS tag specifies the list of the attributes of the # custom filter to add. For more information please see Qt Help Project / Custom -# Filters (see: https://doc.qt.io/archives/qt-4.8/qthelpproject.html#custom- -# filters). +# Filters (see: +# https://doc.qt.io/archives/qt-4.8/qthelpproject.html#custom-filters). # This tag requires that the tag GENERATE_QHP is set to YES. QHP_CUST_FILTER_ATTRS = @@ -1427,9 +1663,9 @@ QHP_CUST_FILTER_ATTRS = QHP_SECT_FILTER_ATTRS = -# The QHG_LOCATION tag can be used to specify the location of Qt's -# qhelpgenerator. If non-empty doxygen will try to run qhelpgenerator on the -# generated .qhp file. +# The QHG_LOCATION tag can be used to specify the location (absolute path +# including file name) of Qt's qhelpgenerator. If non-empty Doxygen will try to +# run qhelpgenerator on the generated .qhp file. # This tag requires that the tag GENERATE_QHP is set to YES. QHG_LOCATION = @@ -1460,7 +1696,7 @@ ECLIPSE_DOC_ID = org.doxygen.Project # of each HTML page. A value of NO enables the index and the value YES disables # it. Since the tabs in the index contain the same information as the navigation # tree, you can set this option to YES if you also set GENERATE_TREEVIEW to YES. -# The default value is: NO. +# The default value is: YES. # This tag requires that the tag GENERATE_HTML is set to YES. DISABLE_INDEX = NO @@ -1472,18 +1708,30 @@ DISABLE_INDEX = NO # to work a browser that supports JavaScript, DHTML, CSS and frames is required # (i.e. any modern browser). Windows users are probably better off using the # HTML help feature. Via custom style sheets (see HTML_EXTRA_STYLESHEET) one can -# further fine-tune the look of the index. As an example, the default style -# sheet generated by doxygen has an example that shows how to put an image at -# the root of the tree instead of the PROJECT_NAME. Since the tree basically has -# the same information as the tab index, you could consider setting -# DISABLE_INDEX to YES when enabling this option. -# The default value is: NO. +# further fine tune the look of the index (see "Fine-tuning the output"). As an +# example, the default style sheet generated by Doxygen has an example that +# shows how to put an image at the root of the tree instead of the PROJECT_NAME. +# Since the tree basically has the same information as the tab index, you could +# consider setting DISABLE_INDEX to YES when enabling this option. +# The default value is: YES. # This tag requires that the tag GENERATE_HTML is set to YES. GENERATE_TREEVIEW = NO +# When both GENERATE_TREEVIEW and DISABLE_INDEX are set to YES, then the +# FULL_SIDEBAR option determines if the side bar is limited to only the treeview +# area (value NO) or if it should extend to the full height of the window (value +# YES). Setting this to YES gives a layout similar to +# https://docs.readthedocs.io with more room for contents, but less room for the +# project logo, title, and description. If either GENERATE_TREEVIEW or +# DISABLE_INDEX is set to NO, this option has no effect. +# The default value is: NO. +# This tag requires that the tag GENERATE_HTML is set to YES. + +FULL_SIDEBAR = NO + # The ENUM_VALUES_PER_LINE tag can be used to set the number of enum values that -# doxygen will group on one line in the generated HTML documentation. +# Doxygen will group on one line in the generated HTML documentation. # # Note that a value of 0 will completely suppress the enum values from appearing # in the overview section. @@ -1492,6 +1740,12 @@ GENERATE_TREEVIEW = NO ENUM_VALUES_PER_LINE = 4 +# When the SHOW_ENUM_VALUES tag is set doxygen will show the specified +# enumeration values besides the enumeration mnemonics. +# The default value is: NO. + +SHOW_ENUM_VALUES = NO + # If the treeview is enabled (see GENERATE_TREEVIEW) then this tag can be used # to set the initial width (in pixels) of the frame in which the tree is shown. # Minimum value: 0, maximum value: 1500, default value: 250. @@ -1499,14 +1753,21 @@ ENUM_VALUES_PER_LINE = 4 TREEVIEW_WIDTH = 250 -# If the EXT_LINKS_IN_WINDOW option is set to YES, doxygen will open links to +# If the EXT_LINKS_IN_WINDOW option is set to YES, Doxygen will open links to # external symbols imported via tag files in a separate window. # The default value is: NO. # This tag requires that the tag GENERATE_HTML is set to YES. EXT_LINKS_IN_WINDOW = NO -# If the HTML_FORMULA_FORMAT option is set to svg, doxygen will use the pdf2svg +# If the OBFUSCATE_EMAILS tag is set to YES, Doxygen will obfuscate email +# addresses. +# The default value is: YES. +# This tag requires that the tag GENERATE_HTML is set to YES. + +OBFUSCATE_EMAILS = YES + +# If the HTML_FORMULA_FORMAT option is set to svg, Doxygen will use the pdf2svg # tool (see https://github.com/dawbarton/pdf2svg) or inkscape (see # https://inkscape.org) to generate formulas as SVG images instead of PNGs for # the HTML output. These images will generally look nicer at scaled resolutions. @@ -1519,24 +1780,13 @@ HTML_FORMULA_FORMAT = png # Use this tag to change the font size of LaTeX formulas included as images in # the HTML documentation. When you change the font size after a successful -# doxygen run you need to manually remove any form_*.png images from the HTML +# Doxygen run you need to manually remove any form_*.png images from the HTML # output directory to force them to be regenerated. # Minimum value: 8, maximum value: 50, default value: 10. # This tag requires that the tag GENERATE_HTML is set to YES. FORMULA_FONTSIZE = 10 -# Use the FORMULA_TRANSPARENT tag to determine whether or not the images -# generated for formulas are transparent PNGs. Transparent PNGs are not -# supported properly for IE 6.0, but are supported on all modern browsers. -# -# Note that when changing this option you need to delete any form_*.png files in -# the HTML output directory before the changes have effect. -# The default value is: YES. -# This tag requires that the tag GENERATE_HTML is set to YES. - -FORMULA_TRANSPARENT = YES - # The FORMULA_MACROFILE can contain LaTeX \newcommand and \renewcommand commands # to create new LaTeX commands to be used in formulas as building blocks. See # the section "Including formulas" for details. @@ -1554,11 +1804,29 @@ FORMULA_MACROFILE = USE_MATHJAX = YES +# With MATHJAX_VERSION it is possible to specify the MathJax version to be used. +# Note that the different versions of MathJax have different requirements with +# regards to the different settings, so it is possible that also other MathJax +# settings have to be changed when switching between the different MathJax +# versions. +# Possible values are: MathJax_2 and MathJax_3. +# The default value is: MathJax_2. +# This tag requires that the tag USE_MATHJAX is set to YES. + +MATHJAX_VERSION = MathJax_2 + # When MathJax is enabled you can set the default output format to be used for -# the MathJax output. See the MathJax site (see: -# http://docs.mathjax.org/en/latest/output.html) for more details. +# the MathJax output. For more details about the output format see MathJax +# version 2 (see: +# http://docs.mathjax.org/en/v2.7-latest/output.html) and MathJax version 3 +# (see: +# http://docs.mathjax.org/en/latest/web/components/output.html). # Possible values are: HTML-CSS (which is slower, but has the best -# compatibility), NativeMML (i.e. MathML) and SVG. +# compatibility. This is the name for Mathjax version 2, for MathJax version 3 +# this will be translated into chtml), NativeMML (i.e. MathML. Only supported +# for MathJax 2. For MathJax version 3 chtml will be used instead.), chtml (This +# is the name for Mathjax version 3, for MathJax version 2 this will be +# translated into HTML-CSS) and SVG. # The default value is: HTML-CSS. # This tag requires that the tag USE_MATHJAX is set to YES. @@ -1571,33 +1839,40 @@ MATHJAX_FORMAT = HTML-CSS # MATHJAX_RELPATH should be ../mathjax. The default value points to the MathJax # Content Delivery Network so you can quickly see the result without installing # MathJax. However, it is strongly recommended to install a local copy of -# MathJax from https://www.mathjax.org before deployment. -# The default value is: https://cdn.jsdelivr.net/npm/mathjax@2. +# MathJax from https://www.mathjax.org before deployment. The default value is: +# - in case of MathJax version 2: https://cdn.jsdelivr.net/npm/mathjax@2 +# - in case of MathJax version 3: https://cdn.jsdelivr.net/npm/mathjax@3 # This tag requires that the tag USE_MATHJAX is set to YES. MATHJAX_RELPATH = http://cdn.mathjax.org/mathjax/latest # The MATHJAX_EXTENSIONS tag can be used to specify one or more MathJax # extension names that should be enabled during MathJax rendering. For example +# for MathJax version 2 (see +# https://docs.mathjax.org/en/v2.7-latest/tex.html#tex-and-latex-extensions): # MATHJAX_EXTENSIONS = TeX/AMSmath TeX/AMSsymbols +# For example for MathJax version 3 (see +# http://docs.mathjax.org/en/latest/input/tex/extensions/index.html): +# MATHJAX_EXTENSIONS = ams # This tag requires that the tag USE_MATHJAX is set to YES. MATHJAX_EXTENSIONS = -# The MATHJAX_CODEFILE tag can be used to specify a file with javascript pieces +# The MATHJAX_CODEFILE tag can be used to specify a file with JavaScript pieces # of code that will be used on startup of the MathJax code. See the MathJax site -# (see: http://docs.mathjax.org/en/latest/output.html) for more details. For an +# (see: +# http://docs.mathjax.org/en/v2.7-latest/output.html) for more details. For an # example see the documentation. # This tag requires that the tag USE_MATHJAX is set to YES. MATHJAX_CODEFILE = -# When the SEARCHENGINE tag is enabled doxygen will generate a search box for -# the HTML output. The underlying search engine uses javascript and DHTML and +# When the SEARCHENGINE tag is enabled Doxygen will generate a search box for +# the HTML output. The underlying search engine uses JavaScript and DHTML and # should work on any modern browser. Note that when using HTML help # (GENERATE_HTMLHELP), Qt help (GENERATE_QHP), or docsets (GENERATE_DOCSET) # there is already a search function so this one should typically be disabled. -# For large projects the javascript based search engine can be slow, then +# For large projects the JavaScript based search engine can be slow, then # enabling SERVER_BASED_SEARCH may provide a better solution. It is possible to # search using the keyboard; to jump to the search box use + S # (what the is depends on the OS and browser, but it is typically @@ -1616,7 +1891,7 @@ SEARCHENGINE = YES # When the SERVER_BASED_SEARCH tag is enabled the search engine will be # implemented using a web server instead of a web client using JavaScript. There # are two flavors of web server based searching depending on the EXTERNAL_SEARCH -# setting. When disabled, doxygen will generate a PHP script for searching and +# setting. When disabled, Doxygen will generate a PHP script for searching and # an index file used by the script. When EXTERNAL_SEARCH is enabled the indexing # and searching needs to be provided by external tools. See the section # "External Indexing and Searching" for details. @@ -1625,7 +1900,7 @@ SEARCHENGINE = YES SERVER_BASED_SEARCH = NO -# When EXTERNAL_SEARCH tag is enabled doxygen will no longer generate the PHP +# When EXTERNAL_SEARCH tag is enabled Doxygen will no longer generate the PHP # script for searching. Instead the search results are written to an XML file # which needs to be processed by an external indexer. Doxygen will invoke an # external search engine pointed to by the SEARCHENGINE_URL option to obtain the @@ -1633,7 +1908,8 @@ SERVER_BASED_SEARCH = NO # # Doxygen ships with an example indexer (doxyindexer) and search engine # (doxysearch.cgi) which are based on the open source search engine library -# Xapian (see: https://xapian.org/). +# Xapian (see: +# https://xapian.org/). # # See the section "External Indexing and Searching" for details. # The default value is: NO. @@ -1646,8 +1922,9 @@ EXTERNAL_SEARCH = NO # # Doxygen ships with an example indexer (doxyindexer) and search engine # (doxysearch.cgi) which are based on the open source search engine library -# Xapian (see: https://xapian.org/). See the section "External Indexing and -# Searching" for details. +# Xapian (see: +# https://xapian.org/). See the section "External Indexing and Searching" for +# details. # This tag requires that the tag SEARCHENGINE is set to YES. SEARCHENGINE_URL = @@ -1668,7 +1945,7 @@ SEARCHDATA_FILE = searchdata.xml EXTERNAL_SEARCH_ID = -# The EXTRA_SEARCH_MAPPINGS tag can be used to enable searching through doxygen +# The EXTRA_SEARCH_MAPPINGS tag can be used to enable searching through Doxygen # projects other than the one defined by this configuration file, but that are # all added to the same external search index. Each project needs to have a # unique id set via EXTERNAL_SEARCH_ID. The search mapping then maps the id of @@ -1682,7 +1959,7 @@ EXTRA_SEARCH_MAPPINGS = # Configuration options related to the LaTeX output #--------------------------------------------------------------------------- -# If the GENERATE_LATEX tag is set to YES, doxygen will generate LaTeX output. +# If the GENERATE_LATEX tag is set to YES, Doxygen will generate LaTeX output. # The default value is: YES. GENERATE_LATEX = NO @@ -1727,7 +2004,7 @@ MAKEINDEX_CMD_NAME = makeindex LATEX_MAKEINDEX_CMD = makeindex -# If the COMPACT_LATEX tag is set to YES, doxygen generates more compact LaTeX +# If the COMPACT_LATEX tag is set to YES, Doxygen generates more compact LaTeX # documents. This may be useful for small projects and may help to save some # trees in general. # The default value is: NO. @@ -1756,36 +2033,38 @@ PAPER_TYPE = a4 EXTRA_PACKAGES = -# The LATEX_HEADER tag can be used to specify a personal LaTeX header for the -# generated LaTeX document. The header should contain everything until the first -# chapter. If it is left blank doxygen will generate a standard header. See -# section "Doxygen usage" for information on how to let doxygen write the -# default header to a separate file. +# The LATEX_HEADER tag can be used to specify a user-defined LaTeX header for +# the generated LaTeX document. The header should contain everything until the +# first chapter. If it is left blank Doxygen will generate a standard header. It +# is highly recommended to start with a default header using +# doxygen -w latex new_header.tex new_footer.tex new_stylesheet.sty +# and then modify the file new_header.tex. See also section "Doxygen usage" for +# information on how to generate the default header that Doxygen normally uses. # -# Note: Only use a user-defined header if you know what you are doing! The -# following commands have a special meaning inside the header: $title, -# $datetime, $date, $doxygenversion, $projectname, $projectnumber, -# $projectbrief, $projectlogo. Doxygen will replace $title with the empty -# string, for the replacement values of the other commands the user is referred -# to HTML_HEADER. +# Note: Only use a user-defined header if you know what you are doing! +# Note: The header is subject to change so you typically have to regenerate the +# default header when upgrading to a newer version of Doxygen. The following +# commands have a special meaning inside the header (and footer): For a +# description of the possible markers and block names see the documentation. # This tag requires that the tag GENERATE_LATEX is set to YES. LATEX_HEADER = -# The LATEX_FOOTER tag can be used to specify a personal LaTeX footer for the -# generated LaTeX document. The footer should contain everything after the last -# chapter. If it is left blank doxygen will generate a standard footer. See +# The LATEX_FOOTER tag can be used to specify a user-defined LaTeX footer for +# the generated LaTeX document. The footer should contain everything after the +# last chapter. If it is left blank Doxygen will generate a standard footer. See # LATEX_HEADER for more information on how to generate a default footer and what -# special commands can be used inside the footer. -# -# Note: Only use a user-defined footer if you know what you are doing! +# special commands can be used inside the footer. See also section "Doxygen +# usage" for information on how to generate the default footer that Doxygen +# normally uses. Note: Only use a user-defined footer if you know what you are +# doing! # This tag requires that the tag GENERATE_LATEX is set to YES. LATEX_FOOTER = # The LATEX_EXTRA_STYLESHEET tag can be used to specify additional user-defined # LaTeX style sheets that are included after the standard style sheets created -# by doxygen. Using this option one can overrule certain style aspects. Doxygen +# by Doxygen. Using this option one can overrule certain style aspects. Doxygen # will copy the style sheet files to the output directory. # Note: The order of the extra style sheet files is of importance (e.g. the last # style sheet in the list overrules the setting of the previous ones in the @@ -1811,7 +2090,7 @@ LATEX_EXTRA_FILES = PDF_HYPERLINKS = YES -# If the USE_PDFLATEX tag is set to YES, doxygen will use the engine as +# If the USE_PDFLATEX tag is set to YES, Doxygen will use the engine as # specified with LATEX_CMD_NAME to generate the PDF file directly from the LaTeX # files. Set this option to YES, to get a higher quality PDF documentation. # @@ -1821,48 +2100,36 @@ PDF_HYPERLINKS = YES USE_PDFLATEX = YES -# If the LATEX_BATCHMODE tag is set to YES, doxygen will add the \batchmode -# command to the generated LaTeX files. This will instruct LaTeX to keep running -# if errors occur, instead of asking the user for help. This option is also used -# when generating formulas in HTML. +# The LATEX_BATCHMODE tag signals the behavior of LaTeX in case of an error. +# Possible values are: NO same as ERROR_STOP, YES same as BATCH, BATCH In batch +# mode nothing is printed on the terminal, errors are scrolled as if is +# hit at every error; missing files that TeX tries to input or request from +# keyboard input (\read on a not open input stream) cause the job to abort, +# NON_STOP In nonstop mode the diagnostic message will appear on the terminal, +# but there is no possibility of user interaction just like in batch mode, +# SCROLL In scroll mode, TeX will stop only for missing files to input or if +# keyboard input is necessary and ERROR_STOP In errorstop mode, TeX will stop at +# each error, asking for user intervention. # The default value is: NO. # This tag requires that the tag GENERATE_LATEX is set to YES. LATEX_BATCHMODE = NO -# If the LATEX_HIDE_INDICES tag is set to YES then doxygen will not include the +# If the LATEX_HIDE_INDICES tag is set to YES then Doxygen will not include the # index chapters (such as File Index, Compound Index, etc.) in the output. # The default value is: NO. # This tag requires that the tag GENERATE_LATEX is set to YES. LATEX_HIDE_INDICES = NO -# If the LATEX_SOURCE_CODE tag is set to YES then doxygen will include source -# code with syntax highlighting in the LaTeX output. -# -# Note that which sources are shown also depends on other settings such as -# SOURCE_BROWSER. -# The default value is: NO. -# This tag requires that the tag GENERATE_LATEX is set to YES. - -LATEX_SOURCE_CODE = NO - # The LATEX_BIB_STYLE tag can be used to specify the style to use for the # bibliography, e.g. plainnat, or ieeetr. See # https://en.wikipedia.org/wiki/BibTeX and \cite for more info. -# The default value is: plain. +# The default value is: plainnat. # This tag requires that the tag GENERATE_LATEX is set to YES. LATEX_BIB_STYLE = plain -# If the LATEX_TIMESTAMP tag is set to YES then the footer of each generated -# page will contain the date and time when the page was generated. Setting this -# to NO can help when comparing the output of multiple runs. -# The default value is: NO. -# This tag requires that the tag GENERATE_LATEX is set to YES. - -LATEX_TIMESTAMP = NO - # The LATEX_EMOJI_DIRECTORY tag is used to specify the (relative or absolute) # path from which the emoji images will be read. If a relative path is entered, # it will be relative to the LATEX_OUTPUT directory. If left blank the @@ -1875,7 +2142,7 @@ LATEX_EMOJI_DIRECTORY = # Configuration options related to the RTF output #--------------------------------------------------------------------------- -# If the GENERATE_RTF tag is set to YES, doxygen will generate RTF output. The +# If the GENERATE_RTF tag is set to YES, Doxygen will generate RTF output. The # RTF output is optimized for Word 97 and may not look too pretty with other RTF # readers/editors. # The default value is: NO. @@ -1890,7 +2157,7 @@ GENERATE_RTF = NO RTF_OUTPUT = rtf -# If the COMPACT_RTF tag is set to YES, doxygen generates more compact RTF +# If the COMPACT_RTF tag is set to YES, Doxygen generates more compact RTF # documents. This may be useful for small projects and may help to save some # trees in general. # The default value is: NO. @@ -1910,38 +2177,36 @@ COMPACT_RTF = NO RTF_HYPERLINKS = NO -# Load stylesheet definitions from file. Syntax is similar to doxygen's +# Load stylesheet definitions from file. Syntax is similar to Doxygen's # configuration file, i.e. a series of assignments. You only have to provide # replacements, missing definitions are set to their default value. # # See also section "Doxygen usage" for information on how to generate the -# default style sheet that doxygen normally uses. +# default style sheet that Doxygen normally uses. # This tag requires that the tag GENERATE_RTF is set to YES. RTF_STYLESHEET_FILE = # Set optional variables used in the generation of an RTF document. Syntax is -# similar to doxygen's configuration file. A template extensions file can be +# similar to Doxygen's configuration file. A template extensions file can be # generated using doxygen -e rtf extensionFile. # This tag requires that the tag GENERATE_RTF is set to YES. RTF_EXTENSIONS_FILE = -# If the RTF_SOURCE_CODE tag is set to YES then doxygen will include source code -# with syntax highlighting in the RTF output. -# -# Note that which sources are shown also depends on other settings such as -# SOURCE_BROWSER. -# The default value is: NO. +# The RTF_EXTRA_FILES tag can be used to specify one or more extra images or +# other source files which should be copied to the RTF_OUTPUT output directory. +# Note that the files will be copied as-is; there are no commands or markers +# available. # This tag requires that the tag GENERATE_RTF is set to YES. -RTF_SOURCE_CODE = NO +RTF_EXTRA_FILES = #--------------------------------------------------------------------------- # Configuration options related to the man page output #--------------------------------------------------------------------------- -# If the GENERATE_MAN tag is set to YES, doxygen will generate man pages for +# If the GENERATE_MAN tag is set to YES, Doxygen will generate man pages for # classes and files. # The default value is: NO. @@ -1972,7 +2237,7 @@ MAN_EXTENSION = .3 MAN_SUBDIR = -# If the MAN_LINKS tag is set to YES and doxygen generates man output, then it +# If the MAN_LINKS tag is set to YES and Doxygen generates man output, then it # will generate one additional man file for each entity documented in the real # man page(s). These additional files only source the real man page, but without # them the man command would be unable to find the correct page. @@ -1985,7 +2250,7 @@ MAN_LINKS = NO # Configuration options related to the XML output #--------------------------------------------------------------------------- -# If the GENERATE_XML tag is set to YES, doxygen will generate an XML file that +# If the GENERATE_XML tag is set to YES, Doxygen will generate an XML file that # captures the structure of the code including all documentation. # The default value is: NO. @@ -1999,7 +2264,7 @@ GENERATE_XML = YES XML_OUTPUT = _xml -# If the XML_PROGRAMLISTING tag is set to YES, doxygen will dump the program +# If the XML_PROGRAMLISTING tag is set to YES, Doxygen will dump the program # listings (including syntax highlighting and cross-referencing information) to # the XML output. Note that enabling this will significantly increase the size # of the XML output. @@ -2008,7 +2273,7 @@ XML_OUTPUT = _xml XML_PROGRAMLISTING = YES -# If the XML_NS_MEMB_FILE_SCOPE tag is set to YES, doxygen will include +# If the XML_NS_MEMB_FILE_SCOPE tag is set to YES, Doxygen will include # namespace members in file scope as well, matching the HTML output. # The default value is: NO. # This tag requires that the tag GENERATE_XML is set to YES. @@ -2019,7 +2284,7 @@ XML_NS_MEMB_FILE_SCOPE = NO # Configuration options related to the DOCBOOK output #--------------------------------------------------------------------------- -# If the GENERATE_DOCBOOK tag is set to YES, doxygen will generate Docbook files +# If the GENERATE_DOCBOOK tag is set to YES, Doxygen will generate Docbook files # that can be used to generate PDF. # The default value is: NO. @@ -2033,32 +2298,49 @@ GENERATE_DOCBOOK = NO DOCBOOK_OUTPUT = docbook -# If the DOCBOOK_PROGRAMLISTING tag is set to YES, doxygen will include the -# program listings (including syntax highlighting and cross-referencing -# information) to the DOCBOOK output. Note that enabling this will significantly -# increase the size of the DOCBOOK output. -# The default value is: NO. -# This tag requires that the tag GENERATE_DOCBOOK is set to YES. - -DOCBOOK_PROGRAMLISTING = NO - #--------------------------------------------------------------------------- # Configuration options for the AutoGen Definitions output #--------------------------------------------------------------------------- -# If the GENERATE_AUTOGEN_DEF tag is set to YES, doxygen will generate an -# AutoGen Definitions (see http://autogen.sourceforge.net/) file that captures +# If the GENERATE_AUTOGEN_DEF tag is set to YES, Doxygen will generate an +# AutoGen Definitions (see https://autogen.sourceforge.net/) file that captures # the structure of the code including all documentation. Note that this feature # is still experimental and incomplete at the moment. # The default value is: NO. GENERATE_AUTOGEN_DEF = NO +#--------------------------------------------------------------------------- +# Configuration options related to Sqlite3 output +#--------------------------------------------------------------------------- + +# If the GENERATE_SQLITE3 tag is set to YES Doxygen will generate a Sqlite3 +# database with symbols found by Doxygen stored in tables. +# The default value is: NO. + +GENERATE_SQLITE3 = NO + +# The SQLITE3_OUTPUT tag is used to specify where the Sqlite3 database will be +# put. If a relative path is entered the value of OUTPUT_DIRECTORY will be put +# in front of it. +# The default directory is: sqlite3. +# This tag requires that the tag GENERATE_SQLITE3 is set to YES. + +SQLITE3_OUTPUT = sqlite3 + +# The SQLITE3_RECREATE_DB tag is set to YES, the existing doxygen_sqlite3.db +# database file will be recreated with each Doxygen run. If set to NO, Doxygen +# will warn if a database file is already found and not modify it. +# The default value is: YES. +# This tag requires that the tag GENERATE_SQLITE3 is set to YES. + +SQLITE3_RECREATE_DB = YES + #--------------------------------------------------------------------------- # Configuration options related to the Perl module output #--------------------------------------------------------------------------- -# If the GENERATE_PERLMOD tag is set to YES, doxygen will generate a Perl module +# If the GENERATE_PERLMOD tag is set to YES, Doxygen will generate a Perl module # file that captures the structure of the code including all documentation. # # Note that this feature is still experimental and incomplete at the moment. @@ -2066,7 +2348,7 @@ GENERATE_AUTOGEN_DEF = NO GENERATE_PERLMOD = NO -# If the PERLMOD_LATEX tag is set to YES, doxygen will generate the necessary +# If the PERLMOD_LATEX tag is set to YES, Doxygen will generate the necessary # Makefile rules, Perl scripts and LaTeX code to be able to generate PDF and DVI # output from the Perl module output. # The default value is: NO. @@ -2096,13 +2378,13 @@ PERLMOD_MAKEVAR_PREFIX = # Configuration options related to the preprocessor #--------------------------------------------------------------------------- -# If the ENABLE_PREPROCESSING tag is set to YES, doxygen will evaluate all +# If the ENABLE_PREPROCESSING tag is set to YES, Doxygen will evaluate all # C-preprocessor directives found in the sources and include files. # The default value is: YES. ENABLE_PREPROCESSING = YES -# If the MACRO_EXPANSION tag is set to YES, doxygen will expand all macro names +# If the MACRO_EXPANSION tag is set to YES, Doxygen will expand all macro names # in the source code. If set to NO, only conditional compilation will be # performed. Macro expansion can be done in a controlled way by setting # EXPAND_ONLY_PREDEF to YES. @@ -2128,7 +2410,8 @@ SEARCH_INCLUDES = YES # The INCLUDE_PATH tag can be used to specify one or more directories that # contain include files that are not input files but should be processed by the -# preprocessor. +# preprocessor. Note that the INCLUDE_PATH is not recursive, so the setting of +# RECURSIVE has no effect here. # This tag requires that the tag SEARCH_INCLUDES is set to YES. INCLUDE_PATH = @@ -2149,7 +2432,8 @@ INCLUDE_FILE_PATTERNS = # recursively expanded use the := operator instead of the = operator. # This tag requires that the tag ENABLE_PREPROCESSING is set to YES. -PREDEFINED = "CUVS_BUILD_MG_ALGOS=1" +PREDEFINED = "CUVS_BUILD_MG_ALGOS=1" \ + "CUVS_BUILD_CAGRA_HNSWLIB=1" # If the MACRO_EXPANSION and EXPAND_ONLY_PREDEF tags are set to YES then this # tag can be used to specify a list of macro names that should be expanded. The @@ -2160,7 +2444,7 @@ PREDEFINED = "CUVS_BUILD_MG_ALGOS=1" EXPAND_AS_DEFINED = -# If the SKIP_FUNCTION_MACROS tag is set to YES then doxygen's preprocessor will +# If the SKIP_FUNCTION_MACROS tag is set to YES then Doxygen's preprocessor will # remove all references to function-like macros that are alone on a line, have # an all uppercase name, and do not end with a semicolon. Such function macros # are typically used for boiler-plate code, and will confuse the parser if not @@ -2184,26 +2468,26 @@ SKIP_FUNCTION_MACROS = YES # section "Linking to external documentation" for more information about the use # of tag files. # Note: Each tag file must have a unique name (where the name does NOT include -# the path). If a tag file is not located in the directory in which doxygen is +# the path). If a tag file is not located in the directory in which Doxygen is # run, you must also specify the path to the tagfile here. TAGFILES = -# When a file name is specified after GENERATE_TAGFILE, doxygen will create a +# When a file name is specified after GENERATE_TAGFILE, Doxygen will create a # tag file that is based on the input files it reads. See section "Linking to # external documentation" for more information about the usage of tag files. GENERATE_TAGFILE = -# If the ALLEXTERNALS tag is set to YES, all external class will be listed in -# the class index. If set to NO, only the inherited external classes will be -# listed. +# If the ALLEXTERNALS tag is set to YES, all external classes and namespaces +# will be listed in the class and namespace index. If set to NO, only the +# inherited external classes will be listed. # The default value is: NO. ALLEXTERNALS = NO # If the EXTERNAL_GROUPS tag is set to YES, all external groups will be listed -# in the modules index. If set to NO, only the current project's groups will be +# in the topic index. If set to NO, only the current project's groups will be # listed. # The default value is: YES. @@ -2217,42 +2501,26 @@ EXTERNAL_GROUPS = YES EXTERNAL_PAGES = YES #--------------------------------------------------------------------------- -# Configuration options related to the dot tool +# Configuration options related to diagram generator tools #--------------------------------------------------------------------------- -# If the CLASS_DIAGRAMS tag is set to YES, doxygen will generate a class diagram -# (in HTML and LaTeX) for classes with base or super classes. Setting the tag to -# NO turns the diagrams off. Note that this option also works with HAVE_DOT -# disabled, but it is recommended to install and use dot, since it yields more -# powerful graphs. -# The default value is: YES. - -CLASS_DIAGRAMS = YES - -# You can include diagrams made with dia in doxygen documentation. Doxygen will -# then run dia to produce the diagram and insert it in the documentation. The -# DIA_PATH tag allows you to specify the directory where the dia binary resides. -# If left empty dia is assumed to be found in the default search path. - -DIA_PATH = - # If set to YES the inheritance and collaboration graphs will hide inheritance # and usage relations if the target is undocumented or is not a class. # The default value is: YES. HIDE_UNDOC_RELATIONS = YES -# If you set the HAVE_DOT tag to YES then doxygen will assume the dot tool is +# If you set the HAVE_DOT tag to YES then Doxygen will assume the dot tool is # available from the path. This tool is part of Graphviz (see: -# http://www.graphviz.org/), a graph visualization toolkit from AT&T and Lucent +# https://www.graphviz.org/), a graph visualization toolkit from AT&T and Lucent # Bell Labs. The other options in this section have no effect if this option is # set to NO # The default value is: NO. -HAVE_DOT = +HAVE_DOT = NO -# The DOT_NUM_THREADS specifies the number of dot invocations doxygen is allowed -# to run in parallel. When set to 0 doxygen will base this on the number of +# The DOT_NUM_THREADS specifies the number of dot invocations Doxygen is allowed +# to run in parallel. When set to 0 Doxygen will base this on the number of # processors available in the system. You can set it explicitly to a value # larger than 0 to get control over the balance between CPU load and processing # speed. @@ -2261,55 +2529,83 @@ HAVE_DOT = DOT_NUM_THREADS = 0 -# When you want a differently looking font in the dot files that doxygen -# generates you can specify the font name using DOT_FONTNAME. You need to make -# sure dot is able to find the font, which can be done by putting it in a -# standard location or by setting the DOTFONTPATH environment variable or by -# setting DOT_FONTPATH to the directory containing the font. -# The default value is: Helvetica. +# DOT_COMMON_ATTR is common attributes for nodes, edges and labels of +# subgraphs. When you want a differently looking font in the dot files that +# Doxygen generates you can specify fontname, fontcolor and fontsize attributes. +# For details please see Node, +# Edge and Graph Attributes specification You need to make sure dot is able +# to find the font, which can be done by putting it in a standard location or by +# setting the DOTFONTPATH environment variable or by setting DOT_FONTPATH to the +# directory containing the font. Default graphviz fontsize is 14. +# The default value is: fontname=Helvetica,fontsize=10. +# This tag requires that the tag HAVE_DOT is set to YES. + +DOT_COMMON_ATTR = "fontname=Helvetica,fontsize=10" + +# DOT_EDGE_ATTR is concatenated with DOT_COMMON_ATTR. For elegant style you can +# add 'arrowhead=open, arrowtail=open, arrowsize=0.5'. Complete documentation about +# arrows shapes. +# The default value is: labelfontname=Helvetica,labelfontsize=10. # This tag requires that the tag HAVE_DOT is set to YES. -DOT_FONTNAME = Helvetica +DOT_EDGE_ATTR = "labelfontname=Helvetica,labelfontsize=10" -# The DOT_FONTSIZE tag can be used to set the size (in points) of the font of -# dot graphs. -# Minimum value: 4, maximum value: 24, default value: 10. +# DOT_NODE_ATTR is concatenated with DOT_COMMON_ATTR. For view without boxes +# around nodes set 'shape=plain' or 'shape=plaintext' Shapes specification +# The default value is: shape=box,height=0.2,width=0.4. # This tag requires that the tag HAVE_DOT is set to YES. -DOT_FONTSIZE = 10 +DOT_NODE_ATTR = "shape=box,height=0.2,width=0.4" -# By default doxygen will tell dot to use the default font as specified with -# DOT_FONTNAME. If you specify a different font using DOT_FONTNAME you can set -# the path where dot can find it using this tag. +# You can set the path where dot can find font specified with fontname in +# DOT_COMMON_ATTR and others dot attributes. # This tag requires that the tag HAVE_DOT is set to YES. DOT_FONTPATH = -# If the CLASS_GRAPH tag is set to YES then doxygen will generate a graph for -# each documented class showing the direct and indirect inheritance relations. -# Setting this tag to YES will force the CLASS_DIAGRAMS tag to NO. +# If the CLASS_GRAPH tag is set to YES or GRAPH or BUILTIN then Doxygen will +# generate a graph for each documented class showing the direct and indirect +# inheritance relations. In case the CLASS_GRAPH tag is set to YES or GRAPH and +# HAVE_DOT is enabled as well, then dot will be used to draw the graph. In case +# the CLASS_GRAPH tag is set to YES and HAVE_DOT is disabled or if the +# CLASS_GRAPH tag is set to BUILTIN, then the built-in generator will be used. +# If the CLASS_GRAPH tag is set to TEXT the direct and indirect inheritance +# relations will be shown as texts / links. Explicit enabling an inheritance +# graph or choosing a different representation for an inheritance graph of a +# specific class, can be accomplished by means of the command \inheritancegraph. +# Disabling an inheritance graph can be accomplished by means of the command +# \hideinheritancegraph. +# Possible values are: NO, YES, TEXT, GRAPH and BUILTIN. # The default value is: YES. -# This tag requires that the tag HAVE_DOT is set to YES. CLASS_GRAPH = YES -# If the COLLABORATION_GRAPH tag is set to YES then doxygen will generate a +# If the COLLABORATION_GRAPH tag is set to YES then Doxygen will generate a # graph for each documented class showing the direct and indirect implementation # dependencies (inheritance, containment, and class references variables) of the -# class with other documented classes. +# class with other documented classes. Explicit enabling a collaboration graph, +# when COLLABORATION_GRAPH is set to NO, can be accomplished by means of the +# command \collaborationgraph. Disabling a collaboration graph can be +# accomplished by means of the command \hidecollaborationgraph. # The default value is: YES. # This tag requires that the tag HAVE_DOT is set to YES. COLLABORATION_GRAPH = YES -# If the GROUP_GRAPHS tag is set to YES then doxygen will generate a graph for -# groups, showing the direct groups dependencies. +# If the GROUP_GRAPHS tag is set to YES then Doxygen will generate a graph for +# groups, showing the direct groups dependencies. Explicit enabling a group +# dependency graph, when GROUP_GRAPHS is set to NO, can be accomplished by means +# of the command \groupgraph. Disabling a directory graph can be accomplished by +# means of the command \hidegroupgraph. See also the chapter Grouping in the +# manual. # The default value is: YES. # This tag requires that the tag HAVE_DOT is set to YES. GROUP_GRAPHS = YES -# If the UML_LOOK tag is set to YES, doxygen will generate inheritance and +# If the UML_LOOK tag is set to YES, Doxygen will generate inheritance and # collaboration diagrams in a style similar to the OMG's Unified Modeling # Language. # The default value is: NO. @@ -2326,10 +2622,32 @@ UML_LOOK = YES # but if the number exceeds 15, the total amount of fields shown is limited to # 10. # Minimum value: 0, maximum value: 100, default value: 10. -# This tag requires that the tag HAVE_DOT is set to YES. +# This tag requires that the tag UML_LOOK is set to YES. UML_LIMIT_NUM_FIELDS = 10 +# If the DOT_UML_DETAILS tag is set to NO, Doxygen will show attributes and +# methods without types and arguments in the UML graphs. If the DOT_UML_DETAILS +# tag is set to YES, Doxygen will add type and arguments for attributes and +# methods in the UML graphs. If the DOT_UML_DETAILS tag is set to NONE, Doxygen +# will not generate fields with class member information in the UML graphs. The +# class diagrams will look similar to the default class diagrams but using UML +# notation for the relationships. +# Possible values are: NO, YES and NONE. +# The default value is: NO. +# This tag requires that the tag UML_LOOK is set to YES. + +DOT_UML_DETAILS = NO + +# The DOT_WRAP_THRESHOLD tag can be used to set the maximum number of characters +# to display on a single line. If the actual line length exceeds this threshold +# significantly it will be wrapped across multiple lines. Some heuristics are +# applied to avoid ugly line breaks. +# Minimum value: 0, maximum value: 1000, default value: 17. +# This tag requires that the tag HAVE_DOT is set to YES. + +DOT_WRAP_THRESHOLD = 17 + # If the TEMPLATE_RELATIONS tag is set to YES then the inheritance and # collaboration graphs will show the relations between templates and their # instances. @@ -2339,24 +2657,29 @@ UML_LIMIT_NUM_FIELDS = 10 TEMPLATE_RELATIONS = NO # If the INCLUDE_GRAPH, ENABLE_PREPROCESSING and SEARCH_INCLUDES tags are set to -# YES then doxygen will generate a graph for each documented file showing the +# YES then Doxygen will generate a graph for each documented file showing the # direct and indirect include dependencies of the file with other documented -# files. +# files. Explicit enabling an include graph, when INCLUDE_GRAPH is is set to NO, +# can be accomplished by means of the command \includegraph. Disabling an +# include graph can be accomplished by means of the command \hideincludegraph. # The default value is: YES. # This tag requires that the tag HAVE_DOT is set to YES. INCLUDE_GRAPH = YES # If the INCLUDED_BY_GRAPH, ENABLE_PREPROCESSING and SEARCH_INCLUDES tags are -# set to YES then doxygen will generate a graph for each documented file showing +# set to YES then Doxygen will generate a graph for each documented file showing # the direct and indirect include dependencies of the file with other documented -# files. +# files. Explicit enabling an included by graph, when INCLUDED_BY_GRAPH is set +# to NO, can be accomplished by means of the command \includedbygraph. Disabling +# an included by graph can be accomplished by means of the command +# \hideincludedbygraph. # The default value is: YES. # This tag requires that the tag HAVE_DOT is set to YES. INCLUDED_BY_GRAPH = YES -# If the CALL_GRAPH tag is set to YES then doxygen will generate a call +# If the CALL_GRAPH tag is set to YES then Doxygen will generate a call # dependency graph for every global function or class method. # # Note that enabling this option will significantly increase the time of a run. @@ -2368,7 +2691,7 @@ INCLUDED_BY_GRAPH = YES CALL_GRAPH = NO -# If the CALLER_GRAPH tag is set to YES then doxygen will generate a caller +# If the CALLER_GRAPH tag is set to YES then Doxygen will generate a caller # dependency graph for every global function or class method. # # Note that enabling this option will significantly increase the time of a run. @@ -2380,44 +2703,59 @@ CALL_GRAPH = NO CALLER_GRAPH = NO -# If the GRAPHICAL_HIERARCHY tag is set to YES then doxygen will graphical +# If the GRAPHICAL_HIERARCHY tag is set to YES then Doxygen will graphical # hierarchy of all classes instead of a textual one. # The default value is: YES. # This tag requires that the tag HAVE_DOT is set to YES. GRAPHICAL_HIERARCHY = YES -# If the DIRECTORY_GRAPH tag is set to YES then doxygen will show the +# If the DIRECTORY_GRAPH tag is set to YES then Doxygen will show the # dependencies a directory has on other directories in a graphical way. The # dependency relations are determined by the #include relations between the -# files in the directories. +# files in the directories. Explicit enabling a directory graph, when +# DIRECTORY_GRAPH is set to NO, can be accomplished by means of the command +# \directorygraph. Disabling a directory graph can be accomplished by means of +# the command \hidedirectorygraph. # The default value is: YES. # This tag requires that the tag HAVE_DOT is set to YES. DIRECTORY_GRAPH = YES +# The DIR_GRAPH_MAX_DEPTH tag can be used to limit the maximum number of levels +# of child directories generated in directory dependency graphs by dot. +# Minimum value: 1, maximum value: 25, default value: 1. +# This tag requires that the tag DIRECTORY_GRAPH is set to YES. + +DIR_GRAPH_MAX_DEPTH = 1 + # The DOT_IMAGE_FORMAT tag can be used to set the image format of the images # generated by dot. For an explanation of the image formats see the section # output formats in the documentation of the dot tool (Graphviz (see: -# http://www.graphviz.org/)). -# Note: If you choose svg you need to set HTML_FILE_EXTENSION to xhtml in order -# to make the SVG files visible in IE 9+ (other browsers do not have this -# requirement). +# https://www.graphviz.org/)). +# +# Note the formats svg:cairo and svg:cairo:cairo cannot be used in combination +# with INTERACTIVE_SVG (the INTERACTIVE_SVG will be set to NO). # Possible values are: png, jpg, gif, svg, png:gd, png:gd:gd, png:cairo, -# png:cairo:gd, png:cairo:cairo, png:cairo:gdiplus, png:gdiplus and -# png:gdiplus:gdiplus. +# png:cairo:gd, png:cairo:cairo, png:cairo:gdiplus, png:gdiplus, +# png:gdiplus:gdiplus, svg:cairo, svg:cairo:cairo, svg:svg, svg:svg:core, +# gif:cairo, gif:cairo:gd, gif:cairo:gdiplus, gif:gdiplus, gif:gdiplus:gdiplus, +# gif:gd, gif:gd:gd, jpg:cairo, jpg:cairo:gd, jpg:cairo:gdiplus, jpg:gd, +# jpg:gd:gd, jpg:gdiplus and jpg:gdiplus:gdiplus. # The default value is: png. # This tag requires that the tag HAVE_DOT is set to YES. DOT_IMAGE_FORMAT = png -# If DOT_IMAGE_FORMAT is set to svg, then this option can be set to YES to -# enable generation of interactive SVG images that allow zooming and panning. +# If DOT_IMAGE_FORMAT is set to svg or svg:svg or svg:svg:core, then this option +# can be set to YES to enable generation of interactive SVG images that allow +# zooming and panning. # # Note that this requires a modern browser other than Internet Explorer. Tested # and working are Firefox, Chrome, Safari, and Opera. -# Note: For IE 9+ you need to set HTML_FILE_EXTENSION to xhtml in order to make -# the SVG files visible. Older versions of IE do not have SVG support. +# +# Note This option will be automatically disabled when DOT_IMAGE_FORMAT is set +# to svg:cairo or svg:cairo:cairo. # The default value is: NO. # This tag requires that the tag HAVE_DOT is set to YES. @@ -2436,11 +2774,12 @@ DOT_PATH = DOTFILE_DIRS = -# The MSCFILE_DIRS tag can be used to specify one or more directories that -# contain msc files that are included in the documentation (see the \mscfile -# command). +# You can include diagrams made with dia in Doxygen documentation. Doxygen will +# then run dia to produce the diagram and insert it in the documentation. The +# DIA_PATH tag allows you to specify the directory where the dia binary resides. +# If left empty dia is assumed to be found in the default search path. -MSCFILE_DIRS = +DIA_PATH = # The DIAFILE_DIRS tag can be used to specify one or more directories that # contain dia files that are included in the documentation (see the \diafile @@ -2448,28 +2787,34 @@ MSCFILE_DIRS = DIAFILE_DIRS = -# When using plantuml, the PLANTUML_JAR_PATH tag should be used to specify the -# path where java can find the plantuml.jar file. If left blank, it is assumed -# PlantUML is not used or called during a preprocessing step. Doxygen will -# generate a warning when it encounters a \startuml command in this case and -# will not generate output for the diagram. +# When using PlantUML, the PLANTUML_JAR_PATH tag should be used to specify the +# path where java can find the plantuml.jar file or to the filename of jar file +# to be used. If left blank, it is assumed PlantUML is not used or called during +# a preprocessing step. Doxygen will generate a warning when it encounters a +# \startuml command in this case and will not generate output for the diagram. PLANTUML_JAR_PATH = -# When using plantuml, the PLANTUML_CFG_FILE tag can be used to specify a -# configuration file for plantuml. +# When using PlantUML, the PLANTUML_CFG_FILE tag can be used to specify a +# configuration file for PlantUML. PLANTUML_CFG_FILE = -# When using plantuml, the specified paths are searched for files specified by -# the !include statement in a plantuml block. +# When using PlantUML, the specified paths are searched for files specified by +# the !include statement in a PlantUML block. PLANTUML_INCLUDE_PATH = +# The PLANTUMLFILE_DIRS tag can be used to specify one or more directories that +# contain PlantUml files that are included in the documentation (see the +# \plantumlfile command). + +PLANTUMLFILE_DIRS = + # The DOT_GRAPH_MAX_NODES tag can be used to set the maximum number of nodes # that will be shown in the graph. If the number of nodes in a graph becomes -# larger than this value, doxygen will truncate the graph, which is visualized -# by representing a node as a red box. Note that doxygen if the number of direct +# larger than this value, Doxygen will truncate the graph, which is visualized +# by representing a node as a red box. Note that if the number of direct # children of the root node in a graph is already larger than # DOT_GRAPH_MAX_NODES then the graph will not be shown at all. Also note that # the size of a graph can be further restricted by MAX_DOT_GRAPH_DEPTH. @@ -2490,18 +2835,6 @@ DOT_GRAPH_MAX_NODES = 100 MAX_DOT_GRAPH_DEPTH = 0 -# Set the DOT_TRANSPARENT tag to YES to generate images with a transparent -# background. This is disabled by default, because dot on Windows does not seem -# to support this out of the box. -# -# Warning: Depending on the platform used, enabling this option may lead to -# badly anti-aliased labels on the edges of a graph (i.e. they become hard to -# read). -# The default value is: NO. -# This tag requires that the tag HAVE_DOT is set to YES. - -DOT_TRANSPARENT = NO - # Set the DOT_MULTI_TARGETS tag to YES to allow dot to generate multiple output # files in one run (i.e. multiple -o and -T options on the command line). This # makes dot run faster, but since only newer versions of dot (>1.8.10) support @@ -2511,17 +2844,37 @@ DOT_TRANSPARENT = NO DOT_MULTI_TARGETS = NO -# If the GENERATE_LEGEND tag is set to YES doxygen will generate a legend page +# If the GENERATE_LEGEND tag is set to YES Doxygen will generate a legend page # explaining the meaning of the various boxes and arrows in the dot generated # graphs. +# Note: This tag requires that UML_LOOK isn't set, i.e. the Doxygen internal +# graphical representation for inheritance and collaboration diagrams is used. # The default value is: YES. # This tag requires that the tag HAVE_DOT is set to YES. GENERATE_LEGEND = YES -# If the DOT_CLEANUP tag is set to YES, doxygen will remove the intermediate dot +# If the DOT_CLEANUP tag is set to YES, Doxygen will remove the intermediate # files that are used to generate the various graphs. +# +# Note: This setting is not only used for dot files but also for msc temporary +# files. # The default value is: YES. -# This tag requires that the tag HAVE_DOT is set to YES. DOT_CLEANUP = YES + +# You can define message sequence charts within Doxygen comments using the \msc +# command. If the MSCGEN_TOOL tag is left empty (the default), then Doxygen will +# use a built-in version of mscgen tool to produce the charts. Alternatively, +# the MSCGEN_TOOL tag can also specify the name an external tool. For instance, +# specifying prog as the value, Doxygen will call the tool as prog -T +# -o . The external tool should support +# output file formats "png", "eps", "svg", and "ismap". + +MSCGEN_TOOL = + +# The MSCFILE_DIRS tag can be used to specify one or more directories that +# contain msc files that are included in the documentation (see the \mscfile +# command). + +MSCFILE_DIRS = diff --git a/cpp/include/cuvs/neighbors/brute_force.h b/cpp/include/cuvs/neighbors/brute_force.h index 91893e7d93..93cdffd7da 100644 --- a/cpp/include/cuvs/neighbors/brute_force.h +++ b/cpp/include/cuvs/neighbors/brute_force.h @@ -1,5 +1,5 @@ /* - * Copyright (c) 2024, NVIDIA CORPORATION. + * Copyright (c) 2024-2025, NVIDIA CORPORATION. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -167,7 +167,7 @@ cuvsError_t cuvsBruteForceSearch(cuvsResources_t res, */ /** - * @defgroup bruteforce_c_serialize BRUTEFORCE C-API serialize functions + * @defgroup bruteforce_c_index_serialize BRUTEFORCE C-API serialize functions * @{ */ /** diff --git a/cpp/include/cuvs/neighbors/cagra.h b/cpp/include/cuvs/neighbors/cagra.h index bd0548d4ba..e892b5304c 100644 --- a/cpp/include/cuvs/neighbors/cagra.h +++ b/cpp/include/cuvs/neighbors/cagra.h @@ -605,7 +605,7 @@ cuvsError_t cuvsCagraSearch(cuvsResources_t res, */ /** - * @defgroup cagra_c_serialize CAGRA C-API serialize functions + * @defgroup cagra_c_index_serialize CAGRA C-API serialize functions * @{ */ /** diff --git a/cpp/include/cuvs/neighbors/cagra.hpp b/cpp/include/cuvs/neighbors/cagra.hpp index fd1fe72b34..ac86492ac6 100644 --- a/cpp/include/cuvs/neighbors/cagra.hpp +++ b/cpp/include/cuvs/neighbors/cagra.hpp @@ -340,11 +340,13 @@ struct index : cuvs::neighbors::index { } // Don't allow copying the index for performance reasons (try avoiding copying data) + /** \cond */ index(const index&) = delete; index(index&&) = default; auto operator=(const index&) -> index& = delete; auto operator=(index&&) -> index& = default; ~index() = default; + /** \endcond */ /** Construct an empty index. */ index(raft::resources const& res, diff --git a/cpp/include/cuvs/neighbors/common.hpp b/cpp/include/cuvs/neighbors/common.hpp index b49b7a4e8a..4ed79dd695 100644 --- a/cpp/include/cuvs/neighbors/common.hpp +++ b/cpp/include/cuvs/neighbors/common.hpp @@ -491,6 +491,7 @@ struct base_filter { /* A filter that filters nothing. This is the default behavior. */ struct none_sample_filter : public base_filter { + /** \cond */ constexpr __forceinline__ _RAFT_HOST_DEVICE bool operator()( // query index const uint32_t query_ix, @@ -504,7 +505,7 @@ struct none_sample_filter : public base_filter { const uint32_t query_ix, // the index of the current sample const uint32_t sample_ix) const; - + /** \endcond */ FilterType get_filter_type() const override { return FilterType::None; } }; @@ -523,6 +524,7 @@ struct ivf_to_sample_filter : public base_filter { ivf_to_sample_filter(const index_t* const* inds_ptrs, const filter_t next_filter); + /** \cond */ /** If the original filter takes three arguments, then don't modify the arguments. * If the original filter takes two arguments, then we are using `inds_ptr_` to obtain the sample * index. @@ -536,6 +538,7 @@ struct ivf_to_sample_filter : public base_filter { const uint32_t sample_ix) const; FilterType get_filter_type() const override { return next_filter_.get_filter_type(); } + /** \endcond */ }; /** @@ -552,11 +555,13 @@ struct bitmap_filter : public base_filter { const view_t bitmap_view_; bitmap_filter(const view_t bitmap_for_filtering); + /** \cond */ inline _RAFT_HOST_DEVICE bool operator()( // query index const uint32_t query_ix, // the index of the current sample const uint32_t sample_ix) const; + /** \endcond */ FilterType get_filter_type() const override { return FilterType::Bitmap; } @@ -579,12 +584,14 @@ struct bitset_filter : public base_filter { // View of the bitset to use as a filter const view_t bitset_view_; + /** \cond */ _RAFT_HOST_DEVICE bitset_filter(const view_t bitset_for_filtering); constexpr __forceinline__ _RAFT_HOST_DEVICE bool operator()( // query index const uint32_t query_ix, // the index of the current sample const uint32_t sample_ix) const; + /** \endcond */ FilterType get_filter_type() const override { return FilterType::Bitset; } diff --git a/cpp/include/cuvs/neighbors/hnsw.h b/cpp/include/cuvs/neighbors/hnsw.h index b5f129e96d..c584f41529 100644 --- a/cpp/include/cuvs/neighbors/hnsw.h +++ b/cpp/include/cuvs/neighbors/hnsw.h @@ -1,5 +1,5 @@ /* - * Copyright (c) 2024, NVIDIA CORPORATION. + * Copyright (c) 2024-2025, NVIDIA CORPORATION. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -369,7 +369,7 @@ cuvsError_t cuvsHnswSearch(cuvsResources_t res, */ /** - * @defgroup hnsw_c_serialize HNSW C-API serialize functions + * @defgroup hnsw_c_index_serialize HNSW C-API serialize functions * @{ */ diff --git a/cpp/include/cuvs/neighbors/hnsw.hpp b/cpp/include/cuvs/neighbors/hnsw.hpp index 7db6d549eb..7be7ccff54 100644 --- a/cpp/include/cuvs/neighbors/hnsw.hpp +++ b/cpp/include/cuvs/neighbors/hnsw.hpp @@ -97,7 +97,9 @@ cuvs::neighbors::cagra::index_params to_cagra_params( int ef_construction, cuvs::distance::DistanceType metric = cuvs::distance::DistanceType::L2Expanded); -/**@}*/ +/** + * @} + */ /** * @defgroup hnsw_cpp_index hnswlib index wrapper @@ -127,7 +129,7 @@ struct index : cuvs::neighbors::index { /** @brief Get underlying index */ - virtual auto get_index() const -> void const* = 0; + virtual void const* get_index() const = 0; auto dim() const -> int const { return dim_; } @@ -146,7 +148,9 @@ struct index : cuvs::neighbors::index { HnswHierarchy hierarchy_; }; -/**@}*/ +/** + * @} + */ /** * @defgroup hnsw_cpp_extend_params HNSW index extend parameters @@ -308,7 +312,9 @@ std::unique_ptr> from_cagra( std::optional> dataset = std::nullopt); -/**@}*/ +/** + * @} + */ /** * @defgroup hnsw_cpp_index_extend Extend HNSW index with additional vectors @@ -342,6 +348,7 @@ std::unique_ptr> from_cagra( * auto additional_dataset = raft::make_host_matrix(res, add_size, index->dim()); * hnsw::extend_params extend_params; * hnsw::extend(res, extend_params, additional_dataset, *hnsw_index.get()); + * @endcode */ void extend(raft::resources const& res, const extend_params& params, @@ -375,6 +382,7 @@ void extend(raft::resources const& res, * auto additional_dataset = raft::make_host_matrix(res, add_size, index->dim()); * hnsw::extend_params extend_params; * hnsw::extend(res, extend_params, additional_dataset, *hnsw_index.get()); + * @endcode */ void extend(raft::resources const& res, const extend_params& params, @@ -408,6 +416,7 @@ void extend(raft::resources const& res, * auto additional_dataset = raft::make_host_matrix(res, add_size, index->dim()); * hnsw::extend_params extend_params; * hnsw::extend(res, extend_params, additional_dataset, *hnsw_index.get()); + * @endcode */ void extend(raft::resources const& res, const extend_params& params, @@ -441,13 +450,16 @@ void extend(raft::resources const& res, * auto additional_dataset = raft::make_host_matrix(res, add_size, index->dim()); * hnsw::extend_params extend_params; * hnsw::extend(res, extend_params, additional_dataset, *hnsw_index.get()); + * @endcode */ void extend(raft::resources const& res, const extend_params& params, raft::host_matrix_view additional_dataset, index& idx); -/**@} */ +/** + * @} + */ /** * @defgroup hnsw_cpp_search_params Build CAGRA index and search with hnswlib @@ -460,7 +472,9 @@ struct search_params : cuvs::neighbors::search_params { // automatically maximizes parallelism }; -/**@}*/ +/** + * @} + */ // TODO: Filtered Search APIs: https://github.com/rapidsai/cuvs/issues/363 @@ -645,7 +659,9 @@ void search(raft::resources const& res, raft::host_matrix_view neighbors, raft::host_matrix_view distances); -/**@}*/ +/** + * @} + */ /** * @defgroup hnsw_cpp_index_serialize Deserialize CAGRA index as hnswlib index @@ -940,7 +956,9 @@ void deserialize(raft::resources const& res, cuvs::distance::DistanceType metric, index** index); -/**@}*/ +/** + * @} + */ } // namespace cuvs::neighbors::hnsw diff --git a/cpp/include/cuvs/neighbors/ivf_flat.h b/cpp/include/cuvs/neighbors/ivf_flat.h index 7d647829ad..9e89ea354d 100644 --- a/cpp/include/cuvs/neighbors/ivf_flat.h +++ b/cpp/include/cuvs/neighbors/ivf_flat.h @@ -300,7 +300,7 @@ cuvsError_t cuvsIvfFlatSearch(cuvsResources_t res, */ /** - * @defgroup ivf_flat_c_serialize IVF-Flat C-API serialize functions + * @defgroup ivf_flat_c_index_serialize IVF-Flat C-API serialize functions * @{ */ /** diff --git a/cpp/include/cuvs/neighbors/ivf_pq.h b/cpp/include/cuvs/neighbors/ivf_pq.h index 5c349be5cb..eef10290c7 100644 --- a/cpp/include/cuvs/neighbors/ivf_pq.h +++ b/cpp/include/cuvs/neighbors/ivf_pq.h @@ -405,7 +405,7 @@ cuvsError_t cuvsIvfPqSearch(cuvsResources_t res, */ /** - * @defgroup ivf_pq_c_serialize IVF-PQ C-API serialize functions + * @defgroup ivf_pq_c_index_serialize IVF-PQ C-API serialize functions * @{ */ /** diff --git a/cpp/include/cuvs/neighbors/vamana.hpp b/cpp/include/cuvs/neighbors/vamana.hpp index 733f9c4e71..25b3d4f7d6 100644 --- a/cpp/include/cuvs/neighbors/vamana.hpp +++ b/cpp/include/cuvs/neighbors/vamana.hpp @@ -38,6 +38,18 @@ namespace cuvs::neighbors::vamana { * @{ */ +/** + * @brief Parameters used to build quantized DiskANN index; to be generated using + * deserialize_codebooks() + */ +template +struct codebook_params { + int pq_codebook_size; + int pq_dim; + std::vector pq_encoding_table; + std::vector rotation_matrix; +}; + /** * @brief Parameters used to build DiskANN index * @@ -53,17 +65,6 @@ namespace cuvs::neighbors::vamana { */ struct index_params : cuvs::neighbors::index_params { - /** - * @brief Parameters used to build quantized DiskANN index; to be generated using - * deserialize_codebooks() - */ - template - struct codebook_params { - int pq_codebook_size; - int pq_dim; - std::vector pq_encoding_table; - std::vector rotation_matrix; - }; /** Maximum degree of output graph corresponds to the R parameter in the original Vamana * literature. */ uint32_t graph_degree = 32; @@ -159,11 +160,13 @@ struct index : cuvs::neighbors::index { [[nodiscard]] inline auto medoid() const noexcept -> IdxT { return medoid_id_; } // Don't allow copying the index for performance reasons (try avoiding copying data) + /** \cond */ index(const index&) = delete; index(index&&) = default; auto operator=(const index&) -> index& = delete; auto operator=(index&&) -> index& = default; ~index() = default; + /** \endcond */ /** Construct an empty index. */ index(raft::resources const& res, @@ -613,7 +616,7 @@ void serialize(raft::resources const& handle, * */ auto deserialize_codebooks(const std::string& codebook_prefix, const int dim) - -> index_params::codebook_params; + -> codebook_params; /** * @} diff --git a/cpp/include/cuvs/preprocessing/quantize/binary.hpp b/cpp/include/cuvs/preprocessing/quantize/binary.hpp index 224c30c8b6..94cdb5ce5b 100644 --- a/cpp/include/cuvs/preprocessing/quantize/binary.hpp +++ b/cpp/include/cuvs/preprocessing/quantize/binary.hpp @@ -234,7 +234,6 @@ void transform(raft::resources const& res, * Usage example: * @code{.cpp} * raft::handle_t handle; - * @param[in] params quantization params * cuvs::preprocessing::quantize::binary::params params; * raft::host_matrix dataset = read_dataset(filename); * int64_t quantized_dim = raft::div_rounding_up_safe(dataset.extent(1), sizeof(uint8_t) * 8); diff --git a/cpp/src/neighbors/detail/vamana/vamana_codebooks.cuh b/cpp/src/neighbors/detail/vamana/vamana_codebooks.cuh index 5657a2840c..ae918af304 100644 --- a/cpp/src/neighbors/detail/vamana/vamana_codebooks.cuh +++ b/cpp/src/neighbors/detail/vamana/vamana_codebooks.cuh @@ -135,10 +135,9 @@ inline std::vector parse_rotation_matrix_file(const std::string& path, co } template -index_params::codebook_params deserialize_codebooks(const std::string& codebook_prefix, - const int dim) +codebook_params deserialize_codebooks(const std::string& codebook_prefix, const int dim) { - index_params::codebook_params codebooks; + codebook_params codebooks; codebooks.pq_encoding_table = parse_pq_pivots_file( codebook_prefix + "_pq_pivots.bin", dim, codebooks.pq_codebook_size, codebooks.pq_dim); codebooks.rotation_matrix = diff --git a/cpp/src/neighbors/vamana.cuh b/cpp/src/neighbors/vamana.cuh index 4cad294854..af8cc432d1 100644 --- a/cpp/src/neighbors/vamana.cuh +++ b/cpp/src/neighbors/vamana.cuh @@ -98,8 +98,7 @@ void serialize(raft::resources const& res, } template -index_params::codebook_params deserialize_codebooks(const std::string& codebook_prefix, - const int dim) +codebook_params deserialize_codebooks(const std::string& codebook_prefix, const int dim) { return cuvs::neighbors::vamana::detail::deserialize_codebooks(codebook_prefix, dim); } diff --git a/cpp/src/neighbors/vamana_codebooks_float.cu b/cpp/src/neighbors/vamana_codebooks_float.cu index e630675908..458c05ebb0 100644 --- a/cpp/src/neighbors/vamana_codebooks_float.cu +++ b/cpp/src/neighbors/vamana_codebooks_float.cu @@ -20,7 +20,7 @@ namespace cuvs::neighbors::vamana { #define CUVS_INST_VAMANA_CODEBOOKS(T) \ auto deserialize_codebooks(const std::string& codebook_prefix, const int dim) \ - -> cuvs::neighbors::vamana::index_params::codebook_params \ + -> cuvs::neighbors::vamana::codebook_params \ { \ return cuvs::neighbors::vamana::deserialize_codebooks(codebook_prefix, dim); \ } diff --git a/docs/source/api_basics.rst b/docs/source/api_basics.rst index d3e2c8a6ff..b085d22de3 100644 --- a/docs/source/api_basics.rst +++ b/docs/source/api_basics.rst @@ -14,7 +14,7 @@ RMM currently has APIs for C++ and Python. C++ ^^^ -Here's an example of configuring RMM to use a pool allocator in C++ (derived from the RMM example `here `_): +Here's an example of configuring RMM to use a pool allocator in C++ (derived from the RMM example `here `__): .. code-block:: c++ @@ -29,7 +29,7 @@ Here's an example of configuring RMM to use a pool allocator in C++ (derived fro Python ^^^^^^ -And the corresponding code in Python (derived from the RMM example `here `_): +And the corresponding code in Python (derived from the RMM example `here `__): .. code-block:: python diff --git a/docs/source/api_interoperability.rst b/docs/source/api_interoperability.rst index 519660036f..92e19de1fa 100644 --- a/docs/source/api_interoperability.rst +++ b/docs/source/api_interoperability.rst @@ -36,7 +36,7 @@ Here's an example on how to represent device memory using `DLManagedTensor`: // free memory after use cuvsRMMFree(dataset_dev); -Please refer to cuVS C API `documentation `_ to learn more. +Please refer to `cuVS C API documentation `_ to learn more. Multi-dimensional span (C++) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^ @@ -99,7 +99,7 @@ Of course, RAFT's `mdspan`/`mdarray` APIs aren't just limited to the `device`. Y auto vector_view = raft::make_host_vector_view(vector.data_handle(), vector.extent(0)); auto matrix_view = raft::make_host_matrix_view(matrix.data_handle(), matrix.extent(0), matrix.extent(1)); -Please refer to RAFT's `mdspan` `documentation `_ to learn more. +Please refer to RAFT's `mdspan documentation `_ to learn more. CUDA array interface (Python) diff --git a/docs/source/c_api/neighbors_cagra_c.rst b/docs/source/c_api/neighbors_cagra_c.rst index a5ffc45b9c..85b34d4947 100644 --- a/docs/source/c_api/neighbors_cagra_c.rst +++ b/docs/source/c_api/neighbors_cagra_c.rst @@ -51,7 +51,7 @@ Index search :content-only: Index serialize ------------- +--------------- .. doxygengroup:: cagra_c_index_serialize :project: cuvs diff --git a/docs/source/conf.py b/docs/source/conf.py index ca7330279d..a7247ed719 100644 --- a/docs/source/conf.py +++ b/docs/source/conf.py @@ -1,4 +1,4 @@ -# Copyright (c) 2018-2024, NVIDIA CORPORATION. +# Copyright (c) 2018-2025, NVIDIA CORPORATION. import os import sys @@ -206,3 +206,5 @@ def setup(app): # backticks`) to be a python object. See # https://www.sphinx-doc.org/en/master/usage/configuration.html#confval-default_role default_role = "py:obj" + +suppress_warnings = ["duplicate_declaration.cpp"] diff --git a/docs/source/cpp_api.rst b/docs/source/cpp_api.rst index 34f48a88f6..131858ff34 100644 --- a/docs/source/cpp_api.rst +++ b/docs/source/cpp_api.rst @@ -13,3 +13,4 @@ C++ API Documentation cpp_api/preprocessing.rst cpp_api/selection.rst cpp_api/stats.rst + cpp_api/core_interop diff --git a/docs/source/cpp_api/distance.rst b/docs/source/cpp_api/distance.rst index 2e83da7b35..6fcf373d48 100644 --- a/docs/source/cpp_api/distance.rst +++ b/docs/source/cpp_api/distance.rst @@ -11,11 +11,11 @@ distances have been highly optimized and support a wide assortment of different Distance Types -------------- -``#include `` +``#include `` namespace *cuvs::distance* -.. doxygenenum:: cuvs::distance::DistanceType +.. doxygenenum:: cuvsDistanceType :project: cuvs diff --git a/docs/source/cpp_api/neighbors_all_neighbors.rst b/docs/source/cpp_api/neighbors_all_neighbors.rst index 50ba35bca2..e11e374f8d 100644 --- a/docs/source/cpp_api/neighbors_all_neighbors.rst +++ b/docs/source/cpp_api/neighbors_all_neighbors.rst @@ -1,5 +1,5 @@ All-neighbors -========== +============= All-neighbors allows building an approximate all-neighbors knn graph. Given a full dataset, it finds nearest neighbors for all the training vectors in the dataset. @@ -12,7 +12,7 @@ All-neighbors allows building an approximate all-neighbors knn graph. Given a fu namespace *cuvs::neighbors::all_neighbors* All neighbors knn graph build parameters ----------------------- +---------------------------------------- .. doxygengroup:: all_neighbors_cpp_params :project: cuvs @@ -21,7 +21,7 @@ All neighbors knn graph build parameters Build ------------ +----- .. doxygengroup:: all_neighbors_cpp_build :project: cuvs diff --git a/docs/source/cpp_api/neighbors_cagra.rst b/docs/source/cpp_api/neighbors_cagra.rst index d9f5038718..b4cef43162 100644 --- a/docs/source/cpp_api/neighbors_cagra.rst +++ b/docs/source/cpp_api/neighbors_cagra.rst @@ -35,14 +35,6 @@ Index extend parameters :members: :content-only: -Index extend memory buffers ---------------------------- - -.. doxygengroup:: cagra_cpp_extend_memory_buffers - :project: cuvs - :members: - :content-only: - Index ----- diff --git a/docs/source/cuvs_bench/index.rst b/docs/source/cuvs_bench/index.rst index cfaa30acb2..0ce5ac8fff 100644 --- a/docs/source/cuvs_bench/index.rst +++ b/docs/source/cuvs_bench/index.rst @@ -181,7 +181,7 @@ All of the datasets above contain ground test datasets with 100 neighbors. Thus End-to-end: large-scale benchmarks (>10M vectors) ------------------------------------------------- -`cuvs_bench.get_dataset` cannot be used to download the `billion-scale datasets`_ due to their size. You should instead use our billion-scale datasets guide to download and prepare them. +`cuvs_bench.get_dataset` cannot be used to download the billion-scale datasets due to their size. You should instead use our billion-scale datasets guide to download and prepare them. All other python commands mentioned below work as intended once the billion-scale dataset has been downloaded. To download billion-scale datasets, visit `big-ann-benchmarks `_ @@ -212,6 +212,7 @@ The steps below demonstrate how to download, install, and run benchmarks on a su The usage of `python -m cuvs_bench.split_groundtruth` is: .. code-block:: bash + usage: split_groundtruth.py [-h] --groundtruth GROUNDTRUTH options: @@ -481,6 +482,7 @@ Implementation of a new algorithm should be a C++ class that inherits `class ANN In addition, it should define two `struct`s for building and searching parameters. The searching parameter class should inherit `struct ANN::AnnSearchParam`. Take `class HnswLib` as an example, its definition is: .. code-block:: c++ + template class HnswLib : public ANN { public: @@ -503,6 +505,7 @@ In addition, it should define two `struct`s for building and searching parameter The benchmark program uses JSON format natively in a configuration file to specify indexes to build, along with the build and search parameters. However the JSON config files are overly verbose and are not meant to be used directly. Instead, the Python scripts parse YAML and create these json files automatically. It's important to realize that these json objects align with the yaml objects for `build_param`, whose value is a JSON object, and `search_param`, whose value is an array of JSON objects. Take the json configuration for `HnswLib` as an example of the json after it's been parsed from yaml: .. code-block:: json + { "name" : "hnswlib.M12.ef500.th32", "algo" : "hnswlib", @@ -546,6 +549,7 @@ The build and search params are ultimately passed to the C++ layer as json objec 2. Next, add corresponding `if` case to functions `create_algo()` (in `cpp/bench/ann/) and `create_search_param()` by calling parsing functions. The string literal in `if` condition statement must be the same as the value of `algo` in configuration file. For example, .. code-block:: c++ + // JSON configuration file contains a line like: "algo" : "hnswlib" if (algo == "hnswlib") { // ... @@ -558,6 +562,7 @@ In `cuvs/cpp/bench/ann/CMakeLists.txt`, we provide a `CMake` function to configu .. code-block:: cmake + ConfigureAnnBench( NAME PATH @@ -580,9 +585,19 @@ This will create an executable called `HNSWLIB_ANN_BENCH`, which can then be use Add a new entry to `algos.yaml` to map the name of the algorithm to its binary executable and specify whether the algorithm requires GPU support. .. code-block:: yaml + cuvs_ivf_pq: executable: CUVS_IVF_PQ_ANN_BENCH requires_gpu: true `executable` : specifies the name of the binary that will build/search the index. It is assumed to be available in `cuvs/cpp/build/`. `requires_gpu` : denotes whether an algorithm requires GPU to run. + + +.. toctree:: + :maxdepth: 4 + + build.rst + datasets.rst + param_tuning.rst + wiki_all_dataset.rst diff --git a/docs/source/cuvs_bench/param_tuning.rst b/docs/source/cuvs_bench/param_tuning.rst index 9e22aed7d1..276fb35a47 100644 --- a/docs/source/cuvs_bench/param_tuning.rst +++ b/docs/source/cuvs_bench/param_tuning.rst @@ -1,6 +1,6 @@ ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ cuVS Bench Parameter Tuning Guide -~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ This guide outlines the various parameter settings that can be specified in :doc:`cuVS Benchmarks ` yaml configuration files and explains the impact they have on corresponding algorithms to help inform their settings for benchmarking across desired levels of recall. @@ -32,6 +32,7 @@ IVF-flat is a simple algorithm which won't save any space, but it provides compe - `build` - Y - Positive integer >0 + - 1024 - Number of clusters to partition the vectors into. Larger values will put less points into each cluster but this will impact index build time as more clusters need to be trained. * - `niter` @@ -48,25 +49,27 @@ IVF-flat is a simple algorithm which won't save any space, but it provides compe - 2 - `1/ratio` is the number of training points which should be used to train the clusters. - * - `dataset_memory_type` - - `build` - - N - - [`device`, `host`, `mmap`] - - `mmap` - - Where should the dataset reside? - - * - `query_memory_type` - - `search` - - [`device`, `host`, `mmap`] - - `device` - - Where should the queries reside? - - * - `nprobe` - - `search` - - Y - - Positive integer >0 - - - - The closest number of clusters to search for each query vector. Larger values will improve recall but will search more points in the index. + * - `dataset_memory_type` + - `build` + - N + - [`device`, `host`, `mmap`] + - `mmap` + - Where should the dataset reside? + + * - `query_memory_type` + - `search` + - N + - [`device`, `host`, `mmap`] + - `device` + - Where should the queries reside? + + * - `nprobe` + - `search` + - Y + - Positive integer >0 + - + - The closest number of clusters to search for each query vector. Larger values will improve recall but will search more points in the index. + cuvs_ivf_pq ----------- @@ -86,6 +89,7 @@ IVF-pq is an inverted-file index, which partitions the vectors into a series of - `build` - Y - Positive integer >0 + - 1024 - Number of clusters to partition the vectors into. Larger values will put less points into each cluster but this will impact index build time as more clusters need to be trained. * - `niter` @@ -132,6 +136,7 @@ IVF-pq is an inverted-file index, which partitions the vectors into a series of * - `query_memory_type` - `search` + - N - [`device`, `host`, `mmap`] - `device` - Where should the queries reside? @@ -140,7 +145,7 @@ IVF-pq is an inverted-file index, which partitions the vectors into a series of - `search` - Y - Positive integer >0 - - + - 20 - The closest number of clusters to search for each query vector. Larger values will improve recall but will search more points in the index. * - `internalDistanceDtype` @@ -195,8 +200,8 @@ CAGRA uses a graph-based index, which creates an intermediate, approximate kNN g * - `graph_build_algo` - `build` - - `N - - [`IVF_PQ`, NN_DESCENT`] + - `N` + - [`IVF_PQ`, `NN_DESCENT`] - `IVF_PQ` - Algorithm to use for building the initial kNN graph, from which CAGRA will optimize into the navigable CAGRA graph @@ -209,6 +214,7 @@ CAGRA uses a graph-based index, which creates an intermediate, approximate kNN g * - `query_memory_type` - `search` + - N - [`device`, `host`, `mmap`] - `device` - Where should the queries reside? @@ -356,7 +362,7 @@ Alternatively, if `graph_build_algo == "NN_DESCENT"`, then we can customize the - 20 - Number of nn-descent iterations - * - `nn_descent_intermediate_graph_degree + * - `nn_descent_intermediate_graph_degree` - `build` - N - Positive integer >0 @@ -463,7 +469,7 @@ IVF-flat is a simple algorithm which won't save any space, but it provides compe - `search` - Y - Positive integer >0 - - + - 20 - The closest number of clusters to search for each query vector. Larger values will improve recall but will search more points in the index. faiss_gpu_ivf_pq diff --git a/docs/source/cuvs_bench/wiki_all_dataset.rst b/docs/source/cuvs_bench/wiki_all_dataset.rst index cedeeb93f5..38b72ae3f5 100644 --- a/docs/source/cuvs_bench/wiki_all_dataset.rst +++ b/docs/source/cuvs_bench/wiki_all_dataset.rst @@ -22,6 +22,7 @@ A version of the dataset is made available in the binary format that can be used The following will download all 10 the parts and untar them to a `wiki_all_88M` directory: .. code-block:: bash + curl -s https://data.rapids.ai/raft/datasets/wiki_all/wiki_all.tar.{00..9} | tar -xf - -C wiki_all_88M/ The above has the unfortunate drawback that if the command should fail for any reason, all the parts need to be re-downloaded. The files can also be downloaded individually and then untarred to the directory. Each file is ~27GB and there are 10 of them. diff --git a/docs/source/filtering.rst b/docs/source/filtering.rst index 35805c5def..7a62270264 100644 --- a/docs/source/filtering.rst +++ b/docs/source/filtering.rst @@ -24,7 +24,7 @@ A bitmap is based on the same principle as a bitset, but in two dimensions. This being searched. Check out RAFT's `bitmap API documentation `. Examples -======= +======== Using a Bitset filter on a CAGRA index -------------------------------------- diff --git a/docs/source/getting_started.rst b/docs/source/getting_started.rst index c4706e5101..f6cd919adc 100644 --- a/docs/source/getting_started.rst +++ b/docs/source/getting_started.rst @@ -43,7 +43,7 @@ Getting Started New to vector search? ===================== -If you are unfamiliar with the basics of vector search or how vector search differs from vector databases, then :doc:`this primer on vector search guide ` should provide some good insight. Another good resource for the uninitiated is our :doc:`vector databases vs vector search ` guide. As outlined in the primer, vector search as used in vector databases is often closer to machine learning than to traditional databases. This means that while traditional databases can often be slow without any performance tuning, they will usually still yield the correct results. Unfortunately, vector search indexes, like other machine learning models, can yield garbage results of not tuned correctly. +If you are unfamiliar with the basics of vector search or how vector search differs from vector databases, then :doc:`this primer on vector search guide ` should provide some good insight. Another good resource for the uninitiated is our :doc:`vector databases vs vector search ` guide. As outlined in the primer, vector search as used in vector databases is often closer to machine learning than to traditional databases. This means that while traditional databases can often be slow without any performance tuning, they will usually still yield the correct results. Unfortunately, vector search indexes, like other machine learning models, can yield garbage results of not tuned correctly. Fortunately, this opens up the whole world of hyperparamer optimization to improve vector search performance and quality. Please see our :doc:`index tuning guide ` for more information. diff --git a/docs/source/index.rst b/docs/source/index.rst index 4dcf59112d..4c7665c162 100644 --- a/docs/source/index.rst +++ b/docs/source/index.rst @@ -88,3 +88,4 @@ Contents cuvs_bench/index.rst api_docs.rst contributing.md + developer_guide.md diff --git a/docs/source/indexes/ivfpq.rst b/docs/source/indexes/ivfpq.rst index ad973adf28..e243a5b562 100644 --- a/docs/source/indexes/ivfpq.rst +++ b/docs/source/indexes/ivfpq.rst @@ -58,7 +58,7 @@ Build parameters Search parameters -~~~~~~~~~~~~~~~~ +~~~~~~~~~~~~~~~~~ .. list-table:: :widths: 25 25 50 diff --git a/docs/source/indexes/vamana.rst b/docs/source/indexes/vamana.rst index 7e0a79fb84..4f5c2eb5f0 100644 --- a/docs/source/indexes/vamana.rst +++ b/docs/source/indexes/vamana.rst @@ -11,7 +11,7 @@ There are many algorithmic details that are outlined in the `paper ` library to perform efficient search. Additional DiskANN functionality, including GPU-accelerated search and 'ssd' index build are planned for future cuVS releases. -[ :doc:`C++ API <../cpp_api/neighbors_vamana>` | :doc:`Python API <../python_api/neighbors_vamana>` ] +[ :doc:`C++ API <../cpp_api/neighbors_vamana>` ] Interoperability with CPU DiskANN --------------------------------- diff --git a/docs/source/python_api/neighbors_cagra.rst b/docs/source/python_api/neighbors_cagra.rst index 76564ce5e7..e7155efb83 100644 --- a/docs/source/python_api/neighbors_cagra.rst +++ b/docs/source/python_api/neighbors_cagra.rst @@ -44,8 +44,3 @@ Index load ########## .. autofunction:: cuvs.neighbors.cagra.load - -Index extend -########### - -.. autofunction:: cuvs.neighbors.cagra.extend diff --git a/docs/source/python_api/neighbors_hnsw.rst b/docs/source/python_api/neighbors_hnsw.rst index fe46a8e4b6..40f3a1de7e 100644 --- a/docs/source/python_api/neighbors_hnsw.rst +++ b/docs/source/python_api/neighbors_hnsw.rst @@ -40,6 +40,6 @@ Index load .. autofunction:: cuvs.neighbors.hnsw.load Index extend -########### +############ .. autofunction:: cuvs.neighbors.hnsw.extend diff --git a/docs/source/python_api/neighbors_nn_decent.rst b/docs/source/python_api/neighbors_nn_decent.rst index df872d7e99..01e9e196c9 100644 --- a/docs/source/python_api/neighbors_nn_decent.rst +++ b/docs/source/python_api/neighbors_nn_decent.rst @@ -1,5 +1,5 @@ NN-Descent -====== +========== .. role:: py(code) :language: python diff --git a/docs/source/python_api/preprocessing.rst b/docs/source/python_api/preprocessing.rst index 679eabc16a..f0e3edac81 100644 --- a/docs/source/python_api/preprocessing.rst +++ b/docs/source/python_api/preprocessing.rst @@ -8,4 +8,15 @@ Preprocessing Scalar Quantizer ################ -.. autofunction:: cuvs.preprocessing.quantize.scalar + +.. autoclass:: cuvs.preprocessing.quantize.scalar.Quantizer + :members: + +.. autoclass:: cuvs.preprocessing.quantize.scalar.QuantizerParams + :members: + +.. autofunction:: cuvs.preprocessing.quantize.scalar.train + +.. autofunction:: cuvs.preprocessing.quantize.scalar.transform + +.. autofunction:: cuvs.preprocessing.quantize.scalar.inverse_transform diff --git a/docs/source/sphinxext/github_link.py b/docs/source/sphinxext/github_link.py index 75acfbd6e1..75a609d454 100644 --- a/docs/source/sphinxext/github_link.py +++ b/docs/source/sphinxext/github_link.py @@ -34,7 +34,7 @@ def isfunction(obj): new_val = hasattr(type(obj), "__code__") - if (orig_val != new_val): + if orig_val != new_val: return new_val return orig_val @@ -42,19 +42,20 @@ def isfunction(obj): inspect.isfunction = isfunction -REVISION_CMD = 'git rev-parse --short HEAD' +REVISION_CMD = "git rev-parse --short HEAD" -source_regex = re.compile(r"^File: (.*?) \(starting at line ([0-9]*?)\)$", - re.MULTILINE) +source_regex = re.compile( + r"^File: (.*?) \(starting at line ([0-9]*?)\)$", re.MULTILINE +) def _get_git_revision(): try: revision = subprocess.check_output(REVISION_CMD.split()).strip() except (subprocess.CalledProcessError, OSError): - print('Failed to execute git to get revision') + print("Failed to execute git to get revision") return None - return revision.decode('utf-8') + return revision.decode("utf-8") def _linkcode_resolve(domain, info, package, url_fmt, revision): @@ -74,14 +75,14 @@ def _linkcode_resolve(domain, info, package, url_fmt, revision): if revision is None: return - if domain not in ('py', 'pyx'): + if domain not in ("py", "pyx"): return - if not info.get('module') or not info.get('fullname'): + if not info.get("module") or not info.get("fullname"): return - class_name = info['fullname'].split('.')[0] - module = __import__(info['module'], fromlist=[class_name]) - obj = attrgetter(info['fullname'])(module) + class_name = info["fullname"].split(".")[0] + module = __import__(info["module"], fromlist=[class_name]) + obj = attrgetter(info["fullname"])(module) # Unwrap the object to get the correct source # file in case that is wrapped by a decorator @@ -104,42 +105,46 @@ def _linkcode_resolve(domain, info, package, url_fmt, revision): # Possibly Cython code. Search docstring for source m = source_regex.search(obj.__doc__) - if (m is not None): + if m is not None: source_file = m.group(1) lineno = m.group(2) # fn is expected to be the absolute path. fn = os.path.relpath(source_file, start=package) - print("{}:{}".format( - os.path.abspath(os.path.join("..", "python", "cuvs", fn)), - lineno)) + print( + "{}:{}".format( + os.path.abspath(os.path.join("..", "python", "cuvs", fn)), + lineno, + ) + ) else: return else: if fn.endswith(".pyx"): - sp_path = next(x for x in sys.path if re.match(".*site-packages$", x)) + sp_path = next( + x for x in sys.path if re.match(".*site-packages$", x) + ) fn = fn.replace("/opt/conda/conda-bld/work/python/cuvs", sp_path) # Convert to relative from module root - fn = os.path.relpath(fn, - start=os.path.dirname( - __import__(package).__file__)) + fn = os.path.relpath( + fn, start=os.path.dirname(__import__(package).__file__) + ) # Get the line number if we need it. (Can work without it) - if (lineno is None): + if lineno is None: try: lineno = inspect.getsourcelines(obj)[1] except Exception: # Can happen if its a cyfunction. See if it has `__code__` - if (hasattr(obj, "__code__")): + if hasattr(obj, "__code__"): lineno = obj.__code__.co_firstlineno else: - lineno = '' - return url_fmt.format(revision=revision, - package=package, - path=fn, - lineno=lineno) + lineno = "" + return url_fmt.format( + revision=revision, package=package, path=fn, lineno=lineno + ) def make_linkcode_resolve(package, url_fmt): @@ -154,7 +159,10 @@ def make_linkcode_resolve(package, url_fmt): '{path}#L{lineno}') """ revision = _get_git_revision() - return partial(_linkcode_resolve, - revision=revision, - package=package, - url_fmt=url_fmt) + + def linkcode_resolve(domain, info): + return _linkcode_resolve( + domain, info, revision=revision, package=package, url_fmt=url_fmt + ) + + return linkcode_resolve diff --git a/docs/source/working_with_ann_indexes_c.rst b/docs/source/working_with_ann_indexes_c.rst index 9d0c17a1e3..1e84141a86 100644 --- a/docs/source/working_with_ann_indexes_c.rst +++ b/docs/source/working_with_ann_indexes_c.rst @@ -3,8 +3,6 @@ Working with ANN Indexes in C - `Building an index`_ - `Searching an index`_ -- `CPU/GPU Interoperability`_ -- `Serializing an index`_ Building an index ----------------- diff --git a/docs/source/working_with_ann_indexes_cpp.rst b/docs/source/working_with_ann_indexes_cpp.rst index ead64107e9..68578bf848 100644 --- a/docs/source/working_with_ann_indexes_cpp.rst +++ b/docs/source/working_with_ann_indexes_cpp.rst @@ -3,8 +3,6 @@ Working with ANN Indexes in C++ - `Building an index`_ - `Searching an index`_ -- `CPU/GPU Interoperability`_ -- `Serializing an index`_ Building an index ----------------- diff --git a/docs/source/working_with_ann_indexes_python.rst b/docs/source/working_with_ann_indexes_python.rst index b2ab47a854..0419c47beb 100644 --- a/docs/source/working_with_ann_indexes_python.rst +++ b/docs/source/working_with_ann_indexes_python.rst @@ -3,8 +3,6 @@ Working with ANN Indexes in Python - `Building an index`_ - `Searching an index`_ -- `CPU/GPU Interoperability`_ -- `Serializing an index`_ Building an index ----------------- diff --git a/docs/source/working_with_ann_indexes_rust.rst b/docs/source/working_with_ann_indexes_rust.rst index ba284d1688..487ad0964b 100644 --- a/docs/source/working_with_ann_indexes_rust.rst +++ b/docs/source/working_with_ann_indexes_rust.rst @@ -1,13 +1,10 @@ Working with ANN Indexes in Rust ================================ -- `Building an index`_ -- `Searching an index`_ -- `CPU/GPU Interoperability`_ -- `Serializing an index`_ +- `Building and Searching an index`_ -Building an index ------------------ +Building and Searching an index +------------------------------- .. code-block:: rust @@ -31,5 +28,35 @@ Building an index let build_params = IndexParams::new()?; let index = Index::build(&res, &build_params, &dataset)?; + // use the first 4 points from the dataset as queries : will test that we get them back + // as their own nearest neighbor + let n_queries = 4; + let queries = dataset.slice(s![0..n_queries, ..]); + + let k = 10; + + // CAGRA search API requires queries and outputs to be on device memory + // copy query data over, and allocate new device memory for the distances/ neighbors + // outputs + let queries = ManagedTensor::from(&queries).to_device(&res)?; + let mut neighbors_host = ndarray::Array::::zeros((n_queries, k)); + let neighbors = ManagedTensor::from(&neighbors_host).to_device(&res)?; + + let mut distances_host = ndarray::Array::::zeros((n_queries, k)); + let distances = ManagedTensor::from(&distances_host).to_device(&res)?; + + let search_params = SearchParams::new()?; + + index.search(&res, &search_params, &queries, &neighbors, &distances)?; + + // Copy back to host memory + distances.to_host(&res, &mut distances_host)?; + neighbors.to_host(&res, &mut neighbors_host)?; + + // nearest neighbors should be themselves, since queries are from the + // dataset + println!("Neighbors {:?}", neighbors_host); + println!("Distances {:?}", distances_host); + Ok(()) } diff --git a/python/cuvs/cuvs/neighbors/brute_force/brute_force.pyx b/python/cuvs/cuvs/neighbors/brute_force/brute_force.pyx index 08690d0932..b271145b6c 100644 --- a/python/cuvs/cuvs/neighbors/brute_force/brute_force.pyx +++ b/python/cuvs/cuvs/neighbors/brute_force/brute_force.pyx @@ -1,5 +1,5 @@ # -# Copyright (c) 2024, NVIDIA CORPORATION. +# Copyright (c) 2024-2025, NVIDIA CORPORATION. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -161,11 +161,12 @@ def search(Index index, - `n_prefilter_rows == 1` when using a bitset prefilter. Each bit in `n_samples` determines whether `queries[i]` should be considered for distance computation with the index. - (default None) + (default None) {resources_docstring} Examples -------- + >>> # Example without pre-filter >>> import cupy as cp >>> from cuvs.neighbors import brute_force @@ -187,8 +188,6 @@ def search(Index index, >>> neighbors = cp.asarray(neighbors) >>> distances = cp.asarray(distances) - Examples - -------- >>> # Example with pre-filter >>> import numpy as np >>> import cupy as cp @@ -223,6 +222,7 @@ def search(Index index, ... prefilter=bitmap_prefilter) >>> neighbors = cp.asarray(neighbors) >>> distances = cp.asarray(distances) + """ if not index.trained: raise ValueError("Index needs to be built before calling search.") diff --git a/python/cuvs/cuvs/neighbors/cagra/cagra.pyx b/python/cuvs/cuvs/neighbors/cagra/cagra.pyx index 7fa73a2cff..0a8fec17b4 100644 --- a/python/cuvs/cuvs/neighbors/cagra/cagra.pyx +++ b/python/cuvs/cuvs/neighbors/cagra/cagra.pyx @@ -136,26 +136,29 @@ cdef class IndexParams: Parameters ---------- - metric : string denoting the metric type, default="sqeuclidean" - Valid values for metric: ["sqeuclidean", "inner_product"], where + + metric : str, default = "sqeuclidean" + String denoting the metric type, valid values for metric are + ["sqeuclidean", "inner_product"], where: + - sqeuclidean is the euclidean distance without the square root operation, i.e.: distance(a,b) = \\sum_i (a_i - b_i)^2 - inner_product distance is defined as distance(a, b) = \\sum_i a_i * b_i. - intermediate_graph_degree : int, default = 128 + intermediate_graph_degree : int, default = 128 graph_degree : int, default = 64 + build_algo: str, default = "ivf_pq" + string denoting the graph building algorithm to use. Valid values for + algo: ["ivf_pq", "nn_descent", "iterative_cagra_search"], where - build_algo: string denoting the graph building algorithm to use, \ - default = "ivf_pq" - Valid values for algo: ["ivf_pq", "nn_descent", - "iterative_cagra_search"], where - ivf_pq will use the IVF-PQ algorithm for building the knn graph - nn_descent (experimental) will use the NN-Descent algorithm for building the knn graph. It is expected to be generally faster than ivf_pq. - iterative_cagra_search will iteratively build the knn graph using CAGRA's search() and optimize() + compression: CompressionParams, optional If compression is desired should be a CompressionParams object. If None compression will be disabled. @@ -166,6 +169,7 @@ cdef class IndexParams: Parameters for IVF-PQ search. If provided, it will be used for searching the graph. refinement_rate: float, default = 1.0 + """ cdef cuvsCagraIndexParams* params @@ -413,6 +417,7 @@ cdef class SearchParams: Parameters ---------- + max_queries: int, default = 0 Maximum number of queries to search at the same time (batch size). Auto select when 0. @@ -422,12 +427,15 @@ cdef class SearchParams: search speed. Higher values improve the search accuracy. max_iterations: int, default = 0 Upper limit of search iterations. Auto select when 0. - algo: string denoting the search algorithm to use, default = "auto" - Valid values for algo: ["auto", "single_cta", "multi_cta"], where + algo: str, default = "auto" + String denoting the search algorithm to use + Valid values for algo: ["auto", "single_cta", "multi_cta"], where: + - auto will automatically select the best value based on query size - single_cta is better when query contains larger number of vectors (e.g >10) - multi_cta is better when query contains only a few vectors + team_size: int, default = 0 Number of threads used to calculate a single distance. 4, 8, 16, or 32. @@ -439,13 +447,15 @@ cdef class SearchParams: thread_block_size: int, default = 0 Thread block size. 0, 64, 128, 256, 512, 1024. Auto selection when 0. - hashmap_mode: string denoting the type of hash map to use. + hashmap_mode: str, default = "auto" + String denoting the type of hash map to use. It's usually better to allow the algorithm to select this value, - default = "auto". - Valid values for hashmap_mode: ["auto", "small", "hash"], where + Valid values for hashmap_mode: ["auto", "small", "hash"], where: + - auto will automatically select the best value based on algo - small will use the small shared memory hash table with resetting. - hash will use a single hash table in global memory. + hashmap_min_bitlen: int, default = 0 Upper limit of hashmap fill rate. More than 0.1, less than 0.9. hashmap_max_fill_rate: float, default = 0.5 @@ -456,12 +466,13 @@ cdef class SearchParams: rand_xor_mask: int, default = 0x128394 Bit mask used for initial random seed node selection. persistent: bool, default = false - Whether to use the persistent version of the kernel + Whether to use the persistent version of the kernel persistent_lifetime: float - Persistent kernel: time in seconds before the kernel stops if no - requests are received. + Persistent kernel: time in seconds before the kernel stops if no + requests are received. persistent_device_usage : float Sets the fraction of maximum grid size used by persistent kernel. + """ cdef cuvsCagraSearchParams * params @@ -608,6 +619,7 @@ def search(SearchParams search_params, Parameters ---------- + search_params : SearchParams index : Index Trained CAGRA index. @@ -623,11 +635,12 @@ def search(SearchParams search_params, neighbors will be written here in-place. (default None) filter: Optional cuvs.neighbors.cuvsFilter can be used to filter neighbors based on a given bitset. - (default None) + (default None) {resources_docstring} Examples -------- + >>> import cupy as cp >>> from cuvs.neighbors import cagra >>> n_samples = 50000 @@ -652,6 +665,7 @@ def search(SearchParams search_params, ... k) >>> neighbors = cp.asarray(neighbors) >>> distances = cp.asarray(distances) + """ if not index.trained: raise ValueError("Index needs to be built before calling search.") diff --git a/python/cuvs/cuvs/neighbors/hnsw/hnsw.pyx b/python/cuvs/cuvs/neighbors/hnsw/hnsw.pyx index 35bf2c8d95..a0f46f59e8 100644 --- a/python/cuvs/cuvs/neighbors/hnsw/hnsw.pyx +++ b/python/cuvs/cuvs/neighbors/hnsw/hnsw.pyx @@ -1,5 +1,5 @@ # -# Copyright (c) 2024, NVIDIA CORPORATION. +# Copyright (c) 2024-2025, NVIDIA CORPORATION. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -303,21 +303,22 @@ def from_cagra(IndexParams index_params, cagra.Index cagra_index, Returns an HNSW index from a CAGRA index. NOTE: When `index_params.hierarchy` is: - 1. `NONE`: This method uses the filesystem to write the CAGRA index - in `/tmp/.bin` before reading it as an - hnswlib index, then deleting the temporary file. The - returned index is immutable and can only be searched by - the hnswlib wrapper in cuVS, as the format is not - compatible with the original hnswlib. - 2. `CPU`: The returned index is mutable and can be extended with - additional vectors. The serialized index is also compatible - with the original hnswlib library. + + 1. `NONE`: This method uses the filesystem to write the CAGRA index in + `/tmp/.bin` before reading it as an hnswlib index, then + deleting the temporary file. The returned index is immutable and can only + be searched by the hnswlib wrapper in cuVS, as the format is not + compatible with the original hnswlib. + 2. `CPU`: The returned index is mutable and can be extended with additional + vectors. The serialized index is also compatible with the original hnswlib + library. Saving / loading the index is experimental. The serialization format is subject to change. Parameters ---------- + index_params : IndexParams Parameters to convert the CAGRA index to HNSW index. cagra_index : cagra.Index @@ -329,6 +330,7 @@ def from_cagra(IndexParams index_params, cagra.Index cagra_index, Examples -------- + >>> import cupy as cp >>> from cuvs.neighbors import cagra >>> from cuvs.neighbors import hnsw @@ -340,6 +342,7 @@ def from_cagra(IndexParams index_params, cagra.Index cagra_index, >>> index = cagra.build(cagra.IndexParams(), dataset) >>> # Serialize the CAGRA index to hnswlib base layer only index format >>> hnsw_index = hnsw.from_cagra(hnsw.IndexParams(), index) + """ cdef Index hnsw_index = Index() diff --git a/python/cuvs/cuvs/neighbors/ivf_flat/ivf_flat.pyx b/python/cuvs/cuvs/neighbors/ivf_flat/ivf_flat.pyx index 6fb4279db1..3a803e5599 100644 --- a/python/cuvs/cuvs/neighbors/ivf_flat/ivf_flat.pyx +++ b/python/cuvs/cuvs/neighbors/ivf_flat/ivf_flat.pyx @@ -61,6 +61,7 @@ cdef class IndexParams: String denoting the metric type. Valid values for metric: ["sqeuclidean", "inner_product", "euclidean", "cosine"], where + - sqeuclidean is the euclidean distance without the square root operation, i.e.: distance(a,b) = \\sum_i (a_i - b_i)^2, - euclidean is the euclidean distance @@ -68,6 +69,7 @@ cdef class IndexParams: distance(a, b) = \\sum_i a_i * b_i. - cosine distance is defined as distance(a, b) = 1 - \\sum_i a_i * b_i / ( ||a||_2 * ||b||_2). + kmeans_n_iters : int, default = 20 The number of iterations searching for kmeans centers during index building. diff --git a/python/cuvs/cuvs/neighbors/ivf_pq/ivf_pq.pyx b/python/cuvs/cuvs/neighbors/ivf_pq/ivf_pq.pyx index 40f22c236b..f784327d33 100644 --- a/python/cuvs/cuvs/neighbors/ivf_pq/ivf_pq.pyx +++ b/python/cuvs/cuvs/neighbors/ivf_pq/ivf_pq.pyx @@ -60,11 +60,13 @@ cdef class IndexParams: String denoting the metric type. Valid values for metric: ["sqeuclidean", "inner_product", "euclidean"], where: + - sqeuclidean is the euclidean distance without the square root operation, i.e.: distance(a,b) = \\sum_i (a_i - b_i)^2, - euclidean is the euclidean distance - inner product distance is defined as distance(a, b) = \\sum_i a_i * b_i. + kmeans_n_iters : int, default = 20 The number of iterations searching for kmeans centers during index building. From 9667a576a3a534a44fec24d7c23d4926eb6a5173 Mon Sep 17 00:00:00 2001 From: Peter Andreas Entschev Date: Tue, 16 Sep 2025 17:01:47 +0200 Subject: [PATCH 64/93] Remove UCX-Py/UCXX from release script (#1331) Contributes to https://github.com/rapidsai/build-planning/issues/198 There are no direct dependencies on either UCX-Py (deprecated) or UCXX but the release script still has code to handle that which may be removed. Authors: - Peter Andreas Entschev (https://github.com/pentschev) Approvers: - James Lamb (https://github.com/jameslamb) URL: https://github.com/rapidsai/cuvs/pull/1331 --- ci/release/update-version.sh | 11 ----------- 1 file changed, 11 deletions(-) diff --git a/ci/release/update-version.sh b/ci/release/update-version.sh index e6056f7332..29cfe0e1a3 100755 --- a/ci/release/update-version.sh +++ b/ci/release/update-version.sh @@ -22,11 +22,9 @@ NEXT_MAJOR=$(echo "$NEXT_FULL_TAG" | awk '{split($0, a, "."); print a[1]}') NEXT_MINOR=$(echo "$NEXT_FULL_TAG" | awk '{split($0, a, "."); print a[2]}') NEXT_PATCH=$(echo "$NEXT_FULL_TAG" | awk '{split($0, a, "."); print a[3]}') NEXT_SHORT_TAG=${NEXT_MAJOR}.${NEXT_MINOR} -NEXT_UCXX_SHORT_TAG="$(curl -sL https://version.gpuci.io/rapids/"${NEXT_SHORT_TAG}")" # Need to distutils-normalize the original version NEXT_SHORT_TAG_PEP440=$(python -c "from packaging.version import Version; print(Version('${NEXT_SHORT_TAG}'))") -NEXT_UCXX_SHORT_TAG_PEP440=$(python -c "from packaging.version import Version; print(Version('${NEXT_UCXX_SHORT_TAG}'))") PATCH_PEP440=$(python -c "from packaging.version import Version; print(Version('${NEXT_PATCH}'))") echo "Preparing release $CURRENT_TAG => $NEXT_FULL_TAG" @@ -51,24 +49,15 @@ DEPENDENCIES=( rmm rapids-dask-dependency ) -UCXX_DEPENDENCIES=( - ucx-py -) for FILE in dependencies.yaml conda/environments/*.yaml; do for DEP in "${DEPENDENCIES[@]}"; do sed_runner "/-.* ${DEP}\(-cu[[:digit:]]\{2\}\)\{0,1\}==/ s/==.*/==${NEXT_SHORT_TAG_PEP440}.*,>=0.0.0a0/g" "${FILE}" done - for DEP in "${UCXX_DEPENDENCIES[@]}"; do - sed_runner "/-.* ${DEP}\(-cu[[:digit:]]\{2\}\)\{0,1\}==/ s/==.*/==${NEXT_UCXX_SHORT_TAG_PEP440}.*,>=0.0.0a0/g" "${FILE}" - done done for FILE in python/*/pyproject.toml; do for DEP in "${DEPENDENCIES[@]}"; do sed_runner "/\"${DEP}==/ s/==.*\"/==${NEXT_SHORT_TAG_PEP440}.*,>=0.0.0a0\"/g" "${FILE}" done - for DEP in "${UCXX_DEPENDENCIES[@]}"; do - sed_runner "/\"${DEP}\(-cu[[:digit:]]\{2\}\)\{0,1\}==/ s/==.*\"/==${NEXT_UCXX_SHORT_TAG_PEP440}.*,>=0.0.0a0\"/g" "${FILE}" - done done for FILE in .github/workflows/*.yaml; do From 439620ad2dec561dbb182cc8d30f671d16aa225f Mon Sep 17 00:00:00 2001 From: "Artem M. Chirkin" <9253178+achirkin@users.noreply.github.com> Date: Tue, 16 Sep 2025 22:34:13 +0200 Subject: [PATCH 65/93] Fix inadvertent uses of copy constructor in mdarrays across cuVS (#1330) As of current date, raft's device mdarray is copy-constructible. This has resulted in cuVS doing accidental mdarray (and its data) copies in a few places across cuVS codebase. In some cases, this leads to segfaults (use-after-destruction of the temporary copy). In other cases, this leads to unnecessary copying of data. This PR fixes the problematic occurrences (found by removing the copy constructor in a local copy of raft and following the compiler errors). Fixes https://github.com/rapidsai/cuvs/pull/1263 Authors: - Artem M. Chirkin (https://github.com/achirkin) Approvers: - Corey J. Nolet (https://github.com/cjnolet) URL: https://github.com/rapidsai/cuvs/pull/1330 --- cpp/bench/ann/src/cuvs/cuvs_ann_bench_utils.h | 8 ++++---- cpp/src/neighbors/ivf_pq_index.cu | 4 ++-- .../preprocessing/spectral/spectral_embedding.cu | 2 +- cpp/tests/distance/masked_nn.cu | 14 +++++++------- cpp/tests/neighbors/ann_ivf_flat.cuh | 6 +++--- examples/cpp/src/cagra_persistent_example.cu | 2 +- 6 files changed, 18 insertions(+), 18 deletions(-) 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 37d1190e66..3aef8a64a4 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 @@ /* - * Copyright (c) 2023-2024, NVIDIA CORPORATION. + * Copyright (c) 2023-2025, NVIDIA CORPORATION. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -192,9 +192,9 @@ inline auto configured_raft_resources::operator=(configured_raft_resources&&) /** A helper to refine the neighbors when the data is on device or on host. */ template void refine_helper(const raft::resources& res, - DatasetT dataset, - QueriesT queries, - CandidatesT candidates, + const DatasetT& dataset, + const QueriesT& queries, + const CandidatesT& candidates, int k, algo_base::index_type* neighbors, float* distances, diff --git a/cpp/src/neighbors/ivf_pq_index.cu b/cpp/src/neighbors/ivf_pq_index.cu index 1bd6b6291d..7711fd63b2 100644 --- a/cpp/src/neighbors/ivf_pq_index.cu +++ b/cpp/src/neighbors/ivf_pq_index.cu @@ -1,5 +1,5 @@ /* - * Copyright (c) 2024, NVIDIA CORPORATION. + * Copyright (c) 2024-2025, NVIDIA CORPORATION. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -302,7 +302,7 @@ uint32_t index::get_list_size_in_bytes(uint32_t label) { RAFT_EXPECTS(label < this->n_lists(), "Expected label to be less than number of lists in the index"); - auto list_data = this->lists()[label]->data; + auto& list_data = this->lists()[label]->data; return list_data.size(); } diff --git a/cpp/src/preprocessing/spectral/spectral_embedding.cu b/cpp/src/preprocessing/spectral/spectral_embedding.cu index 7565ae63b1..0f288d2b13 100644 --- a/cpp/src/preprocessing/spectral/spectral_embedding.cu +++ b/cpp/src/preprocessing/spectral/spectral_embedding.cu @@ -168,7 +168,7 @@ raft::device_csr_matrix create_laplacian( void compute_eigenpairs(raft::resources const& handle, params spectral_embedding_config, const int n_samples, - raft::device_csr_matrix laplacian, + raft::device_csr_matrix& laplacian, raft::device_vector_view diagonal, raft::device_matrix_view embedding) { diff --git a/cpp/tests/distance/masked_nn.cu b/cpp/tests/distance/masked_nn.cu index 63fcee1e31..0d6372f3d9 100644 --- a/cpp/tests/distance/masked_nn.cu +++ b/cpp/tests/distance/masked_nn.cu @@ -1,5 +1,5 @@ /* - * Copyright (c) 2023-2024, NVIDIA CORPORATION. + * Copyright (c) 2023-2025, NVIDIA CORPORATION. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -89,10 +89,10 @@ RAFT_KERNEL init_adj(AdjacencyPattern pattern, template __launch_bounds__(32 * NWARPS, 2) RAFT_KERNEL referenceKernel(raft::KeyValuePair* min, - DataT* x, - DataT* y, - bool* adj, - int* group_idxs, + const DataT* x, + const DataT* y, + const bool* adj, + const int* group_idxs, int m, int n, int k, @@ -196,7 +196,7 @@ struct Inputs { }; template > -auto reference(const raft::handle_t& handle, Inputs inp, const Params& p) +auto reference(const raft::handle_t& handle, const Inputs& inp, const Params& p) -> raft::device_vector { int m = inp.x.extent(0); @@ -244,7 +244,7 @@ auto reference(const raft::handle_t& handle, Inputs inp, const Params& p) } template > -auto run_masked_nn(const raft::handle_t& handle, Inputs inp, const Params& p) +auto run_masked_nn(const raft::handle_t& handle, const Inputs& inp, const Params& p) -> raft::device_vector { // Compute norms: diff --git a/cpp/tests/neighbors/ann_ivf_flat.cuh b/cpp/tests/neighbors/ann_ivf_flat.cuh index bb8c7d1cd7..1bb2609873 100644 --- a/cpp/tests/neighbors/ann_ivf_flat.cuh +++ b/cpp/tests/neighbors/ann_ivf_flat.cuh @@ -314,8 +314,8 @@ class AnnIVFFlatTest : public ::testing::TestWithParam> { if (list_size > 0) { uint32_t padded_list_size = interleaved_group::roundUp(list_size); uint32_t n_elems = padded_list_size * idx.dim(); - auto list_data = lists[label]->data; - auto list_inds = extend_index.lists()[label]->indices; + auto& list_data = lists[label]->data; + auto& list_inds = extend_index.lists()[label]->indices; // fetch the flat codes auto flat_codes = raft::make_device_matrix(handle_, list_size, idx.dim()); @@ -364,7 +364,7 @@ class AnnIVFFlatTest : public ::testing::TestWithParam> { return DataT{0}; }); - auto extend_data = extend_index.lists()[label]->data; + auto& extend_data = extend_index.lists()[label]->data; auto extend_data_filtered = raft::make_device_vector(handle_, n_elems); raft::linalg::map_offset( handle_, diff --git a/examples/cpp/src/cagra_persistent_example.cu b/examples/cpp/src/cagra_persistent_example.cu index e65e10e725..a0234e2472 100644 --- a/examples/cpp/src/cagra_persistent_example.cu +++ b/examples/cpp/src/cagra_persistent_example.cu @@ -30,7 +30,7 @@ // A helper to split the dataset into chunks template -auto slice_matrix(DeviceMatrixOrView source, +auto slice_matrix(DeviceMatrixOrView& source, typename DeviceMatrixOrView::index_type offset_rows, typename DeviceMatrixOrView::index_type count_rows) { From 890929ece6cd03251019fe688b7d4659932d4a12 Mon Sep 17 00:00:00 2001 From: Micka Date: Tue, 16 Sep 2025 22:46:58 +0200 Subject: [PATCH 66/93] Fix debug build (#1258) Closes #462 Fix for the error when building in debug. It will also help downstream libraries that want to build in debug (https://github.com/zilliztech/knowhere/issues/1123) ``` ptxas error : Entry function '_ZN4cuvs9neighbors5cagra6detail17single_cta_search15search_kernel_pILj64ELj64ELj0ENS2_25dataset_descriptor_base_tI6__halfjfEENS0_9filtering18none_sample_filterEEEvPKT2_PNS3_15worker_handle_tEPNS3_10job_desc_tISA_EEPjPKNSA_7INDEX_TEjjmSL_jPSJ_jjjjSI_jjjT3_' with max regcount of 64 calls function '_ZN4cuvs9neighbors5cagra6detail17single_cta_search11search_coreILj64ELj64ELj0ENS2_25dataset_descriptor_base_tI6__halfjfEENS0_9filtering18none_sample_filterEEEvPNT2_7INDEX_TEPNSA_10DISTANCE_TEjPKSA_PKNSA_6DATA_TEPKSB_jjmSL_jSC_jjjjPjjjjjT3_' with regcount of 255 ptxas error : Entry function '_ZN4cuvs9neighbors5cagra6detail17single_cta_search13search_kernelILj64ELj64ELj0ENS2_25dataset_descriptor_base_tI6__halfjfEENS0_9filtering18none_sample_filterEEEvPNT2_7INDEX_TEPNSA_10DISTANCE_TEjPKSA_PKNSA_6DATA_TEPKSB_jjmSL_jSC_jjjjPjjjjT3_' with max regcount of 64 calls function '_ZN4cuvs9neighbors5cagra6detail17single_cta_search11search_coreILj64ELj64ELj0ENS2_25dataset_descriptor_base_tI6__halfjfEENS0_9filtering18none_sample_filterEEEvPNT2_7INDEX_TEPNSA_10DISTANCE_TEjPKSA_PKNSA_6DATA_TEPKSB_jjmSL_jSC_jjjjPjjjjjT3_' with regcount of 255 ptxas error : Entry function '_ZN4cuvs9neighbors5cagra6detail17single_cta_search13search_kernelILj64ELj64ELj1ENS2_25dataset_descriptor_base_tI6__halfjfEENS0_9filtering18none_sample_filterEEEvPNT2_7INDEX_TEPNSA_10DISTANCE_TEjPKSA_PKNSA_6DATA_TEPKSB_jjmSL_jSC_jjjjPjjjjT3_' with max regcount of 64 calls function '_ZN4cuvs9neighbors5cagra6detail17single_cta_search11search_coreILj64ELj64ELj1ENS2_25dataset_descriptor_base_tI6__halfjfEENS0_9filtering18none_sample_filterEEEvPNT2_7INDEX_TEPNSA_10DISTANCE_TEjPKSA_PKNSA_6DATA_TEPKSB_jjmSL_jSC_jjjjPjjjjjT3_' with regcount of 255 ptxas error : Entry function '_ZN4cuvs9neighbors5cagra6detail17single_cta_search15search_kernel_pILj64ELj64ELj1ENS2_25dataset_descriptor_base_tI6__halfjfEENS0_9filtering18none_sample_filterEEEvPKT2_PNS3_15worker_handle_tEPNS3_10job_desc_tISA_EEPjPKNSA_7INDEX_TEjjmSL_jPSJ_jjjjSI_jjjT3_' with max regcount of 64 calls function '_ZN4cuvs9neighbors5cagra6detail17single_cta_search11search_coreILj64ELj64ELj1ENS2_25dataset_descriptor_base_tI6__halfjfEENS0_9filtering18none_sample_filterEEEvPNT2_7INDEX_TEPNSA_10DISTANCE_TEjPKSA_PKNSA_6DATA_TEPKSB_jjmSL_jSC_jjjjPjjjjjT3_' with regcount of 255 ptxas error : Entry function '_ZN4cuvs9neighbors5cagra6detail17single_cta_search15search_kernel_pILj128ELj64ELj1ENS2_25dataset_descriptor_base_tI6__halfjfEENS0_9filtering18none_sample_filterEEEvPKT2_PNS3_15worker_handle_tEPNS3_10job_desc_tISA_EEPjPKNSA_7INDEX_TEjjmSL_jPSJ_jjjjSI_jjjT3_' with max regcount of 64 calls function '_ZN4cuvs9neighbors5cagra6detail17single_cta_search11search_coreILj128ELj64ELj1ENS2_25dataset_descriptor_base_tI6__halfjfEENS0_9filtering18none_sample_filterEEEvPNT2_7INDEX_TEPNSA_10DISTANCE_TEjPKSA_PKNSA_6DATA_TEPKSB_jjmSL_jSC_jjjjPjjjjjT3_' with regcount of 255 ptxas error : Entry function '_ZN4cuvs9neighbors5cagra6detail17single_cta_search13search_kernelILj128ELj64ELj1ENS2_25dataset_descriptor_base_tI6__halfjfEENS0_9filtering18none_sample_filterEEEvPNT2_7INDEX_TEPNSA_10DISTANCE_TEjPKSA_PKNSA_6DATA_TEPKSB_jjmSL_jSC_jjjjPjjjjT3_' with max regcount of 64 calls function '_ZN4cuvs9neighbors5cagra6detail17single_cta_search11search_coreILj128ELj64ELj1ENS2_25dataset_descriptor_base_tI6__halfjfEENS0_9filtering18none_sample_filterEEEvPNT2_7INDEX_TEPNSA_10DISTANCE_TEjPKSA_PKNSA_6DATA_TEPKSB_jjmSL_jSC_jjjjPjjjjjT3_' with regcount of 255 ptxas error : Entry function '_ZN4cuvs9neighbors5cagra6detail17single_cta_search13search_kernelILj256ELj64ELj1ENS2_25dataset_descriptor_base_tI6__halfjfEENS0_9filtering18none_sample_filterEEEvPNT2_7INDEX_TEPNSA_10DISTANCE_TEjPKSA_PKNSA_6DATA_TEPKSB_jjmSL_jSC_jjjjPjjjjT3_' with max regcount of 64 calls function '_ZN4cuvs9neighbors5cagra6detail17single_cta_search11search_coreILj256ELj64ELj1ENS2_25dataset_descriptor_base_tI6__halfjfEENS0_9filtering18none_sample_filterEEEvPNT2_7INDEX_TEPNSA_10DISTANCE_TEjPKSA_PKNSA_6DATA_TEPKSB_jjmSL_jSC_jjjjPjjjjjT3_' with regcount of 255 ptxas error : Entry function '_ZN4cuvs9neighbors5cagra6detail17single_cta_search15search_kernel_pILj256ELj64ELj1ENS2_25dataset_descriptor_base_tI6__halfjfEENS0_9filtering18none_sample_filterEEEvPKT2_PNS3_15worker_handle_tEPNS3_10job_desc_tISA_EEPjPKNSA_7INDEX_TEjjmSL_jPSJ_jjjjSI_jjjT3_' with max regcount of 64 calls function '_ZN4cuvs9neighbors5cagra6detail17single_cta_search11search_coreILj256ELj64ELj1ENS2_25dataset_descriptor_base_tI6__halfjfEENS0_9filtering18none_sample_filterEEEvPNT2_7INDEX_TEPNSA_10DISTANCE_TEjPKSA_PKNSA_6DATA_TEPKSB_jjmSL_jSC_jjjjPjjjjjT3_' with regcount of 255 ptxas error : Entry function '_ZN4cuvs9neighbors5cagra6detail17single_cta_search15search_kernel_pILj512ELj64ELj1ENS2_25dataset_descriptor_base_tI6__halfjfEENS0_9filtering18none_sample_filterEEEvPKT2_PNS3_15worker_handle_tEPNS3_10job_desc_tISA_EEPjPKNSA_7INDEX_TEjjmSL_jPSJ_jjjjSI_jjjT3_' with max regcount of 64 calls function '_ZN4cuvs9neighbors5cagra6detail17single_cta_search11search_coreILj512ELj64ELj1ENS2_25dataset_descriptor_base_tI6__halfjfEENS0_9filtering18none_sample_filterEEEvPNT2_7INDEX_TEPNSA_10DISTANCE_TEjPKSA_PKNSA_6DATA_TEPKSB_jjmSL_jSC_jjjjPjjjjjT3_' with regcount of 255 ptxas error : Entry function '_ZN4cuvs9neighbors5cagra6detail17single_cta_search13search_kernelILj512ELj64ELj1ENS2_25dataset_descriptor_base_tI6__halfjfEENS0_9filtering18none_sample_filterEEEvPNT2_7INDEX_TEPNSA_10DISTANCE_TEjPKSA_PKNSA_6DATA_TEPKSB_jjmSL_jSC_jjjjPjjjjT3_' with max regcount of 64 calls function '_ZN4cuvs9neighbors5cagra6detail17single_cta_search11search_coreILj512ELj64ELj1ENS2_25dataset_descriptor_base_tI6__halfjfEENS0_9filtering18none_sample_filterEEEvPNT2_7INDEX_TEPNSA_10DISTANCE_TEjPKSA_PKNSA_6DATA_TEPKSB_jjmSL_jSC_jjjjPjjjjjT3_' with regcount of 255 ptxas error : Entry function '_ZN4cuvs9neighbors5cagra6detail17single_cta_search13search_kernelILj64ELj128ELj1ENS2_25dataset_descriptor_base_tI6__halfjfEENS0_9filtering18none_sample_filterEEEvPNT2_7INDEX_TEPNSA_10DISTANCE_TEjPKSA_PKNSA_6DATA_TEPKSB_jjmSL_jSC_jjjjPjjjjT3_' with max regcount of 64 calls function '_ZN4cuvs9neighbors5cagra6detail17single_cta_search11search_coreILj64ELj128ELj1ENS2_25dataset_descriptor_base_tI6__halfjfEENS0_9filtering18none_sample_filterEEEvPNT2_7INDEX_TEPNSA_10DISTANCE_TEjPKSA_PKNSA_6DATA_TEPKSB_jjmSL_jSC_jjjjPjjjjjT3_' with regcount of 255 ptxas error : Entry function '_ZN4cuvs9neighbors5cagra6detail17single_cta_search15search_kernel_pILj64ELj128ELj1ENS2_25dataset_descriptor_base_tI6__halfjfEENS0_9filtering18none_sample_filterEEEvPKT2_PNS3_15worker_handle_tEPNS3_10job_desc_tISA_EEPjPKNSA_7INDEX_TEjjmSL_jPSJ_jjjjSI_jjjT3_' with max regcount of 64 calls function '_ZN4cuvs9neighbors5cagra6detail17single_cta_search11search_coreILj64ELj128ELj1ENS2_25dataset_descriptor_base_tI6__halfjfEENS0_9filtering18none_sample_filterEEEvPNT2_7INDEX_TEPNSA_10DISTANCE_TEjPKSA_PKNSA_6DATA_TEPKSB_jjmSL_jSC_jjjjPjjjjjT3_' with regcount of 255 ptxas error : Entry function '_ZN4cuvs9neighbors5cagra6detail17single_cta_search15search_kernel_pILj128ELj128ELj1ENS2_25dataset_descriptor_base_tI6__halfjfEENS0_9filtering18none_sample_filterEEEvPKT2_PNS3_15worker_handle_tEPNS3_10job_desc_tISA_EEPjPKNSA_7INDEX_TEjjmSL_jPSJ_jjjjSI_jjjT3_' with max regcount of 64 calls function '_ZN4cuvs9neighbors5cagra6detail17single_cta_search11search_coreILj128ELj128ELj1ENS2_25dataset_descriptor_base_tI6__halfjfEENS0_9filtering18none_sample_filterEEEvPNT2_7INDEX_TEPNSA_10DISTANCE_TEjPKSA_PKNSA_6DATA_TEPKSB_jjmSL_jSC_jjjjPjjjjjT3_' with regcount of 255 ptxas error : Entry function '_ZN4cuvs9neighbors5cagra6detail17single_cta_search13search_kernelILj128ELj128ELj1ENS2_25dataset_descriptor_base_tI6__halfjfEENS0_9filtering18none_sample_filterEEEvPNT2_7INDEX_TEPNSA_10DISTANCE_TEjPKSA_PKNSA_6DATA_TEPKSB_jjmSL_jSC_jjjjPjjjjT3_' with max regcount of 64 calls function '_ZN4cuvs9neighbors5cagra6detail17single_cta_search11search_coreILj128ELj128ELj1ENS2_25dataset_descriptor_base_tI6__halfjfEENS0_9filtering18none_sample_filterEEEvPNT2_7INDEX_TEPNSA_10DISTANCE_TEjPKSA_PKNSA_6DATA_TEPKSB_jjmSL_jSC_jjjjPjjjjjT3_' with regcount of 255 ptxas error : Entry function '_ZN4cuvs9neighbors5cagra6detail17single_cta_search13search_kernelILj256ELj128ELj1ENS2_25dataset_descriptor_base_tI6__halfjfEENS0_9filtering18none_sample_filterEEEvPNT2_7INDEX_TEPNSA_10DISTANCE_TEjPKSA_PKNSA_6DATA_TEPKSB_jjmSL_jSC_jjjjPjjjjT3_' with max regcount of 64 calls function '_ZN4cuvs9neighbors5cagra6detail17single_cta_search11search_coreILj256ELj128ELj1ENS2_25dataset_descriptor_base_tI6__halfjfEENS0_9filtering18none_sample_filterEEEvPNT2_7INDEX_TEPNSA_10DISTANCE_TEjPKSA_PKNSA_6DATA_TEPKSB_jjmSL_jSC_jjjjPjjjjjT3_' with regcount of 255 ptxas error : Entry function '_ZN4cuvs9neighbors5cagra6detail17single_cta_search15search_kernel_pILj256ELj128ELj1ENS2_25dataset_descriptor_base_tI6__halfjfEENS0_9filtering18none_sample_filterEEEvPKT2_PNS3_15worker_handle_tEPNS3_10job_desc_tISA_EEPjPKNSA_7INDEX_TEjjmSL_jPSJ_jjjjSI_jjjT3_' with max regcount of 64 calls function '_ZN4cuvs9neighbors5cagra6detail17single_cta_search11search_coreILj256ELj128ELj1ENS2_25dataset_descriptor_base_tI6__halfjfEENS0_9filtering18none_sample_filterEEEvPNT2_7INDEX_TEPNSA_10DISTANCE_TEjPKSA_PKNSA_6DATA_TEPKSB_jjmSL_jSC_jjjjPjjjjjT3_' with regcount of 255 ptxas error : Entry function '_ZN4cuvs9neighbors5cagra6detail17single_cta_search15search_kernel_pILj512ELj128ELj1ENS2_25dataset_descriptor_base_tI6__halfjfEENS0_9filtering18none_sample_filterEEEvPKT2_PNS3_15worker_handle_tEPNS3_10job_desc_tISA_EEPjPKNSA_7INDEX_TEjjmSL_jPSJ_jjjjSI_jjjT3_' with max regcount of 64 calls function '_ZN4cuvs9neighbors5cagra6detail17single_cta_search11search_coreILj512ELj128ELj1ENS2_25dataset_descriptor_base_tI6__halfjfEENS0_9filtering18none_sample_filterEEEvPNT2_7INDEX_TEPNSA_10DISTANCE_TEjPKSA_PKNSA_6DATA_TEPKSB_jjmSL_jSC_jjjjPjjjjjT3_' with regcount of 255 ptxas error : Entry function '_ZN4cuvs9neighbors5cagra6detail17single_cta_search13search_kernelILj512ELj128ELj1ENS2_25dataset_descriptor_base_tI6__halfjfEENS0_9filtering18none_sample_filterEEEvPNT2_7INDEX_TEPNSA_10DISTANCE_TEjPKSA_PKNSA_6DATA_TEPKSB_jjmSL_jSC_jjjjPjjjjT3_' with max regcount of 64 calls function '_ZN4cuvs9neighbors5cagra6detail17single_cta_search11search_coreILj512ELj128ELj1ENS2_25dataset_descriptor_base_tI6__halfjfEENS0_9filtering18none_sample_filterEEEvPNT2_7INDEX_TEPNSA_10DISTANCE_TEjPKSA_PKNSA_6DATA_TEPKSB_jjmSL_jSC_jjjjPjjjjjT3_' with regcount of 255 ptxas error : Entry function '_ZN4cuvs9neighbors5cagra6detail17single_cta_search13search_kernelILj64ELj256ELj1ENS2_25dataset_descriptor_base_tI6__halfjfEENS0_9filtering18none_sample_filterEEEvPNT2_7INDEX_TEPNSA_10DISTANCE_TEjPKSA_PKNSA_6DATA_TEPKSB_jjmSL_jSC_jjjjPjjjjT3_' with max regcount of 64 calls function '_ZN4cuvs9neighbors5cagra6detail17single_cta_search11search_coreILj64ELj256ELj1ENS2_25dataset_descriptor_base_tI6__halfjfEENS0_9filtering18none_sample_filterEEEvPNT2_7INDEX_TEPNSA_10DISTANCE_TEjPKSA_PKNSA_6DATA_TEPKSB_jjmSL_jSC_jjjjPjjjjjT3_' with regcount of 255 ptxas error : Entry function '_ZN4cuvs9neighbors5cagra6detail17single_cta_search15search_kernel_pILj64ELj256ELj1ENS2_25dataset_descriptor_base_tI6__halfjfEENS0_9filtering18none_sample_filterEEEvPKT2_PNS3_15worker_handle_tEPNS3_10job_desc_tISA_EEPjPKNSA_7INDEX_TEjjmSL_jPSJ_jjjjSI_jjjT3_' with max regcount of 64 calls function '_ZN4cuvs9neighbors5cagra6detail17single_cta_search11search_coreILj64ELj256ELj1ENS2_25dataset_descriptor_base_tI6__halfjfEENS0_9filtering18none_sample_filterEEEvPNT2_7INDEX_TEPNSA_10DISTANCE_TEjPKSA_PKNSA_6DATA_TEPKSB_jjmSL_jSC_jjjjPjjjjjT3_' with regcount of 255 ptxas error : Entry function '_ZN4cuvs9neighbors5cagra6detail17single_cta_search15search_kernel_pILj128ELj256ELj1ENS2_25dataset_descriptor_base_tI6__halfjfEENS0_9filtering18none_sample_filterEEEvPKT2_PNS3_15worker_handle_tEPNS3_10job_desc_tISA_EEPjPKNSA_7INDEX_TEjjmSL_jPSJ_jjjjSI_jjjT3_' with max regcount of 64 calls function '_ZN4cuvs9neighbors5cagra6detail17single_cta_search11search_coreILj128ELj256ELj1ENS2_25dataset_descriptor_base_tI6__halfjfEENS0_9filtering18none_sample_filterEEEvPNT2_7INDEX_TEPNSA_10DISTANCE_TEjPKSA_PKNSA_6DATA_TEPKSB_jjmSL_jSC_jjjjPjjjjjT3_' with regcount of 255 ptxas error : Entry function '_ZN4cuvs9neighbors5cagra6detail17single_cta_search13search_kernelILj128ELj256ELj1ENS2_25dataset_descriptor_base_tI6__halfjfEENS0_9filtering18none_sample_filterEEEvPNT2_7INDEX_TEPNSA_10DISTANCE_TEjPKSA_PKNSA_6DATA_TEPKSB_jjmSL_jSC_jjjjPjjjjT3_' with max regcount of 64 calls function '_ZN4cuvs9neighbors5cagra6detail17single_cta_search11search_coreILj128ELj256ELj1ENS2_25dataset_descriptor_base_tI6__halfjfEENS0_9filtering18none_sample_filterEEEvPNT2_7INDEX_TEPNSA_10DISTANCE_TEjPKSA_PKNSA_6DATA_TEPKSB_jjmSL_jSC_jjjjPjjjjjT3_' with regcount of 255 ptxas error : Entry function '_ZN4cuvs9neighbors5cagra6detail17single_cta_search13search_kernelILj256ELj256ELj1ENS2_25dataset_descriptor_base_tI6__halfjfEENS0_9filtering18none_sample_filterEEEvPNT2_7INDEX_TEPNSA_10DISTANCE_TEjPKSA_PKNSA_6DATA_TEPKSB_jjmSL_jSC_jjjjPjjjjT3_' with max regcount of 64 calls function '_ZN4cuvs9neighbors5cagra6detail17single_cta_search11search_coreILj256ELj256ELj1ENS2_25dataset_descriptor_base_tI6__halfjfEENS0_9filtering18none_sample_filterEEEvPNT2_7INDEX_TEPNSA_10DISTANCE_TEjPKSA_PKNSA_6DATA_TEPKSB_jjmSL_jSC_jjjjPjjjjjT3_' with regcount of 255 ptxas error : Entry function '_ZN4cuvs9neighbors5cagra6detail17single_cta_search15search_kernel_pILj256ELj256ELj1ENS2_25dataset_descriptor_base_tI6__halfjfEENS0_9filtering18none_sample_filterEEEvPKT2_PNS3_15worker_handle_tEPNS3_10job_desc_tISA_EEPjPKNSA_7INDEX_TEjjmSL_jPSJ_jjjjSI_jjjT3_' with max regcount of 64 calls function '_ZN4cuvs9neighbors5cagra6detail17single_cta_search11search_coreILj256ELj256ELj1ENS2_25dataset_descriptor_base_tI6__halfjfEENS0_9filtering18none_sample_filterEEEvPNT2_7INDEX_TEPNSA_10DISTANCE_TEjPKSA_PKNSA_6DATA_TEPKSB_jjmSL_jSC_jjjjPjjjjjT3_' with regcount of 255 ptxas error : Entry function '_ZN4cuvs9neighbors5cagra6detail17single_cta_search15search_kernel_pILj512ELj256ELj1ENS2_25dataset_descriptor_base_tI6__halfjfEENS0_9filtering18none_sample_filterEEEvPKT2_PNS3_15worker_handle_tEPNS3_10job_desc_tISA_EEPjPKNSA_7INDEX_TEjjmSL_jPSJ_jjjjSI_jjjT3_' with max regcount of 64 calls function '_ZN4cuvs9neighbors5cagra6detail17single_cta_search11search_coreILj512ELj256ELj1ENS2_25dataset_descriptor_base_tI6__halfjfEENS0_9filtering18none_sample_filterEEEvPNT2_7INDEX_TEPNSA_10DISTANCE_TEjPKSA_PKNSA_6DATA_TEPKSB_jjmSL_jSC_jjjjPjjjjjT3_' with regcount of 255 ptxas error : Entry function '_ZN4cuvs9neighbors5cagra6detail17single_cta_search13search_kernelILj512ELj256ELj1ENS2_25dataset_descriptor_base_tI6__halfjfEENS0_9filtering18none_sample_filterEEEvPNT2_7INDEX_TEPNSA_10DISTANCE_TEjPKSA_PKNSA_6DATA_TEPKSB_jjmSL_jSC_jjjjPjjjjT3_' with max regcount of 64 calls function '_ZN4cuvs9neighbors5cagra6detail17single_cta_search11search_coreILj512ELj256ELj1ENS2_25dataset_descriptor_base_tI6__halfjfEENS0_9filtering18none_sample_filterEEEvPNT2_7INDEX_TEPNSA_10DISTANCE_TEjPKSA_PKNSA_6DATA_TEPKSB_jjmSL_jSC_jjjjPjjjjjT3_' with regcount of 255 ptxas error : Entry function '_ZN4cuvs9neighbors5cagra6detail17single_cta_search13search_kernelILj256ELj32ELj0ENS2_25dataset_descriptor_base_tI6__halfjfEENS0_9filtering18none_sample_filterEEEvPNT2_7INDEX_TEPNSA_10DISTANCE_TEjPKSA_PKNSA_6DATA_TEPKSB_jjmSL_jSC_jjjjPjjjjT3_' with max regcount of 64 calls function '_ZN4cuvs9neighbors5cagra6detail17single_cta_search11search_coreILj256ELj32ELj0ENS2_25dataset_descriptor_base_tI6__halfjfEENS0_9filtering18none_sample_filterEEEvPNT2_7INDEX_TEPNSA_10DISTANCE_TEjPKSA_PKNSA_6DATA_TEPKSB_jjmSL_jSC_jjjjPjjjjjT3_' with regcount of 255 ptxas error : Entry function '_ZN4cuvs9neighbors5cagra6detail17single_cta_search15search_kernel_pILj256ELj32ELj0ENS2_25dataset_descriptor_base_tI6__halfjfEENS0_9filtering18none_sample_filterEEEvPKT2_PNS3_15worker_handle_tEPNS3_10job_desc_tISA_EEPjPKNSA_7INDEX_TEjjmSL_jPSJ_jjjjSI_jjjT3_' with max regcount of 64 calls function '_ZN4cuvs9neighbors5cagra6detail17single_cta_search11search_coreILj256ELj32ELj0ENS2_25dataset_descriptor_base_tI6__halfjfEENS0_9filtering18none_sample_filterEEEvPNT2_7INDEX_TEPNSA_10DISTANCE_TEjPKSA_PKNSA_6DATA_TEPKSB_jjmSL_jSC_jjjjPjjjjjT3_' with regcount of 255 ptxas error : Entry function '_ZN4cuvs9neighbors5cagra6detail17single_cta_search15search_kernel_pILj512ELj32ELj0ENS2_25dataset_descriptor_base_tI6__halfjfEENS0_9filtering18none_sample_filterEEEvPKT2_PNS3_15worker_handle_tEPNS3_10job_desc_tISA_EEPjPKNSA_7INDEX_TEjjmSL_jPSJ_jjjjSI_jjjT3_' with max regcount of 64 calls function '_ZN4cuvs9neighbors5cagra6detail17single_cta_search11search_coreILj512ELj32ELj0ENS2_25dataset_descriptor_base_tI6__halfjfEENS0_9filtering18none_sample_filterEEEvPNT2_7INDEX_TEPNSA_10DISTANCE_TEjPKSA_PKNSA_6DATA_TEPKSB_jjmSL_jSC_jjjjPjjjjjT3_' with regcount of 255 ptxas error : Entry function '_ZN4cuvs9neighbors5cagra6detail17single_cta_search13search_kernelILj512ELj32ELj0ENS2_25dataset_descriptor_base_tI6__halfjfEENS0_9filtering18none_sample_filterEEEvPNT2_7INDEX_TEPNSA_10DISTANCE_TEjPKSA_PKNSA_6DATA_TEPKSB_jjmSL_jSC_jjjjPjjjjT3_' with max regcount of 64 calls function '_ZN4cuvs9neighbors5cagra6detail17single_cta_search11search_coreILj512ELj32ELj0ENS2_25dataset_descriptor_base_tI6__halfjfEENS0_9filtering18none_sample_filterEEEvPNT2_7INDEX_TEPNSA_10DISTANCE_TEjPKSA_PKNSA_6DATA_TEPKSB_jjmSL_jSC_jjjjPjjjjjT3_' with regcount of 255 ptxas error : Entry function '_ZN4cuvs9neighbors5cagra6detail17single_cta_search15search_kernel_pILj64ELj64ELj0ENS2_25dataset_descriptor_base_tI6__halfjfEENS2_34CagraSampleFilterWithQueryIdOffsetINS0_9filtering13bitset_filterIjlEEEEEEvPKT2_PNS3_15worker_handle_tEPNS3_10job_desc_tISD_EEPjPKNSD_7INDEX_TEjjmSO_jPSM_jjjjSL_jjjT3_' with max regcount of 64 calls function '_ZN4cuvs9neighbors5cagra6detail17single_cta_search11search_coreILj64ELj64ELj0ENS2_25dataset_descriptor_base_tI6__halfjfEENS2_34CagraSampleFilterWithQueryIdOffsetINS0_9filtering13bitset_filterIjlEEEEEEvPNT2_7INDEX_TEPNSD_10DISTANCE_TEjPKSD_PKNSD_6DATA_TEPKSE_jjmSO_jSF_jjjjPjjjjjT3_' with regcount of 255 ptxas error : Entry function '_ZN4cuvs9neighbors5cagra6detail17single_cta_search13search_kernelILj64ELj64ELj0ENS2_25dataset_descriptor_base_tI6__halfjfEENS2_34CagraSampleFilterWithQueryIdOffsetINS0_9filtering13bitset_filterIjlEEEEEEvPNT2_7INDEX_TEPNSD_10DISTANCE_TEjPKSD_PKNSD_6DATA_TEPKSE_jjmSO_jSF_jjjjPjjjjT3_' with max regcount of 64 calls function '_ZN4cuvs9neighbors5cagra6detail17single_cta_search11search_coreILj64ELj64ELj0ENS2_25dataset_descriptor_base_tI6__halfjfEENS2_34CagraSampleFilterWithQueryIdOffsetINS0_9filtering13bitset_filterIjlEEEEEEvPNT2_7INDEX_TEPNSD_10DISTANCE_TEjPKSD_PKNSD_6DATA_TEPKSE_jjmSO_jSF_jjjjPjjjjjT3_' with regcount of 255 ptxas error : Entry function '_ZN4cuvs9neighbors5cagra6detail17single_cta_search13search_kernelILj64ELj64ELj1ENS2_25dataset_descriptor_base_tI6__halfjfEENS2_34CagraSampleFilterWithQueryIdOffsetINS0_9filtering13bitset_filterIjlEEEEEEvPNT2_7INDEX_TEPNSD_10DISTANCE_TEjPKSD_PKNSD_6DATA_TEPKSE_jjmSO_jSF_jjjjPjjjjT3_' with max regcount of 64 calls function '_ZN4cuvs9neighbors5cagra6detail17single_cta_search11search_coreILj64ELj64ELj1ENS2_25dataset_descriptor_base_tI6__halfjfEENS2_34CagraSampleFilterWithQueryIdOffsetINS0_9filtering13bitset_filterIjlEEEEEEvPNT2_7INDEX_TEPNSD_10DISTANCE_TEjPKSD_PKNSD_6DATA_TEPKSE_jjmSO_jSF_jjjjPjjjjjT3_' with regcount of 255 ptxas error : Entry function '_ZN4cuvs9neighbors5cagra6detail17single_cta_search15search_kernel_pILj64ELj64ELj1ENS2_25dataset_descriptor_base_tI6__halfjfEENS2_34CagraSampleFilterWithQueryIdOffsetINS0_9filtering13bitset_filterIjlEEEEEEvPKT2_PNS3_15worker_handle_tEPNS3_10job_desc_tISD_EEPjPKNSD_7INDEX_TEjjmSO_jPSM_jjjjSL_jjjT3_' with max regcount of 64 calls function '_ZN4cuvs9neighbors5cagra6detail17single_cta_search11search_coreILj64ELj64ELj1ENS2_25dataset_descriptor_base_tI6__halfjfEENS2_34CagraSampleFilterWithQueryIdOffsetINS0_9filtering13bitset_filterIjlEEEEEEvPNT2_7INDEX_TEPNSD_10DISTANCE_TEjPKSD_PKNSD_6DATA_TEPKSE_jjmSO_jSF_jjjjPjjjjjT3_' with regcount of 255 ptxas error : Entry function '_ZN4cuvs9neighbors5cagra6detail17single_cta_search15search_kernel_pILj128ELj64ELj1ENS2_25dataset_descriptor_base_tI6__halfjfEENS2_34CagraSampleFilterWithQueryIdOffsetINS0_9filtering13bitset_filterIjlEEEEEEvPKT2_PNS3_15worker_handle_tEPNS3_10job_desc_tISD_EEPjPKNSD_7INDEX_TEjjmSO_jPSM_jjjjSL_jjjT3_' with max regcount of 64 calls function '_ZN4cuvs9neighbors5cagra6detail17single_cta_search11search_coreILj128ELj64ELj1ENS2_25dataset_descriptor_base_tI6__halfjfEENS2_34CagraSampleFilterWithQueryIdOffsetINS0_9filtering13bitset_filterIjlEEEEEEvPNT2_7INDEX_TEPNSD_10DISTANCE_TEjPKSD_PKNSD_6DATA_TEPKSE_jjmSO_jSF_jjjjPjjjjjT3_' with regcount of 255 ptxas error : Entry function '_ZN4cuvs9neighbors5cagra6detail17single_cta_search13search_kernelILj128ELj64ELj1ENS2_25dataset_descriptor_base_tI6__halfjfEENS2_34CagraSampleFilterWithQueryIdOffsetINS0_9filtering13bitset_filterIjlEEEEEEvPNT2_7INDEX_TEPNSD_10DISTANCE_TEjPKSD_PKNSD_6DATA_TEPKSE_jjmSO_jSF_jjjjPjjjjT3_' with max regcount of 64 calls function '_ZN4cuvs9neighbors5cagra6detail17single_cta_search11search_coreILj128ELj64ELj1ENS2_25dataset_descriptor_base_tI6__halfjfEENS2_34CagraSampleFilterWithQueryIdOffsetINS0_9filtering13bitset_filterIjlEEEEEEvPNT2_7INDEX_TEPNSD_10DISTANCE_TEjPKSD_PKNSD_6DATA_TEPKSE_jjmSO_jSF_jjjjPjjjjjT3_' with regcount of 255 ptxas error : Entry function '_ZN4cuvs9neighbors5cagra6detail17single_cta_search13search_kernelILj256ELj64ELj1ENS2_25dataset_descriptor_base_tI6__halfjfEENS2_34CagraSampleFilterWithQueryIdOffsetINS0_9filtering13bitset_filterIjlEEEEEEvPNT2_7INDEX_TEPNSD_10DISTANCE_TEjPKSD_PKNSD_6DATA_TEPKSE_jjmSO_jSF_jjjjPjjjjT3_' with max regcount of 64 calls function '_ZN4cuvs9neighbors5cagra6detail17single_cta_search11search_coreILj256ELj64ELj1ENS2_25dataset_descriptor_base_tI6__halfjfEENS2_34CagraSampleFilterWithQueryIdOffsetINS0_9filtering13bitset_filterIjlEEEEEEvPNT2_7INDEX_TEPNSD_10DISTANCE_TEjPKSD_PKNSD_6DATA_TEPKSE_jjmSO_jSF_jjjjPjjjjjT3_' with regcount of 255 ptxas error : Entry function '_ZN4cuvs9neighbors5cagra6detail17single_cta_search15search_kernel_pILj256ELj64ELj1ENS2_25dataset_descriptor_base_tI6__halfjfEENS2_34CagraSampleFilterWithQueryIdOffsetINS0_9filtering13bitset_filterIjlEEEEEEvPKT2_PNS3_15worker_handle_tEPNS3_10job_desc_tISD_EEPjPKNSD_7INDEX_TEjjmSO_jPSM_jjjjSL_jjjT3_' with max regcount of 64 calls function '_ZN4cuvs9neighbors5cagra6detail17single_cta_search11search_coreILj256ELj64ELj1ENS2_25dataset_descriptor_base_tI6__halfjfEENS2_34CagraSampleFilterWithQueryIdOffsetINS0_9filtering13bitset_filterIjlEEEEEEvPNT2_7INDEX_TEPNSD_10DISTANCE_TEjPKSD_PKNSD_6DATA_TEPKSE_jjmSO_jSF_jjjjPjjjjjT3_' with regcount of 255 ptxas error : Entry function '_ZN4cuvs9neighbors5cagra6detail17single_cta_search15search_kernel_pILj512ELj64ELj1ENS2_25dataset_descriptor_base_tI6__halfjfEENS2_34CagraSampleFilterWithQueryIdOffsetINS0_9filtering13bitset_filterIjlEEEEEEvPKT2_PNS3_15worker_handle_tEPNS3_10job_desc_tISD_EEPjPKNSD_7INDEX_TEjjmSO_jPSM_jjjjSL_jjjT3_' with max regcount of 64 calls function '_ZN4cuvs9neighbors5cagra6detail17single_cta_search11search_coreILj512ELj64ELj1ENS2_25dataset_descriptor_base_tI6__halfjfEENS2_34CagraSampleFilterWithQueryIdOffsetINS0_9filtering13bitset_filterIjlEEEEEEvPNT2_7INDEX_TEPNSD_10DISTANCE_TEjPKSD_PKNSD_6DATA_TEPKSE_jjmSO_jSF_jjjjPjjjjjT3_' with regcount of 255 ptxas error : Entry function '_ZN4cuvs9neighbors5cagra6detail17single_cta_search13search_kernelILj512ELj64ELj1ENS2_25dataset_descriptor_base_tI6__halfjfEENS2_34CagraSampleFilterWithQueryIdOffsetINS0_9filtering13bitset_filterIjlEEEEEEvPNT2_7INDEX_TEPNSD_10DISTANCE_TEjPKSD_PKNSD_6DATA_TEPKSE_jjmSO_jSF_jjjjPjjjjT3_' with max regcount of 64 calls function '_ZN4cuvs9neighbors5cagra6detail17single_cta_search11search_coreILj512ELj64ELj1ENS2_25dataset_descriptor_base_tI6__halfjfEENS2_34CagraSampleFilterWithQueryIdOffsetINS0_9filtering13bitset_filterIjlEEEEEEvPNT2_7INDEX_TEPNSD_10DISTANCE_TEjPKSD_PKNSD_6DATA_TEPKSE_jjmSO_jSF_jjjjPjjjjjT3_' with regcount of 255 ptxas error : Entry function '_ZN4cuvs9neighbors5cagra6detail17single_cta_search13search_kernelILj64ELj128ELj1ENS2_25dataset_descriptor_base_tI6__halfjfEENS2_34CagraSampleFilterWithQueryIdOffsetINS0_9filtering13bitset_filterIjlEEEEEEvPNT2_7INDEX_TEPNSD_10DISTANCE_TEjPKSD_PKNSD_6DATA_TEPKSE_jjmSO_jSF_jjjjPjjjjT3_' with max regcount of 64 calls function '_ZN4cuvs9neighbors5cagra6detail17single_cta_search11search_coreILj64ELj128ELj1ENS2_25dataset_descriptor_base_tI6__halfjfEENS2_34CagraSampleFilterWithQueryIdOffsetINS0_9filtering13bitset_filterIjlEEEEEEvPNT2_7INDEX_TEPNSD_10DISTANCE_TEjPKSD_PKNSD_6DATA_TEPKSE_jjmSO_jSF_jjjjPjjjjjT3_' with regcount of 255 ptxas error : Entry function '_ZN4cuvs9neighbors5cagra6detail17single_cta_search15search_kernel_pILj64ELj128ELj1ENS2_25dataset_descriptor_base_tI6__halfjfEENS2_34CagraSampleFilterWithQueryIdOffsetINS0_9filtering13bitset_filterIjlEEEEEEvPKT2_PNS3_15worker_handle_tEPNS3_10job_desc_tISD_EEPjPKNSD_7INDEX_TEjjmSO_jPSM_jjjjSL_jjjT3_' with max regcount of 64 calls function '_ZN4cuvs9neighbors5cagra6detail17single_cta_search11search_coreILj64ELj128ELj1ENS2_25dataset_descriptor_base_tI6__halfjfEENS2_34CagraSampleFilterWithQueryIdOffsetINS0_9filtering13bitset_filterIjlEEEEEEvPNT2_7INDEX_TEPNSD_10DISTANCE_TEjPKSD_PKNSD_6DATA_TEPKSE_jjmSO_jSF_jjjjPjjjjjT3_' with regcount of 255 ptxas error : Entry function '_ZN4cuvs9neighbors5cagra6detail17single_cta_search15search_kernel_pILj128ELj128ELj1ENS2_25dataset_descriptor_base_tI6__halfjfEENS2_34CagraSampleFilterWithQueryIdOffsetINS0_9filtering13bitset_filterIjlEEEEEEvPKT2_PNS3_15worker_handle_tEPNS3_10job_desc_tISD_EEPjPKNSD_7INDEX_TEjjmSO_jPSM_jjjjSL_jjjT3_' with max regcount of 64 calls function '_ZN4cuvs9neighbors5cagra6detail17single_cta_search11search_coreILj128ELj128ELj1ENS2_25dataset_descriptor_base_tI6__halfjfEENS2_34CagraSampleFilterWithQueryIdOffsetINS0_9filtering13bitset_filterIjlEEEEEEvPNT2_7INDEX_TEPNSD_10DISTANCE_TEjPKSD_PKNSD_6DATA_TEPKSE_jjmSO_jSF_jjjjPjjjjjT3_' with regcount of 255 ptxas error : Entry function '_ZN4cuvs9neighbors5cagra6detail17single_cta_search13search_kernelILj128ELj128ELj1ENS2_25dataset_descriptor_base_tI6__halfjfEENS2_34CagraSampleFilterWithQueryIdOffsetINS0_9filtering13bitset_filterIjlEEEEEEvPNT2_7INDEX_TEPNSD_10DISTANCE_TEjPKSD_PKNSD_6DATA_TEPKSE_jjmSO_jSF_jjjjPjjjjT3_' with max regcount of 64 calls function '_ZN4cuvs9neighbors5cagra6detail17single_cta_search11search_coreILj128ELj128ELj1ENS2_25dataset_descriptor_base_tI6__halfjfEENS2_34CagraSampleFilterWithQueryIdOffsetINS0_9filtering13bitset_filterIjlEEEEEEvPNT2_7INDEX_TEPNSD_10DISTANCE_TEjPKSD_PKNSD_6DATA_TEPKSE_jjmSO_jSF_jjjjPjjjjjT3_' with regcount of 255 ptxas error : Entry function '_ZN4cuvs9neighbors5cagra6detail17single_cta_search13search_kernelILj256ELj128ELj1ENS2_25dataset_descriptor_base_tI6__halfjfEENS2_34CagraSampleFilterWithQueryIdOffsetINS0_9filtering13bitset_filterIjlEEEEEEvPNT2_7INDEX_TEPNSD_10DISTANCE_TEjPKSD_PKNSD_6DATA_TEPKSE_jjmSO_jSF_jjjjPjjjjT3_' with max regcount of 64 calls function '_ZN4cuvs9neighbors5cagra6detail17single_cta_search11search_coreILj256ELj128ELj1ENS2_25dataset_descriptor_base_tI6__halfjfEENS2_34CagraSampleFilterWithQueryIdOffsetINS0_9filtering13bitset_filterIjlEEEEEEvPNT2_7INDEX_TEPNSD_10DISTANCE_TEjPKSD_PKNSD_6DATA_TEPKSE_jjmSO_jSF_jjjjPjjjjjT3_' with regcount of 255 ptxas error : Entry function '_ZN4cuvs9neighbors5cagra6detail17single_cta_search15search_kernel_pILj256ELj128ELj1ENS2_25dataset_descriptor_base_tI6__halfjfEENS2_34CagraSampleFilterWithQueryIdOffsetINS0_9filtering13bitset_filterIjlEEEEEEvPKT2_PNS3_15worker_handle_tEPNS3_10job_desc_tISD_EEPjPKNSD_7INDEX_TEjjmSO_jPSM_jjjjSL_jjjT3_' with max regcount of 64 calls function '_ZN4cuvs9neighbors5cagra6detail17single_cta_search11search_coreILj256ELj128ELj1ENS2_25dataset_descriptor_base_tI6__halfjfEENS2_34CagraSampleFilterWithQueryIdOffsetINS0_9filtering13bitset_filterIjlEEEEEEvPNT2_7INDEX_TEPNSD_10DISTANCE_TEjPKSD_PKNSD_6DATA_TEPKSE_jjmSO_jSF_jjjjPjjjjjT3_' with regcount of 255 ptxas error : Entry function '_ZN4cuvs9neighbors5cagra6detail17single_cta_search15search_kernel_pILj512ELj128ELj1ENS2_25dataset_descriptor_base_tI6__halfjfEENS2_34CagraSampleFilterWithQueryIdOffsetINS0_9filtering13bitset_filterIjlEEEEEEvPKT2_PNS3_15worker_handle_tEPNS3_10job_desc_tISD_EEPjPKNSD_7INDEX_TEjjmSO_jPSM_jjjjSL_jjjT3_' with max regcount of 64 calls function '_ZN4cuvs9neighbors5cagra6detail17single_cta_search11search_coreILj512ELj128ELj1ENS2_25dataset_descriptor_base_tI6__halfjfEENS2_34CagraSampleFilterWithQueryIdOffsetINS0_9filtering13bitset_filterIjlEEEEEEvPNT2_7INDEX_TEPNSD_10DISTANCE_TEjPKSD_PKNSD_6DATA_TEPKSE_jjmSO_jSF_jjjjPjjjjjT3_' with regcount of 255 ptxas error : Entry function '_ZN4cuvs9neighbors5cagra6detail17single_cta_search13search_kernelILj512ELj128ELj1ENS2_25dataset_descriptor_base_tI6__halfjfEENS2_34CagraSampleFilterWithQueryIdOffsetINS0_9filtering13bitset_filterIjlEEEEEEvPNT2_7INDEX_TEPNSD_10DISTANCE_TEjPKSD_PKNSD_6DATA_TEPKSE_jjmSO_jSF_jjjjPjjjjT3_' with max regcount of 64 calls function '_ZN4cuvs9neighbors5cagra6detail17single_cta_search11search_coreILj512ELj128ELj1ENS2_25dataset_descriptor_base_tI6__halfjfEENS2_34CagraSampleFilterWithQueryIdOffsetINS0_9filtering13bitset_filterIjlEEEEEEvPNT2_7INDEX_TEPNSD_10DISTANCE_TEjPKSD_PKNSD_6DATA_TEPKSE_jjmSO_jSF_jjjjPjjjjjT3_' with regcount of 255 ptxas error : Entry function '_ZN4cuvs9neighbors5cagra6detail17single_cta_search13search_kernelILj64ELj256ELj1ENS2_25dataset_descriptor_base_tI6__halfjfEENS2_34CagraSampleFilterWithQueryIdOffsetINS0_9filtering13bitset_filterIjlEEEEEEvPNT2_7INDEX_TEPNSD_10DISTANCE_TEjPKSD_PKNSD_6DATA_TEPKSE_jjmSO_jSF_jjjjPjjjjT3_' with max regcount of 64 calls function '_ZN4cuvs9neighbors5cagra6detail17single_cta_search11search_coreILj64ELj256ELj1ENS2_25dataset_descriptor_base_tI6__halfjfEENS2_34CagraSampleFilterWithQueryIdOffsetINS0_9filtering13bitset_filterIjlEEEEEEvPNT2_7INDEX_TEPNSD_10DISTANCE_TEjPKSD_PKNSD_6DATA_TEPKSE_jjmSO_jSF_jjjjPjjjjjT3_' with regcount of 255 ptxas error : Entry function '_ZN4cuvs9neighbors5cagra6detail17single_cta_search15search_kernel_pILj64ELj256ELj1ENS2_25dataset_descriptor_base_tI6__halfjfEENS2_34CagraSampleFilterWithQueryIdOffsetINS0_9filtering13bitset_filterIjlEEEEEEvPKT2_PNS3_15worker_handle_tEPNS3_10job_desc_tISD_EEPjPKNSD_7INDEX_TEjjmSO_jPSM_jjjjSL_jjjT3_' with max regcount of 64 calls function '_ZN4cuvs9neighbors5cagra6detail17single_cta_search11search_coreILj64ELj256ELj1ENS2_25dataset_descriptor_base_tI6__halfjfEENS2_34CagraSampleFilterWithQueryIdOffsetINS0_9filtering13bitset_filterIjlEEEEEEvPNT2_7INDEX_TEPNSD_10DISTANCE_TEjPKSD_PKNSD_6DATA_TEPKSE_jjmSO_jSF_jjjjPjjjjjT3_' with regcount of 255 ptxas error : Entry function '_ZN4cuvs9neighbors5cagra6detail17single_cta_search15search_kernel_pILj128ELj256ELj1ENS2_25dataset_descriptor_base_tI6__halfjfEENS2_34CagraSampleFilterWithQueryIdOffsetINS0_9filtering13bitset_filterIjlEEEEEEvPKT2_PNS3_15worker_handle_tEPNS3_10job_desc_tISD_EEPjPKNSD_7INDEX_TEjjmSO_jPSM_jjjjSL_jjjT3_' with max regcount of 64 calls function '_ZN4cuvs9neighbors5cagra6detail17single_cta_search11search_coreILj128ELj256ELj1ENS2_25dataset_descriptor_base_tI6__halfjfEENS2_34CagraSampleFilterWithQueryIdOffsetINS0_9filtering13bitset_filterIjlEEEEEEvPNT2_7INDEX_TEPNSD_10DISTANCE_TEjPKSD_PKNSD_6DATA_TEPKSE_jjmSO_jSF_jjjjPjjjjjT3_' with regcount of 255 ptxas error : Entry function '_ZN4cuvs9neighbors5cagra6detail17single_cta_search13search_kernelILj128ELj256ELj1ENS2_25dataset_descriptor_base_tI6__halfjfEENS2_34CagraSampleFilterWithQueryIdOffsetINS0_9filtering13bitset_filterIjlEEEEEEvPNT2_7INDEX_TEPNSD_10DISTANCE_TEjPKSD_PKNSD_6DATA_TEPKSE_jjmSO_jSF_jjjjPjjjjT3_' with max regcount of 64 calls function '_ZN4cuvs9neighbors5cagra6detail17single_cta_search11search_coreILj128ELj256ELj1ENS2_25dataset_descriptor_base_tI6__halfjfEENS2_34CagraSampleFilterWithQueryIdOffsetINS0_9filtering13bitset_filterIjlEEEEEEvPNT2_7INDEX_TEPNSD_10DISTANCE_TEjPKSD_PKNSD_6DATA_TEPKSE_jjmSO_jSF_jjjjPjjjjjT3_' with regcount of 255 ptxas error : Entry function '_ZN4cuvs9neighbors5cagra6detail17single_cta_search13search_kernelILj256ELj256ELj1ENS2_25dataset_descriptor_base_tI6__halfjfEENS2_34CagraSampleFilterWithQueryIdOffsetINS0_9filtering13bitset_filterIjlEEEEEEvPNT2_7INDEX_TEPNSD_10DISTANCE_TEjPKSD_PKNSD_6DATA_TEPKSE_jjmSO_jSF_jjjjPjjjjT3_' with max regcount of 64 calls function '_ZN4cuvs9neighbors5cagra6detail17single_cta_search11search_coreILj256ELj256ELj1ENS2_25dataset_descriptor_base_tI6__halfjfEENS2_34CagraSampleFilterWithQueryIdOffsetINS0_9filtering13bitset_filterIjlEEEEEEvPNT2_7INDEX_TEPNSD_10DISTANCE_TEjPKSD_PKNSD_6DATA_TEPKSE_jjmSO_jSF_jjjjPjjjjjT3_' with regcount of 255 ptxas error : Entry function '_ZN4cuvs9neighbors5cagra6detail17single_cta_search15search_kernel_pILj256ELj256ELj1ENS2_25dataset_descriptor_base_tI6__halfjfEENS2_34CagraSampleFilterWithQueryIdOffsetINS0_9filtering13bitset_filterIjlEEEEEEvPKT2_PNS3_15worker_handle_tEPNS3_10job_desc_tISD_EEPjPKNSD_7INDEX_TEjjmSO_jPSM_jjjjSL_jjjT3_' with max regcount of 64 calls function '_ZN4cuvs9neighbors5cagra6detail17single_cta_search11search_coreILj256ELj256ELj1ENS2_25dataset_descriptor_base_tI6__halfjfEENS2_34CagraSampleFilterWithQueryIdOffsetINS0_9filtering13bitset_filterIjlEEEEEEvPNT2_7INDEX_TEPNSD_10DISTANCE_TEjPKSD_PKNSD_6DATA_TEPKSE_jjmSO_jSF_jjjjPjjjjjT3_' with regcount of 255 ptxas error : Entry function '_ZN4cuvs9neighbors5cagra6detail17single_cta_search15search_kernel_pILj512ELj256ELj1ENS2_25dataset_descriptor_base_tI6__halfjfEENS2_34CagraSampleFilterWithQueryIdOffsetINS0_9filtering13bitset_filterIjlEEEEEEvPKT2_PNS3_15worker_handle_tEPNS3_10job_desc_tISD_EEPjPKNSD_7INDEX_TEjjmSO_jPSM_jjjjSL_jjjT3_' with max regcount of 64 calls function '_ZN4cuvs9neighbors5cagra6detail17single_cta_search11search_coreILj512ELj256ELj1ENS2_25dataset_descriptor_base_tI6__halfjfEENS2_34CagraSampleFilterWithQueryIdOffsetINS0_9filtering13bitset_filterIjlEEEEEEvPNT2_7INDEX_TEPNSD_10DISTANCE_TEjPKSD_PKNSD_6DATA_TEPKSE_jjmSO_jSF_jjjjPjjjjjT3_' with regcount of 255 ptxas error : Entry function '_ZN4cuvs9neighbors5cagra6detail17single_cta_search13search_kernelILj512ELj256ELj1ENS2_25dataset_descriptor_base_tI6__halfjfEENS2_34CagraSampleFilterWithQueryIdOffsetINS0_9filtering13bitset_filterIjlEEEEEEvPNT2_7INDEX_TEPNSD_10DISTANCE_TEjPKSD_PKNSD_6DATA_TEPKSE_jjmSO_jSF_jjjjPjjjjT3_' with max regcount of 64 calls function '_ZN4cuvs9neighbors5cagra6detail17single_cta_search11search_coreILj512ELj256ELj1ENS2_25dataset_descriptor_base_tI6__halfjfEENS2_34CagraSampleFilterWithQueryIdOffsetINS0_9filtering13bitset_filterIjlEEEEEEvPNT2_7INDEX_TEPNSD_10DISTANCE_TEjPKSD_PKNSD_6DATA_TEPKSE_jjmSO_jSF_jjjjPjjjjjT3_' with regcount of 255 ptxas error : Entry function '_ZN4cuvs9neighbors5cagra6detail17single_cta_search13search_kernelILj256ELj32ELj0ENS2_25dataset_descriptor_base_tI6__halfjfEENS2_34CagraSampleFilterWithQueryIdOffsetINS0_9filtering13bitset_filterIjlEEEEEEvPNT2_7INDEX_TEPNSD_10DISTANCE_TEjPKSD_PKNSD_6DATA_TEPKSE_jjmSO_jSF_jjjjPjjjjT3_' with max regcount of 64 calls function '_ZN4cuvs9neighbors5cagra6detail17single_cta_search11search_coreILj256ELj32ELj0ENS2_25dataset_descriptor_base_tI6__halfjfEENS2_34CagraSampleFilterWithQueryIdOffsetINS0_9filtering13bitset_filterIjlEEEEEEvPNT2_7INDEX_TEPNSD_10DISTANCE_TEjPKSD_PKNSD_6DATA_TEPKSE_jjmSO_jSF_jjjjPjjjjjT3_' with regcount of 255 ptxas error : Entry function '_ZN4cuvs9neighbors5cagra6detail17single_cta_search15search_kernel_pILj256ELj32ELj0ENS2_25dataset_descriptor_base_tI6__halfjfEENS2_34CagraSampleFilterWithQueryIdOffsetINS0_9filtering13bitset_filterIjlEEEEEEvPKT2_PNS3_15worker_handle_tEPNS3_10job_desc_tISD_EEPjPKNSD_7INDEX_TEjjmSO_jPSM_jjjjSL_jjjT3_' with max regcount of 64 calls function '_ZN4cuvs9neighbors5cagra6detail17single_cta_search11search_coreILj256ELj32ELj0ENS2_25dataset_descriptor_base_tI6__halfjfEENS2_34CagraSampleFilterWithQueryIdOffsetINS0_9filtering13bitset_filterIjlEEEEEEvPNT2_7INDEX_TEPNSD_10DISTANCE_TEjPKSD_PKNSD_6DATA_TEPKSE_jjmSO_jSF_jjjjPjjjjjT3_' with regcount of 255 ptxas error : Entry function '_ZN4cuvs9neighbors5cagra6detail17single_cta_search15search_kernel_pILj512ELj32ELj0ENS2_25dataset_descriptor_base_tI6__halfjfEENS2_34CagraSampleFilterWithQueryIdOffsetINS0_9filtering13bitset_filterIjlEEEEEEvPKT2_PNS3_15worker_handle_tEPNS3_10job_desc_tISD_EEPjPKNSD_7INDEX_TEjjmSO_jPSM_jjjjSL_jjjT3_' with max regcount of 64 calls function '_ZN4cuvs9neighbors5cagra6detail17single_cta_search11search_coreILj512ELj32ELj0ENS2_25dataset_descriptor_base_tI6__halfjfEENS2_34CagraSampleFilterWithQueryIdOffsetINS0_9filtering13bitset_filterIjlEEEEEEvPNT2_7INDEX_TEPNSD_10DISTANCE_TEjPKSD_PKNSD_6DATA_TEPKSE_jjmSO_jSF_jjjjPjjjjjT3_' with regcount of 255 ptxas error : Entry function '_ZN4cuvs9neighbors5cagra6detail17single_cta_search13search_kernelILj512ELj32ELj0ENS2_25dataset_descriptor_base_tI6__halfjfEENS2_34CagraSampleFilterWithQueryIdOffsetINS0_9filtering13bitset_filterIjlEEEEEEvPNT2_7INDEX_TEPNSD_10DISTANCE_TEjPKSD_PKNSD_6DATA_TEPKSE_jjmSO_jSF_jjjjPjjjjT3_' with max regcount of 64 calls function '_ZN4cuvs9neighbors5cagra6detail17single_cta_search11search_coreILj512ELj32ELj0ENS2_25dataset_descriptor_base_tI6__halfjfEENS2_34CagraSampleFilterWithQueryIdOffsetINS0_9filtering13bitset_filterIjlEEEEEEvPNT2_7INDEX_TEPNSD_10DISTANCE_TEjPKSD_PKNSD_6DATA_TEPKSE_jjmSO_jSF_jjjjPjjjjjT3_' with regcount of 255 ``` Authors: - Micka (https://github.com/lowener) - MithunR (https://github.com/mythrocks) Approvers: - Tarang Jain (https://github.com/tarang-jain) - Corey J. Nolet (https://github.com/cjnolet) URL: https://github.com/rapidsai/cuvs/pull/1258 --- cpp/cmake/modules/ConfigureCUDA.cmake | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/cpp/cmake/modules/ConfigureCUDA.cmake b/cpp/cmake/modules/ConfigureCUDA.cmake index 5ce9b918b4..0550aa0b11 100644 --- a/cpp/cmake/modules/ConfigureCUDA.cmake +++ b/cpp/cmake/modules/ConfigureCUDA.cmake @@ -64,6 +64,6 @@ endif() # Debug options if(CMAKE_BUILD_TYPE MATCHES Debug) message(VERBOSE "cuVS: Building with debugging flags") - list(APPEND CUVS_CUDA_FLAGS -G -Xcompiler=-rdynamic) + list(APPEND CUVS_CUDA_FLAGS -G -Xcompiler=-rdynamic --maxrregcount=64) list(APPEND CUVS_CUDA_FLAGS -Xptxas --suppress-stack-size-warning) endif() From bcf9025fbf619aa174b688ed5d2545b59271cd74 Mon Sep 17 00:00:00 2001 From: Tarang Jain <40517122+tarang-jain@users.noreply.github.com> Date: Tue, 16 Sep 2025 19:21:10 -0400 Subject: [PATCH 67/93] CAGRA Build + DiskANN Search cuvs-bench Wrapper (#899) cuvs-bench wrapper for CAGRA build + Vamana style serialize. Search the index as a Diskann in-memory index. This PR also brings a bunch of fixes for already existing algos: - `cuvs_vamana` wrapper cleanup via default destructor gives SIGFPE (fixed by declaring raft handle first, so that it is destroyed last) - Low throughput for diskann in-memory and ssd indexes (fixed via removing OpenMP reliance entirely during search and using the default benchmarking methods for multi-threading) - `cuvs_vamana` was using the incorrect template params for device mdspan for the dataset - DiskANN indexes to load the index using the correct number of threads using the `benchmark_n_threads` variable - Faiss CAGRA + HNSW wrapper supports throughput mode search by resetting the CAGRA index `shared_ptr`, which is not used during search (this allows using different CPU threads) Authors: - Tarang Jain (https://github.com/tarang-jain) - Corey J. Nolet (https://github.com/cjnolet) Approvers: - Corey J. Nolet (https://github.com/cjnolet) URL: https://github.com/rapidsai/cuvs/pull/899 --- cpp/bench/ann/CMakeLists.txt | 17 +- .../src/cuvs/cuvs_ann_bench_param_parser.h | 12 +- cpp/bench/ann/src/cuvs/cuvs_cagra_diskann.cu | 88 +++++++ .../ann/src/cuvs/cuvs_cagra_diskann_wrapper.h | 227 ++++++++++++++++++ cpp/bench/ann/src/cuvs/cuvs_vamana.cu | 1 - cpp/bench/ann/src/cuvs/cuvs_vamana_wrapper.h | 6 +- .../ann/src/diskann/diskann_benchmark.cpp | 8 +- cpp/bench/ann/src/diskann/diskann_wrapper.h | 48 ++-- cpp/bench/ann/src/faiss/faiss_gpu_wrapper.h | 14 +- .../cuvs_bench/config/algorithms.yaml | 3 + .../config/algos/cuvs_cagra_diskann.yaml | 11 + .../cuvs_bench/config/algos/cuvs_vamana.yaml | 4 +- .../config/algos/diskann_memory.yaml | 3 +- .../cuvs_bench/config/algos/diskann_ssd.yaml | 1 - 14 files changed, 391 insertions(+), 52 deletions(-) create mode 100644 cpp/bench/ann/src/cuvs/cuvs_cagra_diskann.cu create mode 100644 cpp/bench/ann/src/cuvs/cuvs_cagra_diskann_wrapper.h create mode 100644 python/cuvs_bench/cuvs_bench/config/algos/cuvs_cagra_diskann.yaml diff --git a/cpp/bench/ann/CMakeLists.txt b/cpp/bench/ann/CMakeLists.txt index 7fc49c9e80..9f366e7f83 100644 --- a/cpp/bench/ann/CMakeLists.txt +++ b/cpp/bench/ann/CMakeLists.txt @@ -39,9 +39,13 @@ option(CUVS_ANN_BENCH_USE_HNSWLIB "Include hnsw algorithm in benchmark" ON) option(CUVS_ANN_BENCH_USE_GGNN "Include ggnn algorithm in benchmark" OFF) option(CUVS_ANN_BENCH_USE_DISKANN "Include DISKANN search in benchmark" ON) option(CUVS_ANN_BENCH_USE_CUVS_VAMANA "Include cuVS Vamana with DiskANN search in benchmark" ON) +option(CUVS_ANN_BENCH_USE_CUVS_CAGRA_DISKANN "Include cuVS CAGRA with DISKANN search in benchmark" + ON +) if(CMAKE_SYSTEM_PROCESSOR MATCHES "(ARM|arm|aarch64)") set(CUVS_ANN_BENCH_USE_DISKANN OFF) set(CUVS_ANN_BENCH_USE_CUVS_VAMANA OFF) + set(CUVS_ANN_BENCH_USE_CUVS_CAGRA_DISKANN OFF) endif() option(CUVS_ANN_BENCH_USE_CUVS_MG "Include cuVS ann mg algorithm in benchmark" ${BUILD_MG_ALGOS}) option(CUVS_ANN_BENCH_SINGLE_EXE @@ -69,6 +73,7 @@ if(BUILD_CPU_ONLY) set(CUVS_KNN_BENCH_USE_CUVS_BRUTE_FORCE OFF) set(CUVS_ANN_BENCH_USE_CUVS_MG OFF) set(CUVS_ANN_BENCH_USE_CUVS_VAMANA OFF) + set(CUVS_ANN_BENCH_USE_CUVS_CAGRA_DISKANN OFF) else() set(CUVS_FAISS_ENABLE_GPU ON) endif() @@ -82,6 +87,7 @@ if(CUVS_ANN_BENCH_USE_CUVS_IVF_PQ OR CUVS_KNN_BENCH_USE_CUVS_BRUTE_FORCE OR CUVS_ANN_BENCH_USE_CUVS_MG OR CUVS_ANN_BENCH_USE_CUVS_VAMANA + OR CUVS_ANN_BENCH_USE_CUVS_CAGRA_DISKANN ) set(CUVS_ANN_BENCH_USE_CUVS ON) endif() @@ -99,7 +105,10 @@ if(CUVS_ANN_BENCH_USE_GGNN) include(cmake/thirdparty/get_ggnn) endif() -if(CUVS_ANN_BENCH_USE_DISKANN OR CUVS_ANN_BENCH_USE_CUVS_VAMANA) +if(CUVS_ANN_BENCH_USE_DISKANN + OR CUVS_ANN_BENCH_USE_CUVS_VAMANA + OR CUVS_ANN_BENCH_USE_CUVS_CAGRA_DISKANN +) include(cmake/thirdparty/get_diskann) endif() @@ -341,6 +350,12 @@ if(CUVS_ANN_BENCH_USE_CUVS_VAMANA) ConfigureAnnBench(NAME CUVS_VAMANA PATH src/cuvs/cuvs_vamana.cu LINKS cuvs diskann::diskann aio) endif() +if(CUVS_ANN_BENCH_USE_CUVS_CAGRA_DISKANN) + ConfigureAnnBench( + NAME CUVS_CAGRA_DISKANN PATH src/cuvs/cuvs_cagra_diskann.cu LINKS cuvs diskann::diskann aio + ) +endif() + # ################################################################################################## # * Dynamically-loading ANN_BENCH executable ------------------------------------------------------- if(CUVS_ANN_BENCH_SINGLE_EXE) diff --git a/cpp/bench/ann/src/cuvs/cuvs_ann_bench_param_parser.h b/cpp/bench/ann/src/cuvs/cuvs_ann_bench_param_parser.h index 3ddf026046..3ad01dc2b4 100644 --- a/cpp/bench/ann/src/cuvs/cuvs_ann_bench_param_parser.h +++ b/cpp/bench/ann/src/cuvs/cuvs_ann_bench_param_parser.h @@ -27,7 +27,7 @@ extern template class cuvs::bench::cuvs_ivf_flat; extern template class cuvs::bench::cuvs_ivf_flat; #endif #if defined(CUVS_ANN_BENCH_USE_CUVS_IVF_PQ) || defined(CUVS_ANN_BENCH_USE_CUVS_CAGRA) || \ - defined(CUVS_ANN_BENCH_USE_CUVS_CAGRA_HNSWLIB) + defined(CUVS_ANN_BENCH_USE_CUVS_CAGRA_HNSWLIB) || defined(CUVS_ANN_BENCH_USE_CUVS_CAGRA_DISKANN) #include "cuvs_ivf_pq_wrapper.h" #endif #ifdef CUVS_ANN_BENCH_USE_CUVS_IVF_PQ @@ -35,7 +35,8 @@ extern template class cuvs::bench::cuvs_ivf_pq; extern template class cuvs::bench::cuvs_ivf_pq; extern template class cuvs::bench::cuvs_ivf_pq; #endif -#if defined(CUVS_ANN_BENCH_USE_CUVS_CAGRA) || defined(CUVS_ANN_BENCH_USE_CUVS_CAGRA_HNSWLIB) +#if defined(CUVS_ANN_BENCH_USE_CUVS_CAGRA) || defined(CUVS_ANN_BENCH_USE_CUVS_CAGRA_HNSWLIB) || \ + defined(CUVS_ANN_BENCH_USE_CUVS_CAGRA_DISKANN) #include "cuvs_cagra_wrapper.h" #endif #ifdef CUVS_ANN_BENCH_USE_CUVS_CAGRA @@ -96,8 +97,9 @@ void parse_search_param(const nlohmann::json& conf, } #endif -#if defined(CUVS_ANN_BENCH_USE_CUVS_IVF_PQ) || defined(CUVS_ANN_BENCH_USE_CUVS_CAGRA) || \ - defined(CUVS_ANN_BENCH_USE_CUVS_CAGRA_HNSWLIB) || defined(CUVS_ANN_BENCH_USE_CUVS_MG) +#if defined(CUVS_ANN_BENCH_USE_CUVS_IVF_PQ) || defined(CUVS_ANN_BENCH_USE_CUVS_CAGRA) || \ + defined(CUVS_ANN_BENCH_USE_CUVS_CAGRA_HNSWLIB) || defined(CUVS_ANN_BENCH_USE_CUVS_MG) || \ + defined(CUVS_ANN_BENCH_USE_CUVS_CAGRA_DISKANN) template void parse_build_param(const nlohmann::json& conf, typename cuvs::bench::cuvs_ivf_pq::build_param& param) @@ -188,7 +190,7 @@ void parse_search_param(const nlohmann::json& conf, #endif #if defined(CUVS_ANN_BENCH_USE_CUVS_CAGRA) || defined(CUVS_ANN_BENCH_USE_CUVS_CAGRA_HNSWLIB) || \ - defined(CUVS_ANN_BENCH_USE_CUVS_MG) + defined(CUVS_ANN_BENCH_USE_CUVS_MG) || defined(CUVS_ANN_BENCH_USE_CUVS_CAGRA_DISKANN) template void parse_build_param(const nlohmann::json& conf, cuvs::neighbors::nn_descent::index_params& param) { diff --git a/cpp/bench/ann/src/cuvs/cuvs_cagra_diskann.cu b/cpp/bench/ann/src/cuvs/cuvs_cagra_diskann.cu new file mode 100644 index 0000000000..8719d83867 --- /dev/null +++ b/cpp/bench/ann/src/cuvs/cuvs_cagra_diskann.cu @@ -0,0 +1,88 @@ +/* + * Copyright (c) 2025, NVIDIA CORPORATION. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "../common/ann_types.hpp" +#include "cuvs_ann_bench_param_parser.h" +#include "cuvs_cagra_diskann_wrapper.h" + +#include +#include +#include + +namespace cuvs::bench { + +template +void parse_search_param(const nlohmann::json& conf, + typename cuvs::bench::cuvs_cagra_diskann::search_param& param) +{ + if (conf.contains("L_search")) { param.L_search = conf.at("L_search"); } +} + +template +auto create_algo(const std::string& algo_name, + const std::string& distance, + int dim, + const nlohmann::json& conf) -> std::unique_ptr> +{ + [[maybe_unused]] cuvs::bench::Metric metric = parse_metric(distance); + std::unique_ptr> a; + + if constexpr (std::is_same_v || std::is_same_v || + std::is_same_v) { + if (algo_name == "cuvs_cagra_diskann") { + typename cuvs::bench::cuvs_cagra_diskann::build_param param; + ::parse_build_param(conf, param); + a = std::make_unique>(metric, dim, param); + } + } + + if (!a) { throw std::runtime_error("invalid algo: '" + algo_name + "'"); } + + return a; +} + +template +auto create_search_param(const std::string& algo_name, const nlohmann::json& conf) + -> std::unique_ptr::search_param> +{ + if (algo_name == "cuvs_cagra_diskann") { + auto param = + std::make_unique::search_param>(); + parse_search_param(conf, *param); + return param; + } + + throw std::runtime_error("invalid algo: '" + algo_name + "'"); +} + +} // namespace cuvs::bench + +REGISTER_ALGO_INSTANCE(float); +REGISTER_ALGO_INSTANCE(std::int8_t); +REGISTER_ALGO_INSTANCE(std::uint8_t); + +#ifdef ANN_BENCH_BUILD_MAIN +#include "../common/benchmark.hpp" +/* +[NOTE] Dear developer, + +Please don't modify the content of the `main` function; this will make the behavior of the benchmark +executable differ depending on the cmake flags and will complicate the debugging. In particular, +don't try to setup an RMM memory resource here; it will anyway be modified by the memory resource +set on per-algorithm basis. For example, see `cuvs/cuvs_ann_bench_utils.h`. +*/ +int main(int argc, char** argv) { return cuvs::bench::run_main(argc, argv); } +#endif diff --git a/cpp/bench/ann/src/cuvs/cuvs_cagra_diskann_wrapper.h b/cpp/bench/ann/src/cuvs/cuvs_cagra_diskann_wrapper.h new file mode 100644 index 0000000000..9cf5b50930 --- /dev/null +++ b/cpp/bench/ann/src/cuvs/cuvs_cagra_diskann_wrapper.h @@ -0,0 +1,227 @@ +/* + * Copyright (c) 2023-2025, NVIDIA CORPORATION. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#pragma once + +#include "cuvs_cagra_wrapper.h" +#include +#include + +#include +#include + +#include "../common/ann_types.hpp" +#include "../diskann/diskann_wrapper.h" +#include "cuvs_ann_bench_utils.h" +#include +#include + +#include +#include +#include + +namespace cuvs::bench { + +template +class cuvs_cagra_diskann : public algo, public algo_gpu { + public: + using build_param = typename cuvs_cagra::build_param; + using search_param_base = typename algo::search_param; + using search_param = typename diskann_memory::search_param; + + cuvs_cagra_diskann(Metric metric, int dim, const build_param& param); + + void build(const T* dataset, size_t nrow) final; + + void set_search_param(const search_param_base& param, const void* filter_bitset) override; + + void search(const T* queries, + int batch_size, + int k, + algo_base::index_type* neighbors, + float* distances) const override; + + [[nodiscard]] auto get_sync_stream() const noexcept -> cudaStream_t override + { + return cagra_build_.get_sync_stream(); + } + + // to enable dataset access from GPU memory + [[nodiscard]] auto get_preference() const -> algo_property override + { + algo_property property; + property.dataset_memory_type = MemoryType::kHostMmap; + property.query_memory_type = MemoryType::kHost; + return property; + } + + void save(const std::string& file) const override; + void load(const std::string&) override; + std::unique_ptr> copy() override + { + return std::make_unique>(*this); + } + + private: + raft::resources handle_{}; + build_param build_param_; + search_param search_param_; + cuvs_cagra cagra_build_; + std::shared_ptr> diskann_memory_search_; +}; + +template +cuvs_cagra_diskann::cuvs_cagra_diskann(Metric metric, int dim, const build_param& param) + : algo(metric, dim), build_param_{param}, cagra_build_{metric, dim, param, 1} +{ + diskann_memory_search_ = std::make_shared>( + metric, + dim, + typename diskann_memory::build_param{static_cast(32), static_cast(32)}); +} + +template +void cuvs_cagra_diskann::build(const T* dataset, size_t nrow) +{ + cagra_build_.build(dataset, nrow); +} + +template +void cuvs_cagra_diskann::set_search_param(const search_param_base& param, + const void* filter_bitset) +{ + if (filter_bitset != nullptr) { throw std::runtime_error("Filtering is not supported yet."); } + diskann_memory_search_->set_search_param(param, nullptr); +} + +template +void cuvs_cagra_diskann::save(const std::string& file) const +{ + // Write graph to first index file (format from MSFT DiskANN OSS) + std::ofstream index_of(file, std::ios::out | std::ios::binary); + if (!index_of) { RAFT_FAIL("Cannot open file %s", file.c_str()); } + + size_t file_offset = 0; + index_of.seekp(file_offset, index_of.beg); + uint32_t max_degree = 0; + size_t index_size = 24; + uint32_t start = static_cast(rand() % (cagra_build_.get_index()->graph().extent(0))); + size_t num_frozen_points = 0; + uint32_t max_observed_degree = 0; + + index_of.write((char*)&index_size, sizeof(uint64_t)); + index_of.write((char*)&max_observed_degree, sizeof(uint32_t)); + index_of.write((char*)&start, sizeof(uint32_t)); + index_of.write((char*)&num_frozen_points, sizeof(size_t)); + + auto d_graph = cagra_build_.get_index()->graph(); + auto h_graph = raft::make_host_matrix(d_graph.extent(0), d_graph.extent(1)); + raft::copy(h_graph.data_handle(), + d_graph.data_handle(), + d_graph.size(), + raft::resource::get_cuda_stream(handle_)); + raft::resource::sync_stream(handle_); + + size_t total_edges = 0; + size_t num_sparse = 0; + size_t num_single = 0; + + for (uint32_t i = 0; i < h_graph.extent(0); i++) { + uint32_t node_edges = 0; + for (; node_edges < h_graph.extent(1); node_edges++) { + if (h_graph(i, node_edges) == raft::upper_bound()) { break; } + } + + if (node_edges < 3) num_sparse++; + if (node_edges < 2) num_single++; + total_edges += node_edges; + + index_of.write((char*)&node_edges, sizeof(uint32_t)); + if constexpr (!std::is_same_v) { + RAFT_FAIL("serialization is only implemented for uint32_t graph"); + } + index_of.write((char*)&h_graph(i, 0), node_edges * sizeof(uint32_t)); + + max_degree = node_edges > max_degree ? (uint32_t)node_edges : max_degree; + index_size += (size_t)(sizeof(uint32_t) * (node_edges + 1)); + } + index_of.seekp(file_offset, index_of.beg); + index_of.write((char*)&index_size, sizeof(uint64_t)); + index_of.write((char*)&max_degree, sizeof(uint32_t)); + + RAFT_LOG_DEBUG( + "Wrote file out, index size:%lu, max_degree:%u, num_sparse:%ld, num_single:%ld, total " + "edges:%ld, avg degree:%f", + index_size, + max_degree, + num_sparse, + num_single, + total_edges, + (float)total_edges / (float)h_graph.extent(0)); + + index_of.close(); + if (!index_of) { RAFT_FAIL("Error writing output %s", file.c_str()); } + + // try allocating a buffer for the dataset on host + try { + const cuvs::neighbors::strided_dataset* strided_dataset = + dynamic_cast*>( + const_cast*>(&cagra_build_.get_index()->data())); + if (strided_dataset == nullptr) { + RAFT_LOG_DEBUG("dynamic_cast to strided_dataset failed"); + } else { + auto h_dataset = + raft::make_host_matrix(strided_dataset->n_rows(), strided_dataset->dim()); + raft::copy(h_dataset.data_handle(), + strided_dataset->view().data_handle(), + strided_dataset->n_rows() * strided_dataset->dim(), + raft::resource::get_cuda_stream(handle_)); + std::string dataset_base_file = file + ".data"; + std::ofstream dataset_of(dataset_base_file, std::ios::out | std::ios::binary); + if (!dataset_of) { RAFT_FAIL("Cannot open file %s", dataset_base_file.c_str()); } + size_t dataset_file_offset = 0; + int size = static_cast(cagra_build_.get_index()->size()); + int dim = static_cast(cagra_build_.get_index()->dim()); + dataset_of.seekp(dataset_file_offset, dataset_of.beg); + dataset_of.write((char*)&size, sizeof(int)); + dataset_of.write((char*)&dim, sizeof(int)); + for (int i = 0; i < size; i++) { + dataset_of.write((char*)(h_dataset.data_handle() + i * h_dataset.extent(1)), + dim * sizeof(T)); + } + dataset_of.close(); + if (!dataset_of) { RAFT_FAIL("Error writing output %s", dataset_base_file.c_str()); } + } + } catch (std::bad_alloc& e) { + RAFT_LOG_INFO("Failed to serialize dataset"); + } catch (raft::logic_error& e) { + RAFT_LOG_INFO("Failed to serialize dataset"); + } +} + +template +void cuvs_cagra_diskann::load(const std::string& file) +{ + diskann_memory_search_->load(file); +} + +template +void cuvs_cagra_diskann::search( + const T* queries, int batch_size, int k, algo_base::index_type* neighbors, float* distances) const +{ + diskann_memory_search_->search(queries, batch_size, k, neighbors, distances); +} + +} // namespace cuvs::bench diff --git a/cpp/bench/ann/src/cuvs/cuvs_vamana.cu b/cpp/bench/ann/src/cuvs/cuvs_vamana.cu index 2bc0ca9596..524a61e0b4 100644 --- a/cpp/bench/ann/src/cuvs/cuvs_vamana.cu +++ b/cpp/bench/ann/src/cuvs/cuvs_vamana.cu @@ -37,7 +37,6 @@ void parse_search_param(const nlohmann::json& conf, typename cuvs::bench::cuvs_vamana::search_param& param) { if (conf.contains("L_search")) { param.L_search = conf.at("L_search"); } - if (conf.contains("num_threads")) { param.num_threads = conf.at("num_threads"); } } template diff --git a/cpp/bench/ann/src/cuvs/cuvs_vamana_wrapper.h b/cpp/bench/ann/src/cuvs/cuvs_vamana_wrapper.h index 410f5c2a02..882483853a 100644 --- a/cpp/bench/ann/src/cuvs/cuvs_vamana_wrapper.h +++ b/cpp/bench/ann/src/cuvs/cuvs_vamana_wrapper.h @@ -65,10 +65,10 @@ class cuvs_vamana : public algo, public algo_gpu { std::unique_ptr> copy() override { return std::make_unique>(*this); } private: - std::shared_ptr> vamana_index_; - std::shared_ptr> diskann_memory_search_; configured_raft_resources handle_{}; build_param vamana_index_params_; + std::shared_ptr> vamana_index_; + std::shared_ptr> diskann_memory_search_; }; template @@ -85,7 +85,7 @@ void cuvs_vamana::build(const T* dataset, size_t nrow) { auto dataset_view_host = raft::make_mdspan( dataset, raft::make_extents(nrow, this->dim_)); - auto dataset_view_device = raft::make_mdspan( + auto dataset_view_device = raft::make_mdspan( dataset, raft::make_extents(nrow, this->dim_)); bool dataset_is_on_host = raft::get_device_for_address(dataset) == -1; diff --git a/cpp/bench/ann/src/diskann/diskann_benchmark.cpp b/cpp/bench/ann/src/diskann/diskann_benchmark.cpp index fe47215341..3c1e81c3ad 100644 --- a/cpp/bench/ann/src/diskann/diskann_benchmark.cpp +++ b/cpp/bench/ann/src/diskann/diskann_benchmark.cpp @@ -65,16 +65,14 @@ template void parse_search_param(const nlohmann::json& conf, typename cuvs::bench::diskann_memory::search_param& param) { - param.L_search = conf.at("L_search"); - param.num_threads = conf.at("num_threads"); + param.L_search = conf.at("L_search"); } template void parse_search_param(const nlohmann::json& conf, typename cuvs::bench::diskann_ssd::search_param& param) { - param.L_search = conf.at("L_search"); - param.num_threads = conf.at("num_threads"); + param.L_search = conf.at("L_search"); if (conf.contains("num_nodes_to_cache")) { param.num_nodes_to_cache = conf.at("num_nodes_to_cache"); } @@ -132,6 +130,8 @@ std::unique_ptr::search_param> create_search_param }; // namespace cuvs::bench REGISTER_ALGO_INSTANCE(float); +REGISTER_ALGO_INSTANCE(std::int8_t); +REGISTER_ALGO_INSTANCE(std::uint8_t); #ifdef ANN_BENCH_BUILD_MAIN #include "../common/benchmark.hpp" diff --git a/cpp/bench/ann/src/diskann/diskann_wrapper.h b/cpp/bench/ann/src/diskann/diskann_wrapper.h index c8219c40cb..6e7e0dc1c7 100644 --- a/cpp/bench/ann/src/diskann/diskann_wrapper.h +++ b/cpp/bench/ann/src/diskann/diskann_wrapper.h @@ -16,6 +16,7 @@ #pragma once #include "../common/ann_types.hpp" +#include "../common/util.hpp" #include @@ -58,8 +59,6 @@ class diskann_memory : public algo { using search_param_base = typename algo::search_param; struct search_param : public search_param_base { uint32_t L_search; - uint32_t num_threads = omp_get_max_threads() / 2; - // Mode metric_objective; }; diskann_memory(Metric metric, int dim, const build_param& param); @@ -94,7 +93,6 @@ class diskann_memory : public algo { int num_threads_; uint32_t L_search_; Mode bench_mode_; - int num_search_threads_; std::string index_path_prefix_; std::shared_ptr> mem_index_{nullptr}; void initialize_index_(size_t max_points); @@ -143,24 +141,14 @@ template void diskann_memory::set_search_param(const search_param_base& param, const void* filter_bitset) { if (filter_bitset != nullptr) { throw std::runtime_error("Filtering is not supported yet."); } - auto sp = dynamic_cast(param); - L_search_ = sp.L_search; - num_search_threads_ = sp.num_threads; - - // only latency mode supported. Use the num_threads search param to run search with multiple - // threads - bench_mode_ = Mode::kLatency; - - // Create a pool if multiple query threads have been set and the pool hasn't been created already - initialize_index_(0); - this->mem_index_->load(index_path_prefix_.c_str(), num_search_threads_, L_search_); + auto sp = dynamic_cast(param); + L_search_ = sp.L_search; } template void diskann_memory::search( const T* queries, int batch_size, int k, algo_base::index_type* indices, float* distances) const { -#pragma omp parallel for if (batch_size > 1) schedule(dynamic, 1) num_threads(num_search_threads_) for (int i = 0; i < batch_size; i++) { mem_index_->search(queries + i * this->dim_, static_cast(k), @@ -179,8 +167,14 @@ void diskann_memory::save(const std::string& index_file) const template void diskann_memory::load(const std::string& index_file) { - // only save the index path prefix here index_path_prefix_ = index_file; + + bench_mode_ = (cuvs::bench::benchmark_n_threads > 1) ? Mode::kThroughput : Mode::kLatency; + + initialize_index_(0); + + int load_threads = (bench_mode_ == Mode::kThroughput) ? cuvs::bench::benchmark_n_threads : 1; + this->mem_index_->load(index_path_prefix_.c_str(), load_threads, 2000); } template @@ -200,10 +194,8 @@ class diskann_ssd : public algo { struct search_param : public search_param_base { uint32_t L_search; - uint32_t num_threads = omp_get_max_threads() / 2; uint32_t num_nodes_to_cache = 10000; int beam_width = 2; - // Mode metric_objective; }; diskann_ssd(Metric metric, int dim, const build_param& param); @@ -226,7 +218,7 @@ class diskann_ssd : public algo { [[nodiscard]] auto get_preference() const -> algo_property override { algo_property property; - property.dataset_memory_type = MemoryType::kHost; + property.dataset_memory_type = MemoryType::kHostMmap; property.query_memory_type = MemoryType::kHost; return property; } @@ -239,11 +231,6 @@ class diskann_ssd : public algo { // in-memory index params uint32_t build_pq_bytes_ = 0; uint32_t max_points_; - // for safe scratch space allocs, set the default to half the number of procs for loading the - // index. User must ensure that the number of search threads is less than or equal to this value - int num_search_threads_ = omp_get_max_threads() / 2; - // L_search is hardcoded to the maximum visited list size in the search params. This default is - // for loading the index uint32_t L_search_ = 384; Mode bench_mode_; std::string base_file_; @@ -291,19 +278,14 @@ void diskann_ssd::set_search_param(const search_param_base& param, const void if (filter_bitset != nullptr) { throw std::runtime_error("Filtering is not supported yet."); } auto sp = dynamic_cast(param); L_search_ = sp.L_search; - num_search_threads_ = sp.num_threads; num_nodes_to_cache_ = sp.num_nodes_to_cache; beam_width_ = sp.beam_width; - - // only latency mode supported with thread pool - bench_mode_ = Mode::kLatency; } template void diskann_ssd::search( const T* queries, int batch_size, int k, algo_base::index_type* neighbors, float* distances) const { -#pragma omp parallel for if (batch_size > 1) schedule(dynamic, 1) num_threads(num_search_threads_) for (int64_t i = 0; i < (int64_t)batch_size; i++) { p_flash_index_->cached_beam_search(queries + (i * this->dim_), static_cast(k), @@ -328,10 +310,16 @@ void diskann_ssd::save(const std::string& index_file) const template void diskann_ssd::load(const std::string& index_file) { + index_path_prefix_ = index_file; + + bench_mode_ = (cuvs::bench::benchmark_n_threads > 1) ? Mode::kThroughput : Mode::kLatency; + + int load_threads = (bench_mode_ == Mode::kThroughput) ? cuvs::bench::benchmark_n_threads : 1; + reader.reset(new LinuxAlignedFileReader()); p_flash_index_ = std::make_shared>(reader, parse_metric_to_diskann(this->metric_)); - int result = p_flash_index_->load(num_search_threads_, index_file.c_str()); + int result = p_flash_index_->load(load_threads, index_path_prefix_.c_str()); std::vector node_list; p_flash_index_->cache_bfs_levels(num_nodes_to_cache_, node_list); p_flash_index_->load_cache_list(node_list); diff --git a/cpp/bench/ann/src/faiss/faiss_gpu_wrapper.h b/cpp/bench/ann/src/faiss/faiss_gpu_wrapper.h index 54f685fe79..b4bd384e13 100644 --- a/cpp/bench/ann/src/faiss/faiss_gpu_wrapper.h +++ b/cpp/bench/ann/src/faiss/faiss_gpu_wrapper.h @@ -1,5 +1,5 @@ /* - * Copyright (c) 2023-2024, NVIDIA CORPORATION. + * Copyright (c) 2023-2025, NVIDIA CORPORATION. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -550,7 +550,7 @@ class faiss_gpu_cagra : public faiss_gpu { config.graph_degree = param.graph_degree; config.intermediate_graph_degree = param.intermediate_graph_degree; config.device = this->device_; - config.store_dataset = false; + config.store_dataset = true; if (param.cagra_build_algo == "IVF_PQ") { config.build_algo = faiss::gpu::graph_build_algo::IVF_PQ; this->ivf_pq_build_params_ = param.ivf_pq_build_params; @@ -655,9 +655,17 @@ class faiss_gpu_cagra_hnsw : public faiss_gpu { omp_single_thread_scope omp_single_thread; this->search_index_.reset(static_cast(faiss::read_index(file.c_str()))); } + + [[nodiscard]] auto get_sync_stream() const noexcept -> cudaStream_t override + { + if (this->gpu_resource_ == nullptr) { return 0; } + return this->gpu_resource_->getDefaultStream(this->device_); + } std::unique_ptr> copy() override { - return std::make_unique>(*this); + auto new_instance = std::make_unique>(*this); + new_instance->build_index_ = nullptr; + return new_instance; }; private: diff --git a/python/cuvs_bench/cuvs_bench/config/algorithms.yaml b/python/cuvs_bench/cuvs_bench/config/algorithms.yaml index 191c640c00..fa2195fc61 100644 --- a/python/cuvs_bench/cuvs_bench/config/algorithms.yaml +++ b/python/cuvs_bench/cuvs_bench/config/algorithms.yaml @@ -67,3 +67,6 @@ diskann_ssd: cuvs_vamana: executable: CUVS_VAMANA_ANN_BENCH requires_gpu: true +cuvs_cagra_diskann: + executable: CUVS_CAGRA_DISKANN_ANN_BENCH + requires_gpu: true diff --git a/python/cuvs_bench/cuvs_bench/config/algos/cuvs_cagra_diskann.yaml b/python/cuvs_bench/cuvs_bench/config/algos/cuvs_cagra_diskann.yaml new file mode 100644 index 0000000000..e086e8230b --- /dev/null +++ b/python/cuvs_bench/cuvs_bench/config/algos/cuvs_cagra_diskann.yaml @@ -0,0 +1,11 @@ +name: cuvs_cagra_diskann +constraints: + build: cuvs_bench.config.algos.constraints.cuvs_cagra_build +groups: + base: + build: + graph_degree: [32, 64, 96] + intermediate_graph_degree: [32, 64, 128] + graph_build_algo: ["NN_DESCENT"] + search: + L_search: [10, 20, 30, 40, 50, 100, 200, 300] diff --git a/python/cuvs_bench/cuvs_bench/config/algos/cuvs_vamana.yaml b/python/cuvs_bench/cuvs_bench/config/algos/cuvs_vamana.yaml index 5252f818c5..b840962d92 100644 --- a/python/cuvs_bench/cuvs_bench/config/algos/cuvs_vamana.yaml +++ b/python/cuvs_bench/cuvs_bench/config/algos/cuvs_vamana.yaml @@ -2,8 +2,8 @@ name: cuvs_vamana groups: base: build: - graph_degree: [32, 64] - visited_size: [128, 256] + graph_degree: [32, 64, 96] + visited_size: [128, 256, 512] alpha: [1.2] search: L_search: [10, 20, 30, 40, 50, 100, 200, 300] diff --git a/python/cuvs_bench/cuvs_bench/config/algos/diskann_memory.yaml b/python/cuvs_bench/cuvs_bench/config/algos/diskann_memory.yaml index faf1224656..b62f057084 100644 --- a/python/cuvs_bench/cuvs_bench/config/algos/diskann_memory.yaml +++ b/python/cuvs_bench/cuvs_bench/config/algos/diskann_memory.yaml @@ -4,10 +4,9 @@ constraints: groups: base: build: - R: [64, 96] + R: [32, 64, 96] L_build: [128, 256, 384] alpha: [1.2] num_threads: [32] search: - num_threads: [32] L_search: [10, 20, 30, 40, 50, 100, 200, 300] diff --git a/python/cuvs_bench/cuvs_bench/config/algos/diskann_ssd.yaml b/python/cuvs_bench/cuvs_bench/config/algos/diskann_ssd.yaml index 40afc6fce8..a8eb52abea 100644 --- a/python/cuvs_bench/cuvs_bench/config/algos/diskann_ssd.yaml +++ b/python/cuvs_bench/cuvs_bench/config/algos/diskann_ssd.yaml @@ -10,4 +10,3 @@ groups: num_threads: [32] search: L_search: [10, 20, 30, 40, 50, 100, 200, 300] - num_threads: [32] From 587f975305ebce2c5e904008b423f169a146c48c Mon Sep 17 00:00:00 2001 From: Anupam <54245698+aamijar@users.noreply.github.com> Date: Tue, 16 Sep 2025 20:46:02 -0700 Subject: [PATCH 68/93] Docs for spectral embedding (#1299) Resolves https://github.com/rapidsai/cuvs/issues/1298 Authors: - Anupam (https://github.com/aamijar) Approvers: - Corey J. Nolet (https://github.com/cjnolet) URL: https://github.com/rapidsai/cuvs/pull/1299 --- .../cuvs/preprocessing/spectral_embedding.hpp | 127 +++++++++++++++++- docs/source/cpp_api/preprocessing.rst | 1 + .../preprocessing_spectral_embedding.rst | 108 +++++++++++++++ 3 files changed, 233 insertions(+), 3 deletions(-) create mode 100644 docs/source/cpp_api/preprocessing_spectral_embedding.rst diff --git a/cpp/include/cuvs/preprocessing/spectral_embedding.hpp b/cpp/include/cuvs/preprocessing/spectral_embedding.hpp index 3d935bfd9a..3162f74f02 100644 --- a/cpp/include/cuvs/preprocessing/spectral_embedding.hpp +++ b/cpp/include/cuvs/preprocessing/spectral_embedding.hpp @@ -23,8 +23,18 @@ namespace cuvs::preprocessing::spectral_embedding { +/** + * @defgroup spectral_embedding Spectral Embedding + * @{ + */ + /** * @brief Parameters for spectral embedding algorithm + * + * Spectral embedding is a dimensionality reduction technique that uses the + * eigenvectors of the graph Laplacian to embed data points into a lower-dimensional + * space. This technique is particularly useful for non-linear dimensionality + * reduction and clustering tasks. */ struct params { /** @brief The number of components to reduce the data to. */ @@ -33,24 +43,135 @@ struct params { /** @brief The number of neighbors to use for the nearest neighbors graph. */ int n_neighbors; - /** @brief Whether to normalize the Laplacian matrix. */ + /** + * @brief Whether to normalize the Laplacian matrix. + * + * If true, uses the normalized graph Laplacian (D^(-1/2) L D^(-1/2)). + * If false, uses the unnormalized graph Laplacian (L = D - W). + * Normalized Laplacian often leads to better results for clustering tasks. + */ bool norm_laplacian; - /** @brief Whether to drop the first eigenvector. */ + /** + * @brief Whether to drop the first eigenvector. + * + * The first eigenvector of the normalized Laplacian is constant and + * uninformative. Setting this to true drops it from the embedding. + * This is typically set to true when norm_laplacian is true. + */ bool drop_first; - /** @brief Random seed for reproducibility */ + /** + * @brief Random seed for reproducibility. + * + * Controls the random number generation for k-NN graph construction + * and eigenvalue solver initialization. Use the same seed value to + * ensure reproducible results across runs. + */ uint64_t seed; }; +/** + * @brief Perform spectral embedding on input dataset + * + * This function computes the spectral embedding of the input dataset by: + * 1. Constructing a k-nearest neighbors graph from the input data + * 2. Computing the graph Laplacian (normalized or unnormalized) + * 3. Finding the eigenvectors corresponding to the smallest eigenvalues + * 4. Using these eigenvectors as the embedding coordinates + * + * @code{.cpp} + * #include + * #include + * + * raft::resources handle; + * + * // Set up parameters + * cuvs::preprocessing::spectral_embedding::params params; + * params.n_components = 2; + * params.n_neighbors = 15; + * params.norm_laplacian = true; + * params.drop_first = true; + * params.seed = 42; + * + * // Create input dataset (n_samples x n_features) + * auto dataset = raft::make_device_matrix(handle, n_samples, n_features); + * // ... fill dataset ... + * + * // Create output embedding matrix (n_samples x n_components) + * auto embedding = raft::make_device_matrix( + * handle, n_samples, params.n_components); + * + * // Perform spectral embedding + * cuvs::preprocessing::spectral_embedding::transform( + * handle, params, dataset.view(), embedding.view()); + * @endcode + * + * @param[in] handle RAFT resource handle for managing CUDA resources + * @param[in] config Parameters controlling the spectral embedding algorithm + * @param[in] dataset Input dataset in row-major format [n_samples x n_features] + * @param[out] embedding Output embedding in column-major format [n_samples x n_components] + * + */ void transform(raft::resources const& handle, params config, raft::device_matrix_view dataset, raft::device_matrix_view embedding); +/** + * @brief Perform spectral embedding using a precomputed connectivity graph + * + * This function computes the spectral embedding from a precomputed sparse + * connectivity graph (e.g., from a k-NN search or custom similarity matrix). + * This is useful when you want to use a custom graph construction method + * or when you have a precomputed similarity/affinity matrix. + * + * The function: + * 1. Converts the COO matrix to the graph Laplacian + * 2. Computes eigenvectors of the Laplacian + * 3. Returns the eigenvectors as the embedding + * + * @code{.cpp} + * #include + * #include + * + * raft::resources handle; + * + * // Set up parameters + * cuvs::preprocessing::spectral_embedding::params params; + * params.n_components = 2; + * params.norm_laplacian = true; + * params.drop_first = true; + * params.seed = 42; + * + * // Assume we have a precomputed connectivity graph as COO matrix + * // connectivity_graph represents weighted edges between samples + * raft::device_coo_matrix connectivity_graph(...); + * + * // Create output embedding matrix (n_samples x n_components) + * auto embedding = raft::make_device_matrix( + * handle, n_samples, params.n_components); + * + * // Perform spectral embedding + * cuvs::preprocessing::spectral_embedding::transform( + * handle, params, connectivity_graph.view(), embedding.view()); + * @endcode + * + * @param[in] handle RAFT resource handle for managing CUDA resources + * @param[in] config Parameters controlling the spectral embedding algorithm + * (n_neighbors parameter is ignored when using precomputed graph) + * @param[in] connectivity_graph Precomputed sparse connectivity/affinity graph in COO format + * representing weighted connections between samples + * @param[out] embedding Output embedding in column-major format [n_samples x n_components] + * + */ void transform(raft::resources const& handle, params config, raft::device_coo_matrix_view connectivity_graph, raft::device_matrix_view embedding); +/** + * @} + */ + } // namespace cuvs::preprocessing::spectral_embedding diff --git a/docs/source/cpp_api/preprocessing.rst b/docs/source/cpp_api/preprocessing.rst index 1c2b0f0513..5fd8cfe778 100644 --- a/docs/source/cpp_api/preprocessing.rst +++ b/docs/source/cpp_api/preprocessing.rst @@ -10,3 +10,4 @@ Preprocessing :caption: Contents: preprocessing_quantize.rst + preprocessing_spectral_embedding.rst diff --git a/docs/source/cpp_api/preprocessing_spectral_embedding.rst b/docs/source/cpp_api/preprocessing_spectral_embedding.rst new file mode 100644 index 0000000000..6ab0785836 --- /dev/null +++ b/docs/source/cpp_api/preprocessing_spectral_embedding.rst @@ -0,0 +1,108 @@ +Spectral Embedding +================== + +Spectral embedding is a powerful dimensionality reduction technique that uses the eigenvectors +of the graph Laplacian to embed high-dimensional data into a lower-dimensional space. This +method is particularly effective for discovering non-linear manifold structures in data and +is widely used in clustering, visualization, and feature extraction tasks. + +.. role:: py(code) + :language: c++ + :class: highlight + +Overview +-------- + +The spectral embedding algorithm works by: + +1. **Graph Construction**: Building a k-nearest neighbors graph from the input data +2. **Laplacian Computation**: Computing the graph Laplacian matrix (normalized or unnormalized) +3. **Eigendecomposition**: Finding the eigenvectors corresponding to the smallest eigenvalues +4. **Embedding**: Using these eigenvectors as coordinates in the lower-dimensional space + +Parameters +---------- + +``#include `` + +namespace *cuvs::preprocessing::spectral_embedding* + +.. doxygenstruct:: cuvs::preprocessing::spectral_embedding::params + :project: cuvs + :members: + +Functions +--------- + +``#include `` + +namespace *cuvs::preprocessing::spectral_embedding* + +.. doxygengroup:: spectral_embedding + :project: cuvs + :content-only: + +Example Usage +------------- + +Basic Usage with Dataset +~~~~~~~~~~~~~~~~~~~~~~~~ + +.. code-block:: cpp + + #include + #include + + // Initialize RAFT resources + raft::resources handle; + + // Configure spectral embedding parameters + cuvs::preprocessing::spectral_embedding::params params; + params.n_components = 2; // Reduce to 2D for visualization + params.n_neighbors = 15; // Local neighborhood size + params.norm_laplacian = true; // Use normalized Laplacian + params.drop_first = true; // Drop constant eigenvector + params.seed = 42; // For reproducibility + + // Create input dataset (n_samples x n_features) + int n_samples = 1000; + int n_features = 50; + auto dataset = raft::make_device_matrix(handle, n_samples, n_features); + // ... populate dataset with your data ... + + // Allocate output embedding matrix (n_samples x n_components) + auto embedding = raft::make_device_matrix( + handle, n_samples, params.n_components); + + // Perform spectral embedding + cuvs::preprocessing::spectral_embedding::transform( + handle, params, dataset.view(), embedding.view()); + +Using Precomputed Graph +~~~~~~~~~~~~~~~~~~~~~~~ + +.. code-block:: cpp + + #include + #include + + raft::resources handle; + + // Configure parameters (n_neighbors is ignored with precomputed graph) + cuvs::preprocessing::spectral_embedding::params params; + params.n_components = 3; + params.norm_laplacian = true; + params.drop_first = true; + params.seed = 42; + + // Assume we have a precomputed connectivity graph + // This could be from custom similarity computation or k-NN search + raft::device_coo_matrix connectivity_graph(...); + + // Allocate output embedding + auto embedding = raft::make_device_matrix( + handle, n_samples, params.n_components); + + // Perform spectral embedding with precomputed graph + cuvs::preprocessing::spectral_embedding::transform( + handle, params, connectivity_graph.view(), embedding.view()); From 659b6644a29a596dd77aef6254b859254e483c0c Mon Sep 17 00:00:00 2001 From: Ben Frederickson Date: Tue, 16 Sep 2025 22:18:55 -0700 Subject: [PATCH 69/93] Add cagra.extend bindings for python (#1324) Add the cagra.extend functionality for python. Note that this PR also removes the `return_dataset` parameter from the C-API in favour of the cuvsCagraIndexGetDataset function. This makes it optional to get the extended dataset in c/python, and also doesn't perform a copy. Authors: - Ben Frederickson (https://github.com/benfred) - Corey J. Nolet (https://github.com/cjnolet) Approvers: - Corey J. Nolet (https://github.com/cjnolet) URL: https://github.com/rapidsai/cuvs/pull/1324 --- cpp/include/cuvs/neighbors/cagra.h | 4 +- cpp/src/neighbors/cagra_c.cpp | 41 +++++------ cpp/tests/neighbors/ann_cagra_c.cu | 28 +------- docs/source/python_api/neighbors_cagra.rst | 5 ++ go/cagra/cagra.go | 5 +- python/cuvs/cuvs/neighbors/cagra/__init__.py | 4 ++ python/cuvs/cuvs/neighbors/cagra/cagra.pxd | 13 ++++ python/cuvs/cuvs/neighbors/cagra/cagra.pyx | 72 +++++++++++++++++++- python/cuvs/cuvs/tests/test_cagra.py | 35 +++++----- 9 files changed, 131 insertions(+), 76 deletions(-) diff --git a/cpp/include/cuvs/neighbors/cagra.h b/cpp/include/cuvs/neighbors/cagra.h index e892b5304c..2532bbd3df 100644 --- a/cpp/include/cuvs/neighbors/cagra.h +++ b/cpp/include/cuvs/neighbors/cagra.h @@ -525,14 +525,12 @@ cuvsError_t cuvsCagraBuild(cuvsResources_t res, * @param[in] params cuvsCagraExtendParams_t used to extend CAGRA index * @param[in] additional_dataset DLManagedTensor* additional dataset * @param[in,out] index cuvsCagraIndex_t CAGRA index - * @param[out] return_dataset DLManagedTensor* extended dataset * @return cuvsError_t */ cuvsError_t cuvsCagraExtend(cuvsResources_t res, cuvsCagraExtendParams_t params, DLManagedTensor* additional_dataset, - cuvsCagraIndex_t index, - DLManagedTensor* return_dataset); + cuvsCagraIndex_t index); /** * @} diff --git a/cpp/src/neighbors/cagra_c.cpp b/cpp/src/neighbors/cagra_c.cpp index 7c3061bf93..92d20f85df 100644 --- a/cpp/src/neighbors/cagra_c.cpp +++ b/cpp/src/neighbors/cagra_c.cpp @@ -164,32 +164,24 @@ template void _extend(cuvsResources_t res, cuvsCagraExtendParams params, cuvsCagraIndex index, - DLManagedTensor* additional_dataset_tensor, - DLManagedTensor* return_tensor) + DLManagedTensor* additional_dataset_tensor) { - auto dataset = additional_dataset_tensor->dl_tensor; - auto return_dl_tensor = return_tensor->dl_tensor; - auto index_ptr = reinterpret_cast*>(index.addr); - auto res_ptr = reinterpret_cast(res); + auto dataset = additional_dataset_tensor->dl_tensor; + auto index_ptr = reinterpret_cast*>(index.addr); + auto res_ptr = reinterpret_cast(res); // TODO: use C struct here (see issue #487) auto extend_params = cuvs::neighbors::cagra::extend_params(); extend_params.max_chunk_size = params.max_chunk_size; - if (cuvs::core::is_dlpack_device_compatible(dataset) && - cuvs::core::is_dlpack_device_compatible(return_dl_tensor)) { - using mdspan_type = raft::device_matrix_view; - using mdspan_return_type = raft::device_matrix_view; - auto mds = cuvs::core::from_dlpack(additional_dataset_tensor); - auto return_mds = cuvs::core::from_dlpack(return_tensor); - cuvs::neighbors::cagra::extend(*res_ptr, extend_params, mds, *index_ptr, return_mds); - } else if (cuvs::core::is_dlpack_host_compatible(dataset) && - cuvs::core::is_dlpack_host_compatible(return_dl_tensor)) { - using mdspan_type = raft::host_matrix_view; - using mdspan_return_type = raft::device_matrix_view; - auto mds = cuvs::core::from_dlpack(additional_dataset_tensor); - auto return_mds = cuvs::core::from_dlpack(return_tensor); - cuvs::neighbors::cagra::extend(*res_ptr, extend_params, mds, *index_ptr, return_mds); + if (cuvs::core::is_dlpack_device_compatible(dataset)) { + using mdspan_type = raft::device_matrix_view; + auto mds = cuvs::core::from_dlpack(additional_dataset_tensor); + cuvs::neighbors::cagra::extend(*res_ptr, extend_params, mds, *index_ptr); + } else if (cuvs::core::is_dlpack_host_compatible(dataset)) { + using mdspan_type = raft::host_matrix_view; + auto mds = cuvs::core::from_dlpack(additional_dataset_tensor); + cuvs::neighbors::cagra::extend(*res_ptr, extend_params, mds, *index_ptr); } else { RAFT_FAIL("Unsupported dataset DLtensor dtype: %d and bits: %d", dataset.dtype.code, @@ -540,19 +532,18 @@ extern "C" cuvsError_t cuvsCagraIndexFromArgs(cuvsResources_t res, extern "C" cuvsError_t cuvsCagraExtend(cuvsResources_t res, cuvsCagraExtendParams_t params, DLManagedTensor* additional_dataset_tensor, - cuvsCagraIndex_t index_c_ptr, - DLManagedTensor* return_dataset_tensor) + cuvsCagraIndex_t index_c_ptr) { return cuvs::core::translate_exceptions([=] { auto dataset = additional_dataset_tensor->dl_tensor; auto index = *index_c_ptr; if ((dataset.dtype.code == kDLFloat) && (dataset.dtype.bits == 32)) { - _extend(res, *params, index, additional_dataset_tensor, return_dataset_tensor); + _extend(res, *params, index, additional_dataset_tensor); } else if (dataset.dtype.code == kDLInt && dataset.dtype.bits == 8) { - _extend(res, *params, index, additional_dataset_tensor, return_dataset_tensor); + _extend(res, *params, index, additional_dataset_tensor); } else if (dataset.dtype.code == kDLUInt && dataset.dtype.bits == 8) { - _extend(res, *params, index, additional_dataset_tensor, return_dataset_tensor); + _extend(res, *params, index, additional_dataset_tensor); } else { RAFT_FAIL("Unsupported dataset DLtensor dtype: %d and bits: %d", dataset.dtype.code, diff --git a/cpp/tests/neighbors/ann_cagra_c.cu b/cpp/tests/neighbors/ann_cagra_c.cu index ae80cc8986..3e73f545d4 100644 --- a/cpp/tests/neighbors/ann_cagra_c.cu +++ b/cpp/tests/neighbors/ann_cagra_c.cu @@ -1,5 +1,5 @@ /* - * Copyright (c) 2023-2024, NVIDIA CORPORATION. + * Copyright (c) 2023-2025, NVIDIA CORPORATION. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -210,20 +210,6 @@ TEST(CagraC, BuildExtendSearch) additional_dataset_tensor.dl_tensor.shape = additional_dataset_shape; additional_dataset_tensor.dl_tensor.strides = nullptr; - // create tensor for that points to the extended dataset - rmm::device_uvector extend_return_d((additional_data_size + main_data_size) * dimensions, - stream); - DLManagedTensor additional_dataset_return_tensor; - additional_dataset_return_tensor.dl_tensor.data = extend_return_d.data(); - additional_dataset_return_tensor.dl_tensor.device.device_type = kDLCUDA; - additional_dataset_return_tensor.dl_tensor.ndim = 2; - additional_dataset_return_tensor.dl_tensor.dtype.code = kDLFloat; - additional_dataset_return_tensor.dl_tensor.dtype.bits = 32; - additional_dataset_return_tensor.dl_tensor.dtype.lanes = 1; - int64_t additional_return_dataset_shape[2] = {additional_data_size + main_data_size, dimensions}; - additional_dataset_return_tensor.dl_tensor.shape = additional_return_dataset_shape; - additional_dataset_return_tensor.dl_tensor.strides = nullptr; - // create index cuvsCagraIndex_t index; cuvsCagraIndexCreate(&index); @@ -238,8 +224,7 @@ TEST(CagraC, BuildExtendSearch) // extend index cuvsCagraExtendParams_t extend_params; cuvsCagraExtendParamsCreate(&extend_params); - cuvsCagraExtend( - res, extend_params, &additional_dataset_tensor, index, &additional_dataset_return_tensor); + cuvsCagraExtend(res, extend_params, &additional_dataset_tensor, index); // create queries DLTensor rmm::device_uvector queries_d(num_queries * dimensions, stream); @@ -342,15 +327,6 @@ TEST(CagraC, BuildExtendSearch) cuvsCagraSearch( res, search_params, index, &queries_tensor, &neighbors_tensor, &distances_tensor, filter); - // make sure that extend_return_d points to the extended dataset - ASSERT_TRUE(cuvs::devArrMatch( - main_d.data(), extend_return_d.data(), main_d.size(), cuvs::Compare())); - - ASSERT_TRUE(cuvs::devArrMatch(additional_d.data(), - extend_return_d.data() + main_d.size(), - additional_d.size(), - cuvs::Compare())); - // check neighbors ASSERT_TRUE( cuvs::devArrMatch(min_cols.data_handle(), neighbors_d.data(), 4, cuvs::Compare())); diff --git a/docs/source/python_api/neighbors_cagra.rst b/docs/source/python_api/neighbors_cagra.rst index e7155efb83..42647914f2 100644 --- a/docs/source/python_api/neighbors_cagra.rst +++ b/docs/source/python_api/neighbors_cagra.rst @@ -44,3 +44,8 @@ Index load ########## .. autofunction:: cuvs.neighbors.cagra.load + +Index extend +############ + +.. autofunction:: cuvs.neighbors.cagra.extend diff --git a/go/cagra/cagra.go b/go/cagra/cagra.go index 61b9e71d6e..c68e85e956 100644 --- a/go/cagra/cagra.go +++ b/go/cagra/cagra.go @@ -51,13 +51,12 @@ func BuildIndex[T any](Resources cuvs.Resource, params *IndexParams, dataset *cu // * `Resources` - Resources to use // * `params` - Parameters for extending the index // * `additional_dataset` - A row-major Tensor on the device to extend the index with -// * `return_dataset` - A row-major Tensor on the device that will receive the extended dataset // * `index` - CagraIndex to extend -func ExtendIndex[T any](Resources cuvs.Resource, params *ExtendParams, additional_dataset *cuvs.Tensor[T], return_dataset *cuvs.Tensor[T], index *CagraIndex) error { +func ExtendIndex[T any](Resources cuvs.Resource, params *ExtendParams, additional_dataset *cuvs.Tensor[T], index *CagraIndex) error { if !index.trained { return errors.New("index needs to be built before calling extend") } - err := cuvs.CheckCuvs(cuvs.CuvsError(C.cuvsCagraExtend(C.ulong(Resources.Resource), params.params, (*C.DLManagedTensor)(unsafe.Pointer(additional_dataset.C_tensor)), index.index, (*C.DLManagedTensor)(unsafe.Pointer(return_dataset.C_tensor))))) + err := cuvs.CheckCuvs(cuvs.CuvsError(C.cuvsCagraExtend(C.ulong(Resources.Resource), params.params, (*C.DLManagedTensor)(unsafe.Pointer(additional_dataset.C_tensor)), index.index))) if err != nil { return err } diff --git a/python/cuvs/cuvs/neighbors/cagra/__init__.py b/python/cuvs/cuvs/neighbors/cagra/__init__.py index b32f9ca2d5..8d21e09022 100644 --- a/python/cuvs/cuvs/neighbors/cagra/__init__.py +++ b/python/cuvs/cuvs/neighbors/cagra/__init__.py @@ -15,10 +15,12 @@ from .cagra import ( CompressionParams, + ExtendParams, Index, IndexParams, SearchParams, build, + extend, from_graph, load, save, @@ -27,10 +29,12 @@ __all__ = [ "CompressionParams", + "ExtendParams", "Index", "IndexParams", "SearchParams", "build", + "extend", "from_graph", "load", "save", diff --git a/python/cuvs/cuvs/neighbors/cagra/cagra.pxd b/python/cuvs/cuvs/neighbors/cagra/cagra.pxd index 29e47a5014..bb01c047cd 100644 --- a/python/cuvs/cuvs/neighbors/cagra/cagra.pxd +++ b/python/cuvs/cuvs/neighbors/cagra/cagra.pxd @@ -175,6 +175,19 @@ cdef extern from "cuvs/neighbors/cagra.h" nogil: DLManagedTensor * dataset, cuvsCagraIndex_t index) + ctypedef struct cuvsCagraExtendParams: + uint32_t max_chunk_size + + ctypedef cuvsCagraExtendParams* cuvsCagraExtendParams_t + + cuvsError_t cuvsCagraExtendParamsCreate(cuvsCagraExtendParams_t* params) + cuvsError_t cuvsCagraExtendParamsDestroy(cuvsCagraExtendParams_t params) + cuvsError_t cuvsCagraExtend(cuvsResources_t res, + cuvsCagraExtendParams_t params, + DLManagedTensor* additional_dataset, + cuvsCagraIndex_t index) + + cdef class Index: """ CAGRA index object. This object stores the trained CAGRA index state diff --git a/python/cuvs/cuvs/neighbors/cagra/cagra.pyx b/python/cuvs/cuvs/neighbors/cagra/cagra.pyx index 0a8fec17b4..e09970222f 100644 --- a/python/cuvs/cuvs/neighbors/cagra/cagra.pyx +++ b/python/cuvs/cuvs/neighbors/cagra/cagra.pyx @@ -350,7 +350,7 @@ def build(IndexParams index_params, dataset, resources=None): ---------- index_params : IndexParams object dataset : CUDA array interface compliant matrix shape (n_samples, dim) - Supported dtype [float, int8, uint8] + Supported dtype [float, half, int8, uint8] {resources_docstring} Returns @@ -841,3 +841,73 @@ def from_graph(graph, dataset, metric="sqeuclidean", resources=None): idx.index)) idx.trained = True return idx + + +cdef class ExtendParams: + """ Supplemental parameters to extend CAGRA Index + + Parameters + ---------- + max_chunk_size : int + The additional dataset is divided into chunks and added to the graph. + This is the knob to adjust the tradeoff between the recall and + operation throughput. Large chunk sizes can result in high throughput, + but use more working memory (O(max_chunk_size*degree^2)). This can also + degrade recall because no edges are added between the nodes in the same + chunk. Auto select when 0. + """ + + cdef cuvsCagraExtendParams* params + + def __cinit__(self): + check_cuvs(cuvsCagraExtendParamsCreate(&self.params)) + + def __dealloc__(self): + check_cuvs(cuvsCagraExtendParamsDestroy(self.params)) + + def __init__(self, *, + max_chunk_size=None): + if max_chunk_size is not None: + self.params.max_chunk_size = max_chunk_size + + @property + def max_chunk_size(self): + return self.params.max_chunk_size + + +@auto_sync_resources +def extend(ExtendParams params, Index index, additional_dataset, + resources=None): + """ + Extend a CAGRA index with additional vectors + + Parameters + ---------- + params : ExtendParams object + index: Index + Existing cagra index to extend + additional_dataset : CUDA array interface compliant matrix shape + Supported dtype [float, half, int8, uint8] + {resources_docstring} + + """ + dataset_ai = wrap_array(additional_dataset) + _check_input_array(dataset_ai, [np.dtype("float32"), + np.dtype("float16"), + np.dtype("byte"), + np.dtype("ubyte")]) + + cdef cydlpack.DLManagedTensor* dataset_dlpack = \ + cydlpack.dlpack_c(dataset_ai) + + cdef cuvsResources_t res = resources.get_c_obj() + + with cuda_interruptible(): + check_cuvs(cuvsCagraExtend( + res, + params.params, + dataset_dlpack, + index.index + )) + + return index diff --git a/python/cuvs/cuvs/tests/test_cagra.py b/python/cuvs/cuvs/tests/test_cagra.py index 0b08d48325..359cec02f5 100644 --- a/python/cuvs/cuvs/tests/test_cagra.py +++ b/python/cuvs/cuvs/tests/test_cagra.py @@ -41,7 +41,7 @@ def run_cagra_build_search_test( array_type="device", compare=True, inplace=True, - add_data_on_build=True, + test_extend=False, search_params={}, compression=None, ): @@ -60,26 +60,24 @@ def run_cagra_build_search_test( compression=compression, ) - if array_type == "device": - index = cagra.build(build_params, dataset_device) - else: - index = cagra.build(build_params, dataset) - - if not add_data_on_build: + if test_extend: dataset_1 = dataset[: n_rows // 2, :] dataset_2 = dataset[n_rows // 2 :, :] - indices_1 = np.arange(n_rows // 2, dtype=np.uint32) - indices_2 = np.arange(n_rows // 2, n_rows, dtype=np.uint32) + extend_params = cagra.ExtendParams() if array_type == "device": dataset_1_device = device_ndarray(dataset_1) dataset_2_device = device_ndarray(dataset_2) - indices_1_device = device_ndarray(indices_1) - indices_2_device = device_ndarray(indices_2) - index = cagra.extend(index, dataset_1_device, indices_1_device) - index = cagra.extend(index, dataset_2_device, indices_2_device) + + index = cagra.build(build_params, dataset_1_device) + index = cagra.extend(extend_params, index, dataset_2_device) + else: + index = cagra.build(build_params, dataset_1) + index = cagra.extend(index, dataset_2) + else: + if array_type == "device": + index = cagra.build(build_params, dataset_device) else: - index = cagra.extend(index, dataset_1, indices_1) - index = cagra.extend(index, dataset_2, indices_2) + index = cagra.build(build_params, dataset) queries = generate_data((n_queries, n_cols), dtype) out_idx = np.zeros((n_queries, k), dtype=np.uint32) @@ -185,7 +183,7 @@ def test_filtered_cagra(sparsity): { "intermediate_graph_degree": 64, "graph_degree": 32, - "add_data_on_build": True, + "test_extend": False, "k": 1, "metric": "sqeuclidean", "build_algo": "ivf_pq", @@ -193,7 +191,7 @@ def test_filtered_cagra(sparsity): { "intermediate_graph_degree": 32, "graph_degree": 16, - "add_data_on_build": False, + "test_extend": True, "k": 5, "metric": "sqeuclidean", "build_algo": "ivf_pq", @@ -201,7 +199,7 @@ def test_filtered_cagra(sparsity): { "intermediate_graph_degree": 128, "graph_degree": 32, - "add_data_on_build": True, + "test_extend": False, "k": 10, "metric": "inner_product", "build_algo": "nn_descent", @@ -212,6 +210,7 @@ def test_cagra_index_params(params): # Note that inner_product tests use normalized input which we cannot # represent in int8, therefore we test only sqeuclidean metric here. run_cagra_build_search_test( + test_extend=params["test_extend"], k=params["k"], metric=params["metric"], graph_degree=params["graph_degree"], From b3a4fd80878f0233b74f36b6c91b275f3e980e0e Mon Sep 17 00:00:00 2001 From: Ben Frederickson Date: Tue, 16 Sep 2025 22:37:05 -0700 Subject: [PATCH 70/93] Use int64_t for getters in the cagra/ivf_flat c-api (#1272) Java doesn't support unsigned types well, and this is leading to us only supporting 2^31 items in the cagra index - instead of 2^32. Fix by changing the getters in the c-api to return int64_t values for the size/dim etc, instead of uint32_t. We were also not consistent on returning error codes in ivf_flat/ivf_pq for these getters - with cagra returning an error code, and the ivf* methods just returning the value directly. Change to always return an error code to be consistent. (Note that #1123 is also making this change for ivf-pq - and I'm skipping ivf-pq to avoid merge conflicts). Authors: - Ben Frederickson (https://github.com/benfred) - Corey J. Nolet (https://github.com/cjnolet) Approvers: - MithunR (https://github.com/mythrocks) - Corey J. Nolet (https://github.com/cjnolet) URL: https://github.com/rapidsai/cuvs/pull/1272 --- cpp/include/cuvs/neighbors/cagra.h | 6 +- cpp/include/cuvs/neighbors/ivf_flat.h | 4 +- cpp/include/cuvs/neighbors/ivf_pq.h | 7 +- cpp/src/neighbors/cagra_c.cpp | 6 +- cpp/src/neighbors/ivf_flat_c.cpp | 84 ++++++++++--------- cpp/src/neighbors/ivf_pq_c.cpp | 24 ++++-- cpp/tests/neighbors/ann_cagra_c.cu | 2 +- .../nvidia/cuvs/internal/CagraIndexImpl.java | 1 + .../com/nvidia/cuvs/internal/common/Util.java | 7 -- python/cuvs/cuvs/neighbors/cagra/cagra.pxd | 6 +- python/cuvs/cuvs/neighbors/cagra/cagra.pyx | 6 +- .../cuvs/cuvs/neighbors/ivf_flat/ivf_flat.pxd | 7 +- .../cuvs/cuvs/neighbors/ivf_flat/ivf_flat.pyx | 8 +- python/cuvs/cuvs/neighbors/ivf_pq/ivf_pq.pxd | 9 +- python/cuvs/cuvs/neighbors/ivf_pq/ivf_pq.pyx | 13 ++- 15 files changed, 110 insertions(+), 80 deletions(-) diff --git a/cpp/include/cuvs/neighbors/cagra.h b/cpp/include/cuvs/neighbors/cagra.h index 2532bbd3df..ddf3207aa0 100644 --- a/cpp/include/cuvs/neighbors/cagra.h +++ b/cpp/include/cuvs/neighbors/cagra.h @@ -356,7 +356,7 @@ cuvsError_t cuvsCagraIndexDestroy(cuvsCagraIndex_t index); * @param[out] dim return dimension of the index * @return cuvsError_t */ -cuvsError_t cuvsCagraIndexGetDims(cuvsCagraIndex_t index, int* dim); +cuvsError_t cuvsCagraIndexGetDims(cuvsCagraIndex_t index, int64_t* dim); /** * @brief Get size of the CAGRA index @@ -365,7 +365,7 @@ cuvsError_t cuvsCagraIndexGetDims(cuvsCagraIndex_t index, int* dim); * @param[out] size return number of vectors in the index * @return cuvsError_t */ -cuvsError_t cuvsCagraIndexGetSize(cuvsCagraIndex_t index, uint32_t* size); +cuvsError_t cuvsCagraIndexGetSize(cuvsCagraIndex_t index, int64_t* size); /** * @brief Get graph degree of the CAGRA index @@ -374,7 +374,7 @@ cuvsError_t cuvsCagraIndexGetSize(cuvsCagraIndex_t index, uint32_t* size); * @param[out] graph_degree return graph degree * @return cuvsError_t */ -cuvsError_t cuvsCagraIndexGetGraphDegree(cuvsCagraIndex_t index, uint32_t* graph_degree); +cuvsError_t cuvsCagraIndexGetGraphDegree(cuvsCagraIndex_t index, int64_t* graph_degree); /** * @brief Returns a view of the CAGRA dataset diff --git a/cpp/include/cuvs/neighbors/ivf_flat.h b/cpp/include/cuvs/neighbors/ivf_flat.h index 9e89ea354d..8211f1ef58 100644 --- a/cpp/include/cuvs/neighbors/ivf_flat.h +++ b/cpp/include/cuvs/neighbors/ivf_flat.h @@ -166,10 +166,10 @@ cuvsError_t cuvsIvfFlatIndexCreate(cuvsIvfFlatIndex_t* index); cuvsError_t cuvsIvfFlatIndexDestroy(cuvsIvfFlatIndex_t index); /** Get the number of clusters/inverted lists */ -uint32_t cuvsIvfFlatIndexGetNLists(cuvsIvfFlatIndex_t index); +cuvsError_t cuvsIvfFlatIndexGetNLists(cuvsIvfFlatIndex_t index, int64_t* n_lists); /** Get the dimensionality of the data */ -uint32_t cuvsIvfFlatIndexGetDim(cuvsIvfFlatIndex_t index); +cuvsError_t cuvsIvfFlatIndexGetDim(cuvsIvfFlatIndex_t index, int64_t* dim); /** * @brief Get the cluster centers corresponding to the lists [n_lists, dim] diff --git a/cpp/include/cuvs/neighbors/ivf_pq.h b/cpp/include/cuvs/neighbors/ivf_pq.h index eef10290c7..3ef7b8b7af 100644 --- a/cpp/include/cuvs/neighbors/ivf_pq.h +++ b/cpp/include/cuvs/neighbors/ivf_pq.h @@ -262,10 +262,13 @@ cuvsError_t cuvsIvfPqIndexCreate(cuvsIvfPqIndex_t* index); cuvsError_t cuvsIvfPqIndexDestroy(cuvsIvfPqIndex_t index); /** Get the number of clusters/inverted lists */ -uint32_t cuvsIvfPqIndexGetNLists(cuvsIvfPqIndex_t index); +cuvsError_t cuvsIvfPqIndexGetNLists(cuvsIvfPqIndex_t index, int64_t* n_lists); /** Get the dimensionality */ -uint32_t cuvsIvfPqIndexGetDim(cuvsIvfPqIndex_t index); +cuvsError_t cuvsIvfPqIndexGetDim(cuvsIvfPqIndex_t index, int64_t* dim); + +/** Get the size of the index */ +cuvsError_t cuvsIvfPqIndexGetSize(cuvsIvfPqIndex_t index, int64_t* size); /** * @brief Get the cluster centers corresponding to the lists in the original space diff --git a/cpp/src/neighbors/cagra_c.cpp b/cpp/src/neighbors/cagra_c.cpp index 92d20f85df..b3c8a93a91 100644 --- a/cpp/src/neighbors/cagra_c.cpp +++ b/cpp/src/neighbors/cagra_c.cpp @@ -418,7 +418,7 @@ extern "C" cuvsError_t cuvsCagraIndexDestroy(cuvsCagraIndex_t index_c_ptr) }); } -extern "C" cuvsError_t cuvsCagraIndexGetDims(cuvsCagraIndex_t index, int* dim) +extern "C" cuvsError_t cuvsCagraIndexGetDims(cuvsCagraIndex_t index, int64_t* dim) { return cuvs::core::translate_exceptions([=] { auto index_ptr = reinterpret_cast*>(index->addr); @@ -426,7 +426,7 @@ extern "C" cuvsError_t cuvsCagraIndexGetDims(cuvsCagraIndex_t index, int* dim) }); } -extern "C" cuvsError_t cuvsCagraIndexGetSize(cuvsCagraIndex_t index, uint32_t* size) +extern "C" cuvsError_t cuvsCagraIndexGetSize(cuvsCagraIndex_t index, int64_t* size) { return cuvs::core::translate_exceptions([=] { auto index_ptr = reinterpret_cast*>(index->addr); @@ -434,7 +434,7 @@ extern "C" cuvsError_t cuvsCagraIndexGetSize(cuvsCagraIndex_t index, uint32_t* s }); } -extern "C" cuvsError_t cuvsCagraIndexGetGraphDegree(cuvsCagraIndex_t index, uint32_t* graph_degree) +extern "C" cuvsError_t cuvsCagraIndexGetGraphDegree(cuvsCagraIndex_t index, int64_t* graph_degree) { return cuvs::core::translate_exceptions([=] { auto index_ptr = reinterpret_cast*>(index->addr); diff --git a/cpp/src/neighbors/ivf_flat_c.cpp b/cpp/src/neighbors/ivf_flat_c.cpp index 91ad9abbed..5d8f7462fa 100644 --- a/cpp/src/neighbors/ivf_flat_c.cpp +++ b/cpp/src/neighbors/ivf_flat_c.cpp @@ -373,50 +373,54 @@ extern "C" cuvsError_t cuvsIvfFlatExtend(cuvsResources_t res, }); } -extern "C" uint32_t cuvsIvfFlatIndexGetNLists(cuvsIvfFlatIndex_t index) +extern "C" cuvsError_t cuvsIvfFlatIndexGetNLists(cuvsIvfFlatIndex_t index, int64_t* n_lists) { - if (index->dtype.code == kDLFloat && index->dtype.bits == 32) { - auto index_ptr = - reinterpret_cast*>(index->addr); - return index_ptr->n_lists(); - } else if (index->dtype.code == kDLFloat && index->dtype.bits == 16) { - auto index_ptr = - reinterpret_cast*>(index->addr); - return index_ptr->n_lists(); - } else if (index->dtype.code == kDLInt && index->dtype.bits == 8) { - auto index_ptr = - reinterpret_cast*>(index->addr); - return index_ptr->n_lists(); - } else if (index->dtype.code == kDLUInt && index->dtype.bits == 8) { - auto index_ptr = - reinterpret_cast*>(index->addr); - return index_ptr->n_lists(); - } else { - return 0; - } + return cuvs::core::translate_exceptions([=] { + if (index->dtype.code == kDLFloat && index->dtype.bits == 32) { + auto index_ptr = + reinterpret_cast*>(index->addr); + *n_lists = index_ptr->n_lists(); + } else if (index->dtype.code == kDLFloat && index->dtype.bits == 16) { + auto index_ptr = + reinterpret_cast*>(index->addr); + *n_lists = index_ptr->n_lists(); + } else if (index->dtype.code == kDLInt && index->dtype.bits == 8) { + auto index_ptr = + reinterpret_cast*>(index->addr); + *n_lists = index_ptr->n_lists(); + } else if (index->dtype.code == kDLUInt && index->dtype.bits == 8) { + auto index_ptr = + reinterpret_cast*>(index->addr); + *n_lists = index_ptr->n_lists(); + } else { + RAFT_FAIL("Unsupported index dtype: %d and bits: %d", index->dtype.code, index->dtype.bits); + } + }); } -extern "C" uint32_t cuvsIvfFlatIndexGetDim(cuvsIvfFlatIndex_t index) +extern "C" cuvsError_t cuvsIvfFlatIndexGetDim(cuvsIvfFlatIndex_t index, int64_t* dim) { - if (index->dtype.code == kDLFloat && index->dtype.bits == 32) { - auto index_ptr = - reinterpret_cast*>(index->addr); - return index_ptr->dim(); - } else if (index->dtype.code == kDLFloat && index->dtype.bits == 16) { - auto index_ptr = - reinterpret_cast*>(index->addr); - return index_ptr->dim(); - } else if (index->dtype.code == kDLInt && index->dtype.bits == 8) { - auto index_ptr = - reinterpret_cast*>(index->addr); - return index_ptr->dim(); - } else if (index->dtype.code == kDLUInt && index->dtype.bits == 8) { - auto index_ptr = - reinterpret_cast*>(index->addr); - return index_ptr->dim(); - } else { - return 0; - } + return cuvs::core::translate_exceptions([=] { + if (index->dtype.code == kDLFloat && index->dtype.bits == 32) { + auto index_ptr = + reinterpret_cast*>(index->addr); + *dim = index_ptr->dim(); + } else if (index->dtype.code == kDLFloat && index->dtype.bits == 16) { + auto index_ptr = + reinterpret_cast*>(index->addr); + *dim = index_ptr->dim(); + } else if (index->dtype.code == kDLInt && index->dtype.bits == 8) { + auto index_ptr = + reinterpret_cast*>(index->addr); + *dim = index_ptr->dim(); + } else if (index->dtype.code == kDLUInt && index->dtype.bits == 8) { + auto index_ptr = + reinterpret_cast*>(index->addr); + *dim = index_ptr->dim(); + } else { + RAFT_FAIL("Unsupported index dtype: %d and bits: %d", index->dtype.code, index->dtype.bits); + } + }); } extern "C" cuvsError_t cuvsIvfFlatIndexGetCenters(cuvsIvfFlatIndex_t index, diff --git a/cpp/src/neighbors/ivf_pq_c.cpp b/cpp/src/neighbors/ivf_pq_c.cpp index 4abd080582..5c10d4111b 100644 --- a/cpp/src/neighbors/ivf_pq_c.cpp +++ b/cpp/src/neighbors/ivf_pq_c.cpp @@ -347,16 +347,28 @@ extern "C" cuvsError_t cuvsIvfPqExtend(cuvsResources_t res, }); } -extern "C" uint32_t cuvsIvfPqIndexGetNLists(cuvsIvfPqIndex_t index) +extern "C" cuvsError_t cuvsIvfPqIndexGetNLists(cuvsIvfPqIndex_t index, int64_t* n_lists) { - auto index_ptr = reinterpret_cast*>(index->addr); - return index_ptr->n_lists(); + return cuvs::core::translate_exceptions([=] { + auto index_ptr = reinterpret_cast*>(index->addr); + *n_lists = index_ptr->n_lists(); + }); } -extern "C" uint32_t cuvsIvfPqIndexGetDim(cuvsIvfPqIndex_t index) +extern "C" cuvsError_t cuvsIvfPqIndexGetDim(cuvsIvfPqIndex_t index, int64_t* dim) { - auto index_ptr = reinterpret_cast*>(index->addr); - return index_ptr->dim(); + return cuvs::core::translate_exceptions([=] { + auto index_ptr = reinterpret_cast*>(index->addr); + *dim = index_ptr->dim(); + }); +} + +extern "C" cuvsError_t cuvsIvfPqIndexGetSize(cuvsIvfPqIndex_t index, int64_t* size) +{ + return cuvs::core::translate_exceptions([=] { + auto index_ptr = reinterpret_cast*>(index->addr); + *size = index_ptr->size(); + }); } extern "C" cuvsError_t cuvsIvfPqIndexGetCenters(cuvsIvfPqIndex_t index, DLManagedTensor* centers) diff --git a/cpp/tests/neighbors/ann_cagra_c.cu b/cpp/tests/neighbors/ann_cagra_c.cu index 3e73f545d4..3123e2b3b3 100644 --- a/cpp/tests/neighbors/ann_cagra_c.cu +++ b/cpp/tests/neighbors/ann_cagra_c.cu @@ -517,7 +517,7 @@ TEST(CagraC, BuildMergeSearch) cuvsCagraIndex_t index_array[2] = {index_main, index_add}; ASSERT_EQ(cuvsCagraMerge(res, merge_params, index_array, 2, index_merged), CUVS_SUCCESS); - int merged_dim = -1; + int64_t merged_dim = -1; ASSERT_EQ(cuvsCagraIndexGetDims(index_merged, &merged_dim), CUVS_SUCCESS); EXPECT_EQ(merged_dim, 2); diff --git a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CagraIndexImpl.java b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CagraIndexImpl.java index aa082482dc..99b0ecc71e 100644 --- a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CagraIndexImpl.java +++ b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CagraIndexImpl.java @@ -36,6 +36,7 @@ import com.nvidia.cuvs.internal.common.CloseableRMMAllocation; import com.nvidia.cuvs.internal.common.CompositeCloseableHandle; import com.nvidia.cuvs.internal.panama.*; + import java.io.FileInputStream; import java.io.InputStream; import java.io.OutputStream; diff --git a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/common/Util.java b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/common/Util.java index 28dfc5c020..1ffb5ad559 100644 --- a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/common/Util.java +++ b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/common/Util.java @@ -106,13 +106,6 @@ public static void checkCudaError(int value, String caller) { } } - private static final long UNSIGNED_INT_MASK = 0xFFFFFFFFL; - - public static long dereferenceUnsignedInt(MemorySegment ptr) { - assert ptr.byteSize() == 4; - return ptr.get(uint32_t, 0) & UNSIGNED_INT_MASK; - } - /** * Java analog to CUDA's cudaMemcpyKind, used for cudaMemcpy() calls. * @see CUDA Runtime API diff --git a/python/cuvs/cuvs/neighbors/cagra/cagra.pxd b/python/cuvs/cuvs/neighbors/cagra/cagra.pxd index bb01c047cd..b498cf1681 100644 --- a/python/cuvs/cuvs/neighbors/cagra/cagra.pxd +++ b/python/cuvs/cuvs/neighbors/cagra/cagra.pxd @@ -134,10 +134,10 @@ cdef extern from "cuvs/neighbors/cagra.h" nogil: cuvsError_t cuvsCagraIndexDestroy(cuvsCagraIndex_t index) - cuvsError_t cuvsCagraIndexGetDims(cuvsCagraIndex_t index, int32_t* dim) - cuvsError_t cuvsCagraIndexGetSize(cuvsCagraIndex_t index, uint32_t* size) + cuvsError_t cuvsCagraIndexGetDims(cuvsCagraIndex_t index, int64_t* dim) + cuvsError_t cuvsCagraIndexGetSize(cuvsCagraIndex_t index, int64_t* size) cuvsError_t cuvsCagraIndexGetGraphDegree(cuvsCagraIndex_t index, - uint32_t* degree) + int64_t* degree) cuvsError_t cuvsCagraIndexGetGraph(cuvsCagraIndex_t index, DLManagedTensor * graph) cuvsError_t cuvsCagraIndexGetDataset(cuvsCagraIndex_t index, diff --git a/python/cuvs/cuvs/neighbors/cagra/cagra.pyx b/python/cuvs/cuvs/neighbors/cagra/cagra.pyx index e09970222f..d82ea99ffa 100644 --- a/python/cuvs/cuvs/neighbors/cagra/cagra.pyx +++ b/python/cuvs/cuvs/neighbors/cagra/cagra.pyx @@ -275,18 +275,18 @@ cdef class Index: @property def dim(self): - cdef int32_t dim + cdef int64_t dim check_cuvs(cuvsCagraIndexGetDims(self.index, &dim)) return dim @property def graph_degree(self): - cdef uint32_t degree + cdef int64_t degree check_cuvs(cuvsCagraIndexGetGraphDegree(self.index, °ree)) return degree def __len__(self): - cdef uint32_t size + cdef int64_t size check_cuvs(cuvsCagraIndexGetSize(self.index, &size)) return size diff --git a/python/cuvs/cuvs/neighbors/ivf_flat/ivf_flat.pxd b/python/cuvs/cuvs/neighbors/ivf_flat/ivf_flat.pxd index 6e30ac0a07..470234d9cf 100644 --- a/python/cuvs/cuvs/neighbors/ivf_flat/ivf_flat.pxd +++ b/python/cuvs/cuvs/neighbors/ivf_flat/ivf_flat.pxd @@ -15,7 +15,7 @@ # # cython: language_level=3 -from libc.stdint cimport uint32_t, uintptr_t +from libc.stdint cimport int64_t, uint32_t, uintptr_t from libcpp cimport bool from cuvs.common.c_api cimport cuvsError_t, cuvsResources_t @@ -62,9 +62,10 @@ cdef extern from "cuvs/neighbors/ivf_flat.h" nogil: cuvsError_t cuvsIvfFlatIndexDestroy(cuvsIvfFlatIndex_t index) - uint32_t cuvsIvfFlatIndexGetNLists(cuvsIvfFlatIndex_t index) + cuvsError_t cuvsIvfFlatIndexGetNLists(cuvsIvfFlatIndex_t index, + int64_t * n_lists) - uint32_t cuvsIvfFlatIndexGetDim(cuvsIvfFlatIndex_t index) + cuvsError_t cuvsIvfFlatIndexGetDim(cuvsIvfFlatIndex_t index, int64_t * dim) cuvsError_t cuvsIvfFlatIndexGetCenters(cuvsIvfFlatIndex_t index, DLManagedTensor * centers) diff --git a/python/cuvs/cuvs/neighbors/ivf_flat/ivf_flat.pyx b/python/cuvs/cuvs/neighbors/ivf_flat/ivf_flat.pyx index 3a803e5599..996a39d4fe 100644 --- a/python/cuvs/cuvs/neighbors/ivf_flat/ivf_flat.pyx +++ b/python/cuvs/cuvs/neighbors/ivf_flat/ivf_flat.pyx @@ -188,12 +188,16 @@ cdef class Index: @property def n_lists(self): """ The number of inverted lists (clusters) """ - return cuvsIvfFlatIndexGetNLists(self.index) + cdef int64_t n_lists = 0 + cuvsIvfFlatIndexGetNLists(self.index, &n_lists) + return n_lists @property def dim(self): """ dimensionality of the cluster centers """ - return cuvsIvfFlatIndexGetDim(self.index) + cdef int64_t dim = 0 + cuvsIvfFlatIndexGetDim(self.index, &dim) + return dim @property def centers(self): diff --git a/python/cuvs/cuvs/neighbors/ivf_pq/ivf_pq.pxd b/python/cuvs/cuvs/neighbors/ivf_pq/ivf_pq.pxd index 897e4951a9..5ca7b97602 100644 --- a/python/cuvs/cuvs/neighbors/ivf_pq/ivf_pq.pxd +++ b/python/cuvs/cuvs/neighbors/ivf_pq/ivf_pq.pxd @@ -15,7 +15,7 @@ # # cython: language_level=3 -from libc.stdint cimport uint32_t, uintptr_t +from libc.stdint cimport int64_t, uint32_t, uintptr_t from libcpp cimport bool from cuvs.common.c_api cimport cuvsError_t, cuvsResources_t @@ -84,9 +84,12 @@ cdef extern from "cuvs/neighbors/ivf_pq.h" nogil: cuvsError_t cuvsIvfPqIndexDestroy(cuvsIvfPqIndex_t index) - uint32_t cuvsIvfPqIndexGetNLists(cuvsIvfPqIndex_t index) + cuvsError_t cuvsIvfPqIndexGetNLists(cuvsIvfPqIndex_t index, + int64_t * n_lists) - uint32_t cuvsIvfPqIndexGetDim(cuvsIvfPqIndex_t index) + cuvsError_t cuvsIvfPqIndexGetDim(cuvsIvfPqIndex_t index, int64_t * dim) + + cuvsError_t cuvsIvfPqIndexGetSize(cuvsIvfPqIndex_t index, int64_t * size) cuvsError_t cuvsIvfPqIndexGetCenters(cuvsIvfPqIndex_t index, DLManagedTensor * centers) diff --git a/python/cuvs/cuvs/neighbors/ivf_pq/ivf_pq.pyx b/python/cuvs/cuvs/neighbors/ivf_pq/ivf_pq.pyx index f784327d33..5ffa0b0c88 100644 --- a/python/cuvs/cuvs/neighbors/ivf_pq/ivf_pq.pyx +++ b/python/cuvs/cuvs/neighbors/ivf_pq/ivf_pq.pyx @@ -249,12 +249,21 @@ cdef class Index: @property def n_lists(self): """ The number of inverted lists (clusters) """ - return cuvsIvfPqIndexGetNLists(self.index) + cdef int64_t n_lists + check_cuvs(cuvsIvfPqIndexGetNLists(self.index, &n_lists)) + return n_lists @property def dim(self): """ dimensionality of the cluster centers """ - return cuvsIvfPqIndexGetDim(self.index) + cdef int64_t dim + check_cuvs(cuvsIvfPqIndexGetDim(self.index, &dim)) + return dim + + def __len__(self): + cdef int64_t size + check_cuvs(cuvsIvfPqIndexGetSize(self.index, &size)) + return size @property def centers(self): From d5edf1298ed2ab2b01d5941ed7d35b22ff61e65f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Lorenzo=20Dematt=C3=A9?= Date: Wed, 17 Sep 2025 16:36:38 +0200 Subject: [PATCH 71/93] [Java] Uniform toHost/toDevice to work across all CuVSMatrix classes (#1328) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit This PR moves `toHost`/`toDevice` operations to `CuVSMatrix` in order to uniform them; implementations then do not need to worry or perform checks on the input type, but simply accept any `CuVSMatrix` and then "request" the type they want or need (no ifs, custom cudaAlloc or cudaMemcpy operations needed): ``` CuVSDeviceMatrix matrix = input.toDevice(resources); ``` If the source matrix is a `CuVSHostMatrix`, device memory will be allocated and data will be copied; if the source matrix is a already a `CuVSDeviceMatrix`, a (wrapped) reference to the original data is returned (very fast, no copies, no additional allocations). The wrapper is currently implemented to follow C++ `weak_ptr` semantics: it delegates everything to the wrapped matrix, the only difference is on `close()` (which is a no-op in the wrapper, as you'd expect). The lifecycle of the data follows the one of the original matrix, the "weak" delegate does not influence it in any way. This way the caller can handle it uniformly (e.g. with try-with-resources), without side effects on the original matrix (which lifecycle will be already handled by its owner). An alternative would be to give this C++ `shared_ptr` semantics, making it reference counted. I think this would be more complex and not necessary, but let me know. Authors: - Lorenzo Dematté (https://github.com/ldematte) - MithunR (https://github.com/mythrocks) Approvers: - MithunR (https://github.com/mythrocks) URL: https://github.com/rapidsai/cuvs/pull/1328 --- .../com/nvidia/cuvs/CuVSDeviceMatrix.java | 74 +++++++++++-- .../java/com/nvidia/cuvs/CuVSHostMatrix.java | 73 ++++++++++++ .../main/java/com/nvidia/cuvs/CuVSMatrix.java | 36 ++++++ .../cuvs/internal/CuVSDeviceMatrixImpl.java | 5 + .../cuvs/internal/CuVSHostMatrixImpl.java | 15 +++ .../cuvs/internal/CuVSMatrixBaseImpl.java | 24 ++++ .../java/com/nvidia/cuvs/CuVSMatrixIT.java | 104 +++++++++++++++++- 7 files changed, 317 insertions(+), 14 deletions(-) diff --git a/java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSDeviceMatrix.java b/java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSDeviceMatrix.java index b654163a63..ec34a3de1c 100644 --- a/java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSDeviceMatrix.java +++ b/java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSDeviceMatrix.java @@ -21,16 +21,7 @@ public interface CuVSDeviceMatrix extends CuVSMatrix { /** - * Fills the provided, pre-allocated host matrix with data from this device matrix. - * The content of the provided host matrix will be overwritten; the 2 matrices must have the - * same element type and dimension. - * - * @param hostMatrix the host-memory-backed matrix to fill. - */ - void toHost(CuVSHostMatrix hostMatrix); - - /** - * Returns a new, host matrix with data from this device matrix. + * Returns a new host matrix with data from this device matrix. * The returned host matrix will need to be managed by the caller, which will be * responsible to call {@link CuVSMatrix#close()} to free its resources when done. */ @@ -39,4 +30,67 @@ default CuVSHostMatrix toHost() { toHost(hostMatrix); return hostMatrix; } + + default CuVSDeviceMatrix toDevice(CuVSResources resources) { + return new CuVSDeviceMatrixDelegate(this); + } + + class CuVSDeviceMatrixDelegate implements CuVSDeviceMatrix { + + private final CuVSDeviceMatrix deviceMatrix; + + private CuVSDeviceMatrixDelegate(CuVSDeviceMatrix deviceMatrix) { + this.deviceMatrix = deviceMatrix; + } + + @Override + public long size() { + return deviceMatrix.size(); + } + + @Override + public long columns() { + return deviceMatrix.columns(); + } + + @Override + public DataType dataType() { + return deviceMatrix.dataType(); + } + + @Override + public RowView getRow(long row) { + return deviceMatrix.getRow(row); + } + + @Override + public void toArray(int[][] array) { + deviceMatrix.toArray(array); + } + + @Override + public void toArray(float[][] array) { + deviceMatrix.toArray(array); + } + + @Override + public void toArray(byte[][] array) { + deviceMatrix.toArray(array); + } + + @Override + public void toHost(CuVSHostMatrix hostMatrix) { + deviceMatrix.toHost(hostMatrix); + } + + @Override + public void toDevice(CuVSDeviceMatrix deviceMatrix, CuVSResources cuVSResources) { + this.deviceMatrix.toDevice(deviceMatrix, cuVSResources); + } + + @Override + public void close() { + // Do nothing + } + } } diff --git a/java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSHostMatrix.java b/java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSHostMatrix.java index 237dfa39f7..91ff9012f3 100644 --- a/java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSHostMatrix.java +++ b/java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSHostMatrix.java @@ -20,4 +20,77 @@ */ public interface CuVSHostMatrix extends CuVSMatrix { int get(int row, int col); + + default CuVSHostMatrix toHost() { + return new CuVSHostMatrixDelegate(this); + } + + default CuVSDeviceMatrix toDevice(CuVSResources resources) { + var deviceMatrix = CuVSMatrix.deviceBuilder(resources, size(), columns(), dataType()).build(); + toDevice(deviceMatrix, resources); + return deviceMatrix; + } + + class CuVSHostMatrixDelegate implements CuVSHostMatrix { + private final CuVSHostMatrix hostMatrix; + + public CuVSHostMatrixDelegate(CuVSHostMatrix cuVSHostMatrix) { + this.hostMatrix = cuVSHostMatrix; + } + + @Override + public int get(int row, int col) { + return hostMatrix.get(row, col); + } + + @Override + public long size() { + return hostMatrix.size(); + } + + @Override + public long columns() { + return hostMatrix.columns(); + } + + @Override + public DataType dataType() { + return hostMatrix.dataType(); + } + + @Override + public RowView getRow(long row) { + return hostMatrix.getRow(row); + } + + @Override + public void toArray(int[][] array) { + hostMatrix.toArray(array); + } + + @Override + public void toArray(float[][] array) { + hostMatrix.toArray(array); + } + + @Override + public void toArray(byte[][] array) { + hostMatrix.toArray(array); + } + + @Override + public void toHost(CuVSHostMatrix hostMatrix) { + this.hostMatrix.toHost(hostMatrix); + } + + @Override + public void toDevice(CuVSDeviceMatrix deviceMatrix, CuVSResources cuVSResources) { + hostMatrix.toDevice(deviceMatrix, cuVSResources); + } + + @Override + public void close() { + // Do nothing + } + } } diff --git a/java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSMatrix.java b/java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSMatrix.java index 708009be02..3ccdb9ae28 100644 --- a/java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSMatrix.java +++ b/java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSMatrix.java @@ -177,6 +177,42 @@ static Builder deviceBuilder( */ void toArray(byte[][] array); + /** + * Fills the provided, pre-allocated host matrix with data from this matrix. + * The content of the provided host matrix will be overwritten; the 2 matrices must have the + * same element type and dimension. + * + * @param hostMatrix the host-memory-backed matrix to fill. + */ + void toHost(CuVSHostMatrix hostMatrix); + + /** + * Returns a host matrix; if the matrix is already a host matrix, a "weak" reference to the same host memory + * is returned. If the matrix is a device matrix, a newly allocated matrix will be populated with data from + * the device matrix. + * The returned host matrix will need to be managed by the caller, which will be + * responsible to call {@link CuVSMatrix#close()} to free its resources when done. + */ + CuVSHostMatrix toHost(); + + /** + * Fills the provided, pre-allocated device matrix with data from this matrix. + * The content of the provided device matrix will be overwritten; the 2 matrices must have the + * same element type and dimension. + * + * @param deviceMatrix the device-memory-backed matrix to fill. + */ + void toDevice(CuVSDeviceMatrix deviceMatrix, CuVSResources cuVSResources); + + /** + * Returns a device matrix; if this matrix is already a device matrix, a "weak" reference to the same host memory + * is returned. If the matrix is a host matrix, a newly allocated matrix will be populated with data from + * the host matrix. + * The returned device matrix will need to be managed by the caller, which will be + * responsible to call {@link CuVSMatrix#close()} to free its resources when done. + */ + CuVSDeviceMatrix toDevice(CuVSResources cuVSResources); + @Override void close(); } diff --git a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CuVSDeviceMatrixImpl.java b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CuVSDeviceMatrixImpl.java index 50b914c76b..0f97608ae1 100644 --- a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CuVSDeviceMatrixImpl.java +++ b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CuVSDeviceMatrixImpl.java @@ -236,6 +236,11 @@ public void toHost(CuVSHostMatrix hostMatrix) { } } + @Override + public void toDevice(CuVSDeviceMatrix targetMatrix, CuVSResources cuVSResources) { + copyMatrix(this, (CuVSMatrixBaseImpl) targetMatrix, cuVSResources); + } + @Override public void close() { if (hostBuffer != MemorySegment.NULL) { diff --git a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CuVSHostMatrixImpl.java b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CuVSHostMatrixImpl.java index 317b315c77..408feda2bf 100644 --- a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CuVSHostMatrixImpl.java +++ b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CuVSHostMatrixImpl.java @@ -18,7 +18,9 @@ import static com.nvidia.cuvs.internal.common.Util.prepareTensor; import static com.nvidia.cuvs.internal.panama.headers_h.*; +import com.nvidia.cuvs.CuVSDeviceMatrix; import com.nvidia.cuvs.CuVSHostMatrix; +import com.nvidia.cuvs.CuVSResources; import com.nvidia.cuvs.RowView; import java.lang.foreign.*; import java.lang.invoke.VarHandle; @@ -142,6 +144,19 @@ public void toArray(byte[][] array) { } } + @Override + public void toHost(CuVSHostMatrix hostMatrix) { + var targetMatrix = (CuVSHostMatrixImpl) hostMatrix; + var valueByteSize = valueLayout.byteSize(); + MemorySegment.copy( + this.memorySegment, 0L, targetMatrix.memorySegment, 0L, size * columns * valueByteSize); + } + + @Override + public void toDevice(CuVSDeviceMatrix deviceMatrix, CuVSResources cuVSResources) { + copyMatrix(this, (CuVSMatrixBaseImpl) deviceMatrix, cuVSResources); + } + @Override public void close() {} diff --git a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CuVSMatrixBaseImpl.java b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CuVSMatrixBaseImpl.java index d629645e61..b46ab63c9f 100644 --- a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CuVSMatrixBaseImpl.java +++ b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CuVSMatrixBaseImpl.java @@ -18,6 +18,7 @@ import static com.nvidia.cuvs.internal.common.LinkerHelper.C_CHAR; import static com.nvidia.cuvs.internal.common.LinkerHelper.C_FLOAT; import static com.nvidia.cuvs.internal.common.LinkerHelper.C_INT; +import static com.nvidia.cuvs.internal.common.Util.checkCuVSError; import static com.nvidia.cuvs.internal.panama.headers_h.*; import com.nvidia.cuvs.CuVSMatrix; @@ -49,6 +50,29 @@ protected CuVSMatrixBaseImpl( this.columns = columns; } + protected static void copyMatrix( + CuVSMatrixBaseImpl sourceMatrix, CuVSMatrixBaseImpl targetMatrix, CuVSResources resources) { + if (targetMatrix.columns() != sourceMatrix.columns + || targetMatrix.size() != sourceMatrix.size) { + throw new IllegalArgumentException( + "Source and target matrices must have the same dimensions"); + } + if (targetMatrix.dataType() != sourceMatrix.dataType) { + throw new IllegalArgumentException("Source and target matrices must have the same dataType"); + } + + try (var localArena = Arena.ofConfined()) { + var targetTensor = targetMatrix.toTensor(localArena); + + try (var resourceAccess = resources.access()) { + var cuvsRes = resourceAccess.handle(); + var sourceTensor = sourceMatrix.toTensor(localArena); + checkCuVSError(cuvsMatrixCopy(cuvsRes, sourceTensor, targetTensor), "cuvsMatrixCopy"); + checkCuVSError(cuvsStreamSync(cuvsRes), "cuvsStreamSync"); + } + } + } + @Override public long size() { return size; 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 0eaf666d15..499c7d366f 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 @@ -262,7 +262,7 @@ public void testFloatDatasetCopy() { } } - private void testFloatDatasetBuilder(int rows, int cols, CuVSMatrix.Builder builder) { + private void testFloatDatasetBuilder(int rows, int cols, CuVSMatrix.Builder builder) { float[][] data = new float[rows][cols]; for (int r = 0; r < rows; ++r) { @@ -306,7 +306,7 @@ public void testFloatDatasetDeviceBuilder() throws Throwable { } } - private void testIntDatasetBuilder(int rows, int cols, CuVSMatrix.Builder builder) { + private void testIntDatasetBuilder(int rows, int cols, CuVSMatrix.Builder builder) { var data = new int[rows][cols]; for (int r = 0; r < rows; ++r) { @@ -349,7 +349,7 @@ public void testIntDatasetDeviceBuilder() throws Throwable { } } - private void testByteDatasetBuilder(int rows, int cols, CuVSMatrix.Builder builder) { + private void testByteDatasetBuilder(int rows, int cols, CuVSMatrix.Builder builder) { var data = new byte[rows][cols]; for (int r = 0; r < rows; ++r) { @@ -408,7 +408,7 @@ public void testDeviceToHost() throws Throwable { builder.addVector(array); } - try (var deviceMatrix = (CuVSDeviceMatrix) builder.build(); + try (var deviceMatrix = builder.build(); var hostMatrix = deviceMatrix.toHost()) { assertEquals(data.length, deviceMatrix.size()); @@ -429,4 +429,100 @@ public void testDeviceToHost() throws Throwable { } } } + + @Test + public void testHostToDevice() throws Throwable { + + final int size = 16 * 1024; + final int columns = 2048; + final float[][] data = createFloatMatrix(size, columns); + + try (var resources = CuVSResources.create()) { + + try (var hostMatrix = CuVSMatrix.ofArray(data); + var deviceMatrix = hostMatrix.toDevice(resources)) { + + assertEquals(data.length, deviceMatrix.size()); + assertEquals(data[0].length, deviceMatrix.columns()); + + assertEquals(deviceMatrix.size(), hostMatrix.size()); + assertEquals(deviceMatrix.columns(), hostMatrix.columns()); + + var roundTripData = new float[size][columns]; + + deviceMatrix.toArray(roundTripData); + + for (int n = 0; n < deviceMatrix.size(); ++n) { + for (int i = 0; i < deviceMatrix.columns(); ++i) { + assertEquals(data[n][i], roundTripData[n][i], 1e-9); + } + } + } + } + } + + @Test + public void testHostToHostReturnsWeakReferenceSameData() { + + final int size = 16 * 1024; + final int columns = 2048; + final float[][] data = createFloatMatrix(size, columns); + + try (var hostMatrix = CuVSMatrix.ofArray(data); + var hostMatrix2 = hostMatrix.toHost()) { + + assertEquals(data.length, hostMatrix2.size()); + assertEquals(data[0].length, hostMatrix2.columns()); + + assertEquals(hostMatrix2.size(), hostMatrix.size()); + assertEquals(hostMatrix2.columns(), hostMatrix.columns()); + + var roundTripData = new float[size][columns]; + + hostMatrix2.toArray(roundTripData); + + for (int n = 0; n < hostMatrix2.size(); ++n) { + for (int i = 0; i < hostMatrix2.columns(); ++i) { + assertEquals(data[n][i], roundTripData[n][i], 1e-9); + } + } + } + } + + @Test + public void testDeviceToDeviceReturnsWeakReferenceSameData() throws Throwable { + + final int size = 16 * 1024; + final int columns = 2048; + final float[][] data = createFloatMatrix(size, columns); + + try (var resources = CuVSResources.create()) { + + var builder = CuVSMatrix.deviceBuilder(resources, size, columns, CuVSMatrix.DataType.FLOAT); + for (int i = 0; i < size; ++i) { + var array = data[i]; + builder.addVector(array); + } + + try (var deviceMatrix = builder.build(); + var deviceMatrix2 = deviceMatrix.toDevice(resources)) { + + assertEquals(data.length, deviceMatrix.size()); + assertEquals(data[0].length, deviceMatrix.columns()); + + assertEquals(deviceMatrix.size(), deviceMatrix2.size()); + assertEquals(deviceMatrix.columns(), deviceMatrix2.columns()); + + var roundTripData = new float[size][columns]; + + deviceMatrix2.toArray(roundTripData); + + for (int n = 0; n < deviceMatrix2.size(); ++n) { + for (int i = 0; i < deviceMatrix2.columns(); ++i) { + assertEquals(data[n][i], roundTripData[n][i], 1e-9); + } + } + } + } + } } From 5492a061f7124485c146fbbf4173d8b6d8afc072 Mon Sep 17 00:00:00 2001 From: Robert Maynard Date: Wed, 17 Sep 2025 10:58:10 -0400 Subject: [PATCH 72/93] Refactor CMakeLists to support static only builds of libcuvs (#1317) Fixes: https://github.com/rapidsai/cuvs/issues/1312 Fixes: https://github.com/rapidsai/cuvs/issues/1322 Supports: https://github.com/rapidsai/cuvs/issues/1241 If you want a shared libcuvs but static RAPIDS dependencies ( raft, rmm, logger, .... ). You can use `CUVS_STATIC_RAPIDS_LIBRARIES=ON` with `CUVS_COMPILE_DYNAMIC_ONLY=ON`. If you want a fully static cuvs and dependencies you can use `BUILD_SHARED_LIBS=OFF` ( `CUVS_STATIC_RAPIDS_LIBRARIES` is not needed but can also be set to `ON` ) Authors: - Robert Maynard (https://github.com/robertmaynard) - Corey J. Nolet (https://github.com/cjnolet) Approvers: - Corey J. Nolet (https://github.com/cjnolet) URL: https://github.com/rapidsai/cuvs/pull/1317 --- cpp/CMakeLists.txt | 281 ++++++++++++++-------------- cpp/cmake/thirdparty/get_raft.cmake | 7 +- python/libcuvs/CMakeLists.txt | 2 +- 3 files changed, 143 insertions(+), 147 deletions(-) diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt index 04621aaffd..d21f130ee8 100644 --- a/cpp/CMakeLists.txt +++ b/cpp/CMakeLists.txt @@ -63,7 +63,7 @@ option(CUDA_ENABLE_LINEINFO ) option(CUDA_STATIC_RUNTIME "Statically link the CUDA runtime" OFF) option(CUDA_STATIC_MATH_LIBRARIES "Statically link the CUDA math libraries" OFF) -option(CUVS_USE_RAFT_STATIC "Build and statically link the RAFT libraries" OFF) +option(CUVS_STATIC_RAPIDS_LIBRARIES "Build and statically link RAPIDS libraries" OFF) option(CUDA_LOG_COMPILE_TIME "Write a log of compilation times to nvcc_compile_log.csv" OFF) option(DETECT_CONDA_ENV "Enable detection of conda environment for dependencies" ON) option(DISABLE_DEPRECATION_WARNINGS "Disable deprecaction warnings " ON) @@ -83,6 +83,23 @@ elseif(NOT BUILD_SHARED_LIBS) set(BUILD_CAGRA_HNSWLIB OFF) endif() +if(NOT BUILD_SHARED_LIBS AND CUVS_COMPILE_DYNAMIC_ONLY) + message( + VERBOSE + "cuVS: Ignoring `BUILD_SHARED_LIBS` and forcing shared library due to `CUVS_COMPILE_DYNAMIC_ONLY`" + ) +endif() + +set(cuvs_compile_mode "both") +set(_cuvs_lib_targets cuvs cuvs_static) +if(CUVS_COMPILE_DYNAMIC_ONLY) + set(cuvs_compile_mode "shared_only") + set(_cuvs_lib_targets cuvs) +elseif(NOT BUILD_SHARED_LIBS) + set(cuvs_compile_mode "static_only") + set(_cuvs_lib_targets cuvs_static) +endif() + # Needed because GoogleBenchmark changes the state of FindThreads.cmake, causing subsequent runs to # have different values for the `Threads::Threads` target. Setting this flag ensures # `Threads::Threads` is the same value across all builds so that cache hits occur @@ -104,7 +121,7 @@ message(VERBOSE "cuVS: Enable lineinfo in nvcc: ${CUDA_ENABLE_LINEINFO}") message(VERBOSE "cuVS: Enable nvtx markers: ${CUVS_NVTX}") message(VERBOSE "cuVS: Statically link the CUDA runtime: ${CUDA_STATIC_RUNTIME}") message(VERBOSE "cuVS: Statically link the CUDA math libraries: ${CUDA_STATIC_MATH_LIBRARIES}") -message(VERBOSE "cuVS: Build and statically link RAFT libraries: ${CUVS_USE_RAFT_STATIC}") +message(VERBOSE "cuVS: Build and statically link RAPIDS libraries: ${CUVS_STATIC_RAPIDS_LIBRARIES}") # Set RMM logging level set(RMM_LOGGING_LEVEL @@ -200,7 +217,7 @@ endif() # ################################################################################################## # * cuvs --------------------------------------------------------------------- -if(BUILD_SHARED_LIBS) +if(NOT BUILD_CPU_ONLY) add_library( cuvs-cagra-search STATIC src/neighbors/cagra_search_float.cu @@ -519,25 +536,116 @@ if(BUILD_SHARED_LIBS) cuvs_objs PRIVATE "$<$:${CUVS_CXX_FLAGS}>" "$<$:${CUVS_CUDA_FLAGS}>" ) + + target_compile_definitions( + cuvs_objs PRIVATE $<$:CUVS_BUILD_CAGRA_HNSWLIB> + $<$:NVTX_ENABLED> + ) + target_link_libraries( - cuvs_objs PUBLIC raft::raft rmm::rmm ${CUVS_CTK_MATH_DEPENDENCIES} - $ + cuvs_objs + PRIVATE raft::raft rmm::rmm ${CUVS_CTK_MATH_DEPENDENCIES} + $ $ ) - add_library(cuvs SHARED $) + target_include_directories( + cuvs_objs + PUBLIC "$" + "$" + INTERFACE "$" + ) - if(NOT CUVS_COMPILE_DYNAMIC_ONLY) - add_library(cuvs_static STATIC $) + # Endian detection + include(TestBigEndian) + test_big_endian(BIG_ENDIAN) + if(BIG_ENDIAN) + target_compile_definitions(cuvs_objs PRIVATE CUVS_SYSTEM_LITTLE_ENDIAN=0) + else() + target_compile_definitions(cuvs_objs PRIVATE CUVS_SYSTEM_LITTLE_ENDIAN=1) endif() - target_compile_options( - cuvs INTERFACE $<$:--expt-extended-lambda - --expt-relaxed-constexpr> + if(BUILD_MG_ALGOS) + rapids_find_generate_module( + NCCL + HEADER_NAMES nccl.h + LIBRARY_NAMES nccl + ) + find_package(NCCL REQUIRED) + target_link_libraries(cuvs_objs PUBLIC $) + + target_compile_definitions(cuvs_objs PUBLIC CUVS_BUILD_MG_ALGOS) + target_compile_definitions(cuvs-cagra-search PUBLIC CUVS_BUILD_MG_ALGOS) + endif() + + set(CUVS_CUSOLVER_DEPENDENCY CUDA::cusolver${_ctk_static_suffix}) + set(CUVS_CUBLAS_DEPENDENCY CUDA::cublas${_ctk_static_suffix}) + set(CUVS_CURAND_DEPENDENCY CUDA::curand${_ctk_static_suffix}) + set(CUVS_CUSPARSE_DEPENDENCY CUDA::cusparse${_ctk_static_suffix}) + + set(CUVS_CTK_MATH_DEPENDENCIES ${CUVS_CUBLAS_DEPENDENCY} ${CUVS_CUSOLVER_DEPENDENCY} + ${CUVS_CUSPARSE_DEPENDENCY} ${CUVS_CURAND_DEPENDENCY} ) - add_library(cuvs::cuvs ALIAS cuvs) + if(NOT cuvs_compile_mode STREQUAL "static_only") + add_library(cuvs SHARED $) + add_library(cuvs::cuvs ALIAS cuvs) + set_target_properties( + cuvs + PROPERTIES BUILD_RPATH "\$ORIGIN" + INSTALL_RPATH "\$ORIGIN" + CXX_STANDARD 17 + CXX_STANDARD_REQUIRED ON + CUDA_STANDARD 17 + CUDA_STANDARD_REQUIRED ON + INTERFACE_POSITION_INDEPENDENT_CODE ON + POSITION_INDEPENDENT_CODE ON + ) + target_compile_options( + cuvs + INTERFACE $<$:--expt-extended-lambda + --expt-relaxed-constexpr> + PRIVATE "$<$:${CUVS_CXX_FLAGS}>" + "$<$:${CUVS_CUDA_FLAGS}>" + ) + target_compile_definitions( + cuvs PUBLIC $<$:CUVS_BUILD_CAGRA_HNSWLIB> + $<$:NVTX_ENABLED> + ) + + target_include_directories( + cuvs PUBLIC "$" + "$" + ) - if(NOT CUVS_COMPILE_DYNAMIC_ONLY) + target_link_libraries( + cuvs + PUBLIC rmm::rmm + raft::raft + ${CUVS_CTK_MATH_DEPENDENCIES} + $> + $> + $<$:CUDA::nvtx3> + PRIVATE nvidia::cutlass::cutlass $ + cuvs-cagra-search + ) + + # ensure CUDA symbols aren't relocated to the middle of the debug build binaries + file( + WRITE "${CMAKE_CURRENT_BINARY_DIR}/fatbin.ld" + [=[ +SECTIONS +{ +.nvFatBinSegment : { *(.nvFatBinSegment) } +.nv_fatbin : { *(.nv_fatbin) } +} +]=] + ) + target_link_options(cuvs PRIVATE $) + + endif() + + if(NOT cuvs_compile_mode STREQUAL "shared_only") + add_library(cuvs_static STATIC $) add_library(cuvs::cuvs_static ALIAS cuvs_static) set_target_properties( @@ -552,16 +660,11 @@ if(BUILD_SHARED_LIBS) ) target_compile_options(cuvs_static PRIVATE "$<$:${CUVS_CXX_FLAGS}>") - endif() - - target_include_directories( - cuvs_objs - PUBLIC "$" - "$" - INTERFACE "$" - ) + target_compile_definitions( + cuvs_static PUBLIC $<$:CUVS_BUILD_CAGRA_HNSWLIB> + $<$:NVTX_ENABLED> + ) - if(NOT CUVS_COMPILE_DYNAMIC_ONLY) target_include_directories( cuvs_static PUBLIC "$" @@ -575,102 +678,21 @@ if(BUILD_SHARED_LIBS) cuvs_static PUBLIC "$" "$" ) - endif() - - target_include_directories( - cuvs PUBLIC "$" - "$" - ) - - rapids_find_package( - OpenMP REQUIRED - BUILD_EXPORT_SET cuvs-exports - INSTALL_EXPORT_SET cuvs-exports - ) - - if(NOT BUILD_CPU_ONLY) - - set(CUVS_CUSOLVER_DEPENDENCY CUDA::cusolver${_ctk_static_suffix}) - set(CUVS_CUBLAS_DEPENDENCY CUDA::cublas${_ctk_static_suffix}) - set(CUVS_CURAND_DEPENDENCY CUDA::curand${_ctk_static_suffix}) - set(CUVS_CUSPARSE_DEPENDENCY CUDA::cusparse${_ctk_static_suffix}) - - set(CUVS_CTK_MATH_DEPENDENCIES ${CUVS_CUBLAS_DEPENDENCY} ${CUVS_CUSOLVER_DEPENDENCY} - ${CUVS_CUSPARSE_DEPENDENCY} ${CUVS_CURAND_DEPENDENCY} - ) - - if(BUILD_MG_ALGOS) - rapids_find_generate_module( - NCCL - HEADER_NAMES nccl.h - LIBRARY_NAMES nccl - ) - find_package(NCCL REQUIRED) - target_link_libraries(cuvs_objs PRIVATE NCCL::NCCL) - target_link_libraries(cuvs PRIVATE NCCL::NCCL) - endif() - - # Keep cuVS as lightweight as possible. Only CUDA libs and rmm should be used in global target. target_link_libraries( - cuvs - PUBLIC rmm::rmm raft::raft ${CUVS_CTK_MATH_DEPENDENCIES} + cuvs_static + PUBLIC rmm::rmm + raft::raft + ${CUVS_CTK_MATH_DEPENDENCIES} + $> + $> + $<$:CUDA::nvtx3> PRIVATE nvidia::cutlass::cutlass $ - cuvs-cagra-search ) - - if(NOT CUVS_COMPILE_DYNAMIC_ONLY) - target_link_libraries( - cuvs_static - PUBLIC rmm::rmm raft::raft ${CUVS_CTK_MATH_DEPENDENCIES} - PRIVATE nvidia::cutlass::cutlass $ - ) - endif() - endif() - - if(BUILD_MG_ALGOS) - target_compile_definitions(cuvs PUBLIC CUVS_BUILD_MG_ALGOS) - target_compile_definitions(cuvs_objs PUBLIC CUVS_BUILD_MG_ALGOS) - target_compile_definitions(cuvs-cagra-search PUBLIC CUVS_BUILD_MG_ALGOS) - endif() - - if(BUILD_CAGRA_HNSWLIB) - target_link_libraries(cuvs_objs PRIVATE hnswlib::hnswlib) - target_compile_definitions(cuvs PUBLIC CUVS_BUILD_CAGRA_HNSWLIB) - target_compile_definitions(cuvs_objs PUBLIC CUVS_BUILD_CAGRA_HNSWLIB) - endif() - - # Endian detection - include(TestBigEndian) - test_big_endian(BIG_ENDIAN) - if(BIG_ENDIAN) - target_compile_definitions(cuvs PRIVATE CUVS_SYSTEM_LITTLE_ENDIAN=0) - else() - target_compile_definitions(cuvs PRIVATE CUVS_SYSTEM_LITTLE_ENDIAN=1) endif() - file( - WRITE "${CMAKE_CURRENT_BINARY_DIR}/fatbin.ld" - [=[ -SECTIONS -{ -.nvFatBinSegment : { *(.nvFatBinSegment) } -.nv_fatbin : { *(.nv_fatbin) } -} -]=] - ) - # ################################################################################################ # * NVTX support in cuvs ----------------------------------------------------- - - if(CUVS_NVTX) - # This enables NVTX within the project with no option to disable it downstream. - target_link_libraries(cuvs PUBLIC CUDA::nvtx3) - target_compile_definitions(cuvs_objs PUBLIC NVTX_ENABLED) - target_compile_definitions(cuvs PUBLIC NVTX_ENABLED) - - target_link_libraries(cuvs-cagra-search PUBLIC CUDA::nvtx3) - target_compile_definitions(cuvs-cagra-search PUBLIC NVTX_ENABLED) - else() + if(NOT CUVS_NVTX) # Allow enable NVTX downstream if not set here. This creates a new option at build/install time, # which is set by default to OFF, but can be enabled in the dependent project. get_property( @@ -693,25 +715,6 @@ target_compile_definitions(cuvs::cuvs INTERFACE $<$:NVTX_ENAB ) endif() - set_target_properties( - cuvs - PROPERTIES BUILD_RPATH "\$ORIGIN" - INSTALL_RPATH "\$ORIGIN" - CXX_STANDARD 17 - CXX_STANDARD_REQUIRED ON - CUDA_STANDARD 17 - CUDA_STANDARD_REQUIRED ON - INTERFACE_POSITION_INDEPENDENT_CODE ON - POSITION_INDEPENDENT_CODE ON - ) - - target_compile_options( - cuvs PRIVATE "$<$:${CUVS_CXX_FLAGS}>" - "$<$:${CUVS_CUDA_FLAGS}>" - ) - # ensure CUDA symbols aren't relocated to the middle of the debug build binaries - target_link_options(cuvs PRIVATE $) - # ################################################################################################ # * cuvs_c ------------------------------------------------------------------------------- @@ -739,14 +742,7 @@ target_compile_definitions(cuvs::cuvs INTERFACE $<$:NVTX_ENAB src/distance/pairwise_distance_c.cpp ${CUVS_C_MG_ALGOS} ) - - if(BUILD_CAGRA_HNSWLIB) - target_link_libraries(cuvs_c PRIVATE hnswlib::hnswlib) - target_compile_definitions(cuvs_c PUBLIC CUVS_BUILD_CAGRA_HNSWLIB) - endif() - add_library(cuvs::c_api ALIAS cuvs_c) - set_target_properties( cuvs_c PROPERTIES BUILD_RPATH "\$ORIGIN" @@ -758,6 +754,10 @@ target_compile_definitions(cuvs::cuvs INTERFACE $<$:NVTX_ENAB EXPORT_NAME c_api ) + target_compile_definitions( + cuvs_c PUBLIC $<$:CUVS_BUILD_CAGRA_HNSWLIB> + ) + target_compile_options(cuvs_c PRIVATE "$<$:${CUVS_CXX_FLAGS}>") target_include_directories( @@ -770,7 +770,7 @@ target_compile_definitions(cuvs::cuvs INTERFACE $<$:NVTX_ENAB target_link_libraries( cuvs_c PUBLIC cuvs::cuvs ${CUVS_CTK_MATH_DEPENDENCIES} - PRIVATE raft::raft + PRIVATE raft::raft $ ) # ensure CUDA symbols aren't relocated to the middle of the debug build binaries @@ -783,11 +783,6 @@ target_compile_definitions(cuvs::cuvs INTERFACE $<$:NVTX_ENAB include(GNUInstallDirs) include(CPack) - set(_cuvs_lib_targets cuvs) - if(NOT CUVS_COMPILE_DYNAMIC_ONLY) - list(APPEND _cuvs_lib_targets cuvs_static) - endif() - install( TARGETS ${_cuvs_lib_targets} DESTINATION ${lib_dir} diff --git a/cpp/cmake/thirdparty/get_raft.cmake b/cpp/cmake/thirdparty/get_raft.cmake index 424e829cb6..6a08edb2ab 100644 --- a/cpp/cmake/thirdparty/get_raft.cmake +++ b/cpp/cmake/thirdparty/get_raft.cmake @@ -17,16 +17,17 @@ set(RAFT_FORK "rapidsai") set(RAFT_PINNED_TAG "${rapids-cmake-checkout-tag}") function(find_and_configure_raft) - set(oneValueArgs VERSION FORK PINNED_TAG USE_RAFT_STATIC ENABLE_NVTX ENABLE_MNMG_DEPENDENCIES CLONE_ON_PIN) + set(oneValueArgs VERSION FORK PINNED_TAG BUILD_STATIC_DEPS ENABLE_NVTX ENABLE_MNMG_DEPENDENCIES CLONE_ON_PIN) cmake_parse_arguments(PKG "${options}" "${oneValueArgs}" "${multiValueArgs}" ${ARGN} ) if(PKG_CLONE_ON_PIN AND NOT PKG_PINNED_TAG STREQUAL "${rapids-cmake-checkout-tag}") message(STATUS "cuVS: RAFT pinned tag found: ${PKG_PINNED_TAG}. Cloning raft locally.") set(CPM_DOWNLOAD_raft ON) - elseif(PKG_USE_RAFT_STATIC AND (NOT CPM_raft_SOURCE)) + elseif(PKG_BUILD_STATIC_DEPS AND (NOT CPM_raft_SOURCE)) message(STATUS "cuVS: Cloning raft locally to build static libraries.") set(CPM_DOWNLOAD_raft ON) + set(BUILD_SHARED_LIBS OFF) endif() set(RAFT_COMPONENTS "") @@ -65,7 +66,7 @@ find_and_configure_raft(VERSION ${RAFT_VERSION}.00 PINNED_TAG ${RAFT_PINNED_TAG} ENABLE_MNMG_DEPENDENCIES OFF ENABLE_NVTX OFF - USE_RAFT_STATIC ${CUVS_USE_RAFT_STATIC} + BUILD_STATIC_DEPS ${CUVS_STATIC_RAPIDS_LIBRARIES} # When PINNED_TAG above doesn't match the default rapids branch, # force local raft clone in build directory # even if it's already installed. diff --git a/python/libcuvs/CMakeLists.txt b/python/libcuvs/CMakeLists.txt index 02574ac104..aec24b9ca0 100644 --- a/python/libcuvs/CMakeLists.txt +++ b/python/libcuvs/CMakeLists.txt @@ -42,7 +42,7 @@ set(CUDA_STATIC_RUNTIME ON) set(CUDA_STATIC_MATH_LIBRARIES OFF) # --- RAFT ---# -set(CUVS_USE_RAFT_STATIC OFF) +set(CUVS_STATIC_RAPIDS_LIBRARIES OFF) # --- cuVS ---# set(BUILD_TESTS OFF) From f33de2a8b499ca805082e3dc8e9617fc3d7e92a5 Mon Sep 17 00:00:00 2001 From: Victor Lafargue Date: Thu, 18 Sep 2025 18:43:31 +0200 Subject: [PATCH 73/93] All neighbors C and Python bindings (#1282) Answers https://github.com/rapidsai/cuvs/issues/1279 Authors: - Victor Lafargue (https://github.com/viclafargue) - Corey J. Nolet (https://github.com/cjnolet) Approvers: - Kyle Edwards (https://github.com/KyleFromNVIDIA) - Corey J. Nolet (https://github.com/cjnolet) - Jinsol Park (https://github.com/jinsolp) URL: https://github.com/rapidsai/cuvs/pull/1282 --- cpp/CMakeLists.txt | 1 + cpp/include/cuvs/core/c_api.h | 28 ++ cpp/include/cuvs/neighbors/all_neighbors.h | 122 ++++++ cpp/src/core/c_api.cpp | 28 ++ cpp/src/neighbors/all_neighbors_c.cpp | 287 ++++++++++++++ cpp/tests/CMakeLists.txt | 5 + cpp/tests/neighbors/all_neighbors_c.cu | 301 ++++++++++++++ cpp/tests/neighbors/c_api.c | 16 + cpp/tests/neighbors/run_all_neighbors_c.c | 124 ++++++ docs/source/c_api/neighbors.rst | 1 + .../c_api/neighbors_all_neighbors_c.rst | 18 + docs/source/cuvs_bench/param_tuning.rst | 4 +- docs/source/getting_started.rst | 6 +- docs/source/neighbors/all_neighbors.rst | 71 ++++ .../{indexes => neighbors}/bruteforce.rst | 0 docs/source/{indexes => neighbors}/cagra.rst | 0 .../source/{indexes => neighbors}/ivfflat.rst | 0 docs/source/{indexes => neighbors}/ivfpq.rst | 0 .../indexes.rst => neighbors/neighbors.rst} | 5 +- docs/source/{indexes => neighbors}/vamana.rst | 0 docs/source/python_api/neighbors.rst | 1 + .../python_api/neighbors_all_neighbors.rst | 19 + docs/source/tuning_guide.rst | 2 +- python/cuvs/cuvs/common/CMakeLists.txt | 6 +- python/cuvs/cuvs/common/__init__.py | 10 +- python/cuvs/cuvs/common/c_api.pxd | 6 + python/cuvs/cuvs/common/mg_resources.pxd | 24 ++ python/cuvs/cuvs/common/mg_resources.pyx | 153 +++++++ python/cuvs/cuvs/neighbors/CMakeLists.txt | 3 +- python/cuvs/cuvs/neighbors/__init__.py | 4 +- .../neighbors/all_neighbors/CMakeLists.txt | 24 ++ .../cuvs/neighbors/all_neighbors/__init__.pxd | 0 .../cuvs/neighbors/all_neighbors/__init__.py | 17 + .../neighbors/all_neighbors/all_neighbors.pxd | 52 +++ .../neighbors/all_neighbors/all_neighbors.pyx | 373 ++++++++++++++++++ .../cuvs/neighbors/nn_descent/nn_descent.pyx | 4 + python/cuvs/cuvs/tests/test_all_neighbors.py | 219 ++++++++++ 37 files changed, 1920 insertions(+), 14 deletions(-) create mode 100644 cpp/include/cuvs/neighbors/all_neighbors.h create mode 100644 cpp/src/neighbors/all_neighbors_c.cpp create mode 100644 cpp/tests/neighbors/all_neighbors_c.cu create mode 100644 cpp/tests/neighbors/run_all_neighbors_c.c create mode 100644 docs/source/c_api/neighbors_all_neighbors_c.rst create mode 100644 docs/source/neighbors/all_neighbors.rst rename docs/source/{indexes => neighbors}/bruteforce.rst (100%) rename docs/source/{indexes => neighbors}/cagra.rst (100%) rename docs/source/{indexes => neighbors}/ivfflat.rst (100%) rename docs/source/{indexes => neighbors}/ivfpq.rst (100%) rename docs/source/{indexes/indexes.rst => neighbors/neighbors.rst} (80%) rename docs/source/{indexes => neighbors}/vamana.rst (100%) create mode 100644 docs/source/python_api/neighbors_all_neighbors.rst create mode 100644 python/cuvs/cuvs/common/mg_resources.pxd create mode 100644 python/cuvs/cuvs/common/mg_resources.pyx create mode 100644 python/cuvs/cuvs/neighbors/all_neighbors/CMakeLists.txt create mode 100644 python/cuvs/cuvs/neighbors/all_neighbors/__init__.pxd create mode 100644 python/cuvs/cuvs/neighbors/all_neighbors/__init__.py create mode 100644 python/cuvs/cuvs/neighbors/all_neighbors/all_neighbors.pxd create mode 100644 python/cuvs/cuvs/neighbors/all_neighbors/all_neighbors.pyx create mode 100644 python/cuvs/cuvs/tests/test_all_neighbors.py diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt index d21f130ee8..85e1c171ab 100644 --- a/cpp/CMakeLists.txt +++ b/cpp/CMakeLists.txt @@ -737,6 +737,7 @@ target_compile_definitions(cuvs::cuvs INTERFACE $<$:NVTX_ENAB src/neighbors/nn_descent_c.cpp src/neighbors/refine/refine_c.cpp src/neighbors/tiered_index_c.cpp + src/neighbors/all_neighbors_c.cpp src/preprocessing/quantize/binary_c.cpp src/preprocessing/quantize/scalar_c.cpp src/distance/pairwise_distance_c.cpp diff --git a/cpp/include/cuvs/core/c_api.h b/cpp/include/cuvs/core/c_api.h index e2fc4ebc94..b6319fe3b0 100644 --- a/cpp/include/cuvs/core/c_api.h +++ b/cpp/include/cuvs/core/c_api.h @@ -127,6 +127,34 @@ cuvsError_t cuvsStreamSync(cuvsResources_t res); * @return cuvsError_t */ cuvsError_t cuvsDeviceIdGet(cuvsResources_t res, int* device_id); + +/** + * @brief Create an Initialized opaque C handle for C++ type `raft::device_resources_snmg` + * for multi-GPU operations + * + * @param[in] res cuvsResources_t opaque C handle + * @return cuvsError_t + */ +cuvsError_t cuvsMultiGpuResourcesCreate(cuvsResources_t* res); + +/** + * @brief Create an Initialized opaque C handle for C++ type `raft::device_resources_snmg` + * for multi-GPU operations with specific device IDs + * + * @param[in] res cuvsResources_t opaque C handle + * @param[in] device_ids DLManagedTensor* containing device IDs to use + * @return cuvsError_t + */ +cuvsError_t cuvsMultiGpuResourcesCreateWithDeviceIds(cuvsResources_t* res, + DLManagedTensor* device_ids); + +/** + * @brief Destroy and de-allocate opaque C handle for C++ type `raft::device_resources_snmg` + * + * @param[in] res cuvsResources_t opaque C handle + * @return cuvsError_t + */ +cuvsError_t cuvsMultiGpuResourcesDestroy(cuvsResources_t res); /** @} */ /** diff --git a/cpp/include/cuvs/neighbors/all_neighbors.h b/cpp/include/cuvs/neighbors/all_neighbors.h new file mode 100644 index 0000000000..08d509c825 --- /dev/null +++ b/cpp/include/cuvs/neighbors/all_neighbors.h @@ -0,0 +1,122 @@ +/* + * Copyright (c) 2025, NVIDIA CORPORATION. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include +#include +#include +#include +#include +#include + +#ifdef __cplusplus +extern "C" { +#endif + +/** + * @defgroup all_neighbors_c_build All-neighbors C-API build + * @{ + * + * All-neighbors constructs an approximate k-NN graph for all vectors in a dataset. + * The All-neighbors API can either be used in single-GPU or multi-GPU mode. + * For multi-GPU deployment, please pass a multi-GPU resources handle and + * provide the dataset on host. + * + * Notes: + * - Outputs (indices, distances, core_distances) are expected to be on device memory. + * - Host variant accepts host-resident dataset; device variant accepts device-resident dataset. + * - For batching, `overlap_factor < n_clusters` must hold. + * - When `core_distances` is provided, mutual-reachability distances are produced (see alpha). + */ + +/** + * @brief Graph build algorithm selection. + */ +typedef enum { + CUVS_ALL_NEIGHBORS_ALGO_BRUTE_FORCE = 0, ///< Use Brute Force for local kNN subgraphs + CUVS_ALL_NEIGHBORS_ALGO_IVF_PQ = 1, ///< Use IVF-PQ for local kNN subgraphs (host dataset only) + CUVS_ALL_NEIGHBORS_ALGO_NN_DESCENT = 2 ///< Use NN-Descent for local kNN subgraphs +} cuvsAllNeighborsAlgo; + +/** + * @brief Parameters controlling SNMG all-neighbors build. + */ +struct cuvsAllNeighborsIndexParams { + cuvsAllNeighborsAlgo algo; ///< Local kNN graph build algorithm + size_t overlap_factor; ///< Number of clusters each point is assigned to (must be < n_clusters) + size_t + n_clusters; ///< Number of clusters/batches to partition the dataset into (> overlap_factor) + cuvsDistanceType metric; ///< Distance metric used for graph construction + + // Algorithm-specific parameters (only one should be set based on algo) + cuvsIvfPqIndexParams_t ivf_pq_params; ///< Parameters for IVF-PQ algorithm (when algo == + ///< CUVS_ALL_NEIGHBORS_ALGO_IVF_PQ) + cuvsNNDescentIndexParams_t nn_descent_params; ///< Parameters for NN-Descent algorithm (when algo + ///< == CUVS_ALL_NEIGHBORS_ALGO_NN_DESCENT) +}; + +typedef struct cuvsAllNeighborsIndexParams* cuvsAllNeighborsIndexParams_t; + +/** + * @brief Create a default all-neighbors index parameters struct. + * + * @param[out] index_params Pointer to allocated index_params struct + * + * @return cuvsError_t + */ +cuvsError_t cuvsAllNeighborsIndexParamsCreate(cuvsAllNeighborsIndexParams_t* index_params); + +/** + * @brief Destroy an all-neighbors index parameters struct. + * + * @param[in] index_params Index parameters struct to destroy + * + * @return cuvsError_t + */ +cuvsError_t cuvsAllNeighborsIndexParamsDestroy(cuvsAllNeighborsIndexParams_t index_params); + +/** + * @brief Build an all-neighbors k-NN graph automatically detecting host vs device dataset. + * + * @param[in] res Can be a SNMG multi-GPU resources (`cuvsResources_t`) or single-GPU + * resources + * @param[in] params Build parameters (see cuvsAllNeighborsIndexParams) + * @param[in] dataset 2D tensor [num_rows x dim] on host or device (auto-detected) + * @param[out] indices 2D tensor [num_rows x k] on device (int64) + * @param[out] distances Optional 2D tensor [num_rows x k] on device (float32); can be NULL + * @param[out] core_distances Optional 1D tensor [num_rows] on device (float32); can be NULL + * @param[in] alpha Mutual-reachability scaling; used only when core_distances is provided + * + * The function automatically detects whether the dataset is host-resident or device-resident + * and calls the appropriate implementation. For host datasets, it partitions data into + * `n_clusters` clusters and assigns each row to `overlap_factor` nearest clusters. For device + * datasets, `n_clusters` must be 1 (no batching); `overlap_factor` is ignored. + * Outputs always reside in device memory. + */ +cuvsError_t cuvsAllNeighborsBuild(cuvsResources_t res, + cuvsAllNeighborsIndexParams_t params, + DLManagedTensor* dataset, + DLManagedTensor* indices, + DLManagedTensor* distances, + DLManagedTensor* core_distances, + float alpha); + +/** @} */ + +#ifdef __cplusplus +} +#endif diff --git a/cpp/src/core/c_api.cpp b/cpp/src/core/c_api.cpp index a36bb472b6..3cf6db457b 100644 --- a/cpp/src/core/c_api.cpp +++ b/cpp/src/core/c_api.cpp @@ -60,6 +60,34 @@ extern "C" cuvsError_t cuvsMultiGpuResourcesCreate(cuvsResources_t* res) }); } +extern "C" cuvsError_t cuvsMultiGpuResourcesCreateWithDeviceIds(cuvsResources_t* res, + DLManagedTensor* device_ids) +{ + return cuvs::core::translate_exceptions([=] { + // Basic validation + if (device_ids == nullptr || device_ids->dl_tensor.data == nullptr) { + throw std::invalid_argument("device_ids cannot be null"); + } + + // Check data type is int32 + if (device_ids->dl_tensor.dtype.code != kDLInt || device_ids->dl_tensor.dtype.bits != 32) { + throw std::invalid_argument("device_ids must be int32"); + } + + // Check data is on host + if (device_ids->dl_tensor.device.device_type != kDLCPU) { + throw std::invalid_argument("device_ids must be on host memory"); + } + + // Cast void* to int* to perform pointer arithmetic + int* data_ptr = static_cast(device_ids->dl_tensor.data); + std::vector ids(data_ptr, data_ptr + device_ids->dl_tensor.shape[0]); + + auto res_ptr = new raft::device_resources_snmg{ids}; + *res = reinterpret_cast(res_ptr); + }); +} + extern "C" cuvsError_t cuvsMultiGpuResourcesDestroy(cuvsResources_t res) { return cuvs::core::translate_exceptions([=] { diff --git a/cpp/src/neighbors/all_neighbors_c.cpp b/cpp/src/neighbors/all_neighbors_c.cpp new file mode 100644 index 0000000000..38f6b5c2a2 --- /dev/null +++ b/cpp/src/neighbors/all_neighbors_c.cpp @@ -0,0 +1,287 @@ +/* + * Copyright (c) 2025, NVIDIA CORPORATION. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include +#include + +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +#include "ivf_pq_c.hpp" + +namespace { + +void convert_nn_descent_params(cuvsNNDescentIndexParams params, + cuvs::neighbors::nn_descent::index_params* out) +{ + out->metric = static_cast((int)params.metric); + out->metric_arg = params.metric_arg; + out->graph_degree = params.graph_degree; + out->intermediate_graph_degree = params.intermediate_graph_degree; + out->max_iterations = params.max_iterations; + out->termination_threshold = params.termination_threshold; + out->return_distances = params.return_distances; +} + +static cuvs::neighbors::all_neighbors::all_neighbors_params convert_params( + cuvsAllNeighborsIndexParams_t params_ptr, int64_t n_rows, int64_t n_cols) +{ + using namespace cuvs::neighbors; + cuvs::neighbors::all_neighbors::all_neighbors_params out{}; + auto& p = *params_ptr; + out.metric = static_cast((int)p.metric); + out.overlap_factor = p.overlap_factor; + out.n_clusters = p.n_clusters; + + switch (p.algo) { + case CUVS_ALL_NEIGHBORS_ALGO_BRUTE_FORCE: { + graph_build_params::brute_force_params b{}; + b.build_params.metric = out.metric; + out.graph_build_params = b; + break; + } + case CUVS_ALL_NEIGHBORS_ALGO_NN_DESCENT: { + graph_build_params::nn_descent_params n{}; + n.metric = out.metric; + // Use nn_descent_params if provided, otherwise use defaults + if (p.nn_descent_params != nullptr) { convert_nn_descent_params(*p.nn_descent_params, &n); } + out.graph_build_params = n; + break; + } + case CUVS_ALL_NEIGHBORS_ALGO_IVF_PQ: { + auto dataset_extents = raft::matrix_extent{n_rows, n_cols}; + graph_build_params::ivf_pq_params ivf(dataset_extents, out.metric); + // Use ivf_pq_params if provided, otherwise use defaults + if (p.ivf_pq_params != nullptr) { + cuvs::neighbors::ivf_pq::convert_c_index_params(*p.ivf_pq_params, &ivf.build_params); + } + out.graph_build_params = ivf; + break; + } + default: RAFT_FAIL("Unsupported all-neighbors algo"); + } + + return out; +} + +static void ensure_indices_dtype_and_device_compatibility(DLManagedTensor* indices) +{ + auto dtype = indices->dl_tensor.dtype; + RAFT_EXPECTS(dtype.code == kDLInt && dtype.bits == 64, "indices must be int64 output tensor"); + RAFT_EXPECTS(cuvs::core::is_dlpack_device_compatible(indices->dl_tensor), + "indices tensor must be device-compatible"); +} + +static void ensure_optional_distance_dtype_and_device_compatibility(DLManagedTensor* distances) +{ + if (distances == nullptr) { return; } + auto dtype = distances->dl_tensor.dtype; + RAFT_EXPECTS(dtype.code == kDLFloat && dtype.bits == 32, + "distances must be float32 output tensor"); + RAFT_EXPECTS(cuvs::core::is_dlpack_device_compatible(distances->dl_tensor), + "distances tensor must be device-compatible"); +} + +static void ensure_optional_core_distance_dtype_and_device_compatibility( + DLManagedTensor* core_distances) +{ + if (core_distances == nullptr) { return; } + auto dtype = core_distances->dl_tensor.dtype; + RAFT_EXPECTS(dtype.code == kDLFloat && dtype.bits == 32, + "core_distances must be float32 output tensor"); + RAFT_EXPECTS(cuvs::core::is_dlpack_device_compatible(core_distances->dl_tensor), + "core_distances tensor must be device-compatible"); +} + +template +void _build_host(cuvsResources_t res, + cuvsAllNeighborsIndexParams_t params, + DLManagedTensor* dataset_tensor, + DLManagedTensor* indices_tensor, + DLManagedTensor* distances_tensor, + DLManagedTensor* core_distances_tensor, + float alpha) +{ + auto& cpp_res = *reinterpret_cast(res); + + auto& dlt = dataset_tensor->dl_tensor; + RAFT_EXPECTS(cuvs::core::is_dlpack_host_compatible(dlt), + "Host build expects host-compatible dataset tensor"); + + ensure_indices_dtype_and_device_compatibility(indices_tensor); + ensure_optional_distance_dtype_and_device_compatibility(distances_tensor); + ensure_optional_core_distance_dtype_and_device_compatibility(core_distances_tensor); + + // Check dependencies between parameters + if (core_distances_tensor != nullptr && distances_tensor == nullptr) { + RAFT_FAIL("distances tensor must be provided when core_distances tensor is provided"); + } + + int64_t n_rows = dlt.shape[0]; + int64_t n_cols = dlt.shape[1]; + + auto cpp_params = convert_params(params, n_rows, n_cols); + + using dataset_mdspan_t = raft::host_matrix_view; + using indices_mdspan_t = raft::device_matrix_view; + using distances_mdspan_t = raft::device_matrix_view; + using core_mdspan_t = raft::device_vector_view; + + auto dataset = cuvs::core::from_dlpack(dataset_tensor); + auto indices = cuvs::core::from_dlpack(indices_tensor); + + std::optional distances = std::nullopt; + if (distances_tensor) { + distances = cuvs::core::from_dlpack(distances_tensor); + } + + std::optional core_distances = std::nullopt; + if (core_distances_tensor) { + core_distances = cuvs::core::from_dlpack(core_distances_tensor); + } + + cuvs::neighbors::all_neighbors::build( + cpp_res, cpp_params, dataset, indices, distances, core_distances, alpha); +} + +template +void _build_device(cuvsResources_t device_res, + cuvsAllNeighborsIndexParams_t params, + DLManagedTensor* dataset_tensor, + DLManagedTensor* indices_tensor, + DLManagedTensor* distances_tensor, + DLManagedTensor* core_distances_tensor, + float alpha) +{ + auto& cpp_res = *reinterpret_cast(device_res); + + auto& dlt = dataset_tensor->dl_tensor; + RAFT_EXPECTS(cuvs::core::is_dlpack_device_compatible(dlt), + "Device build expects device-compatible dataset tensor"); + + ensure_indices_dtype_and_device_compatibility(indices_tensor); + ensure_optional_distance_dtype_and_device_compatibility(distances_tensor); + ensure_optional_core_distance_dtype_and_device_compatibility(core_distances_tensor); + + // Check dependencies between parameters + if (core_distances_tensor != nullptr && distances_tensor == nullptr) { + RAFT_FAIL("distances tensor must be provided when core_distances tensor is provided"); + } + + int64_t n_rows = dlt.shape[0]; + int64_t n_cols = dlt.shape[1]; + + auto cpp_params = convert_params(params, n_rows, n_cols); + + using dataset_mdspan_t = raft::device_matrix_view; + using indices_mdspan_t = raft::device_matrix_view; + using distances_mdspan_t = raft::device_matrix_view; + using core_mdspan_t = raft::device_vector_view; + + auto dataset = cuvs::core::from_dlpack(dataset_tensor); + auto indices = cuvs::core::from_dlpack(indices_tensor); + + std::optional distances = std::nullopt; + if (distances_tensor) { + distances = cuvs::core::from_dlpack(distances_tensor); + } + + std::optional core_distances = std::nullopt; + if (core_distances_tensor) { + core_distances = cuvs::core::from_dlpack(core_distances_tensor); + } + + cuvs::neighbors::all_neighbors::build( + cpp_res, cpp_params, dataset, indices, distances, core_distances, alpha); +} + +} // namespace + +extern "C" cuvsError_t cuvsAllNeighborsIndexParamsCreate(cuvsAllNeighborsIndexParams_t* params) +{ + return cuvs::core::translate_exceptions([=] { + *params = new cuvsAllNeighborsIndexParams{}; + (*params)->algo = CUVS_ALL_NEIGHBORS_ALGO_BRUTE_FORCE; + (*params)->overlap_factor = 1; + (*params)->n_clusters = 1; + (*params)->metric = L2Expanded; + (*params)->ivf_pq_params = nullptr; + (*params)->nn_descent_params = nullptr; + }); +} + +extern "C" cuvsError_t cuvsAllNeighborsIndexParamsDestroy(cuvsAllNeighborsIndexParams_t params) +{ + return cuvs::core::translate_exceptions([=] { + if (params != nullptr) { + if (params->ivf_pq_params != nullptr) { cuvsIvfPqIndexParamsDestroy(params->ivf_pq_params); } + if (params->nn_descent_params != nullptr) { + cuvsNNDescentIndexParamsDestroy(params->nn_descent_params); + } + delete params; + } + }); +} + +extern "C" cuvsError_t cuvsAllNeighborsBuild(cuvsResources_t res, + cuvsAllNeighborsIndexParams_t params, + DLManagedTensor* dataset_tensor, + DLManagedTensor* indices_tensor, + DLManagedTensor* distances_tensor, + DLManagedTensor* core_distances_tensor, + float alpha) +{ + return cuvs::core::translate_exceptions([=] { + auto dataset = dataset_tensor->dl_tensor; + + if (dataset.dtype.code == kDLFloat && dataset.dtype.bits == 32) { + // Check if dataset is host-compatible or device-compatible + if (cuvs::core::is_dlpack_host_compatible(dataset)) { + _build_host(res, + params, + dataset_tensor, + indices_tensor, + distances_tensor, + core_distances_tensor, + alpha); + } else if (cuvs::core::is_dlpack_device_compatible(dataset)) { + _build_device(res, + params, + dataset_tensor, + indices_tensor, + distances_tensor, + core_distances_tensor, + alpha); + } else { + RAFT_FAIL("Dataset tensor must be either host-compatible or device-compatible"); + } + } else { + RAFT_FAIL("Unsupported dataset DLtensor dtype: %d and bits: %d", + dataset.dtype.code, + dataset.dtype.bits); + } + }); +} diff --git a/cpp/tests/CMakeLists.txt b/cpp/tests/CMakeLists.txt index c6bb7ce8e4..33f9a12193 100644 --- a/cpp/tests/CMakeLists.txt +++ b/cpp/tests/CMakeLists.txt @@ -317,6 +317,11 @@ if(TARGET cuvs::c_api) ConfigureTest(NAME MG_C_TEST PATH neighbors/run_mg_c.c neighbors/ann_mg_c.cu C_LIB) + ConfigureTest( + NAME ALL_NEIGHBORS_C_TEST PATH neighbors/run_all_neighbors_c.c neighbors/all_neighbors_c.cu + C_LIB + ) + if(BUILD_CAGRA_HNSWLIB) ConfigureTest(NAME HNSW_C_TEST PATH neighbors/ann_hnsw_c.cu C_LIB) target_link_libraries(NEIGHBORS_HNSW_TEST PRIVATE hnswlib::hnswlib) diff --git a/cpp/tests/neighbors/all_neighbors_c.cu b/cpp/tests/neighbors/all_neighbors_c.cu new file mode 100644 index 0000000000..a4c79cc6f2 --- /dev/null +++ b/cpp/tests/neighbors/all_neighbors_c.cu @@ -0,0 +1,301 @@ +/* + * Copyright (c) 2025, NVIDIA CORPORATION. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include +#include +#include +#include +#include + +#include "ann_utils.cuh" +#include "naive_knn.cuh" +#include + +#include +#include +#include +#include +#include +#include + +extern "C" void run_all_neighbors(int64_t n_rows, + int64_t n_dim, + int64_t k, + float* dataset_data, + int64_t* indices_data, + float* distances_data, + float* core_distances_data, + cuvsDistanceType metric, + cuvsAllNeighborsAlgo algo, + float alpha); + +template +void generate_random_data(T* devPtr, size_t size) +{ + raft::handle_t handle; + raft::random::RngState r(1234ULL); + // Generate more diverse data with larger range to ensure different neighbors + raft::random::uniform(handle, r, devPtr, size, T(-10.0), T(10.0)); +} + +struct AllNeighborsCInputs { + int64_t n_rows; + int64_t n_dim; + int64_t k; + cuvsDistanceType metric; + cuvsAllNeighborsAlgo algo; + float alpha; + bool include_distances; + bool include_core_distances; +}; + +class AllNeighborsCTest : public ::testing::TestWithParam { + public: + AllNeighborsCTest() + : stream_(0), + data_(0, stream_), + indices_(0, stream_), + distances_(0, stream_), + core_distances_(0, stream_) + { + } + + protected: + void SetUp() override + { + params_ = ::testing::TestWithParam::GetParam(); + RAFT_CUDA_TRY(cudaStreamCreate(&stream_)); + } + + void TearDown() override { RAFT_CUDA_TRY(cudaStreamDestroy(stream_)); } + + void run() + { + size_t data_size = params_.n_rows * params_.n_dim; + size_t indices_size = params_.n_rows * params_.k; + + data_.resize(data_size, stream_); + indices_.resize(indices_size, stream_); + + if (params_.include_distances) { distances_.resize(indices_size, stream_); } + + if (params_.include_core_distances) { core_distances_.resize(params_.n_rows, stream_); } + + generate_random_data(data_.data(), data_size); + + float* distances_ptr = params_.include_distances ? distances_.data() : nullptr; + float* core_distances_ptr = params_.include_core_distances ? core_distances_.data() : nullptr; + + // Test the C API + run_all_neighbors(params_.n_rows, + params_.n_dim, + params_.k, + data_.data(), + indices_.data(), + distances_ptr, + core_distances_ptr, + params_.metric, + params_.algo, + params_.alpha); + + RAFT_CUDA_TRY(cudaStreamSynchronize(stream_)); + + // Basic sanity checks + verify_results(); + } + + private: + template + void recall_eval(T* query_data, + T* index_data, + IdxT* neighbors, + T* distances, + size_t n_queries, + size_t n_rows, + size_t n_dim, + size_t n_neighbors, + cuvsDistanceType metric) + { + raft::handle_t handle; + auto distances_ref = raft::make_device_matrix(handle, n_queries, n_neighbors); + auto neighbors_ref = raft::make_device_matrix(handle, n_queries, n_neighbors); + cuvs::neighbors::naive_knn( + handle, + distances_ref.data_handle(), + neighbors_ref.data_handle(), + query_data, + index_data, + n_queries, + n_rows, + n_dim, + n_neighbors, + static_cast((uint16_t)metric)); + + size_t size = n_queries * n_neighbors; + std::vector neighbors_h(size); + std::vector distances_h(size); + std::vector neighbors_ref_h(size); + std::vector distances_ref_h(size); + + auto stream = raft::resource::get_cuda_stream(handle); + raft::copy(neighbors_h.data(), neighbors, size, stream); + raft::copy(distances_h.data(), distances, size, stream); + raft::copy(neighbors_ref_h.data(), neighbors_ref.data_handle(), size, stream); + raft::copy(distances_ref_h.data(), distances_ref.data_handle(), size, stream); + + // verify output with algorithm-specific minimum recall + double min_recall = 0.8; // Default for approximate algorithms + + if (params_.algo == CUVS_ALL_NEIGHBORS_ALGO_BRUTE_FORCE) { + min_recall = 0.95; // Brute force should be nearly perfect + } else if (params_.algo == CUVS_ALL_NEIGHBORS_ALGO_NN_DESCENT) { + min_recall = 0.7; // NN Descent is approximate + } + + // Lower recall threshold when using mutual reachability distance (core_distances) + if (params_.include_core_distances) { + min_recall = std::min(min_recall, 0.8); // Mutual reachability is more approximate + } + + ASSERT_TRUE(cuvs::neighbors::eval_neighbours(neighbors_ref_h, + neighbors_h, + distances_ref_h, + distances_h, + n_queries, + n_neighbors, + 0.001, + min_recall)); + } + + void verify_results() + { + // Copy results to host for verification + std::vector indices_host(params_.n_rows * params_.k); + RAFT_CUDA_TRY(cudaMemcpy(indices_host.data(), + indices_.data(), + indices_host.size() * sizeof(int64_t), + cudaMemcpyDeviceToHost)); + + // Verify that all indices are valid (within bounds) or indicate missing neighbors + for (size_t i = 0; i < indices_host.size(); ++i) { + // Accept either valid indices or max value (indicating missing neighbors) + EXPECT_TRUE(indices_host[i] >= 0 && indices_host[i] < params_.n_rows || + indices_host[i] == std::numeric_limits::max()) + << "Index " << i << " has value " << indices_host[i] << " which is out of bounds"; + } + + // Perform recall evaluation using the same approach as IVF-PQ tests + if (params_.include_distances) { + // We have distances, so we can do full recall evaluation + std::vector distances_host(params_.n_rows * params_.k); + RAFT_CUDA_TRY(cudaMemcpy(distances_host.data(), + distances_.data(), + distances_host.size() * sizeof(float), + cudaMemcpyDeviceToHost)); + + recall_eval(data_.data(), // query_data (same as index for all_neighbors) + data_.data(), // index_data + indices_.data(), // neighbors + distances_.data(), // distances + params_.n_rows, // n_queries + params_.n_rows, // n_rows + params_.n_dim, // n_dim + params_.k, // n_neighbors + params_.metric); // metric + + // Verify distances are non-negative for valid neighbors + for (size_t i = 0; i < distances_host.size(); ++i) { + int64_t corresponding_index = indices_host[i]; + if (corresponding_index != std::numeric_limits::max()) { + // Allow tiny negative distances due to floating-point precision errors + EXPECT_GE(distances_host[i], -1e-6f) + << "Distance at index " << i << " is significantly negative: " << distances_host[i]; + } + } + } else { + // No distances available, do basic neighbor validation only + for (int64_t row = 0; row < params_.n_rows; ++row) { + std::unordered_set unique_neighbors; + int64_t valid_neighbors = 0; + for (int64_t neighbor_idx = 0; neighbor_idx < params_.k; ++neighbor_idx) { + int64_t neighbor = indices_host[row * params_.k + neighbor_idx]; + if (neighbor != std::numeric_limits::max() && neighbor >= 0 && + neighbor < params_.n_rows) { + unique_neighbors.insert(neighbor); + valid_neighbors++; + } + } + + // For all_neighbors, we should have at least some valid neighbors + EXPECT_GT(valid_neighbors, 0) << "Row " << row << " has no valid neighbors"; + + // The number of unique neighbors should be reasonable + EXPECT_GE(unique_neighbors.size(), 1) << "Row " << row << " has no unique neighbors"; + + // We should have at most min(k, n_rows) neighbors + EXPECT_LE(unique_neighbors.size(), static_cast(std::min(params_.k, params_.n_rows))) + << "Row " << row << " has more unique neighbors than expected"; + } + } + + // If core distances are computed, verify they are non-negative + if (params_.include_core_distances) { + std::vector core_distances_host(params_.n_rows); + RAFT_CUDA_TRY(cudaMemcpy(core_distances_host.data(), + core_distances_.data(), + core_distances_host.size() * sizeof(float), + cudaMemcpyDeviceToHost)); + + for (size_t i = 0; i < core_distances_host.size(); ++i) { + EXPECT_GE(core_distances_host[i], 0.0f) << "Core distance at index " << i << " is negative"; + } + } + } + + cudaStream_t stream_; + AllNeighborsCInputs params_; + rmm::device_uvector data_; + rmm::device_uvector indices_; + rmm::device_uvector distances_; + rmm::device_uvector core_distances_; +}; + +const std::vector inputs = { + // Basic brute force tests + {100, 16, 10, L2Expanded, CUVS_ALL_NEIGHBORS_ALGO_BRUTE_FORCE, 1.0f, true, false}, + {100, 16, 10, L2Expanded, CUVS_ALL_NEIGHBORS_ALGO_BRUTE_FORCE, 1.0f, false, false}, + {100, 16, 10, L2Expanded, CUVS_ALL_NEIGHBORS_ALGO_BRUTE_FORCE, 1.0f, true, true}, + + // Different metrics + {50, 8, 5, InnerProduct, CUVS_ALL_NEIGHBORS_ALGO_BRUTE_FORCE, 1.0f, true, false}, + {50, 8, 5, CosineExpanded, CUVS_ALL_NEIGHBORS_ALGO_BRUTE_FORCE, 1.0f, true, false}, + + // NN Descent tests + {200, 32, 15, L2Expanded, CUVS_ALL_NEIGHBORS_ALGO_NN_DESCENT, 1.0f, true, false}, + {200, 32, 15, L2Expanded, CUVS_ALL_NEIGHBORS_ALGO_NN_DESCENT, 1.0f, false, false}, + + // Different k values + {80, 12, 20, L2Expanded, CUVS_ALL_NEIGHBORS_ALGO_BRUTE_FORCE, 1.0f, true, false}, + {120, 24, 30, L2Expanded, CUVS_ALL_NEIGHBORS_ALGO_BRUTE_FORCE, 1.0f, true, false}, + + // Different alpha values for mutual reachability + {60, 10, 8, L2Expanded, CUVS_ALL_NEIGHBORS_ALGO_BRUTE_FORCE, 0.5f, true, true}, + {60, 10, 8, L2Expanded, CUVS_ALL_NEIGHBORS_ALGO_BRUTE_FORCE, 2.0f, true, true}, +}; + +TEST_P(AllNeighborsCTest, AllNeighborsC) { run(); } + +INSTANTIATE_TEST_CASE_P(AllNeighborsCTest, AllNeighborsCTest, ::testing::ValuesIn(inputs)); diff --git a/cpp/tests/neighbors/c_api.c b/cpp/tests/neighbors/c_api.c index cc71847eaf..75d232cd7f 100644 --- a/cpp/tests/neighbors/c_api.c +++ b/cpp/tests/neighbors/c_api.c @@ -15,6 +15,7 @@ */ #include +#include #include #include @@ -57,12 +58,27 @@ void test_compile_tiered_index() cuvsTieredIndexExtend(resources, &dataset, tiered_index); } +void test_compile_all_neighbors() +{ + // Smoke test to ensure that the all_neighbors.h API compiles correctly + // using a c compiler. Not a full test. + assert(!"test_compile_all_neighbors is not meant to be run"); + + cuvsAllNeighborsIndexParams_t params; + cuvsResources_t resources; + DLManagedTensor dataset, indices, distances, core_distances; + cuvsAllNeighborsIndexParamsCreate(¶ms); + cuvsAllNeighborsIndexParamsDestroy(params); + cuvsAllNeighborsBuild(resources, params, &dataset, &indices, &distances, &core_distances, 1.0f); +} + int main() { // These are smoke tests that check that the C-APIs compile with a C compiler. // These are not meant to be run. test_compile_cagra(); test_compile_tiered_index(); + test_compile_all_neighbors(); return 0; } diff --git a/cpp/tests/neighbors/run_all_neighbors_c.c b/cpp/tests/neighbors/run_all_neighbors_c.c new file mode 100644 index 0000000000..796fbb74cb --- /dev/null +++ b/cpp/tests/neighbors/run_all_neighbors_c.c @@ -0,0 +1,124 @@ +/* + * Copyright (c) 2025, NVIDIA CORPORATION. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include +#include + +void run_all_neighbors(int64_t n_rows, + int64_t n_dim, + int64_t k, + float* dataset_data, + int64_t* indices_data, + float* distances_data, + float* core_distances_data, + cuvsDistanceType metric, + cuvsAllNeighborsAlgo algo, + float alpha) +{ + // create cuvsResources_t + cuvsResources_t res; + cuvsResourcesCreate(&res); + + // create all neighbors index parameters + cuvsAllNeighborsIndexParams_t params; + cuvsAllNeighborsIndexParamsCreate(¶ms); + + // configure parameters + params->algo = algo; + params->metric = metric; + params->overlap_factor = 1; + params->n_clusters = 1; + + // create dataset DLTensor + DLManagedTensor dataset_tensor; + dataset_tensor.dl_tensor.data = dataset_data; + dataset_tensor.dl_tensor.device.device_type = kDLCUDA; + dataset_tensor.dl_tensor.device.device_id = 0; + dataset_tensor.dl_tensor.ndim = 2; + dataset_tensor.dl_tensor.dtype.code = kDLFloat; + dataset_tensor.dl_tensor.dtype.bits = 32; + dataset_tensor.dl_tensor.dtype.lanes = 1; + int64_t dataset_shape[2] = {n_rows, n_dim}; + dataset_tensor.dl_tensor.shape = dataset_shape; + dataset_tensor.dl_tensor.strides = NULL; + dataset_tensor.dl_tensor.byte_offset = 0; + dataset_tensor.manager_ctx = NULL; + dataset_tensor.deleter = NULL; + + // create indices DLTensor + DLManagedTensor indices_tensor; + indices_tensor.dl_tensor.data = (void*)indices_data; + indices_tensor.dl_tensor.device.device_type = kDLCUDA; + indices_tensor.dl_tensor.device.device_id = 0; + indices_tensor.dl_tensor.ndim = 2; + indices_tensor.dl_tensor.dtype.code = kDLInt; + indices_tensor.dl_tensor.dtype.bits = 64; + indices_tensor.dl_tensor.dtype.lanes = 1; + int64_t indices_shape[2] = {n_rows, k}; + indices_tensor.dl_tensor.shape = indices_shape; + indices_tensor.dl_tensor.strides = NULL; + indices_tensor.dl_tensor.byte_offset = 0; + indices_tensor.manager_ctx = NULL; + indices_tensor.deleter = NULL; + + // create distances DLTensor (optional) + DLManagedTensor* distances_ptr = NULL; + DLManagedTensor distances_tensor; + int64_t distances_shape[2] = {n_rows, k}; // Moved outside if block + if (distances_data != NULL) { + distances_tensor.dl_tensor.data = (void*)distances_data; + distances_tensor.dl_tensor.device.device_type = kDLCUDA; + distances_tensor.dl_tensor.device.device_id = 0; + distances_tensor.dl_tensor.ndim = 2; + distances_tensor.dl_tensor.dtype.code = kDLFloat; + distances_tensor.dl_tensor.dtype.bits = 32; + distances_tensor.dl_tensor.dtype.lanes = 1; + distances_tensor.dl_tensor.shape = distances_shape; + distances_tensor.dl_tensor.strides = NULL; + distances_tensor.dl_tensor.byte_offset = 0; + distances_tensor.manager_ctx = NULL; + distances_tensor.deleter = NULL; + distances_ptr = &distances_tensor; + } + + // create core_distances DLTensor (optional) + DLManagedTensor* core_distances_ptr = NULL; + DLManagedTensor core_distances_tensor; + int64_t core_distances_shape[1] = {n_rows}; // Moved outside if block + if (core_distances_data != NULL) { + core_distances_tensor.dl_tensor.data = (void*)core_distances_data; + core_distances_tensor.dl_tensor.device.device_type = kDLCUDA; + core_distances_tensor.dl_tensor.device.device_id = 0; + core_distances_tensor.dl_tensor.ndim = 1; + core_distances_tensor.dl_tensor.dtype.code = kDLFloat; + core_distances_tensor.dl_tensor.dtype.bits = 32; + core_distances_tensor.dl_tensor.dtype.lanes = 1; + core_distances_tensor.dl_tensor.shape = core_distances_shape; + core_distances_tensor.dl_tensor.strides = NULL; + core_distances_tensor.dl_tensor.byte_offset = 0; + core_distances_tensor.manager_ctx = NULL; + core_distances_tensor.deleter = NULL; + core_distances_ptr = &core_distances_tensor; + } + + // build all neighbors graph + cuvsError_t result = cuvsAllNeighborsBuild( + res, params, &dataset_tensor, &indices_tensor, distances_ptr, core_distances_ptr, alpha); + + // cleanup + cuvsAllNeighborsIndexParamsDestroy(params); + cuvsResourcesDestroy(res); +} diff --git a/docs/source/c_api/neighbors.rst b/docs/source/c_api/neighbors.rst index 9de9c40a96..bb6dbec5c2 100644 --- a/docs/source/c_api/neighbors.rst +++ b/docs/source/c_api/neighbors.rst @@ -15,3 +15,4 @@ Nearest Neighbors neighbors_cagra_c.rst neighbors_hnsw_c.rst neighbors_mg.rst + neighbors_all_neighbors_c.rst diff --git a/docs/source/c_api/neighbors_all_neighbors_c.rst b/docs/source/c_api/neighbors_all_neighbors_c.rst new file mode 100644 index 0000000000..f8936384a6 --- /dev/null +++ b/docs/source/c_api/neighbors_all_neighbors_c.rst @@ -0,0 +1,18 @@ +All-Neighbors +============= + +The all-neighbors method constructs a k-NN graph for all vectors in a dataset. It supports multiple algorithms including brute force, IVF-PQ (approximate), and NN-Descent (approximate) for building local k-NN subgraphs. The API automatically detects whether the dataset is host-resident or device-resident and applies appropriate optimizations. + +.. role:: py(code) + :language: c + :class: highlight + +``#include `` + +Build +----- + +.. doxygengroup:: all_neighbors_c_build + :project: cuvs + :members: + :content-only: diff --git a/docs/source/cuvs_bench/param_tuning.rst b/docs/source/cuvs_bench/param_tuning.rst index 276fb35a47..a0de6daf09 100644 --- a/docs/source/cuvs_bench/param_tuning.rst +++ b/docs/source/cuvs_bench/param_tuning.rst @@ -125,7 +125,7 @@ IVF-pq is an inverted-file index, which partitions the vectors into a series of - N - [`cluster`, `subspace`] - `subspace` - - Type of codebook. See :doc:`IVF-PQ index overview <../indexes/ivfpq>` for more detail + - Type of codebook. See :doc:`IVF-PQ index overview <../neighbors/ivfpq>` for more detail * - `dataset_memory_type` - `build` @@ -314,7 +314,7 @@ To fine tune CAGRA index building we can customize IVF-PQ index builder options - N - [`cluster`, `subspace`] - `subspace` - - Type of codebook. See :doc:`IVF-PQ index overview <../indexes/ivfpq>` for more detail + - Type of codebook. See :doc:`IVF-PQ index overview <../neighbors/ivfpq>` for more detail * - `ivf_pq_build_nprobe` - `search` diff --git a/docs/source/getting_started.rst b/docs/source/getting_started.rst index f6cd919adc..89be0f1e5b 100644 --- a/docs/source/getting_started.rst +++ b/docs/source/getting_started.rst @@ -14,7 +14,7 @@ Getting Started - `Supported indexes`_ - * :doc:`Vector search index guide ` + * :doc:`Vector search index guide ` - `Using cuVS APIs`_ @@ -58,7 +58,7 @@ Please see the :doc:`primer on comparing vector search index performance ` for to learn more about each individual index type, when they can be useful on the GPU, the tuning knobs they offer to trade off performance and quality. +cuVS supports many of the standard index types with the list continuing to grow and stay current with the state-of-the-art. Please refer to our :doc:`vector search index guide ` for to learn more about each individual index type, when they can be useful on the GPU, the tuning knobs they offer to trade off performance and quality. The primary goal of cuVS is to enable speed, scale, and flexibility (in that order)- and one of the important value propositions is to enhance existing software deployments with extensible GPU capabilities to improve pain points while not interrupting parts of the system that work well today with CPU. @@ -115,7 +115,7 @@ We always welcome patches for new features and bug fixes. Please read our `contr vector_databases_vs_vector_search.rst tuning_guide.rst comparing_indexes.rst - indexes/indexes.rst + neighbors/neighbors.rst api_basics.rst api_interoperability.rst working_with_ann_indexes.rst diff --git a/docs/source/neighbors/all_neighbors.rst b/docs/source/neighbors/all_neighbors.rst new file mode 100644 index 0000000000..c95d07488e --- /dev/null +++ b/docs/source/neighbors/all_neighbors.rst @@ -0,0 +1,71 @@ +All-neighbors +============= + +All-neighbors is a specialized algorithm for building approximate all-neighbors k-NN graphs. Unlike traditional nearest neighbor indexes that are designed for searching, all-neighbors focuses on constructing complete k-NN graphs for entire datasets. + +This algorithm is particularly useful for: +- Graph construction for visualization algorithms (UMAP, t-SNE) +- Building connectivity graphs for clustering algorithms +- Creating similarity graphs for graph-based machine learning +- Batch processing of large datasets across multiple GPUs + +All-neighbors supports multiple underlying algorithms: +- **Brute Force**: Exact nearest neighbors computation +- **IVF-PQ**: Approximate nearest neighbors using inverted file with product quantization +- **NN-Descent**: Approximate nearest neighbors using graph-based descent + +The algorithm partitions the dataset into clusters and distributes the work across multiple GPUs when possible, making it suitable for large-scale graph construction tasks. + +[ :doc:`C API <../c_api/neighbors_all_neighbors_c>` |:doc:`C++ API <../cpp_api/neighbors_all_neighbors>` | :doc:`Python API <../python_api/neighbors_all_neighbors>` ] + +Algorithm Overview +------------------ + +All-neighbors works by: + +1. **Partitioning**: Dividing the dataset into `n_clusters` clusters/batches +2. **Assignment**: Assigning each point to `overlap_factor` clusters (must be < n_clusters) +3. **Local Computation**: Building k-NN graphs within each cluster using the specified algorithm +4. **Aggregation**: Combining results from all clusters to form the complete graph + +This approach enables: +- **Scalability**: Work distribution across multiple GPUs +- **Memory Efficiency**: Processing large datasets that don't fit in single GPU memory +- **Flexibility**: Choice of underlying algorithm based on accuracy vs. speed requirements + +Use Cases +--------- + +**Data Mining and Machine Learning** +- Clustering algorithms (K-means, HDBSCAN) +- Visualization algorithms (UMAP, t-SNE) +- Sampling and ensemble methods + +**Graph Construction** +- Building similarity graphs for graph neural networks +- Creating connectivity matrices for spectral clustering +- Constructing neighborhood graphs for manifold learning + +**Large-scale Processing** +- Processing datasets that exceed single GPU memory +- Batch processing for distributed computing environments +- Building graphs for graph databases and analytics + +Parameters +---------- + +- **algo**: Underlying algorithm (brute_force, ivf_pq, nn_descent) +- **overlap_factor**: Number of clusters each point is assigned to +- **n_clusters**: Total number of clusters for work distribution +- **metric**: Distance metric for graph construction +- **algorithm-specific parameters**: IVF-PQ or NN-Descent specific settings + +Performance Characteristics +--------------------------- + +- **Build Time**: Scales with dataset size and chosen algorithm +- **Memory Usage**: Depends on cluster size and overlap factor +- **Accuracy**: Varies by algorithm (brute_force, others) and also according to the number of clusters (n_clusters). n_clusters>1 will result in an approximation. +- **Scalability**: Linear scaling with number of GPUs when n_clusters > 1 + +The algorithm automatically chooses between single-GPU and multi-GPU execution based on the n_clusters parameter and available resources. diff --git a/docs/source/indexes/bruteforce.rst b/docs/source/neighbors/bruteforce.rst similarity index 100% rename from docs/source/indexes/bruteforce.rst rename to docs/source/neighbors/bruteforce.rst diff --git a/docs/source/indexes/cagra.rst b/docs/source/neighbors/cagra.rst similarity index 100% rename from docs/source/indexes/cagra.rst rename to docs/source/neighbors/cagra.rst diff --git a/docs/source/indexes/ivfflat.rst b/docs/source/neighbors/ivfflat.rst similarity index 100% rename from docs/source/indexes/ivfflat.rst rename to docs/source/neighbors/ivfflat.rst diff --git a/docs/source/indexes/ivfpq.rst b/docs/source/neighbors/ivfpq.rst similarity index 100% rename from docs/source/indexes/ivfpq.rst rename to docs/source/neighbors/ivfpq.rst diff --git a/docs/source/indexes/indexes.rst b/docs/source/neighbors/neighbors.rst similarity index 80% rename from docs/source/indexes/indexes.rst rename to docs/source/neighbors/neighbors.rst index 8746b84acc..f66b73f867 100644 --- a/docs/source/indexes/indexes.rst +++ b/docs/source/neighbors/neighbors.rst @@ -1,5 +1,5 @@ -Nearest Neighbor Indexes -======================== +Nearest Neighbor +================ .. toctree:: :maxdepth: 3 @@ -10,6 +10,7 @@ Nearest Neighbor Indexes ivfflat.rst ivfpq.rst vamana.rst + all_neighbors.rst Indices and tables diff --git a/docs/source/indexes/vamana.rst b/docs/source/neighbors/vamana.rst similarity index 100% rename from docs/source/indexes/vamana.rst rename to docs/source/neighbors/vamana.rst diff --git a/docs/source/python_api/neighbors.rst b/docs/source/python_api/neighbors.rst index 8e2bad9f14..47e4b2044e 100644 --- a/docs/source/python_api/neighbors.rst +++ b/docs/source/python_api/neighbors.rst @@ -15,3 +15,4 @@ Nearest Neighbors neighbors_ivf_flat.rst neighbors_ivf_pq.rst neighbors_nn_decent.rst + neighbors_all_neighbors.rst diff --git a/docs/source/python_api/neighbors_all_neighbors.rst b/docs/source/python_api/neighbors_all_neighbors.rst new file mode 100644 index 0000000000..241848b7f7 --- /dev/null +++ b/docs/source/python_api/neighbors_all_neighbors.rst @@ -0,0 +1,19 @@ +All-neighbors KNN +================= + +.. role:: py(code) + :language: python + :class: highlight + +All-neighbors allows building an approximate all-neighbors knn graph. Given a full dataset, it finds nearest neighbors for all the training vectors in the dataset. + +Parameters +########## + +.. autoclass:: cuvs.neighbors.all_neighbors.AllNeighborsParams + :members: + +Build +##### + +.. autofunction:: cuvs.neighbors.all_neighbors.build diff --git a/docs/source/tuning_guide.rst b/docs/source/tuning_guide.rst index 26cb2a7010..0f49802c86 100644 --- a/docs/source/tuning_guide.rst +++ b/docs/source/tuning_guide.rst @@ -12,7 +12,7 @@ For more information on the various different types of vector search indexes, pl Why automated tuning? ===================== -As much as 75% of users have told us they will not be able to tune a vector database beyond one or two simple knobs and we suggest that an ideal “knob” would be to balance training time and search time with search quality. The more time, the higher the quality, and the more needed to find an acceptable search performance. Even the 25% of users that want to tune are still asking for simple tools for doing so. These users also ask for some simple guidelines for setting tuning parameters, like :doc:`this guide `. +As much as 75% of users have told us they will not be able to tune a vector database beyond one or two simple knobs and we suggest that an ideal “knob” would be to balance training time and search time with search quality. The more time, the higher the quality, and the more needed to find an acceptable search performance. Even the 25% of users that want to tune are still asking for simple tools for doing so. These users also ask for some simple guidelines for setting tuning parameters, like :doc:`this guide `. Since vector search indexes are more closely related to machine learning models than traditional databases indexes, one option for easing the parameter tuning burden is to use hyper-parameter optimization tools like `Ray Tune `_ and `Optuna `_. to verify this. diff --git a/python/cuvs/cuvs/common/CMakeLists.txt b/python/cuvs/cuvs/common/CMakeLists.txt index ba3ee12bcb..13306d0133 100644 --- a/python/cuvs/cuvs/common/CMakeLists.txt +++ b/python/cuvs/cuvs/common/CMakeLists.txt @@ -1,5 +1,5 @@ # ============================================================================= -# Copyright (c) 2024-2025, NVIDIA CORPORATION. +# Copyright (c) 2025, NVIDIA CORPORATION. # # Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except # in compliance with the License. You may obtain a copy of the License at @@ -13,7 +13,9 @@ # ============================================================================= # Set the list of Cython files to build -set(cython_sources cydlpack.pyx exceptions.pyx resources.pyx device_tensor_view.pyx) +set(cython_sources cydlpack.pyx exceptions.pyx resources.pyx mg_resources.pyx + device_tensor_view.pyx +) set(linked_libraries cuvs::cuvs cuvs::c_api) # Build all of the Cython targets diff --git a/python/cuvs/cuvs/common/__init__.py b/python/cuvs/cuvs/common/__init__.py index e66418306b..52cff6dd9a 100644 --- a/python/cuvs/cuvs/common/__init__.py +++ b/python/cuvs/cuvs/common/__init__.py @@ -1,4 +1,4 @@ -# Copyright (c) 2024, NVIDIA CORPORATION. +# Copyright (c) 2025, NVIDIA CORPORATION. # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -13,6 +13,12 @@ # limitations under the License. +from .mg_resources import MultiGpuResources, auto_sync_multi_gpu_resources from .resources import Resources, auto_sync_resources -__all__ = ["auto_sync_resources", "Resources"] +__all__ = [ + "auto_sync_resources", + "Resources", + "MultiGpuResources", + "auto_sync_multi_gpu_resources", +] diff --git a/python/cuvs/cuvs/common/c_api.pxd b/python/cuvs/cuvs/common/c_api.pxd index 0f9e93fc7c..1d125ba851 100644 --- a/python/cuvs/cuvs/common/c_api.pxd +++ b/python/cuvs/cuvs/common/c_api.pxd @@ -35,6 +35,12 @@ cdef extern from "cuvs/core/c_api.h": cuvsError_t cuvsStreamSync(cuvsResources_t res) const char * cuvsGetLastErrorText() + cuvsError_t cuvsMultiGpuResourcesCreate(cuvsResources_t* res) + cuvsError_t cuvsMultiGpuResourcesCreateWithDeviceIds( + cuvsResources_t* res, + DLManagedTensor* device_ids) + cuvsError_t cuvsMultiGpuResourcesDestroy(cuvsResources_t res) + cuvsError_t cuvsMatrixCopy(cuvsResources_t res, DLManagedTensor * src, DLManagedTensor * dst) diff --git a/python/cuvs/cuvs/common/mg_resources.pxd b/python/cuvs/cuvs/common/mg_resources.pxd new file mode 100644 index 0000000000..ffd9357a9e --- /dev/null +++ b/python/cuvs/cuvs/common/mg_resources.pxd @@ -0,0 +1,24 @@ +# +# Copyright (c) 2025, NVIDIA CORPORATION. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +# cython: language_level=3 + +from cuda.bindings.cyruntime cimport cudaStream_t + +from cuvs.common.c_api cimport cuvsError_t, cuvsResources_t + + +cdef class MultiGpuResources: + cdef cuvsResources_t c_obj diff --git a/python/cuvs/cuvs/common/mg_resources.pyx b/python/cuvs/cuvs/common/mg_resources.pyx new file mode 100644 index 0000000000..d1dc34b623 --- /dev/null +++ b/python/cuvs/cuvs/common/mg_resources.pyx @@ -0,0 +1,153 @@ +# +# Copyright (c) 2025, NVIDIA CORPORATION. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +# cython: language_level=3 + +from cuda.bindings.cyruntime cimport cudaStream_t + +from cuvs.common.c_api cimport ( + cuvsMultiGpuResourcesCreate, + cuvsMultiGpuResourcesCreateWithDeviceIds, + cuvsMultiGpuResourcesDestroy, + cuvsResources_t, + cuvsStreamSet, + cuvsStreamSync, +) +from cuvs.common.cydlpack cimport DLManagedTensor, dlpack_c + +import numpy as np +from pylibraft.common.cai_wrapper import wrap_array + +from cuvs.common.exceptions import check_cuvs + + +cdef class MultiGpuResources: + """ + Multi-GPU Resources is a lightweight python wrapper around the + corresponding C++ class of multi-GPU resources exposed by RAFT's C++ + interface. This class provides a handle for multi-GPU operations across + all available GPUs. + + Parameters + ---------- + stream : int, optional + A CUDA stream pointer to use for this resource handle. If None, a + default stream will be used. + device_ids : list of int, optional + A list of device IDs to use for multi-GPU operations. If None, all + available GPUs will be used. + + Examples + -------- + + Basic usage: + + >>> from cuvs.common import MultiGpuResources + >>> handle = MultiGpuResources() + >>> + >>> # call multi-GPU algos here + >>> + >>> # final sync of all work launched in the stream of this handle + >>> handle.sync() + + Using a cuPy stream with cuVS Multi-GPU Resources: + + >>> import cupy + >>> from cuvs.common import MultiGpuResources + >>> + >>> cupy_stream = cupy.cuda.Stream() + >>> handle = MultiGpuResources(stream=cupy_stream.ptr) + + Using specific device IDs: + + >>> from cuvs.common import MultiGpuResources + >>> handle = MultiGpuResources(device_ids=[0]) + >>> + >>> # call multi-GPU algos here + >>> + >>> handle.sync() + """ + + def __cinit__(self, stream=None, device_ids=None): + cdef DLManagedTensor* device_ids_dlpack = NULL + + if device_ids is not None: + # Convert device_ids list to DLManagedTensor (keep on host/CPU) + # NumPy arrays are naturally on the host, which is what we need + device_ids_array = np.asarray(device_ids, dtype=np.int32) + ai = wrap_array(device_ids_array) + device_ids_dlpack = dlpack_c(ai) + check_cuvs(cuvsMultiGpuResourcesCreateWithDeviceIds( + &self.c_obj, device_ids_dlpack)) + else: + check_cuvs(cuvsMultiGpuResourcesCreate(&self.c_obj)) + + if stream: + check_cuvs(cuvsStreamSet(self.c_obj, stream)) + + def sync(self): + check_cuvs(cuvsStreamSync(self.c_obj)) + + def get_c_obj(self): + """ + Return the pointer to the underlying c_obj as a size_t + """ + return self.c_obj + + def __dealloc__(self): + check_cuvs(cuvsMultiGpuResourcesDestroy(self.c_obj)) + + +_multi_gpu_resources_param_string = """ + resources : Optional cuVS Multi-GPU Resource handle for reusing CUDA \ +resources. + If Multi-GPU Resources aren't supplied, CUDA resources will be + allocated inside this function and synchronized before the + function exits. If resources are supplied, you will need to + explicitly synchronize yourself by calling `resources.sync()` + before accessing the output. +""".strip() + + +def auto_sync_multi_gpu_resources(f): + """Decorator to automatically call sync on a cuVS Multi-GPU Resources + object when it isn't passed to a function. + + When a resources=None is passed to the wrapped function, this decorator + will automatically create a default multi-GPU resources for the function, + and call sync on that resources when the function exits. + + This will also insert the appropriate docstring for the resources + parameter + """ + import functools + + @functools.wraps(f) + def wrapper(*args, resources=None, **kwargs): + sync_resources = resources is None + resources = (resources if resources is not None + else MultiGpuResources()) + + ret_value = f(*args, resources=resources, **kwargs) + + if sync_resources: + resources.sync() + + return ret_value + + wrapper.__doc__ = wrapper.__doc__.format( + resources_docstring=_multi_gpu_resources_param_string + ) + return wrapper diff --git a/python/cuvs/cuvs/neighbors/CMakeLists.txt b/python/cuvs/cuvs/neighbors/CMakeLists.txt index 1b08a32544..0c9196dc43 100644 --- a/python/cuvs/cuvs/neighbors/CMakeLists.txt +++ b/python/cuvs/cuvs/neighbors/CMakeLists.txt @@ -1,5 +1,5 @@ # ============================================================================= -# Copyright (c) 2024, NVIDIA CORPORATION. +# Copyright (c) 2024-2025, NVIDIA CORPORATION. # # Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except # in compliance with the License. You may obtain a copy of the License at @@ -20,6 +20,7 @@ add_subdirectory(ivf_pq) add_subdirectory(filters) add_subdirectory(nn_descent) add_subdirectory(tiered_index) +add_subdirectory(all_neighbors) # Set the list of Cython files to build set(cython_sources refine.pyx) diff --git a/python/cuvs/cuvs/neighbors/__init__.py b/python/cuvs/cuvs/neighbors/__init__.py index 1ba7f79ec8..8ae8c5678b 100644 --- a/python/cuvs/cuvs/neighbors/__init__.py +++ b/python/cuvs/cuvs/neighbors/__init__.py @@ -1,4 +1,4 @@ -# Copyright (c) 2024, NVIDIA CORPORATION. +# Copyright (c) 2024-2025, NVIDIA CORPORATION. # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -14,6 +14,7 @@ from cuvs.neighbors import ( + all_neighbors, brute_force, cagra, filters, @@ -31,5 +32,6 @@ "ivf_flat", "ivf_pq", "nn_descent", + "all_neighbors", "refine", ] diff --git a/python/cuvs/cuvs/neighbors/all_neighbors/CMakeLists.txt b/python/cuvs/cuvs/neighbors/all_neighbors/CMakeLists.txt new file mode 100644 index 0000000000..e38b4ac879 --- /dev/null +++ b/python/cuvs/cuvs/neighbors/all_neighbors/CMakeLists.txt @@ -0,0 +1,24 @@ +# ============================================================================= +# Copyright (c) 2025, NVIDIA CORPORATION. +# +# Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except +# in compliance with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software distributed under the License +# is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express +# or implied. See the License for the specific language governing permissions and limitations under +# the License. +# ============================================================================= + +# Set the list of Cython files to build +set(cython_sources all_neighbors.pyx) +set(linked_libraries cuvs::cuvs cuvs::c_api) + +# Build all of the Cython targets +rapids_cython_create_modules( + CXX + SOURCE_FILES "${cython_sources}" + LINKED_LIBRARIES "${linked_libraries}" MODULE_PREFIX neighbors_all_neighbors_ +) diff --git a/python/cuvs/cuvs/neighbors/all_neighbors/__init__.pxd b/python/cuvs/cuvs/neighbors/all_neighbors/__init__.pxd new file mode 100644 index 0000000000..e69de29bb2 diff --git a/python/cuvs/cuvs/neighbors/all_neighbors/__init__.py b/python/cuvs/cuvs/neighbors/all_neighbors/__init__.py new file mode 100644 index 0000000000..511bb013e2 --- /dev/null +++ b/python/cuvs/cuvs/neighbors/all_neighbors/__init__.py @@ -0,0 +1,17 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. + +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at + +# http://www.apache.org/licenses/LICENSE-2.0 + +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from .all_neighbors import AllNeighborsParams, build + +__all__ = ["AllNeighborsParams", "build"] diff --git a/python/cuvs/cuvs/neighbors/all_neighbors/all_neighbors.pxd b/python/cuvs/cuvs/neighbors/all_neighbors/all_neighbors.pxd new file mode 100644 index 0000000000..8eef428151 --- /dev/null +++ b/python/cuvs/cuvs/neighbors/all_neighbors/all_neighbors.pxd @@ -0,0 +1,52 @@ +# +# Copyright (c) 2025, NVIDIA CORPORATION. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +# cython: language_level=3 + +from libc.stdint cimport int64_t + +from cuvs.common.c_api cimport cuvsError_t, cuvsResources_t +from cuvs.common.cydlpack cimport DLManagedTensor +from cuvs.distance_type cimport cuvsDistanceType +from cuvs.neighbors.ivf_pq.ivf_pq cimport cuvsIvfPqIndexParams_t +from cuvs.neighbors.nn_descent.nn_descent cimport cuvsNNDescentIndexParams_t + + +cdef extern from "cuvs/neighbors/all_neighbors.h" nogil: + ctypedef enum cuvsAllNeighborsAlgo: + CUVS_ALL_NEIGHBORS_ALGO_BRUTE_FORCE + CUVS_ALL_NEIGHBORS_ALGO_IVF_PQ + CUVS_ALL_NEIGHBORS_ALGO_NN_DESCENT + + ctypedef struct cuvsAllNeighborsIndexParams: + cuvsAllNeighborsAlgo algo + size_t overlap_factor + size_t n_clusters + cuvsDistanceType metric + + cuvsIvfPqIndexParams_t ivf_pq_params + cuvsNNDescentIndexParams_t nn_descent_params + + ctypedef cuvsAllNeighborsIndexParams* cuvsAllNeighborsIndexParams_t + + cuvsError_t cuvsAllNeighborsBuild( + cuvsResources_t res, + cuvsAllNeighborsIndexParams_t params, + DLManagedTensor* dataset, + DLManagedTensor* indices, + DLManagedTensor* distances, + DLManagedTensor* core_distances, + float alpha + ) diff --git a/python/cuvs/cuvs/neighbors/all_neighbors/all_neighbors.pyx b/python/cuvs/cuvs/neighbors/all_neighbors/all_neighbors.pyx new file mode 100644 index 0000000000..a824c1a512 --- /dev/null +++ b/python/cuvs/cuvs/neighbors/all_neighbors/all_neighbors.pyx @@ -0,0 +1,373 @@ +# +# Copyright (c) 2025, NVIDIA CORPORATION. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +# cython: language_level=3 + +import numpy as np + +cimport cuvs.common.cydlpack + +from cuvs.common.mg_resources import MultiGpuResources +from cuvs.common.resources import Resources + +from cython.operator cimport dereference as deref + +from cuvs.common cimport cydlpack +from cuvs.common.c_api cimport cuvsResources_t +from cuvs.distance_type cimport cuvsDistanceType + +from pylibraft.common import auto_convert_output, device_ndarray +from pylibraft.common.cai_wrapper import wrap_array +from pylibraft.common.interruptible import cuda_interruptible + +from cuvs.distance import DISTANCE_TYPES +from cuvs.neighbors.common import _check_input_array + +from libc.stdint cimport int64_t + +from cuvs.common.exceptions import check_cuvs + +from cuvs.neighbors.ivf_pq.ivf_pq cimport cuvsIvfPqIndexParams_t +from cuvs.neighbors.nn_descent.nn_descent cimport cuvsNNDescentIndexParams_t + +from .all_neighbors cimport ( + CUVS_ALL_NEIGHBORS_ALGO_BRUTE_FORCE, + CUVS_ALL_NEIGHBORS_ALGO_IVF_PQ, + CUVS_ALL_NEIGHBORS_ALGO_NN_DESCENT, + cuvsAllNeighborsAlgo, + cuvsAllNeighborsBuild, + cuvsAllNeighborsIndexParams, + cuvsAllNeighborsIndexParams_t, +) + +from cuvs.neighbors.ivf_pq.ivf_pq import IndexParams as IvfPqIndexParams +from cuvs.neighbors.nn_descent.nn_descent import ( + IndexParams as NNDescentIndexParams, +) + + +cdef inline cuvsAllNeighborsAlgo _algo_from_str(object algo): + if isinstance(algo, str): + if algo == "brute_force": + return CUVS_ALL_NEIGHBORS_ALGO_BRUTE_FORCE + elif algo == "ivf_pq": + return CUVS_ALL_NEIGHBORS_ALGO_IVF_PQ + elif algo == "nn_descent": + return CUVS_ALL_NEIGHBORS_ALGO_NN_DESCENT + else: + raise ValueError(f"Invalid algo: {algo}") + elif isinstance(algo, int): + return algo + else: + raise ValueError(f"Invalid algo type: {type(algo)}") + + +cdef class AllNeighborsParams: + """ + Parameters for all-neighbors k-NN graph building. + + Parameters + ---------- + algo : str or cuvsAllNeighborsAlgo + Algorithm to use for local k-NN graph building. + Options: "brute_force", "ivf_pq", "nn_descent" + overlap_factor : int, default=2 + Number of clusters each point is assigned to (must be < n_clusters) + n_clusters : int, default=1 + Number of clusters/batches to partition the dataset into + (> overlap_factor). Use n_clusters>1 to distribute the work + across GPUs. + metric : str or cuvsDistanceType, default="sqeuclidean" + Distance metric to use for graph construction + ivf_pq_params : cuvs.neighbors.ivf_pq.IndexParams, optional + IVF-PQ specific parameters (used when algo="ivf_pq") + nn_descent_params : cuvs.neighbors.nn_descent.IndexParams, optional + NN-Descent specific parameters (used when algo="nn_descent") + """ + + cdef cuvsAllNeighborsIndexParams params + cdef object _ivf_pq_params + cdef object _nn_descent_params + + def __init__(self, *, + algo="nn_descent", + overlap_factor=2, + n_clusters=1, + metric="sqeuclidean", + ivf_pq_params=None, + nn_descent_params=None): + + self.params.algo = _algo_from_str(algo) + self.params.overlap_factor = overlap_factor + self.params.n_clusters = n_clusters + self.params.metric = DISTANCE_TYPES[metric] + + if ivf_pq_params is not None: + if not isinstance(ivf_pq_params, IvfPqIndexParams): + raise TypeError( + "ivf_pq_params must be an instance of " + "cuvs.neighbors.ivf_pq.IndexParams" + ) + + # Check metric consistency + ivf_pq_metric = ivf_pq_params.metric + if ivf_pq_metric != metric: + raise ValueError( + f"Metric conflict: AllNeighborsParams metric '{metric}' " + f"does not match IVF-PQ metric '{ivf_pq_metric}'. Please " + f"ensure both use the same metric." + ) + + if nn_descent_params is not None: + if not isinstance(nn_descent_params, NNDescentIndexParams): + raise TypeError( + "nn_descent_params must be an instance of " + "cuvs.neighbors.nn_descent.IndexParams" + ) + + # Check metric consistency + nn_descent_metric = nn_descent_params.metric + if nn_descent_metric != metric: + raise ValueError( + f"Metric conflict: AllNeighborsParams metric '{metric}' " + f"does not match NN-Descent metric '{nn_descent_metric}'. " + f"Please ensure both use the same metric." + ) + + # Store references to prevent garbage collection + self._ivf_pq_params = ivf_pq_params + self._nn_descent_params = nn_descent_params + + # Set algorithm-specific parameter pointers + if ivf_pq_params is not None: + self.params.ivf_pq_params = ( + ivf_pq_params.get_handle() + ) + else: + self.params.ivf_pq_params = NULL + + if nn_descent_params is not None: + self.params.nn_descent_params = ( + ( + nn_descent_params.get_handle() + ) + ) + else: + self.params.nn_descent_params = NULL + + def get_handle(self): + """Get a pointer to the underlying C object.""" + return &self.params + + @property + def algo(self): + """Algorithm used for local k-NN graph building.""" + if self.params.algo == CUVS_ALL_NEIGHBORS_ALGO_BRUTE_FORCE: + return "brute_force" + elif self.params.algo == CUVS_ALL_NEIGHBORS_ALGO_IVF_PQ: + return "ivf_pq" + elif self.params.algo == CUVS_ALL_NEIGHBORS_ALGO_NN_DESCENT: + return "nn_descent" + else: + return self.params.algo + + @property + def overlap_factor(self): + """Number of clusters each point is assigned to.""" + return self.params.overlap_factor + + @property + def n_clusters(self): + """Number of clusters/batches to partition the dataset into.""" + return self.params.n_clusters + + @property + def metric(self): + """Distance metric used for graph construction.""" + # Reverse lookup in DISTANCE_TYPES + for name, value in DISTANCE_TYPES.items(): + if value == self.params.metric: + return name + return self.params.metric + + +@auto_convert_output +def build(dataset, k, params, *, + indices=None, + distances=None, + core_distances=None, + alpha=1.0, + resources=None): + """ + All-neighbors allows building an approximate all-neighbors knn graph. + Given a full dataset, it finds nearest neighbors for all the training + vectors in the dataset. + + Parameters + ---------- + dataset : array_like + Training dataset to build the k-NN graph for. Can be provided + on host (for multi-GPU build) or device (for single-GPU build). + Host vs device location is automatically detected. + Supported dtype: float32 + k : int + Number of nearest neighbors to find for each point + params : AllNeighborsParams + Parameters object containing all build settings including algorithm + choice and algorithm-specific parameters. + indices : array_like, optional + Optional output buffer for indices [num_rows x k] on device + (int64). If not provided, will be allocated automatically. + distances : array_like, optional + Optional output buffer for distances [num_rows x k] on device + (float32) + core_distances : array_like, optional + Optional output buffer for core distances [num_rows] on device + (float32). Requires distances parameter to be provided. + alpha : float, default=1.0 + Mutual-reachability scaling; used only when core_distances is + provided + resources : Resources or MultiGpuResources, optional + CUDA resources to use for the operation. If not provided, a default + Resources object will be created. Use MultiGpuResources to enable + multi-GPU execution across multiple devices. + + Returns + ------- + indices : array_like + k-NN indices for each point [num_rows x k], always on device. + If indices buffer was provided, returns the same array filled + with results. + distances : array_like or None + k-NN distances if distances buffer was provided, None otherwise + core_distances : array_like or None + Core distances if core_distances buffer was provided, None otherwise + """ + if not isinstance(params, AllNeighborsParams): + raise TypeError("params must be an instance of AllNeighborsParams") + + # Check if data is on device for validation purposes + on_device = hasattr(dataset, "__cuda_array_interface__") + + if on_device and params.n_clusters > 1: + raise ValueError( + "Batched all-neighbors build is not supported with data on " + "device. Put data on host for batch build." + ) + + if not isinstance(resources, (Resources, MultiGpuResources)): + resources = Resources() + + resources.sync() + + dataset_ai = wrap_array(dataset) + _check_input_array(dataset_ai, [np.dtype('float32')]) + n_rows, n_cols = dataset_ai.shape + + # Check dependencies between parameters + if core_distances is not None and distances is None: + raise ValueError( + "distances must be provided when core_distances is provided" + ) + + # Validate user-provided outputs (must be device arrays if provided) + if indices is not None and not hasattr( + indices, "__cuda_array_interface__" + ): + raise ValueError( + "indices must be a device array (CUDA array interface)" + ) + if distances is not None and not hasattr( + distances, "__cuda_array_interface__" + ): + raise ValueError( + "distances must be a device array (CUDA array interface)" + ) + if core_distances is not None and not hasattr( + core_distances, "__cuda_array_interface__" + ): + raise ValueError( + "core_distances must be a device array (CUDA array interface)" + ) + + # Handle indices array (create if not provided) + if indices is None: + indices = device_ndarray.empty((n_rows, k), dtype="int64") + + indices_out = wrap_array(indices) + _check_input_array( + indices_out, [np.dtype("int64")], exp_rows=n_rows, exp_cols=k + ) + + distances_out = None + if distances is not None: + distances_out = wrap_array(distances) + _check_input_array( + distances_out, [np.dtype("float32")], exp_rows=n_rows, exp_cols=k + ) + + core_out = None + if core_distances is not None: + core_out = wrap_array(core_distances) + _check_input_array( + core_out, [np.dtype("float32")], exp_rows=n_rows, exp_cols=None + ) + + cdef cydlpack.DLManagedTensor* indices_dlpack = cydlpack.dlpack_c( + indices_out + ) + + cdef cydlpack.DLManagedTensor* distances_dlpack = NULL + if distances_out is not None: + distances_dlpack = cydlpack.dlpack_c(distances_out) + + cdef cydlpack.DLManagedTensor* core_dlpack = NULL + if core_out is not None: + core_dlpack = cydlpack.dlpack_c(core_out) + + cdef cydlpack.DLManagedTensor* dataset_dlpack = cydlpack.dlpack_c( + dataset_ai + ) + + cdef cuvsAllNeighborsIndexParams_t params_ptr = ( + params.get_handle() + ) + + cdef cuvsResources_t res = resources.get_c_obj() + + with cuda_interruptible(): + # Use unified function that auto-detects host vs device + check_cuvs(cuvsAllNeighborsBuild( + res, + params_ptr, + dataset_dlpack, + indices_dlpack, + distances_dlpack, + core_dlpack, + alpha, + )) + + # Build return tuple based on provided parameters + result = [indices] + if distances is not None: + result.append(distances) + if core_distances is not None: + result.append(core_distances) + + # Return single element if only indices, otherwise return tuple + if len(result) == 1: + return result[0] + else: + return tuple(result) diff --git a/python/cuvs/cuvs/neighbors/nn_descent/nn_descent.pyx b/python/cuvs/cuvs/neighbors/nn_descent/nn_descent.pyx index 5c9d0d8155..f9edf7cdf2 100644 --- a/python/cuvs/cuvs/neighbors/nn_descent/nn_descent.pyx +++ b/python/cuvs/cuvs/neighbors/nn_descent/nn_descent.pyx @@ -131,6 +131,10 @@ cdef class IndexParams: def termination_threshold(self): return self.params.termination_threshold + def get_handle(self): + """Get a pointer to the underlying C object.""" + return self.params + cdef class Index: """ NN-Descent index object. This object stores the trained NN-Descent index, diff --git a/python/cuvs/cuvs/tests/test_all_neighbors.py b/python/cuvs/cuvs/tests/test_all_neighbors.py new file mode 100644 index 0000000000..8fff7b8ed9 --- /dev/null +++ b/python/cuvs/cuvs/tests/test_all_neighbors.py @@ -0,0 +1,219 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + + +import cupy +import numpy as np +import pytest +from pylibraft.common import device_ndarray +from sklearn.datasets import make_blobs + +from cuvs.common import MultiGpuResources, Resources +from cuvs.neighbors import all_neighbors, brute_force, ivf_pq, nn_descent +from cuvs.tests.ann_utils import calc_recall + + +def make_cosine( + n_samples=100, + n_features=2, + x_range=(0, 2 * np.pi), + noise=0.0, + random_state=None, +): + r = np.random.default_rng(random_state) + x = r.uniform(x_range[0], x_range[1], n_samples) + y = np.cos(x) + r.normal(0, noise, n_samples) + X = ( + y.reshape(-1, 1) + if n_features == 1 + else np.column_stack( + (x, y, r.normal(size=(n_samples, max(0, n_features - 2)))) + ) + ) + return X, y + + +@pytest.mark.parametrize("algo", ["nn_descent", "brute_force", "ivf_pq"]) +@pytest.mark.parametrize("cluster", ["single_cluster", "multi_cluster"]) +@pytest.mark.parametrize("metric", ["sqeuclidean", "cosine"]) +def test_all_neighbors_device_build_quality(algo, cluster, metric): + """Test device build with quality validation against brute force ground + truth. + """ + n_rows, n_cols, k = 7151, 64, 16 + + if algo == "ivf_pq" and metric == "cosine": + pytest.skip("Skipping IVF-PQ with cosine distance") + + if cluster == "single_cluster": + overlap_factor = 0 + else: + overlap_factor = 3 + + np.random.seed(42) + + if metric == "cosine": + X, _ = make_cosine( + n_samples=n_rows, n_features=n_cols, random_state=42 + ) + else: + X, _ = make_blobs( + n_samples=n_rows, + n_features=n_cols, + centers=10, + cluster_std=1.0, + center_box=(-10.0, 10.0), + random_state=42, + ) + X = X.astype(np.float32) + X_device = device_ndarray(X) + + ivf_pq_params = None + nn_descent_params = None + if algo == "ivf_pq": + ivf_pq_params = ivf_pq.IndexParams( + metric=metric, + n_lists=8 if cluster == "multi_cluster" else 4, + pq_bits=8, + pq_dim=0, + add_data_on_build=True, + ) + elif algo == "nn_descent": + nn_descent_params = nn_descent.IndexParams( + metric=metric, + graph_degree=k, + intermediate_graph_degree=k * 2, + max_iterations=100, + termination_threshold=0.001, + ) + + params = all_neighbors.AllNeighborsParams( + algo=algo, + overlap_factor=overlap_factor, + n_clusters=1, + metric=metric, + ivf_pq_params=ivf_pq_params, + nn_descent_params=nn_descent_params, + ) + + res = Resources() + indices, distances = all_neighbors.build( + X_device, + k, + params, + distances=cupy.empty((n_rows, k), dtype=cupy.float32), + resources=res, + ) + + bf_index = brute_force.build(X_device, metric=metric) + bf_distances, bf_indices = brute_force.search(bf_index, X_device, k=k) + + indices_host = cupy.asnumpy(indices) + bf_indices_host = cupy.asnumpy(bf_indices) + + assert indices.shape == (n_rows, k) + assert indices.dtype == cupy.int64 + assert distances.shape == (n_rows, k) + assert distances.dtype == cupy.float32 + + recall = calc_recall(indices_host, bf_indices_host) + assert recall > 0.85 + + +@pytest.mark.parametrize("algo", ["nn_descent", "brute_force", "ivf_pq"]) +@pytest.mark.parametrize("cluster", ["single_cluster", "multi_cluster"]) +@pytest.mark.parametrize("snmg", [False, True]) +def test_all_neighbors_host_build_quality(algo, cluster, snmg): + """Test host build with quality validation against brute force ground + truth. + """ + n_rows, n_cols, k = 7151, 64, 16 + + if cluster == "single_cluster": + n_clusters = 1 + overlap_factor = 0 + else: + n_clusters = 8 + overlap_factor = 3 + + np.random.seed(42) + + X_host, _ = make_blobs( + n_samples=n_rows, + n_features=n_cols, + centers=10, + cluster_std=1.0, + center_box=(-10.0, 10.0), + random_state=42, + ) + X_host = X_host.astype(np.float32) + X_device = device_ndarray(X_host) + + ivf_pq_params = None + nn_descent_params = None + + if algo == "ivf_pq": + ivf_pq_params = ivf_pq.IndexParams( + metric="sqeuclidean", + n_lists=8 if cluster == "multi_cluster" else 4, + pq_bits=8, + pq_dim=0, + add_data_on_build=True, + ) + elif algo == "nn_descent": + nn_descent_params = nn_descent.IndexParams( + metric="sqeuclidean", + graph_degree=k, + intermediate_graph_degree=k * 2, + max_iterations=100, + termination_threshold=0.001, + ) + + params = all_neighbors.AllNeighborsParams( + algo=algo, + overlap_factor=overlap_factor, + n_clusters=n_clusters, + metric="sqeuclidean", + ivf_pq_params=ivf_pq_params, + nn_descent_params=nn_descent_params, + ) + + if snmg: + res = MultiGpuResources() + else: + res = Resources() + + indices, distances = all_neighbors.build( + X_host, + k, + params, + distances=cupy.empty((n_rows, k), dtype=cupy.float32), + resources=res, + ) + + bf_index = brute_force.build(X_device, metric="sqeuclidean") + bf_distances, bf_indices = brute_force.search(bf_index, X_device, k=k) + + indices_host = cupy.asnumpy(indices) + bf_indices_host = cupy.asnumpy(bf_indices) + + assert indices.shape == (n_rows, k) + assert indices.dtype == cupy.int64 + assert distances.shape == (n_rows, k) + assert distances.dtype == cupy.float32 + + recall = calc_recall(indices_host, bf_indices_host) + + assert recall > 0.85 From be6024d0eafcdb4f3c0795028a1e3893dc2749f2 Mon Sep 17 00:00:00 2001 From: Jinsol Park Date: Thu, 18 Sep 2025 09:44:38 -0700 Subject: [PATCH 74/93] Improve performance of assigning clusters in batched all_neighbors (#1336) Turns out we were using the wrong resource to create a per-device vector. Authors: - Jinsol Park (https://github.com/jinsolp) Approvers: - Corey J. Nolet (https://github.com/cjnolet) URL: https://github.com/rapidsai/cuvs/pull/1336 --- .../neighbors/all_neighbors/all_neighbors_batched.cuh | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/cpp/src/neighbors/all_neighbors/all_neighbors_batched.cuh b/cpp/src/neighbors/all_neighbors/all_neighbors_batched.cuh index 75eec6d50f..2649babdc2 100644 --- a/cpp/src/neighbors/all_neighbors/all_neighbors_batched.cuh +++ b/cpp/src/neighbors/all_neighbors/all_neighbors_batched.cuh @@ -126,13 +126,15 @@ 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); raft::copy(dataset_batch_d.data_handle(), dataset.data_handle() + row_offset * num_cols, n_rows_of_current_batch * num_cols, - resource::get_cuda_stream(res)); + stream); // n_clusters is usually not large, so okay to do this brute-force cuvs::neighbors::brute_force::search(res, @@ -143,7 +145,7 @@ void single_gpu_assign_clusters( raft::copy(global_nearest_cluster.data_handle() + row_offset * overlap_factor, nearest_clusters_idx_d.data_handle(), n_rows_of_current_batch * overlap_factor, - resource::get_cuda_stream(res)); + stream); } } @@ -185,7 +187,8 @@ void assign_clusters(raft::resources const& res, for (int rank = 0; rank < num_ranks; rank++) { auto dev_res = raft::resource::set_current_device_to_rank(res, rank); - auto centroids_matrix = raft::make_device_matrix(res, params.n_clusters, num_cols); + auto centroids_matrix = + raft::make_device_matrix(dev_res, params.n_clusters, num_cols); raft::copy(centroids_matrix.data_handle(), centroids_h.data_handle(), params.n_clusters * num_cols, From cb2fafda82f1de23a6fbe098d8fbc344610d1157 Mon Sep 17 00:00:00 2001 From: Victor Lafargue Date: Fri, 19 Sep 2025 18:12:53 +0200 Subject: [PATCH 75/93] Update SNMG ANN API testing (#1023) Enabling all tests that are able to be ran through CI and refers to issues otherwise. Authors: - Victor Lafargue (https://github.com/viclafargue) - Corey J. Nolet (https://github.com/cjnolet) Approvers: - Corey J. Nolet (https://github.com/cjnolet) URL: https://github.com/rapidsai/cuvs/pull/1023 --- cpp/tests/neighbors/mg.cuh | 16 +--------------- 1 file changed, 1 insertion(+), 15 deletions(-) diff --git a/cpp/tests/neighbors/mg.cuh b/cpp/tests/neighbors/mg.cuh index 046968ed2e..9818e0c856 100644 --- a/cpp/tests/neighbors/mg.cuh +++ b/cpp/tests/neighbors/mg.cuh @@ -1,5 +1,5 @@ /* - * Copyright (c) 2023, NVIDIA CORPORATION. + * Copyright (c) 2023-2025, NVIDIA CORPORATION. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -686,8 +686,6 @@ const std::vector inputs = { 1024, cuvs::distance::DistanceType::L2Expanded, true}, - - /* {7000, 10000, 8, @@ -699,9 +697,6 @@ const std::vector inputs = { 1024, cuvs::distance::DistanceType::L2Expanded, true}, - */ - - /* {7000, 10000, 8, @@ -768,8 +763,6 @@ const std::vector inputs = { 1024, cuvs::distance::DistanceType::L2Expanded, true}, - */ - {7000, 10000, 8, @@ -792,8 +785,6 @@ const std::vector inputs = { 1024, cuvs::distance::DistanceType::L2Expanded, true}, - - /* {7000, 10000, 8, @@ -805,8 +796,6 @@ const std::vector inputs = { 1024, cuvs::distance::DistanceType::L2Expanded, true}, - */ - {3, 10000, 8, @@ -829,8 +818,6 @@ const std::vector inputs = { 1024, cuvs::distance::DistanceType::L2Expanded, true}, - - /* {3, 10000, 8, @@ -842,6 +829,5 @@ const std::vector inputs = { 1024, cuvs::distance::DistanceType::L2Expanded, true}, - */ }; } // namespace cuvs::neighbors::mg From d02e0fc1aabe4e72cc5b0383b1385a0c992cf4f3 Mon Sep 17 00:00:00 2001 From: Victor Lafargue Date: Sat, 20 Sep 2025 00:56:39 +0200 Subject: [PATCH 76/93] MG Python API (#1307) Authors: - Victor Lafargue (https://github.com/viclafargue) - Corey J. Nolet (https://github.com/cjnolet) Approvers: - Ben Frederickson (https://github.com/benfred) URL: https://github.com/rapidsai/cuvs/pull/1307 --- cpp/include/cuvs/core/c_api.h | 17 - cpp/scripts/gitutils.py | 4 +- cpp/src/neighbors/mg_cagra_c.cpp | 50 +- cpp/src/neighbors/mg_ivf_flat_c.cpp | 50 +- cpp/src/neighbors/mg_ivf_pq_c.cpp | 76 +- docs/source/python_api/neighbors.rst | 14 +- docs/source/python_api/neighbors_mg_cagra.rst | 55 ++ .../python_api/neighbors_mg_ivf_flat.rst | 60 ++ .../source/python_api/neighbors_mg_ivf_pq.rst | 60 ++ .../source/python_api/neighbors_multi_gpu.rst | 116 +++ python/cuvs/cuvs/neighbors/CMakeLists.txt | 4 +- python/cuvs/cuvs/neighbors/__init__.py | 2 + python/cuvs/cuvs/neighbors/cagra/cagra.pxd | 10 + python/cuvs/cuvs/neighbors/cagra/cagra.pyx | 15 +- python/cuvs/cuvs/neighbors/common.py | 44 +- .../cuvs/cuvs/neighbors/ivf_flat/ivf_flat.pxd | 7 + .../cuvs/cuvs/neighbors/ivf_flat/ivf_flat.pyx | 10 +- python/cuvs/cuvs/neighbors/ivf_pq/ivf_pq.pxd | 8 + python/cuvs/cuvs/neighbors/ivf_pq/ivf_pq.pyx | 11 +- python/cuvs/cuvs/neighbors/mg/CMakeLists.txt | 17 + python/cuvs/cuvs/neighbors/mg/__init__.py | 21 + .../cuvs/neighbors/mg/cagra/CMakeLists.txt | 24 + .../cuvs/cuvs/neighbors/mg/cagra/__init__.py | 37 + python/cuvs/cuvs/neighbors/mg/cagra/cagra.pxd | 126 ++++ python/cuvs/cuvs/neighbors/mg/cagra/cagra.pyx | 571 +++++++++++++++ .../cuvs/neighbors/mg/ivf_flat/CMakeLists.txt | 24 + .../cuvs/neighbors/mg/ivf_flat/__init__.py | 39 + .../cuvs/neighbors/mg/ivf_flat/ivf_flat.pxd | 128 ++++ .../cuvs/neighbors/mg/ivf_flat/ivf_flat.pyx | 575 +++++++++++++++ .../cuvs/neighbors/mg/ivf_pq/CMakeLists.txt | 24 + .../cuvs/cuvs/neighbors/mg/ivf_pq/__init__.py | 39 + .../cuvs/cuvs/neighbors/mg/ivf_pq/ivf_pq.pxd | 125 ++++ .../cuvs/cuvs/neighbors/mg/ivf_pq/ivf_pq.pyx | 572 +++++++++++++++ python/cuvs/cuvs/tests/test_mg_cagra.py | 608 ++++++++++++++++ python/cuvs/cuvs/tests/test_mg_ivf_flat.py | 650 +++++++++++++++++ python/cuvs/cuvs/tests/test_mg_ivf_pq.py | 682 ++++++++++++++++++ 36 files changed, 4799 insertions(+), 76 deletions(-) create mode 100644 docs/source/python_api/neighbors_mg_cagra.rst create mode 100644 docs/source/python_api/neighbors_mg_ivf_flat.rst create mode 100644 docs/source/python_api/neighbors_mg_ivf_pq.rst create mode 100644 docs/source/python_api/neighbors_multi_gpu.rst create mode 100644 python/cuvs/cuvs/neighbors/mg/CMakeLists.txt create mode 100644 python/cuvs/cuvs/neighbors/mg/__init__.py create mode 100644 python/cuvs/cuvs/neighbors/mg/cagra/CMakeLists.txt create mode 100644 python/cuvs/cuvs/neighbors/mg/cagra/__init__.py create mode 100644 python/cuvs/cuvs/neighbors/mg/cagra/cagra.pxd create mode 100644 python/cuvs/cuvs/neighbors/mg/cagra/cagra.pyx create mode 100644 python/cuvs/cuvs/neighbors/mg/ivf_flat/CMakeLists.txt create mode 100644 python/cuvs/cuvs/neighbors/mg/ivf_flat/__init__.py create mode 100644 python/cuvs/cuvs/neighbors/mg/ivf_flat/ivf_flat.pxd create mode 100644 python/cuvs/cuvs/neighbors/mg/ivf_flat/ivf_flat.pyx create mode 100644 python/cuvs/cuvs/neighbors/mg/ivf_pq/CMakeLists.txt create mode 100644 python/cuvs/cuvs/neighbors/mg/ivf_pq/__init__.py create mode 100644 python/cuvs/cuvs/neighbors/mg/ivf_pq/ivf_pq.pxd create mode 100644 python/cuvs/cuvs/neighbors/mg/ivf_pq/ivf_pq.pyx create mode 100644 python/cuvs/cuvs/tests/test_mg_cagra.py create mode 100644 python/cuvs/cuvs/tests/test_mg_ivf_flat.py create mode 100644 python/cuvs/cuvs/tests/test_mg_ivf_pq.py diff --git a/cpp/include/cuvs/core/c_api.h b/cpp/include/cuvs/core/c_api.h index b6319fe3b0..b47af2c773 100644 --- a/cpp/include/cuvs/core/c_api.h +++ b/cpp/include/cuvs/core/c_api.h @@ -75,23 +75,6 @@ cuvsError_t cuvsResourcesCreate(cuvsResources_t* res); */ cuvsError_t cuvsResourcesDestroy(cuvsResources_t res); -/** - * @brief Create an Initialized opaque C handle for C++ type `raft::device_resources_snmg` - * for multi-GPU operations - * - * @param[in] res cuvsResources_t opaque C handle - * @return cuvsError_t - */ -cuvsError_t cuvsMultiGpuResourcesCreate(cuvsResources_t* res); - -/** - * @brief Destroy and de-allocate opaque C handle for C++ type `raft::device_resources_snmg` - * - * @param[in] res cuvsResources_t opaque C handle - * @return cuvsError_t - */ -cuvsError_t cuvsMultiGpuResourcesDestroy(cuvsResources_t res); - /** * @brief Set cudaStream_t on cuvsResources_t to queue CUDA kernels on APIs * that accept a cuvsResources_t handle diff --git a/cpp/scripts/gitutils.py b/cpp/scripts/gitutils.py index a7337ed4df..800d7797e8 100644 --- a/cpp/scripts/gitutils.py +++ b/cpp/scripts/gitutils.py @@ -1,4 +1,4 @@ -# Copyright (c) 2020-2023, NVIDIA CORPORATION. +# Copyright (c) 2020-2025, NVIDIA CORPORATION. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -155,7 +155,7 @@ def uncommittedFiles(): ret = [] for f in files.splitlines(): f = f.strip(" ") - f = re.sub("\s+", " ", f) # noqa: W605 + f = re.sub(r"\s+", " ", f) # noqa: W605 tmp = f.split(" ", 1) # only consider staged files or uncommitted files # in other words, ignore untracked files diff --git a/cpp/src/neighbors/mg_cagra_c.cpp b/cpp/src/neighbors/mg_cagra_c.cpp index c6d05605d7..e661297b97 100644 --- a/cpp/src/neighbors/mg_cagra_c.cpp +++ b/cpp/src/neighbors/mg_cagra_c.cpp @@ -267,7 +267,12 @@ extern "C" cuvsError_t cuvsMultiGpuCagraBuild(cuvsResources_t res, cuvsMultiGpuCagraIndex_t index) { return cuvs::core::translate_exceptions([=] { - auto dataset = dataset_tensor->dl_tensor; + auto dataset = dataset_tensor->dl_tensor; + + // Multi-GPU CAGRA requires dataset to be in host memory + RAFT_EXPECTS(cuvs::core::is_dlpack_host_compatible(dataset), + "Multi-GPU CAGRA build requires dataset to have host compatible memory"); + index->dtype.code = dataset.dtype.code; index->dtype.bits = dataset.dtype.bits; @@ -295,7 +300,29 @@ extern "C" cuvsError_t cuvsMultiGpuCagraSearch(cuvsResources_t res, DLManagedTensor* distances_tensor) { return cuvs::core::translate_exceptions([=] { - auto queries = queries_tensor->dl_tensor; + auto queries = queries_tensor->dl_tensor; + auto neighbors = neighbors_tensor->dl_tensor; + auto distances = distances_tensor->dl_tensor; + + // Multi-GPU CAGRA requires all tensors to be in host memory + RAFT_EXPECTS(cuvs::core::is_dlpack_host_compatible(queries), + "Multi-GPU CAGRA search requires queries to have host compatible memory"); + RAFT_EXPECTS(cuvs::core::is_dlpack_host_compatible(neighbors), + "Multi-GPU CAGRA search requires neighbors to have host compatible memory"); + RAFT_EXPECTS(cuvs::core::is_dlpack_host_compatible(distances), + "Multi-GPU CAGRA search requires distances to have host compatible memory"); + + // Validate data types + RAFT_EXPECTS(neighbors.dtype.code == kDLInt && neighbors.dtype.bits == 64, + "neighbors should be of type int64_t"); + RAFT_EXPECTS(distances.dtype.code == kDLFloat && distances.dtype.bits == 32, + "distances should be of type float32"); + + // Check type compatibility between index and queries + RAFT_EXPECTS(queries.dtype.code == index->dtype.code, + "type mismatch between index and queries"); + RAFT_EXPECTS(queries.dtype.bits == index->dtype.bits, + "type mismatch between index and queries"); if (queries.dtype.code == kDLFloat && queries.dtype.bits == 32) { _mg_search(res, *params, *index, queries_tensor, neighbors_tensor, distances_tensor); @@ -321,6 +348,25 @@ extern "C" cuvsError_t cuvsMultiGpuCagraExtend(cuvsResources_t res, return cuvs::core::translate_exceptions([=] { auto vectors = new_vectors_tensor->dl_tensor; + // Multi-GPU CAGRA requires vectors to be in host memory + RAFT_EXPECTS(cuvs::core::is_dlpack_host_compatible(vectors), + "Multi-GPU CAGRA extend requires new_vectors to have host compatible memory"); + + // Check type compatibility between index and vectors + RAFT_EXPECTS(vectors.dtype.code == index->dtype.code, + "type mismatch between index and new_vectors"); + RAFT_EXPECTS(vectors.dtype.bits == index->dtype.bits, + "type mismatch between index and new_vectors"); + + // If indices are provided, they should also be in host memory + if (new_indices_tensor != nullptr) { + auto indices = new_indices_tensor->dl_tensor; + RAFT_EXPECTS(cuvs::core::is_dlpack_host_compatible(indices), + "Multi-GPU CAGRA extend requires new_indices to have host compatible memory"); + RAFT_EXPECTS(indices.dtype.code == kDLUInt && indices.dtype.bits == 32, + "new_indices should be of type uint32_t"); + } + if (vectors.dtype.code == kDLFloat && vectors.dtype.bits == 32) { _mg_extend(res, *index, new_vectors_tensor, new_indices_tensor); } else if (vectors.dtype.code == kDLFloat && vectors.dtype.bits == 16) { diff --git a/cpp/src/neighbors/mg_ivf_flat_c.cpp b/cpp/src/neighbors/mg_ivf_flat_c.cpp index c012cb4c7c..bec2fe8149 100644 --- a/cpp/src/neighbors/mg_ivf_flat_c.cpp +++ b/cpp/src/neighbors/mg_ivf_flat_c.cpp @@ -264,7 +264,12 @@ extern "C" cuvsError_t cuvsMultiGpuIvfFlatBuild(cuvsResources_t res, cuvsMultiGpuIvfFlatIndex_t index) { return cuvs::core::translate_exceptions([=] { - auto dataset = dataset_tensor->dl_tensor; + auto dataset = dataset_tensor->dl_tensor; + + // Multi-GPU IVF-Flat requires dataset to be in host memory + RAFT_EXPECTS(cuvs::core::is_dlpack_host_compatible(dataset), + "Multi-GPU IVF-Flat build requires dataset to have host compatible memory"); + index->dtype.code = dataset.dtype.code; index->dtype.bits = dataset.dtype.bits; @@ -292,7 +297,29 @@ extern "C" cuvsError_t cuvsMultiGpuIvfFlatSearch(cuvsResources_t res, DLManagedTensor* distances_tensor) { return cuvs::core::translate_exceptions([=] { - auto queries = queries_tensor->dl_tensor; + auto queries = queries_tensor->dl_tensor; + auto neighbors = neighbors_tensor->dl_tensor; + auto distances = distances_tensor->dl_tensor; + + // Multi-GPU IVF-Flat requires all tensors to be in host memory + RAFT_EXPECTS(cuvs::core::is_dlpack_host_compatible(queries), + "Multi-GPU IVF-Flat search requires queries to have host compatible memory"); + RAFT_EXPECTS(cuvs::core::is_dlpack_host_compatible(neighbors), + "Multi-GPU IVF-Flat search requires neighbors to have host compatible memory"); + RAFT_EXPECTS(cuvs::core::is_dlpack_host_compatible(distances), + "Multi-GPU IVF-Flat search requires distances to have host compatible memory"); + + // Validate data types + RAFT_EXPECTS(neighbors.dtype.code == kDLInt && neighbors.dtype.bits == 64, + "neighbors should be of type int64_t"); + RAFT_EXPECTS(distances.dtype.code == kDLFloat && distances.dtype.bits == 32, + "distances should be of type float32"); + + // Check type compatibility between index and queries + RAFT_EXPECTS(queries.dtype.code == index->dtype.code, + "type mismatch between index and queries"); + RAFT_EXPECTS(queries.dtype.bits == index->dtype.bits, + "type mismatch between index and queries"); if (queries.dtype.code == kDLFloat && queries.dtype.bits == 32) { _mg_search(res, *params, *index, queries_tensor, neighbors_tensor, distances_tensor); @@ -318,6 +345,25 @@ extern "C" cuvsError_t cuvsMultiGpuIvfFlatExtend(cuvsResources_t res, return cuvs::core::translate_exceptions([=] { auto vectors = new_vectors_tensor->dl_tensor; + // Multi-GPU IVF-Flat requires vectors to be in host memory + RAFT_EXPECTS(cuvs::core::is_dlpack_host_compatible(vectors), + "Multi-GPU IVF-Flat extend requires new_vectors to have host compatible memory"); + + // Check type compatibility between index and vectors + RAFT_EXPECTS(vectors.dtype.code == index->dtype.code, + "type mismatch between index and new_vectors"); + RAFT_EXPECTS(vectors.dtype.bits == index->dtype.bits, + "type mismatch between index and new_vectors"); + + // If indices are provided, they should also be in host memory + if (new_indices_tensor != nullptr) { + auto indices = new_indices_tensor->dl_tensor; + RAFT_EXPECTS(cuvs::core::is_dlpack_host_compatible(indices), + "Multi-GPU IVF-Flat extend requires new_indices to have host compatible memory"); + RAFT_EXPECTS(indices.dtype.code == kDLInt && indices.dtype.bits == 64, + "new_indices should be of type int64_t"); + } + if (vectors.dtype.code == kDLFloat && vectors.dtype.bits == 32) { _mg_extend(res, *index, new_vectors_tensor, new_indices_tensor); } else if (vectors.dtype.code == kDLFloat && vectors.dtype.bits == 16) { diff --git a/cpp/src/neighbors/mg_ivf_pq_c.cpp b/cpp/src/neighbors/mg_ivf_pq_c.cpp index 57d11f5264..0307a659d1 100644 --- a/cpp/src/neighbors/mg_ivf_pq_c.cpp +++ b/cpp/src/neighbors/mg_ivf_pq_c.cpp @@ -256,7 +256,12 @@ extern "C" cuvsError_t cuvsMultiGpuIvfPqBuild(cuvsResources_t res, cuvsMultiGpuIvfPqIndex_t index) { return cuvs::core::translate_exceptions([=] { - auto dataset = dataset_tensor->dl_tensor; + auto dataset = dataset_tensor->dl_tensor; + + // Multi-GPU IVF-PQ requires dataset to be in host memory + RAFT_EXPECTS(cuvs::core::is_dlpack_host_compatible(dataset), + "Multi-GPU IVF-PQ build requires dataset to have host compatible memory"); + index->dtype.code = dataset.dtype.code; index->dtype.bits = dataset.dtype.bits; @@ -284,7 +289,29 @@ extern "C" cuvsError_t cuvsMultiGpuIvfPqSearch(cuvsResources_t res, DLManagedTensor* distances_tensor) { return cuvs::core::translate_exceptions([=] { - auto queries = queries_tensor->dl_tensor; + auto queries = queries_tensor->dl_tensor; + auto neighbors = neighbors_tensor->dl_tensor; + auto distances = distances_tensor->dl_tensor; + + // Multi-GPU IVF-PQ requires all tensors to be in host memory + RAFT_EXPECTS(cuvs::core::is_dlpack_host_compatible(queries), + "Multi-GPU IVF-PQ search requires queries to have host compatible memory"); + RAFT_EXPECTS(cuvs::core::is_dlpack_host_compatible(neighbors), + "Multi-GPU IVF-PQ search requires neighbors to have host compatible memory"); + RAFT_EXPECTS(cuvs::core::is_dlpack_host_compatible(distances), + "Multi-GPU IVF-PQ search requires distances to have host compatible memory"); + + // Validate data types + RAFT_EXPECTS(neighbors.dtype.code == kDLInt && neighbors.dtype.bits == 64, + "neighbors should be of type int64_t"); + RAFT_EXPECTS(distances.dtype.code == kDLFloat && distances.dtype.bits == 32, + "distances should be of type float32"); + + // Check type compatibility between index and queries + RAFT_EXPECTS(queries.dtype.code == index->dtype.code, + "type mismatch between index and queries"); + RAFT_EXPECTS(queries.dtype.bits == index->dtype.bits, + "type mismatch between index and queries"); if (queries.dtype.code == kDLFloat && queries.dtype.bits == 32) { _mg_search(res, *params, *index, queries_tensor, neighbors_tensor, distances_tensor); @@ -310,6 +337,25 @@ extern "C" cuvsError_t cuvsMultiGpuIvfPqExtend(cuvsResources_t res, return cuvs::core::translate_exceptions([=] { auto vectors = new_vectors_tensor->dl_tensor; + // Multi-GPU IVF-PQ requires vectors to be in host memory + RAFT_EXPECTS(cuvs::core::is_dlpack_host_compatible(vectors), + "Multi-GPU IVF-PQ extend requires new_vectors to have host compatible memory"); + + // Check type compatibility between index and vectors + RAFT_EXPECTS(vectors.dtype.code == index->dtype.code, + "type mismatch between index and new_vectors"); + RAFT_EXPECTS(vectors.dtype.bits == index->dtype.bits, + "type mismatch between index and new_vectors"); + + // If indices are provided, they should also be in host memory + if (new_indices_tensor != nullptr) { + auto indices = new_indices_tensor->dl_tensor; + RAFT_EXPECTS(cuvs::core::is_dlpack_host_compatible(indices), + "Multi-GPU IVF-PQ extend requires new_indices to have host compatible memory"); + RAFT_EXPECTS(indices.dtype.code == kDLInt && indices.dtype.bits == 64, + "new_indices should be of type int64_t"); + } + if (vectors.dtype.code == kDLFloat && vectors.dtype.bits == 32) { _mg_extend(res, *index, new_vectors_tensor, new_indices_tensor); } else if (vectors.dtype.code == kDLFloat && vectors.dtype.bits == 16) { @@ -381,28 +427,8 @@ extern "C" cuvsError_t cuvsMultiGpuIvfPqDistribute(cuvsResources_t res, cuvsMultiGpuIvfPqIndex_t index) { return cuvs::core::translate_exceptions([=] { - std::ifstream is(filename, std::ios::in | std::ios::binary); - if (!is) { RAFT_FAIL("Cannot open file %s", filename); } - char dtype_string[4]; - is.read(dtype_string, 4); - auto dtype = raft::detail::numpy_serializer::parse_descr(std::string(dtype_string, 4)); - is.close(); - - index->dtype.bits = dtype.itemsize * 8; - if (dtype.kind == 'f' && dtype.itemsize == 4) { - index->dtype.code = kDLFloat; - index->addr = reinterpret_cast(_mg_distribute(res, filename)); - } else if (dtype.kind == 'f' && dtype.itemsize == 2) { - index->dtype.code = kDLFloat; - index->addr = reinterpret_cast(_mg_distribute(res, filename)); - } else if (dtype.kind == 'i' && dtype.itemsize == 1) { - index->dtype.code = kDLInt; - index->addr = reinterpret_cast(_mg_distribute(res, filename)); - } else if (dtype.kind == 'u' && dtype.itemsize == 1) { - index->dtype.code = kDLUInt; - index->addr = reinterpret_cast(_mg_distribute(res, filename)); - } else { - RAFT_FAIL("Unsupported index dtype"); - } + index->dtype.code = kDLFloat; + index->dtype.bits = 32; + index->addr = reinterpret_cast(_mg_distribute(res, filename)); }); } diff --git a/docs/source/python_api/neighbors.rst b/docs/source/python_api/neighbors.rst index 47e4b2044e..909f2013ad 100644 --- a/docs/source/python_api/neighbors.rst +++ b/docs/source/python_api/neighbors.rst @@ -5,9 +5,12 @@ Nearest Neighbors :language: python :class: highlight +Single-GPU Algorithms +##################### + .. toctree:: :maxdepth: 2 - :caption: Contents: + :caption: Single-GPU ANN Algorithms: neighbors_brute_force.rst neighbors_cagra.rst @@ -15,4 +18,13 @@ Nearest Neighbors neighbors_ivf_flat.rst neighbors_ivf_pq.rst neighbors_nn_decent.rst + +Multi-GPU Algorithms +#################### + +.. toctree:: + :maxdepth: 2 + :caption: Multi-GPU Distributed ANN: + + neighbors_multi_gpu.rst neighbors_all_neighbors.rst diff --git a/docs/source/python_api/neighbors_mg_cagra.rst b/docs/source/python_api/neighbors_mg_cagra.rst new file mode 100644 index 0000000000..2f03fa389b --- /dev/null +++ b/docs/source/python_api/neighbors_mg_cagra.rst @@ -0,0 +1,55 @@ +Multi-GPU CAGRA +=============== + +Multi-GPU CAGRA extends the graph-based CAGRA algorithm to work across multiple GPUs, providing improved scalability and performance for large-scale vector search. It supports both replicated and sharded distribution modes. + +.. role:: py(code) + :language: python + :class: highlight + +.. note:: + **IMPORTANT**: Multi-GPU CAGRA requires all data (datasets, queries, output arrays) to be in host memory (CPU). + If using CuPy/device arrays, transfer to host with ``array.get()`` or ``cp.asnumpy(array)`` before use. + +Index build parameters +###################### + +.. autoclass:: cuvs.neighbors.mg_cagra.IndexParams + :members: + +Index search parameters +####################### + +.. autoclass:: cuvs.neighbors.mg_cagra.SearchParams + :members: + +Index +##### + +.. autoclass:: cuvs.neighbors.mg_cagra.Index + :members: + +Index build +########### + +.. autofunction:: cuvs.neighbors.mg_cagra.build + +Index search +############ + +.. autofunction:: cuvs.neighbors.mg_cagra.search + +Index save +########## + +.. autofunction:: cuvs.neighbors.mg_cagra.save + +Index load +########## + +.. autofunction:: cuvs.neighbors.mg_cagra.load + +Index distribute +################ + +.. autofunction:: cuvs.neighbors.mg_cagra.distribute diff --git a/docs/source/python_api/neighbors_mg_ivf_flat.rst b/docs/source/python_api/neighbors_mg_ivf_flat.rst new file mode 100644 index 0000000000..37515ce546 --- /dev/null +++ b/docs/source/python_api/neighbors_mg_ivf_flat.rst @@ -0,0 +1,60 @@ +Multi-GPU IVF-Flat +================== + +Multi-GPU IVF-Flat extends the IVF-Flat algorithm to work across multiple GPUs, providing improved scalability and performance for large-scale vector search. It supports both replicated and sharded distribution modes. + +.. role:: py(code) + :language: python + :class: highlight + +.. note:: + **IMPORTANT**: Multi-GPU IVF-Flat requires all data (datasets, queries, output arrays) to be in host memory (CPU). + If using CuPy/device arrays, transfer to host with ``array.get()`` or ``cp.asnumpy(array)`` before use. + +Index build parameters +###################### + +.. autoclass:: cuvs.neighbors.mg_ivf_flat.IndexParams + :members: + +Index search parameters +####################### + +.. autoclass:: cuvs.neighbors.mg_ivf_flat.SearchParams + :members: + +Index +##### + +.. autoclass:: cuvs.neighbors.mg_ivf_flat.Index + :members: + +Index build +########### + +.. autofunction:: cuvs.neighbors.mg_ivf_flat.build + +Index search +############ + +.. autofunction:: cuvs.neighbors.mg_ivf_flat.search + +Index extend +############ + +.. autofunction:: cuvs.neighbors.mg_ivf_flat.extend + +Index save +########## + +.. autofunction:: cuvs.neighbors.mg_ivf_flat.save + +Index load +########## + +.. autofunction:: cuvs.neighbors.mg_ivf_flat.load + +Index distribute +################ + +.. autofunction:: cuvs.neighbors.mg_ivf_flat.distribute diff --git a/docs/source/python_api/neighbors_mg_ivf_pq.rst b/docs/source/python_api/neighbors_mg_ivf_pq.rst new file mode 100644 index 0000000000..d7d13b4734 --- /dev/null +++ b/docs/source/python_api/neighbors_mg_ivf_pq.rst @@ -0,0 +1,60 @@ +Multi-GPU IVF-PQ +================ + +Multi-GPU IVF-PQ extends the IVF-PQ (Inverted File with Product Quantization) algorithm to work across multiple GPUs, providing improved scalability and performance for large-scale vector search. It supports both replicated and sharded distribution modes. + +.. role:: py(code) + :language: python + :class: highlight + +.. note:: + **IMPORTANT**: Multi-GPU IVF-PQ requires all data (datasets, queries, output arrays) to be in host memory (CPU). + If using CuPy/device arrays, transfer to host with ``array.get()`` or ``cp.asnumpy(array)`` before use. + +Index build parameters +###################### + +.. autoclass:: cuvs.neighbors.mg_ivf_pq.IndexParams + :members: + +Index search parameters +####################### + +.. autoclass:: cuvs.neighbors.mg_ivf_pq.SearchParams + :members: + +Index +##### + +.. autoclass:: cuvs.neighbors.mg_ivf_pq.Index + :members: + +Index build +########### + +.. autofunction:: cuvs.neighbors.mg_ivf_pq.build + +Index search +############ + +.. autofunction:: cuvs.neighbors.mg_ivf_pq.search + +Index extend +############ + +.. autofunction:: cuvs.neighbors.mg_ivf_pq.extend + +Index save +########## + +.. autofunction:: cuvs.neighbors.mg_ivf_pq.save + +Index load +########## + +.. autofunction:: cuvs.neighbors.mg_ivf_pq.load + +Index distribute +################ + +.. autofunction:: cuvs.neighbors.mg_ivf_pq.distribute diff --git a/docs/source/python_api/neighbors_multi_gpu.rst b/docs/source/python_api/neighbors_multi_gpu.rst new file mode 100644 index 0000000000..e8230d7695 --- /dev/null +++ b/docs/source/python_api/neighbors_multi_gpu.rst @@ -0,0 +1,116 @@ +Multi-GPU Nearest Neighbors +=========================== + +Multi-GPU support in cuVS enables scaling ANN (Approximate Nearest Neighbors) algorithms across multiple GPUs on a single node, providing improved performance and the ability to handle larger datasets. + +.. role:: py(code) + :language: python + :class: highlight + +Overview +-------- + +The multi-GPU implementations extend the single-GPU algorithms to work across multiple GPUs using two main distribution strategies: + +- **Replicated Mode**: The entire index is replicated across all GPUs. This mode provides higher query throughput by distributing queries across GPUs while maintaining the full index on each GPU. + +- **Sharded Mode**: The index is partitioned (sharded) across GPUs. This mode allows handling larger datasets that don't fit on a single GPU by distributing the data across multiple GPUs. + +Important Notes +--------------- + +.. warning:: + **Memory Requirements**: Multi-GPU algorithms require all data to be in host memory (CPU). This is different from single-GPU algorithms that typically work with device memory. + +.. note:: + **Supported Algorithms**: Currently, multi-GPU support is available for: + + - CAGRA (Graph-based ANN) + - IVF-Flat (Inverted File with Flat storage) + - IVF-PQ (Inverted File with Product Quantization) + +Configuration Options +--------------------- + +Distribution Modes +^^^^^^^^^^^^^^^^^^ + +- **Replicated Mode** + + In replicated mode, the complete index is stored on each GPU. This approach: + + - Maximizes query throughput by processing queries in parallel across all GPUs + - Requires each GPU to have enough memory to store the entire index + - Is ideal for scenarios where query throughput is more important than index size limitations + +- **Sharded Mode** + + In sharded mode, the index is distributed across GPUs. This approach: + + - Enables handling of larger datasets by partitioning across GPUs + - Requires coordination between GPUs during search operations + - Is ideal for scenarios where the dataset is too large for a single GPU + +Search Modes +^^^^^^^^^^^^ + +- **Load Balancer** + + Divides each query across multiple GPUs, distributing workload efficiently to maximize performance and throughput. + +- **Round Robin** + + Distributes queries evenly across GPUs in a rotating sequence, ensuring balanced workload allocation. This mode is best suited for frequent, small-scale search operations. + +Merge Modes +^^^^^^^^^^^ + +- **Merge on Root Rank** + + Results from all GPUs are collected and merged on the root rank (typically GPU 0). + +- **Tree Merge** + + Results are merged in a tree-like fashion across GPUs to reduce communication overhead. + +Usage Examples +-------------- + +Basic Multi-GPU Usage +^^^^^^^^^^^^^^^^^^^^^^ + +.. code-block:: python + + import numpy as np + from cuvs.neighbors import mg_cagra + + # Create dataset in host memory + n_samples = 100000 + n_features = 128 + dataset = np.random.random_sample((n_samples, n_features), dtype=np.float32) + + # Build multi-GPU index + build_params = mg_cagra.IndexParams( + distribution_mode="sharded", + metric="sqeuclidean" + ) + index = mg_cagra.build(build_params, dataset) + + # Search with multi-GPU + queries = np.random.random_sample((1000, n_features), dtype=np.float32) + search_params = mg_cagra.SearchParams( + search_mode="load_balancer", + merge_mode="merge_on_root_rank" + ) + distances, neighbors = mg_cagra.search(search_params, index, queries, k=10) + +Algorithm-Specific Documentation +-------------------------------- + +.. toctree:: + :maxdepth: 2 + :caption: Multi-GPU Algorithms: + + neighbors_mg_cagra.rst + neighbors_mg_ivf_flat.rst + neighbors_mg_ivf_pq.rst diff --git a/python/cuvs/cuvs/neighbors/CMakeLists.txt b/python/cuvs/cuvs/neighbors/CMakeLists.txt index 0c9196dc43..ee48687d69 100644 --- a/python/cuvs/cuvs/neighbors/CMakeLists.txt +++ b/python/cuvs/cuvs/neighbors/CMakeLists.txt @@ -10,8 +10,7 @@ # is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express # or implied. See the License for the specific language governing permissions and limitations under # the License. -# ============================================================================= - +# add_subdirectory(brute_force) add_subdirectory(cagra) add_subdirectory(hnsw) @@ -20,6 +19,7 @@ add_subdirectory(ivf_pq) add_subdirectory(filters) add_subdirectory(nn_descent) add_subdirectory(tiered_index) +add_subdirectory(mg) add_subdirectory(all_neighbors) # Set the list of Cython files to build diff --git a/python/cuvs/cuvs/neighbors/__init__.py b/python/cuvs/cuvs/neighbors/__init__.py index 8ae8c5678b..b34a3b3598 100644 --- a/python/cuvs/cuvs/neighbors/__init__.py +++ b/python/cuvs/cuvs/neighbors/__init__.py @@ -20,6 +20,7 @@ filters, ivf_flat, ivf_pq, + mg, nn_descent, ) @@ -31,6 +32,7 @@ "filters", "ivf_flat", "ivf_pq", + "mg", "nn_descent", "all_neighbors", "refine", diff --git a/python/cuvs/cuvs/neighbors/cagra/cagra.pxd b/python/cuvs/cuvs/neighbors/cagra/cagra.pxd index b498cf1681..b142a4b33b 100644 --- a/python/cuvs/cuvs/neighbors/cagra/cagra.pxd +++ b/python/cuvs/cuvs/neighbors/cagra/cagra.pxd @@ -197,3 +197,13 @@ cdef class Index: cdef cuvsCagraIndex_t index cdef bool trained cdef str active_index_type + + +cdef class IndexParams: + cdef cuvsCagraIndexParams* params + cdef public object compression + cdef public object ivf_pq_build_params + cdef public object ivf_pq_search_params + +cdef class SearchParams: + cdef cuvsCagraSearchParams * params diff --git a/python/cuvs/cuvs/neighbors/cagra/cagra.pyx b/python/cuvs/cuvs/neighbors/cagra/cagra.pyx index d82ea99ffa..8738e035c3 100644 --- a/python/cuvs/cuvs/neighbors/cagra/cagra.pyx +++ b/python/cuvs/cuvs/neighbors/cagra/cagra.pyx @@ -172,13 +172,6 @@ cdef class IndexParams: """ - cdef cuvsCagraIndexParams* params - - # hold on to a reference to the compression, to keep from being GC'ed - cdef public object compression - cdef public object ivf_pq_build_params - cdef public object ivf_pq_search_params - def __cinit__(self): check_cuvs(cuvsCagraIndexParamsCreate(&self.params)) self.compression = None @@ -186,7 +179,8 @@ cdef class IndexParams: self.ivf_pq_search_params = None def __dealloc__(self): - check_cuvs(cuvsCagraIndexParamsDestroy(self.params)) + if self.params != NULL: + check_cuvs(cuvsCagraIndexParamsDestroy(self.params)) def __init__(self, *, metric="sqeuclidean", @@ -475,13 +469,12 @@ cdef class SearchParams: """ - cdef cuvsCagraSearchParams * params - def __cinit__(self): check_cuvs(cuvsCagraSearchParamsCreate(&self.params)) def __dealloc__(self): - check_cuvs(cuvsCagraSearchParamsDestroy(self.params)) + if self.params != NULL: + check_cuvs(cuvsCagraSearchParamsDestroy(self.params)) def __init__(self, *, max_queries=0, diff --git a/python/cuvs/cuvs/neighbors/common.py b/python/cuvs/cuvs/neighbors/common.py index f49d9eb1f0..4eaaf91d65 100644 --- a/python/cuvs/cuvs/neighbors/common.py +++ b/python/cuvs/cuvs/neighbors/common.py @@ -1,5 +1,5 @@ # -# Copyright (c) 2024, NVIDIA CORPORATION. +# Copyright (c) 2024-2025, NVIDIA CORPORATION. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -13,6 +13,8 @@ # See the License for the specific language governing permissions and # limitations under the License. +import numpy as np + def _check_input_array( cai, exp_dt, exp_rows=None, exp_cols=None, exp_row_major=True @@ -36,3 +38,43 @@ def _check_input_array( exp_rows, cai.shape[0] ) ) + + +def _check_memory_location(array_like, expected_host=True, name="array"): + """ + Check if array is in expected memory location for multi-GPU operations. + + Parameters + ---------- + array_like : array-like + Array to check memory location of + expected_host : bool, default=True + If True, expects host memory. If False, expects device memory. + name : str + Name of the array for error messages + + Raises + ------ + ValueError + If array is not in expected memory location + """ + # Check if array has __cuda_array_interface__ (device memory indicator) + has_cuda_interface = hasattr(array_like, "__cuda_array_interface__") + + # Check if array is NumPy array (host memory indicator) + is_numpy = isinstance(array_like, np.ndarray) + + if expected_host: + if has_cuda_interface and not is_numpy: + raise ValueError( + f"Multi-GPU IVF-PQ requires {name} to be in host memory " + f"(CPU), but received device memory (GPU). Please use " + f"array.get() or cp.asnumpy(array) to transfer to host memory." + ) + else: + if is_numpy and not has_cuda_interface: + raise ValueError( + f"Expected {name} to be in device memory (GPU), but received " + f"host memory (CPU). Please use cp.asarray(array) to transfer " + f"to device memory." + ) diff --git a/python/cuvs/cuvs/neighbors/ivf_flat/ivf_flat.pxd b/python/cuvs/cuvs/neighbors/ivf_flat/ivf_flat.pxd index 470234d9cf..2078210d30 100644 --- a/python/cuvs/cuvs/neighbors/ivf_flat/ivf_flat.pxd +++ b/python/cuvs/cuvs/neighbors/ivf_flat/ivf_flat.pxd @@ -95,3 +95,10 @@ cdef extern from "cuvs/neighbors/ivf_flat.h" nogil: DLManagedTensor* new_vectors, DLManagedTensor* new_indices, cuvsIvfFlatIndex_t index) + + +cdef class IndexParams: + cdef cuvsIvfFlatIndexParams* params + +cdef class SearchParams: + cdef cuvsIvfFlatSearchParams* params diff --git a/python/cuvs/cuvs/neighbors/ivf_flat/ivf_flat.pyx b/python/cuvs/cuvs/neighbors/ivf_flat/ivf_flat.pyx index 996a39d4fe..c5f5fce5a4 100644 --- a/python/cuvs/cuvs/neighbors/ivf_flat/ivf_flat.pyx +++ b/python/cuvs/cuvs/neighbors/ivf_flat/ivf_flat.pyx @@ -99,13 +99,12 @@ cdef class IndexParams: distribution of the newly added data. """ - cdef cuvsIvfFlatIndexParams* params - def __cinit__(self): cuvsIvfFlatIndexParamsCreate(&self.params) def __dealloc__(self): - check_cuvs(cuvsIvfFlatIndexParamsDestroy(self.params)) + if self.params != NULL: + check_cuvs(cuvsIvfFlatIndexParamsDestroy(self.params)) def __init__(self, *, n_lists=1024, @@ -284,13 +283,12 @@ cdef class SearchParams: The number of clusters to search. """ - cdef cuvsIvfFlatSearchParams* params - def __cinit__(self): cuvsIvfFlatSearchParamsCreate(&self.params) def __dealloc__(self): - check_cuvs(cuvsIvfFlatSearchParamsDestroy(self.params)) + if self.params != NULL: + check_cuvs(cuvsIvfFlatSearchParamsDestroy(self.params)) def __init__(self, *, n_probes=20): self.params.n_probes = n_probes diff --git a/python/cuvs/cuvs/neighbors/ivf_pq/ivf_pq.pxd b/python/cuvs/cuvs/neighbors/ivf_pq/ivf_pq.pxd index 5ca7b97602..928a0cba1b 100644 --- a/python/cuvs/cuvs/neighbors/ivf_pq/ivf_pq.pxd +++ b/python/cuvs/cuvs/neighbors/ivf_pq/ivf_pq.pxd @@ -121,3 +121,11 @@ cdef extern from "cuvs/neighbors/ivf_pq.h" nogil: DLManagedTensor* new_vectors, DLManagedTensor* new_indices, cuvsIvfPqIndex_t index) + + +cdef class IndexParams: + cdef cuvsIvfPqIndexParams* params + cdef object _metric + +cdef class SearchParams: + cdef cuvsIvfPqSearchParams* params diff --git a/python/cuvs/cuvs/neighbors/ivf_pq/ivf_pq.pyx b/python/cuvs/cuvs/neighbors/ivf_pq/ivf_pq.pyx index 5ffa0b0c88..dd3b17f949 100644 --- a/python/cuvs/cuvs/neighbors/ivf_pq/ivf_pq.pyx +++ b/python/cuvs/cuvs/neighbors/ivf_pq/ivf_pq.pyx @@ -123,14 +123,12 @@ cdef class IndexParams: train each codebook. """ - cdef cuvsIvfPqIndexParams* params - cdef object _metric - def __cinit__(self): cuvsIvfPqIndexParamsCreate(&self.params) def __dealloc__(self): - check_cuvs(cuvsIvfPqIndexParamsDestroy(self.params)) + if self.params != NULL: + check_cuvs(cuvsIvfPqIndexParamsDestroy(self.params)) def __init__(self, *, n_lists=1024, @@ -399,13 +397,12 @@ cdef class SearchParams: of larger memory footprint. """ - cdef cuvsIvfPqSearchParams* params - def __cinit__(self): cuvsIvfPqSearchParamsCreate(&self.params) def __dealloc__(self): - check_cuvs(cuvsIvfPqSearchParamsDestroy(self.params)) + if self.params != NULL: + check_cuvs(cuvsIvfPqSearchParamsDestroy(self.params)) def __init__(self, *, n_probes=20, lut_dtype=np.float32, internal_distance_dtype=np.float32, diff --git a/python/cuvs/cuvs/neighbors/mg/CMakeLists.txt b/python/cuvs/cuvs/neighbors/mg/CMakeLists.txt new file mode 100644 index 0000000000..24a2ae01a4 --- /dev/null +++ b/python/cuvs/cuvs/neighbors/mg/CMakeLists.txt @@ -0,0 +1,17 @@ +# +# Copyright (c) 2025, NVIDIA CORPORATION. +# +# Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except +# in compliance with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software distributed under the License +# is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express +# or implied. See the License for the specific language governing permissions and limitations under +# the License. +# + +add_subdirectory(cagra) +add_subdirectory(ivf_flat) +add_subdirectory(ivf_pq) diff --git a/python/cuvs/cuvs/neighbors/mg/__init__.py b/python/cuvs/cuvs/neighbors/mg/__init__.py new file mode 100644 index 0000000000..a36b96d653 --- /dev/null +++ b/python/cuvs/cuvs/neighbors/mg/__init__.py @@ -0,0 +1,21 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from . import cagra, ivf_flat, ivf_pq + +__all__ = [ + "cagra", + "ivf_flat", + "ivf_pq", +] diff --git a/python/cuvs/cuvs/neighbors/mg/cagra/CMakeLists.txt b/python/cuvs/cuvs/neighbors/mg/cagra/CMakeLists.txt new file mode 100644 index 0000000000..4a9ffb13e0 --- /dev/null +++ b/python/cuvs/cuvs/neighbors/mg/cagra/CMakeLists.txt @@ -0,0 +1,24 @@ +# +# Copyright (c) 2025, NVIDIA CORPORATION. +# +# Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except +# in compliance with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software distributed under the License +# is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express +# or implied. See the License for the specific language governing permissions and limitations under +# the License. +# + +# Set the list of Cython files to build +set(cython_sources cagra.pyx) +set(linked_libraries cuvs::cuvs cuvs::c_api) + +# Build all of the Cython targets +rapids_cython_create_modules( + CXX + SOURCE_FILES "${cython_sources}" + LINKED_LIBRARIES "${linked_libraries}" MODULE_PREFIX neighbors_mg_cagra_ +) diff --git a/python/cuvs/cuvs/neighbors/mg/cagra/__init__.py b/python/cuvs/cuvs/neighbors/mg/cagra/__init__.py new file mode 100644 index 0000000000..597f6317fd --- /dev/null +++ b/python/cuvs/cuvs/neighbors/mg/cagra/__init__.py @@ -0,0 +1,37 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from .cagra import ( + Index, + IndexParams, + SearchParams, + build, + distribute, + extend, + load, + save, + search, +) + +__all__ = [ + "Index", + "IndexParams", + "SearchParams", + "build", + "extend", + "search", + "save", + "load", + "distribute", +] diff --git a/python/cuvs/cuvs/neighbors/mg/cagra/cagra.pxd b/python/cuvs/cuvs/neighbors/mg/cagra/cagra.pxd new file mode 100644 index 0000000000..bb42c07d4a --- /dev/null +++ b/python/cuvs/cuvs/neighbors/mg/cagra/cagra.pxd @@ -0,0 +1,126 @@ +# +# Copyright (c) 2025, NVIDIA CORPORATION. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +# cython: language_level=3 + +from libc.stdint cimport uint32_t +from libcpp cimport bool + +# Import base single-GPU extension module for subclassing +cimport cuvs.neighbors.cagra.cagra as _cagra +from cuvs.common.c_api cimport cuvsError_t, cuvsResources_t +from cuvs.common.cydlpack cimport DLManagedTensor +from cuvs.neighbors.cagra.cagra cimport ( + IndexParams as SingleGpuIndexParams, + SearchParams as SingleGpuSearchParams, + cuvsCagraIndexParams_t, + cuvsCagraSearchParams_t, +) + + +# Multi-GPU distribution modes +cdef extern from "cuvs/neighbors/mg_common.h" nogil: + ctypedef enum cuvsMultiGpuDistributionMode: + CUVS_NEIGHBORS_MG_REPLICATED + CUVS_NEIGHBORS_MG_SHARDED + + ctypedef enum cuvsMultiGpuReplicatedSearchMode: + CUVS_NEIGHBORS_MG_LOAD_BALANCER + CUVS_NEIGHBORS_MG_ROUND_ROBIN + + ctypedef enum cuvsMultiGpuShardedMergeMode: + CUVS_NEIGHBORS_MG_MERGE_ON_ROOT_RANK + CUVS_NEIGHBORS_MG_TREE_MERGE + +# Multi-GPU CAGRA structures and functions +cdef extern from "cuvs/neighbors/mg_cagra.h" nogil: + + cdef struct cuvsMultiGpuCagraIndexParams: + cuvsCagraIndexParams_t base_params + cuvsMultiGpuDistributionMode mode + + cdef struct cuvsMultiGpuCagraSearchParams: + cuvsCagraSearchParams_t base_params + cuvsMultiGpuReplicatedSearchMode search_mode + cuvsMultiGpuShardedMergeMode merge_mode + uint32_t n_rows_per_batch + + cdef struct cuvsMultiGpuCagraIndex: + pass + + ctypedef cuvsMultiGpuCagraIndexParams* cuvsMultiGpuCagraIndexParams_t + ctypedef cuvsMultiGpuCagraSearchParams* cuvsMultiGpuCagraSearchParams_t + ctypedef cuvsMultiGpuCagraIndex* cuvsMultiGpuCagraIndex_t + + cuvsError_t cuvsMultiGpuCagraIndexParamsCreate( + cuvsMultiGpuCagraIndexParams_t* index_params) + + cuvsError_t cuvsMultiGpuCagraIndexParamsDestroy( + cuvsMultiGpuCagraIndexParams_t index_params) + + cuvsError_t cuvsMultiGpuCagraSearchParamsCreate( + cuvsMultiGpuCagraSearchParams_t* params) + + cuvsError_t cuvsMultiGpuCagraSearchParamsDestroy( + cuvsMultiGpuCagraSearchParams_t params) + + cuvsError_t cuvsMultiGpuCagraIndexCreate(cuvsMultiGpuCagraIndex_t* index) + + cuvsError_t cuvsMultiGpuCagraIndexDestroy(cuvsMultiGpuCagraIndex_t index) + + cuvsError_t cuvsMultiGpuCagraBuild(cuvsResources_t res, + cuvsMultiGpuCagraIndexParams_t params, + DLManagedTensor* dataset_tensor, + cuvsMultiGpuCagraIndex_t index) except + + + cuvsError_t cuvsMultiGpuCagraSearch( + cuvsResources_t res, + cuvsMultiGpuCagraSearchParams_t params, + cuvsMultiGpuCagraIndex_t index, + DLManagedTensor* queries_tensor, + DLManagedTensor* neighbors_tensor, + DLManagedTensor* distances_tensor) except + + + cuvsError_t cuvsMultiGpuCagraSerialize( + cuvsResources_t res, + cuvsMultiGpuCagraIndex_t index, + const char* filename) except + + + cuvsError_t cuvsMultiGpuCagraDeserialize( + cuvsResources_t res, + const char* filename, + cuvsMultiGpuCagraIndex_t index) except + + + cuvsError_t cuvsMultiGpuCagraDistribute( + cuvsResources_t res, + const char* filename, + cuvsMultiGpuCagraIndex_t index) except + + + cuvsError_t cuvsMultiGpuCagraExtend( + cuvsResources_t res, + cuvsMultiGpuCagraIndex_t index, + DLManagedTensor* new_vectors_tensor, + DLManagedTensor* new_indices_tensor) except + + + +cdef class IndexParams(SingleGpuIndexParams): + cdef cuvsMultiGpuCagraIndexParams_t mg_params + +cdef class SearchParams(SingleGpuSearchParams): + cdef cuvsMultiGpuCagraSearchParams_t mg_params + +cdef class Index: + cdef cuvsMultiGpuCagraIndex_t mg_index + cdef bool mg_trained diff --git a/python/cuvs/cuvs/neighbors/mg/cagra/cagra.pyx b/python/cuvs/cuvs/neighbors/mg/cagra/cagra.pyx new file mode 100644 index 0000000000..6efcd0cd24 --- /dev/null +++ b/python/cuvs/cuvs/neighbors/mg/cagra/cagra.pyx @@ -0,0 +1,571 @@ +# +# Copyright (c) 2025, NVIDIA CORPORATION. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +# cython: language_level=3 + +import numpy as np + +from libc.stdint cimport uint32_t +from libcpp.string cimport string + +from pylibraft.common import auto_convert_output +from pylibraft.common.cai_wrapper import wrap_array +from pylibraft.common.interruptible import cuda_interruptible + +from cuvs.common.exceptions import check_cuvs +from cuvs.common.mg_resources import auto_sync_multi_gpu_resources +from cuvs.neighbors.common import _check_input_array, _check_memory_location + +from cuvs.common cimport cydlpack +from cuvs.common.c_api cimport cuvsResources_t +from cuvs.neighbors.cagra.cagra cimport ( + IndexParams as SingleGpuIndexParams, + SearchParams as SingleGpuSearchParams, + cuvsCagraIndexParams_t, + cuvsCagraIndexParamsDestroy, + cuvsCagraSearchParams_t, + cuvsCagraSearchParamsDestroy, +) + +from .cagra cimport ( + cuvsMultiGpuCagraBuild, + cuvsMultiGpuCagraDeserialize, + cuvsMultiGpuCagraDistribute, + cuvsMultiGpuCagraExtend, + cuvsMultiGpuCagraIndex_t, + cuvsMultiGpuCagraIndexCreate, + cuvsMultiGpuCagraIndexDestroy, + cuvsMultiGpuCagraIndexParams_t, + cuvsMultiGpuCagraIndexParamsCreate, + cuvsMultiGpuCagraIndexParamsDestroy, + cuvsMultiGpuCagraSearch, + cuvsMultiGpuCagraSearchParams_t, + cuvsMultiGpuCagraSearchParamsCreate, + cuvsMultiGpuCagraSearchParamsDestroy, + cuvsMultiGpuCagraSerialize, + cuvsMultiGpuDistributionMode, + cuvsMultiGpuReplicatedSearchMode, + cuvsMultiGpuShardedMergeMode, +) + + +cdef class IndexParams(SingleGpuIndexParams): + """ + Parameters to build multi-GPU CAGRA index for efficient search. + Extends single-GPU IndexParams with multi-GPU specific parameters. + + Parameters + ---------- + distribution_mode : str, default = "sharded" + Distribution mode for multi-GPU setup. + Valid values: ["replicated", "sharded"] + **kwargs : Additional parameters passed to single-GPU IndexParams + + Note + ---- + CAGRA currently only supports "sqeuclidean" and "inner_product" metrics. + """ + + def __cinit__(self): + # Base class __cinit__ has already created self.params + # We need to destroy it and use our embedded params instead + if self.params != NULL: + check_cuvs(cuvsCagraIndexParamsDestroy(self.params)) + + # Create multi-GPU params which includes embedded base params + check_cuvs(cuvsMultiGpuCagraIndexParamsCreate(&self.mg_params)) + # Replace base pointer with embedded base params + self.params = self.mg_params.base_params + + def __dealloc__(self): + # Only destroy the mg_params, which will handle base_params cleanup + check_cuvs(cuvsMultiGpuCagraIndexParamsDestroy(self.mg_params)) + self.mg_params = NULL + self.params = NULL + + def __init__(self, *, distribution_mode="sharded", **kwargs): + super().__init__(**kwargs) + if distribution_mode == "replicated": + self.mg_params.mode = CUVS_NEIGHBORS_MG_REPLICATED + elif distribution_mode == "sharded": + self.mg_params.mode = CUVS_NEIGHBORS_MG_SHARDED + else: + raise ValueError( + "distribution_mode must be 'replicated' or 'sharded'") + + def get_handle(self): + return self.mg_params + + @property + def distribution_mode(self): + return ("replicated" if self.mg_params.mode == + CUVS_NEIGHBORS_MG_REPLICATED else "sharded") + + +cdef class Index: + """ + Multi-GPU CAGRA index object. Stores the trained multi-GPU CAGRA index + state which can be used to perform nearest neighbors searches across + multiple GPUs. + """ + + def __cinit__(self): + # Initialize multi-GPU index + check_cuvs(cuvsMultiGpuCagraIndexCreate(&self.mg_index)) + # Initialize multi-GPU trained state + self.mg_trained = False + + def __dealloc__(self): + check_cuvs(cuvsMultiGpuCagraIndexDestroy(self.mg_index)) + + def __repr__(self): + return "Index(type=MultiGpuCagra)" + + @property + def trained(self): + return self.mg_trained + + +@auto_sync_multi_gpu_resources +def build(IndexParams index_params, dataset, resources=None): + """ + Build the multi-GPU CAGRA index from the dataset for efficient search. + + Parameters + ---------- + index_params : :py:class:`cuvs.neighbors.cagra.IndexParams` + dataset : Array interface compliant matrix shape (n_samples, dim) + Supported dtype [float32, float16, int8, uint8] + **IMPORTANT**: For multi-GPU CAGRA, the dataset MUST be in host + memory (CPU). If using CuPy/device arrays, transfer to host with + array.get() or cp.asnumpy(array). + {resources_docstring} + + Returns + ------- + index: py:class:`cuvs.neighbors.cagra.Index` + + Examples + -------- + + >>> import numpy as np + >>> from cuvs.neighbors.mg import cagra + >>> n_samples = 50000 + >>> n_features = 50 + >>> n_queries = 1000 + >>> k = 10 + >>> # For multi-GPU CAGRA, use host (NumPy) arrays + >>> dataset = np.random.random_sample((n_samples, n_features)).astype( + ... np.float32) + >>> build_params = cagra.IndexParams(metric="sqeuclidean") + >>> index = cagra.build(build_params, dataset) + >>> distances, neighbors = cagra.search(cagra.SearchParams(), + ... index, dataset, k) + >>> # Results are already in host memory (NumPy arrays) + """ + + dataset_ai = wrap_array(dataset) + _check_input_array(dataset_ai, [np.dtype('float32'), np.dtype('float16'), + np.dtype('byte'), np.dtype('ubyte')]) + + # Multi-GPU CAGRA requires dataset in host memory + _check_memory_location(dataset, expected_host=True, name="dataset") + + cdef Index idx = Index() + cdef cydlpack.DLManagedTensor* dataset_dlpack = ( + cydlpack.dlpack_c(dataset_ai)) + cdef cuvsMultiGpuCagraIndexParams_t params = index_params.mg_params + + cdef cuvsResources_t res = resources.get_c_obj() + + # Build the multi-GPU index + with cuda_interruptible(): + check_cuvs(cuvsMultiGpuCagraBuild( + res, params, dataset_dlpack, idx.mg_index)) + idx.mg_trained = True + + return idx + + +cdef class SearchParams(SingleGpuSearchParams): + """ + Parameters to search multi-GPU CAGRA index. + """ + + def __cinit__(self): + # Base class __cinit__ has already created self.params + # We need to destroy it and use our embedded params instead + if self.params != NULL: + check_cuvs(cuvsCagraSearchParamsDestroy(self.params)) + + # Create multi-GPU search params which includes embedded base params + check_cuvs(cuvsMultiGpuCagraSearchParamsCreate(&self.mg_params)) + # Replace base pointer with embedded base params + self.params = self.mg_params.base_params + + def __dealloc__(self): + # Only destroy the mg_params, which will handle base_params cleanup + check_cuvs(cuvsMultiGpuCagraSearchParamsDestroy(self.mg_params)) + self.mg_params = NULL + self.params = NULL + + def __init__(self, *, search_mode="load_balancer", + merge_mode="merge_on_root_rank", + n_rows_per_batch=1000, **kwargs): + super().__init__(**kwargs) + # Use the property setters for consistent validation + self.search_mode = search_mode + self.merge_mode = merge_mode + self.n_rows_per_batch = n_rows_per_batch + + def get_handle(self): + return self.mg_params + + @property + def search_mode(self): + """Get the search mode for multi-GPU search.""" + return ("load_balancer" if self.mg_params.search_mode == + CUVS_NEIGHBORS_MG_LOAD_BALANCER else "round_robin") + + @search_mode.setter + def search_mode(self, value): + """Set the search mode for multi-GPU search.""" + if value == "load_balancer": + self.mg_params.search_mode = CUVS_NEIGHBORS_MG_LOAD_BALANCER + elif value == "round_robin": + self.mg_params.search_mode = CUVS_NEIGHBORS_MG_ROUND_ROBIN + else: + raise ValueError( + "search_mode must be 'load_balancer' or 'round_robin'") + + @property + def merge_mode(self): + """Get the merge mode for multi-GPU search.""" + return ("merge_on_root_rank" if self.mg_params.merge_mode == + CUVS_NEIGHBORS_MG_MERGE_ON_ROOT_RANK else "tree_merge") + + @merge_mode.setter + def merge_mode(self, value): + """Set the merge mode for multi-GPU search.""" + if value == "merge_on_root_rank": + self.mg_params.merge_mode = CUVS_NEIGHBORS_MG_MERGE_ON_ROOT_RANK + elif value == "tree_merge": + self.mg_params.merge_mode = CUVS_NEIGHBORS_MG_TREE_MERGE + else: + raise ValueError( + "merge_mode must be 'merge_on_root_rank' or 'tree_merge'") + + @property + def n_rows_per_batch(self): + """Get the number of rows per batch for multi-GPU search.""" + return self.mg_params.n_rows_per_batch + + @n_rows_per_batch.setter + def n_rows_per_batch(self, value): + """Set the number of rows per batch for multi-GPU search.""" + if not isinstance(value, int) or value <= 0: + raise ValueError("n_rows_per_batch must be a positive integer") + self.mg_params.n_rows_per_batch = value + + +@auto_sync_multi_gpu_resources +@auto_convert_output +def search(SearchParams search_params, Index index, queries, + k, neighbors=None, distances=None, resources=None): + """ + Search the multi-GPU CAGRA index for the k-nearest neighbors of each query. + + Parameters + ---------- + search_params : :py:class:`cuvs.neighbors.cagra.SearchParams` + index : :py:class:`cuvs.neighbors.cagra.Index` + queries : Array interface compliant matrix shape (n_queries, dim) + Supported dtype [float32, float16, int8, uint8] + **IMPORTANT**: For multi-GPU CAGRA, queries MUST be in host memory + (CPU). If using CuPy/device arrays, transfer to host with + array.get() or cp.asnumpy(array). + k : int + The number of neighbors to search for each query. + neighbors : Array interface compliant matrix shape (n_queries, k), optional + If provided, this array will be filled with the indices of + the k-nearest neighbors. + If not provided, a new host array will be allocated. + **IMPORTANT**: Must be in host memory (CPU) for multi-GPU CAGRA. + Expected dtype: int64 + distances : Array interface compliant matrix shape (n_queries, k), optional + If provided, this array will be filled with the distances + to the k-nearest neighbors. + If not provided, a new host array will be allocated. + **IMPORTANT**: Must be in host memory (CPU) for multi-GPU CAGRA. + {resources_docstring} + + Returns + ------- + distances : numpy.ndarray + The distances to the k-nearest neighbors for each query + (in host memory). + neighbors : numpy.ndarray + The indices of the k-nearest neighbors for each query + (in host memory). + + Examples + -------- + + >>> import numpy as np + >>> from cuvs.neighbors.mg import cagra + >>> n_samples = 50000 + >>> n_features = 50 + >>> n_queries = 1000 + >>> k = 10 + >>> # For multi-GPU CAGRA, use host (NumPy) arrays + >>> dataset = np.random.random_sample((n_samples, n_features)).astype( + ... np.float32) + >>> queries = np.random.random_sample((n_queries, n_features)).astype( + ... np.float32) + >>> build_params = cagra.IndexParams(metric="sqeuclidean") + >>> index = cagra.build(build_params, dataset) + >>> distances, neighbors = cagra.search(cagra.SearchParams(), + ... index, queries, k) + >>> # Results are already in host memory (NumPy arrays) + """ + + if not index.trained: + raise ValueError("Index needs to be built before searching") + + queries_ai = wrap_array(queries) + _check_input_array(queries_ai, [np.dtype('float32'), np.dtype('float16'), + np.dtype('byte'), np.dtype('ubyte')]) + + # Multi-GPU CAGRA requires queries in host memory + _check_memory_location(queries, expected_host=True, name="queries") + + # Get resources + cdef cuvsResources_t res = resources.get_c_obj() + + # Prepare output arrays + cdef uint32_t n_queries = queries.shape[0] + if neighbors is None: + # For multi-GPU, create host arrays instead of device arrays + neighbors = np.empty((n_queries, k), dtype='int64') + if distances is None: + # For multi-GPU, create host arrays instead of device arrays + distances = np.empty((n_queries, k), dtype='float32') + + neighbors_ai = wrap_array(neighbors) + _check_input_array(neighbors_ai, [np.dtype('int64')], + exp_rows=n_queries, exp_cols=k) + distances_ai = wrap_array(distances) + _check_input_array(distances_ai, [np.dtype('float32')], + exp_rows=n_queries, exp_cols=k) + + # Multi-GPU CAGRA requires output arrays in host memory + _check_memory_location(neighbors, expected_host=True, + name="neighbors") + _check_memory_location(distances, expected_host=True, + name="distances") + + cdef cydlpack.DLManagedTensor* queries_dlpack = ( + cydlpack.dlpack_c(queries_ai)) + cdef cydlpack.DLManagedTensor* neighbors_dlpack = ( + cydlpack.dlpack_c(neighbors_ai)) + cdef cydlpack.DLManagedTensor* distances_dlpack = ( + cydlpack.dlpack_c(distances_ai)) + + # Perform search + with cuda_interruptible(): + check_cuvs(cuvsMultiGpuCagraSearch( + res, search_params.mg_params, index.mg_index, queries_dlpack, + neighbors_dlpack, distances_dlpack)) + + return (distances, neighbors) + + +@auto_sync_multi_gpu_resources +def extend(Index index, new_vectors, new_indices=None, resources=None): + """ + Extend the multi-GPU CAGRA index with new vectors. + + Parameters + ---------- + index : :py:class:`cuvs.neighbors.cagra.Index` + new_vectors : Array interface compliant matrix shape (n_new_vectors, dim) + Supported dtype [float32, float16, int8, uint8] + **IMPORTANT**: For multi-GPU CAGRA, new_vectors MUST be in host + memory (CPU). If using CuPy/device arrays, transfer to host with + array.get() or cp.asnumpy(array). + new_indices : Array interface compliant matrix shape (n_new_vectors,), + optional + If provided, these indices will be used for the new vectors. + If not provided, indices will be automatically assigned. + **IMPORTANT**: Must be in host memory (CPU) for multi-GPU CAGRA. + Expected dtype: uint32 + {resources_docstring} + + Examples + -------- + + >>> import numpy as np + >>> from cuvs.neighbors.mg import cagra + >>> n_samples = 50000 + >>> n_features = 50 + >>> n_new_vectors = 1000 + >>> # For multi-GPU CAGRA, use host (NumPy) arrays + >>> dataset = np.random.random_sample((n_samples, n_features)).astype( + ... np.float32) + >>> new_vectors = np.random.random_sample( + ... (n_new_vectors, n_features)).astype(np.float32) + >>> new_indices = np.arange(n_samples, n_samples + n_new_vectors, + ... dtype=np.uint32) + >>> build_params = cagra.IndexParams(metric="sqeuclidean") + >>> index = cagra.build(build_params, dataset) + >>> cagra.extend(index, new_vectors, new_indices) # doctest: +SKIP + """ + + if not index.trained: + raise ValueError("Index needs to be built before extending") + + new_vectors_ai = wrap_array(new_vectors) + _check_input_array(new_vectors_ai, + [np.dtype('float32'), np.dtype('float16'), + np.dtype('byte'), np.dtype('ubyte')]) + + # Multi-GPU CAGRA requires new_vectors in host memory + _check_memory_location(new_vectors, expected_host=True, name="new_vectors") + + # Get resources + cdef cuvsResources_t res = resources.get_c_obj() + + cdef cydlpack.DLManagedTensor* new_vectors_dlpack = \ + cydlpack.dlpack_c(new_vectors_ai) + cdef cydlpack.DLManagedTensor* new_indices_dlpack = NULL + + if new_indices is not None: + new_indices_ai = wrap_array(new_indices) + _check_input_array(new_indices_ai, [np.dtype('uint32')]) + # Multi-GPU CAGRA requires new_indices in host memory + _check_memory_location(new_indices, expected_host=True, + name="new_indices") + new_indices_dlpack = cydlpack.dlpack_c(new_indices_ai) + + with cuda_interruptible(): + check_cuvs(cuvsMultiGpuCagraExtend(res, index.mg_index, + new_vectors_dlpack, + new_indices_dlpack)) + + +@auto_sync_multi_gpu_resources +def save(Index index, filename, resources=None): + """ + Serialize the multi-GPU CAGRA index to a file. + + Parameters + ---------- + index : :py:class:`cuvs.neighbors.cagra.Index` + filename : str + The filename to serialize the index to. + {resources_docstring} + + Examples + -------- + + >>> import numpy as np + >>> from cuvs.neighbors.mg import cagra + >>> n_samples = 50000 + >>> n_features = 50 + >>> # For multi-GPU CAGRA, use host (NumPy) arrays + >>> dataset = np.random.random_sample((n_samples, n_features)).astype( + ... np.float32) + >>> build_params = cagra.IndexParams(metric="sqeuclidean") + >>> index = cagra.build(build_params, dataset) + >>> cagra.save(index, "index.bin") + """ + + if not index.trained: + raise ValueError("Index needs to be built before serializing") + + # Get resources + cdef cuvsResources_t res = resources.get_c_obj() + + cdef string filename_str = filename.encode('utf-8') + check_cuvs(cuvsMultiGpuCagraSerialize( + res, index.mg_index, filename_str.c_str())) + + +@auto_sync_multi_gpu_resources +def load(filename, resources=None): + """ + Deserialize the multi-GPU CAGRA index from a file. + + Parameters + ---------- + filename : str + The filename to deserialize the index from. + {resources_docstring} + + Returns + ------- + index : Index + The deserialized index. + + Examples + -------- + + >>> from cuvs.neighbors.mg import cagra + >>> index = cagra.load("index.bin") # doctest: +SKIP + """ + + cdef Index index = Index() + cdef cuvsResources_t res = resources.get_c_obj() + + cdef string filename_str = filename.encode('utf-8') + check_cuvs(cuvsMultiGpuCagraDeserialize( + res, filename_str.c_str(), index.mg_index)) + index.mg_trained = True + + return index + + +@auto_sync_multi_gpu_resources +def distribute(filename, resources=None): + """ + Distribute a single-GPU CAGRA index across multiple GPUs from a file. + + Parameters + ---------- + filename : str + The filename to distribute the index from. + {resources_docstring} + + Returns + ------- + index : Index + The distributed index. + + Examples + -------- + + >>> from cuvs.neighbors.mg import cagra + >>> index = cagra.distribute("single_gpu_index.bin") # doctest: +SKIP + """ + + cdef Index index = Index() + cdef cuvsResources_t res = resources.get_c_obj() + + cdef string filename_str = filename.encode('utf-8') + check_cuvs(cuvsMultiGpuCagraDistribute( + res, filename_str.c_str(), index.mg_index)) + index.mg_trained = True + + return index diff --git a/python/cuvs/cuvs/neighbors/mg/ivf_flat/CMakeLists.txt b/python/cuvs/cuvs/neighbors/mg/ivf_flat/CMakeLists.txt new file mode 100644 index 0000000000..faacbfe2ea --- /dev/null +++ b/python/cuvs/cuvs/neighbors/mg/ivf_flat/CMakeLists.txt @@ -0,0 +1,24 @@ +# +# Copyright (c) 2025, NVIDIA CORPORATION. +# +# Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except +# in compliance with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software distributed under the License +# is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express +# or implied. See the License for the specific language governing permissions and limitations under +# the License. +# + +# Set the list of Cython files to build +set(cython_sources ivf_flat.pyx) +set(linked_libraries cuvs::cuvs cuvs::c_api) + +# Build all of the Cython targets +rapids_cython_create_modules( + CXX + SOURCE_FILES "${cython_sources}" + LINKED_LIBRARIES "${linked_libraries}" MODULE_PREFIX neighbors_mg_ivf_flat_ +) diff --git a/python/cuvs/cuvs/neighbors/mg/ivf_flat/__init__.py b/python/cuvs/cuvs/neighbors/mg/ivf_flat/__init__.py new file mode 100644 index 0000000000..e4ea5ce643 --- /dev/null +++ b/python/cuvs/cuvs/neighbors/mg/ivf_flat/__init__.py @@ -0,0 +1,39 @@ +# +# Copyright (c) 2025, NVIDIA CORPORATION. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +from .ivf_flat import ( + Index, + IndexParams, + SearchParams, + build, + distribute, + extend, + load, + save, + search, +) + +__all__ = [ + "Index", + "IndexParams", + "SearchParams", + "build", + "extend", + "search", + "save", + "load", + "distribute", +] diff --git a/python/cuvs/cuvs/neighbors/mg/ivf_flat/ivf_flat.pxd b/python/cuvs/cuvs/neighbors/mg/ivf_flat/ivf_flat.pxd new file mode 100644 index 0000000000..65dfe0db15 --- /dev/null +++ b/python/cuvs/cuvs/neighbors/mg/ivf_flat/ivf_flat.pxd @@ -0,0 +1,128 @@ +# +# Copyright (c) 2025, NVIDIA CORPORATION. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +# cython: language_level=3 + +from libc.stdint cimport int64_t, uintptr_t +from libcpp cimport bool + +# Import base single-GPU extension module for subclassing +from cuvs.common.c_api cimport cuvsError_t, cuvsResources_t +from cuvs.common.cydlpack cimport DLDataType, DLManagedTensor +from cuvs.neighbors.ivf_flat.ivf_flat cimport ( + IndexParams as SingleGpuIndexParams, + SearchParams as SingleGpuSearchParams, + cuvsIvfFlatIndexParams_t, + cuvsIvfFlatSearchParams_t, +) + + +# Multi-GPU distribution modes +cdef extern from "cuvs/neighbors/mg_common.h" nogil: + ctypedef enum cuvsMultiGpuDistributionMode: + CUVS_NEIGHBORS_MG_REPLICATED + CUVS_NEIGHBORS_MG_SHARDED + + ctypedef enum cuvsMultiGpuReplicatedSearchMode: + CUVS_NEIGHBORS_MG_LOAD_BALANCER + CUVS_NEIGHBORS_MG_ROUND_ROBIN + + ctypedef enum cuvsMultiGpuShardedMergeMode: + CUVS_NEIGHBORS_MG_MERGE_ON_ROOT_RANK + CUVS_NEIGHBORS_MG_TREE_MERGE + +# Multi-GPU IVF-Flat structures and functions +cdef extern from "cuvs/neighbors/mg_ivf_flat.h" nogil: + cdef struct cuvsMultiGpuIvfFlatIndexParams: + cuvsIvfFlatIndexParams_t base_params + cuvsMultiGpuDistributionMode mode + + cdef struct cuvsMultiGpuIvfFlatSearchParams: + cuvsIvfFlatSearchParams_t base_params + cuvsMultiGpuReplicatedSearchMode search_mode + cuvsMultiGpuShardedMergeMode merge_mode + int64_t n_rows_per_batch + + cdef struct cuvsMultiGpuIvfFlatIndex: + uintptr_t addr + DLDataType dtype + + ctypedef cuvsMultiGpuIvfFlatIndexParams* cuvsMultiGpuIvfFlatIndexParams_t + ctypedef cuvsMultiGpuIvfFlatSearchParams* cuvsMultiGpuIvfFlatSearchParams_t + ctypedef cuvsMultiGpuIvfFlatIndex* cuvsMultiGpuIvfFlatIndex_t + + cuvsError_t cuvsMultiGpuIvfFlatIndexParamsCreate( + cuvsMultiGpuIvfFlatIndexParams_t* index_params) + + cuvsError_t cuvsMultiGpuIvfFlatIndexParamsDestroy( + cuvsMultiGpuIvfFlatIndexParams_t index_params) + + cuvsError_t cuvsMultiGpuIvfFlatSearchParamsCreate( + cuvsMultiGpuIvfFlatSearchParams_t* params) + + cuvsError_t cuvsMultiGpuIvfFlatSearchParamsDestroy( + cuvsMultiGpuIvfFlatSearchParams_t params) + + cuvsError_t cuvsMultiGpuIvfFlatIndexCreate( + cuvsMultiGpuIvfFlatIndex_t* index) + + cuvsError_t cuvsMultiGpuIvfFlatIndexDestroy( + cuvsMultiGpuIvfFlatIndex_t index) + + cuvsError_t cuvsMultiGpuIvfFlatBuild( + cuvsResources_t res, + cuvsMultiGpuIvfFlatIndexParams_t params, + DLManagedTensor* dataset_tensor, + cuvsMultiGpuIvfFlatIndex_t index) except + + + cuvsError_t cuvsMultiGpuIvfFlatSearch( + cuvsResources_t res, + cuvsMultiGpuIvfFlatSearchParams_t params, + cuvsMultiGpuIvfFlatIndex_t index, + DLManagedTensor* queries_tensor, + DLManagedTensor* neighbors_tensor, + DLManagedTensor* distances_tensor) except + + + cuvsError_t cuvsMultiGpuIvfFlatExtend( + cuvsResources_t res, + cuvsMultiGpuIvfFlatIndex_t index, + DLManagedTensor* new_vectors_tensor, + DLManagedTensor* new_indices_tensor) except + + + cuvsError_t cuvsMultiGpuIvfFlatSerialize( + cuvsResources_t res, + cuvsMultiGpuIvfFlatIndex_t index, + const char* filename) except + + + cuvsError_t cuvsMultiGpuIvfFlatDeserialize( + cuvsResources_t res, + const char* filename, + cuvsMultiGpuIvfFlatIndex_t index) except + + + cuvsError_t cuvsMultiGpuIvfFlatDistribute( + cuvsResources_t res, + const char* filename, + cuvsMultiGpuIvfFlatIndex_t index) except + + + +cdef class IndexParams(SingleGpuIndexParams): + cdef cuvsMultiGpuIvfFlatIndexParams_t mg_params + +cdef class SearchParams(SingleGpuSearchParams): + cdef cuvsMultiGpuIvfFlatSearchParams_t mg_params + +cdef class Index: + cdef cuvsMultiGpuIvfFlatIndex_t mg_index + cdef bool mg_trained diff --git a/python/cuvs/cuvs/neighbors/mg/ivf_flat/ivf_flat.pyx b/python/cuvs/cuvs/neighbors/mg/ivf_flat/ivf_flat.pyx new file mode 100644 index 0000000000..e40b6b82ea --- /dev/null +++ b/python/cuvs/cuvs/neighbors/mg/ivf_flat/ivf_flat.pyx @@ -0,0 +1,575 @@ +# +# Copyright (c) 2025, NVIDIA CORPORATION. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +# cython: language_level=3 + +import numpy as np + +from libc.stdint cimport uint32_t +from libcpp.string cimport string + +from pylibraft.common import auto_convert_output +from pylibraft.common.cai_wrapper import wrap_array +from pylibraft.common.interruptible import cuda_interruptible + +from cuvs.common.exceptions import check_cuvs +from cuvs.common.mg_resources import auto_sync_multi_gpu_resources +from cuvs.neighbors.common import _check_input_array, _check_memory_location + +from cuvs.common cimport cydlpack +from cuvs.common.c_api cimport cuvsResources_t +from cuvs.neighbors.ivf_flat.ivf_flat cimport ( + IndexParams as SingleGpuIndexParams, + SearchParams as SingleGpuSearchParams, + cuvsIvfFlatIndexParams_t, + cuvsIvfFlatIndexParamsDestroy, + cuvsIvfFlatSearchParams_t, + cuvsIvfFlatSearchParamsDestroy, +) + +from .ivf_flat cimport ( + cuvsMultiGpuDistributionMode, + cuvsMultiGpuIvfFlatBuild, + cuvsMultiGpuIvfFlatDeserialize, + cuvsMultiGpuIvfFlatDistribute, + cuvsMultiGpuIvfFlatExtend, + cuvsMultiGpuIvfFlatIndex, + cuvsMultiGpuIvfFlatIndex_t, + cuvsMultiGpuIvfFlatIndexCreate, + cuvsMultiGpuIvfFlatIndexDestroy, + cuvsMultiGpuIvfFlatIndexParams, + cuvsMultiGpuIvfFlatIndexParams_t, + cuvsMultiGpuIvfFlatIndexParamsCreate, + cuvsMultiGpuIvfFlatIndexParamsDestroy, + cuvsMultiGpuIvfFlatSearch, + cuvsMultiGpuIvfFlatSearchParams, + cuvsMultiGpuIvfFlatSearchParams_t, + cuvsMultiGpuIvfFlatSearchParamsCreate, + cuvsMultiGpuIvfFlatSearchParamsDestroy, + cuvsMultiGpuIvfFlatSerialize, + cuvsMultiGpuReplicatedSearchMode, + cuvsMultiGpuShardedMergeMode, +) + + +cdef class IndexParams(SingleGpuIndexParams): + """ + Parameters to build multi-GPU IVF-Flat index for efficient search. + Extends single-GPU IndexParams with multi-GPU specific parameters. + + Parameters + ---------- + distribution_mode : str, default = "sharded" + Distribution mode for multi-GPU setup. + Valid values: ["replicated", "sharded"] + **kwargs : Additional parameters passed to single-GPU IndexParams + """ + + def __cinit__(self): + # Base class __cinit__ has already created self.params + # We need to destroy it and use our embedded params instead + if self.params != NULL: + check_cuvs(cuvsIvfFlatIndexParamsDestroy(self.params)) + + # Create multi-GPU params which includes embedded base params + check_cuvs(cuvsMultiGpuIvfFlatIndexParamsCreate(&self.mg_params)) + # Replace base pointer with embedded base params + self.params = self.mg_params.base_params + + def __dealloc__(self): + # Only destroy the mg_params, which will handle base_params cleanup + check_cuvs(cuvsMultiGpuIvfFlatIndexParamsDestroy(self.mg_params)) + self.mg_params = NULL + self.params = NULL + + def __init__(self, *, distribution_mode="sharded", **kwargs): + super().__init__(**kwargs) + if distribution_mode == "replicated": + self.mg_params.mode = CUVS_NEIGHBORS_MG_REPLICATED + elif distribution_mode == "sharded": + self.mg_params.mode = CUVS_NEIGHBORS_MG_SHARDED + else: + raise ValueError( + "distribution_mode must be 'replicated' or 'sharded'") + + def get_handle(self): + return self.mg_params + + @property + def distribution_mode(self): + return ("replicated" if self.mg_params.mode == + CUVS_NEIGHBORS_MG_REPLICATED else "sharded") + + +cdef class Index: + """ + Multi-GPU IVF-Flat index object. Stores the trained multi-GPU IVF-Flat + index state which can be used to perform nearest neighbors searches + across multiple GPUs. + """ + + def __cinit__(self): + # Initialize multi-GPU index + check_cuvs(cuvsMultiGpuIvfFlatIndexCreate(&self.mg_index)) + # Initialize multi-GPU trained state + self.mg_trained = False + + def __dealloc__(self): + check_cuvs(cuvsMultiGpuIvfFlatIndexDestroy(self.mg_index)) + + def __repr__(self): + return "Index(type=MultiGpuIvfFlat)" + + @property + def trained(self): + return self.mg_trained + + +@auto_sync_multi_gpu_resources +def build(IndexParams index_params, dataset, resources=None): + """ + Build the multi-GPU IVF-Flat index from the dataset for efficient search. + + Parameters + ---------- + index_params : :py:class:`cuvs.neighbors.ivf_flat.\ +IndexParams` + dataset : Array interface compliant matrix shape (n_samples, dim) + Supported dtype [float32, float16, int8, uint8] + **IMPORTANT**: For multi-GPU IVF-Flat, the dataset MUST be in host + memory (CPU). If using CuPy/device arrays, transfer to host with + array.get() or cp.asnumpy(array). + {resources_docstring} + + Returns + ------- + index: py:class:`cuvs.neighbors.ivf_flat.Index` + + Examples + -------- + + >>> import numpy as np + >>> from cuvs.neighbors.mg import ivf_flat + >>> n_samples = 50000 + >>> n_features = 50 + >>> n_queries = 1000 + >>> k = 10 + >>> # For multi-GPU IVF-Flat, use host (NumPy) arrays + >>> dataset = np.random.random_sample((n_samples, n_features)).astype( + ... np.float32) + >>> build_params = ivf_flat.IndexParams(metric="sqeuclidean") + >>> index = ivf_flat.build(build_params, dataset) + >>> distances, neighbors = ivf_flat.search( + ... ivf_flat.SearchParams(), + ... index, dataset, k) + >>> # Results are already in host memory (NumPy arrays) + """ + + dataset_ai = wrap_array(dataset) + _check_input_array(dataset_ai, [np.dtype('float32'), np.dtype('float16'), + np.dtype('byte'), np.dtype('ubyte')]) + + # Multi-GPU IVF-Flat requires dataset in host memory + _check_memory_location(dataset, expected_host=True, name="dataset") + + cdef Index idx = Index() + cdef cydlpack.DLManagedTensor* dataset_dlpack = ( + cydlpack.dlpack_c(dataset_ai)) + cdef cuvsMultiGpuIvfFlatIndexParams_t params = index_params.mg_params + + cdef cuvsResources_t res = resources.get_c_obj() + + # Build the multi-GPU index + with cuda_interruptible(): + check_cuvs(cuvsMultiGpuIvfFlatBuild( + res, params, dataset_dlpack, idx.mg_index)) + idx.mg_trained = True + + return idx + + +cdef class SearchParams(SingleGpuSearchParams): + """ + Parameters to search multi-GPU IVF-Flat index. + """ + + def __cinit__(self): + # Base class __cinit__ has already created self.params + # We need to destroy it and use our embedded params instead + if self.params != NULL: + check_cuvs(cuvsIvfFlatSearchParamsDestroy(self.params)) + + # Create multi-GPU search params which includes embedded base params + check_cuvs(cuvsMultiGpuIvfFlatSearchParamsCreate(&self.mg_params)) + # Replace base pointer with embedded base params + self.params = self.mg_params.base_params + + def __dealloc__(self): + # Only destroy the mg_params, which will handle base_params cleanup + check_cuvs(cuvsMultiGpuIvfFlatSearchParamsDestroy(self.mg_params)) + self.mg_params = NULL + self.params = NULL + + def __init__(self, *, n_probes=1, search_mode="load_balancer", + merge_mode="merge_on_root_rank", + n_rows_per_batch=1000, **kwargs): + super().__init__(n_probes=n_probes, **kwargs) + # Use the property setters for consistent validation + self.search_mode = search_mode + self.merge_mode = merge_mode + self.n_rows_per_batch = n_rows_per_batch + + def get_handle(self): + return self.mg_params + + @property + def search_mode(self): + """Get the search mode for multi-GPU search.""" + return ("load_balancer" if self.mg_params.search_mode == + CUVS_NEIGHBORS_MG_LOAD_BALANCER else "round_robin") + + @search_mode.setter + def search_mode(self, value): + """Set the search mode for multi-GPU search.""" + if value == "load_balancer": + self.mg_params.search_mode = CUVS_NEIGHBORS_MG_LOAD_BALANCER + elif value == "round_robin": + self.mg_params.search_mode = CUVS_NEIGHBORS_MG_ROUND_ROBIN + else: + raise ValueError( + "search_mode must be 'load_balancer' or 'round_robin'") + + @property + def merge_mode(self): + """Get the merge mode for multi-GPU search.""" + return ("merge_on_root_rank" if self.mg_params.merge_mode == + CUVS_NEIGHBORS_MG_MERGE_ON_ROOT_RANK else "tree_merge") + + @merge_mode.setter + def merge_mode(self, value): + """Set the merge mode for multi-GPU search.""" + if value == "merge_on_root_rank": + self.mg_params.merge_mode = CUVS_NEIGHBORS_MG_MERGE_ON_ROOT_RANK + elif value == "tree_merge": + self.mg_params.merge_mode = CUVS_NEIGHBORS_MG_TREE_MERGE + else: + raise ValueError( + "merge_mode must be 'merge_on_root_rank' or 'tree_merge'") + + @property + def n_rows_per_batch(self): + """Get the number of rows per batch for multi-GPU search.""" + return self.mg_params.n_rows_per_batch + + @n_rows_per_batch.setter + def n_rows_per_batch(self, value): + """Set the number of rows per batch for multi-GPU search.""" + if not isinstance(value, int) or value <= 0: + raise ValueError("n_rows_per_batch must be a positive integer") + self.mg_params.n_rows_per_batch = value + + +@auto_sync_multi_gpu_resources +@auto_convert_output +def search(SearchParams search_params, Index index, queries, + k, neighbors=None, distances=None, resources=None): + """ + Search the multi-GPU IVF-Flat index for the k-nearest neighbors + of each query. + + Parameters + ---------- + search_params : :py:class:`cuvs.neighbors.ivf_flat.SearchParams` + index : :py:class:`cuvs.neighbors.ivf_flat.Index` + queries : Array interface compliant matrix shape (n_queries, dim) + Supported dtype [float32, float16, int8, uint8] + **IMPORTANT**: For multi-GPU IVF-Flat, queries MUST be + in host memory (CPU). + If using CuPy/device arrays, transfer to host with array.get() + or cp.asnumpy(array). + k : int + The number of neighbors to search for each query. + neighbors : Array interface compliant matrix shape (n_queries, k), optional + If provided, this array will be filled with the indices of + the k-nearest neighbors. + If not provided, a new host array will be allocated. + **IMPORTANT**: Must be in host memory (CPU) for multi-GPU IVF-Flat. + distances : Array interface compliant matrix shape (n_queries, k), optional + If provided, this array will be filled with the distances to + the k-nearest neighbors. + If not provided, a new host array will be allocated. + **IMPORTANT**: Must be in host memory (CPU) for multi-GPU IVF-Flat. + {resources_docstring} + + Returns + ------- + distances : numpy.ndarray + The distances to the k-nearest neighbors for each query + (in host memory). + neighbors : numpy.ndarray + The indices of the k-nearest neighbors for each query + (in host memory). + + Examples + -------- + + >>> import numpy as np + >>> from cuvs.neighbors.mg import ivf_flat + >>> n_samples = 50000 + >>> n_features = 50 + >>> n_queries = 1000 + >>> k = 10 + >>> # For multi-GPU IVF-Flat, use host (NumPy) arrays + >>> dataset = np.random.random_sample((n_samples, n_features)).astype( + ... np.float32) + >>> queries = np.random.random_sample((n_queries, n_features)).astype( + ... np.float32) + >>> build_params = ivf_flat.IndexParams(metric="sqeuclidean") + >>> index = ivf_flat.build(build_params, dataset) + >>> distances, neighbors = ivf_flat.search( + ... ivf_flat.SearchParams(), + ... index, queries, k) + >>> # Results are already in host memory (NumPy arrays) + """ + + if not index.trained: + raise ValueError("Index needs to be built before searching") + + queries_ai = wrap_array(queries) + _check_input_array(queries_ai, [np.dtype('float32'), np.dtype('float16'), + np.dtype('byte'), np.dtype('ubyte')]) + + # Multi-GPU IVF-Flat requires queries in host memory + _check_memory_location(queries, expected_host=True, name="queries") + + # Get resources + cdef cuvsResources_t res = resources.get_c_obj() + + # Prepare output arrays + cdef uint32_t n_queries = queries.shape[0] + if neighbors is None: + # For multi-GPU, create host arrays instead of device arrays + neighbors = np.empty((n_queries, k), dtype='int64') + if distances is None: + # For multi-GPU, create host arrays instead of device arrays + distances = np.empty((n_queries, k), dtype='float32') + + neighbors_ai = wrap_array(neighbors) + _check_input_array(neighbors_ai, [np.dtype('int64')], + exp_rows=n_queries, exp_cols=k) + distances_ai = wrap_array(distances) + _check_input_array(distances_ai, [np.dtype('float32')], + exp_rows=n_queries, exp_cols=k) + + # Multi-GPU IVF-Flat requires output arrays in host memory + _check_memory_location(neighbors, expected_host=True, name="neighbors") + _check_memory_location(distances, expected_host=True, name="distances") + + cdef cydlpack.DLManagedTensor* queries_dlpack = \ + cydlpack.dlpack_c(queries_ai) + cdef cydlpack.DLManagedTensor* neighbors_dlpack = \ + cydlpack.dlpack_c(neighbors_ai) + cdef cydlpack.DLManagedTensor* distances_dlpack = \ + cydlpack.dlpack_c(distances_ai) + + # Perform search + with cuda_interruptible(): + check_cuvs(cuvsMultiGpuIvfFlatSearch(res, search_params.mg_params, + index.mg_index, queries_dlpack, + neighbors_dlpack, + distances_dlpack)) + + return (distances, neighbors) + + +@auto_sync_multi_gpu_resources +def extend(Index index, new_vectors, new_indices=None, + resources=None): + """ + Extend the multi-GPU IVF-Flat index with new vectors. + + Parameters + ---------- + index : :py:class:`cuvs.neighbors.ivf_flat.Index` + new_vectors : Array interface compliant matrix shape (n_new_vectors, dim) + Supported dtype [float32, float16, int8, uint8] + **IMPORTANT**: For multi-GPU IVF-Flat, new_vectors MUST be + in host memory (CPU). If using CuPy/device arrays, transfer + to host with array.get() or cp.asnumpy(array). + new_indices : Array interface compliant matrix shape (n_new_vectors,) + , optional + If provided, these indices will be used for the new vectors. + If not provided, indices will be automatically assigned. + **IMPORTANT**: Must be in host memory (CPU) for multi-GPU IVF-Flat. + {resources_docstring} + + Examples + -------- + + >>> import numpy as np + >>> from cuvs.neighbors.mg import ivf_flat + >>> n_samples = 50000 + >>> n_features = 50 + >>> n_new_vectors = 1000 + >>> # For multi-GPU IVF-Flat, use host (NumPy) arrays + >>> dataset = np.random.random_sample((n_samples, n_features)).astype( + ... np.float32) + >>> new_vectors = np.random.random_sample( + ... (n_new_vectors, n_features)).astype(np.float32) + >>> new_indices = np.arange(n_samples, n_new_vectors, dtype=np.int64) + >>> build_params = ivf_flat.IndexParams(metric="sqeuclidean") + >>> index = ivf_flat.build(build_params, dataset) + >>> ivf_flat.extend(index, new_vectors, new_indices) + """ + + if not index.trained: + raise ValueError("Index needs to be built before extending") + + new_vectors_ai = wrap_array(new_vectors) + _check_input_array(new_vectors_ai, + [np.dtype('float32'), np.dtype('float16'), + np.dtype('byte'), np.dtype('ubyte')]) + + # Multi-GPU IVF-Flat requires new_vectors in host memory + _check_memory_location(new_vectors, expected_host=True, name="new_vectors") + + # Get resources + cdef cuvsResources_t res = resources.get_c_obj() + + cdef cydlpack.DLManagedTensor* new_vectors_dlpack = \ + cydlpack.dlpack_c(new_vectors_ai) + cdef cydlpack.DLManagedTensor* new_indices_dlpack = NULL + + if new_indices is not None: + new_indices_ai = wrap_array(new_indices) + _check_input_array(new_indices_ai, [np.dtype('int64')]) + # Multi-GPU IVF-Flat requires new_indices in host memory + _check_memory_location(new_indices, expected_host=True, + name="new_indices") + new_indices_dlpack = cydlpack.dlpack_c(new_indices_ai) + + with cuda_interruptible(): + check_cuvs(cuvsMultiGpuIvfFlatExtend(res, index.mg_index, + new_vectors_dlpack, + new_indices_dlpack)) + + +@auto_sync_multi_gpu_resources +def save(Index index, filename, resources=None): + """ + Serialize the multi-GPU IVF-Flat index to a file. + + Parameters + ---------- + index : :py:class:`cuvs.neighbors.ivf_flat.Index` + filename : str + The filename to serialize the index to. + {resources_docstring} + + Examples + -------- + + >>> import numpy as np + >>> from cuvs.neighbors.mg import ivf_flat + >>> n_samples = 50000 + >>> n_features = 50 + >>> # For multi-GPU IVF-Flat, use host (NumPy) arrays + >>> dataset = np.random.random_sample((n_samples, n_features)).astype( + ... np.float32) + >>> build_params = ivf_flat.IndexParams(metric="sqeuclidean") + >>> index = ivf_flat.build(build_params, dataset) + >>> ivf_flat.save(index, "index.bin") + """ + + if not index.trained: + raise ValueError("Index needs to be built before serializing") + + # Get resources + cdef cuvsResources_t res = resources.get_c_obj() + + cdef string filename_str = filename.encode('utf-8') + check_cuvs(cuvsMultiGpuIvfFlatSerialize(res, + index.mg_index, + filename_str.c_str())) + + +@auto_sync_multi_gpu_resources +def load(filename, resources=None): + """ + Deserialize the multi-GPU IVF-Flat index from a file. + + Parameters + ---------- + filename : str + The filename to deserialize the index from. + {resources_docstring} + + Returns + ------- + index : Index + The deserialized index. + + Examples + -------- + + >>> from cuvs.neighbors.mg import ivf_flat + >>> index = ivf_flat.load("index.bin") # doctest: +SKIP + """ + + cdef Index index = Index() + cdef cuvsResources_t res = resources.get_c_obj() + + cdef string filename_str = filename.encode('utf-8') + check_cuvs(cuvsMultiGpuIvfFlatDeserialize(res, + filename_str.c_str(), + index.mg_index)) + index.mg_trained = True + + return index + + +@auto_sync_multi_gpu_resources +def distribute(filename, resources=None): + """ + Distribute a single-GPU IVF-Flat index across multiple GPUs from a file. + + Parameters + ---------- + filename : str + The filename to distribute the index from. + {resources_docstring} + + Returns + ------- + index : Index + The distributed index. + + Examples + -------- + + >>> from cuvs.neighbors.mg import ivf_flat + >>> index = ivf_flat.distribute("single_gpu_index.bin") # doctest: +SKIP + """ + + cdef Index index = Index() + cdef cuvsResources_t res = resources.get_c_obj() + + cdef string filename_str = filename.encode('utf-8') + check_cuvs(cuvsMultiGpuIvfFlatDistribute(res, + filename_str.c_str(), + index.mg_index)) + index.mg_trained = True + + return index diff --git a/python/cuvs/cuvs/neighbors/mg/ivf_pq/CMakeLists.txt b/python/cuvs/cuvs/neighbors/mg/ivf_pq/CMakeLists.txt new file mode 100644 index 0000000000..2b5c5a18c8 --- /dev/null +++ b/python/cuvs/cuvs/neighbors/mg/ivf_pq/CMakeLists.txt @@ -0,0 +1,24 @@ +# +# Copyright (c) 2025, NVIDIA CORPORATION. +# +# Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except +# in compliance with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software distributed under the License +# is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express +# or implied. See the License for the specific language governing permissions and limitations under +# the License. +# + +# Set the list of Cython files to build +set(cython_sources ivf_pq.pyx) +set(linked_libraries cuvs::cuvs cuvs::c_api) + +# Build all of the Cython targets +rapids_cython_create_modules( + CXX + SOURCE_FILES "${cython_sources}" + LINKED_LIBRARIES "${linked_libraries}" MODULE_PREFIX neighbors_mg_ivf_pq_ +) diff --git a/python/cuvs/cuvs/neighbors/mg/ivf_pq/__init__.py b/python/cuvs/cuvs/neighbors/mg/ivf_pq/__init__.py new file mode 100644 index 0000000000..c75cedc267 --- /dev/null +++ b/python/cuvs/cuvs/neighbors/mg/ivf_pq/__init__.py @@ -0,0 +1,39 @@ +# +# Copyright (c) 2025, NVIDIA CORPORATION. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +from .ivf_pq import ( + Index, + IndexParams, + SearchParams, + build, + distribute, + extend, + load, + save, + search, +) + +__all__ = [ + "Index", + "IndexParams", + "SearchParams", + "build", + "extend", + "search", + "save", + "load", + "distribute", +] diff --git a/python/cuvs/cuvs/neighbors/mg/ivf_pq/ivf_pq.pxd b/python/cuvs/cuvs/neighbors/mg/ivf_pq/ivf_pq.pxd new file mode 100644 index 0000000000..b0a635eb52 --- /dev/null +++ b/python/cuvs/cuvs/neighbors/mg/ivf_pq/ivf_pq.pxd @@ -0,0 +1,125 @@ +# +# Copyright (c) 2025, NVIDIA CORPORATION. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +# cython: language_level=3 + +from libc.stdint cimport int64_t, uintptr_t +from libcpp cimport bool + +from cuvs.common.c_api cimport cuvsError_t, cuvsResources_t +from cuvs.common.cydlpack cimport DLDataType, DLManagedTensor +from cuvs.neighbors.ivf_pq.ivf_pq cimport ( + IndexParams as SingleGpuIndexParams, + SearchParams as SingleGpuSearchParams, + cuvsIvfPqIndexParams_t, + cuvsIvfPqSearchParams_t, +) + +# Import base single-GPU extension module for subclassing + +# Multi-GPU distribution modes +cdef extern from "cuvs/neighbors/mg_common.h" nogil: + ctypedef enum cuvsMultiGpuDistributionMode: + CUVS_NEIGHBORS_MG_REPLICATED + CUVS_NEIGHBORS_MG_SHARDED + + ctypedef enum cuvsMultiGpuReplicatedSearchMode: + CUVS_NEIGHBORS_MG_LOAD_BALANCER + CUVS_NEIGHBORS_MG_ROUND_ROBIN + + ctypedef enum cuvsMultiGpuShardedMergeMode: + CUVS_NEIGHBORS_MG_MERGE_ON_ROOT_RANK + CUVS_NEIGHBORS_MG_TREE_MERGE + +# Multi-GPU IVF-PQ structures and functions +cdef extern from "cuvs/neighbors/mg_ivf_pq.h" nogil: + cdef struct cuvsMultiGpuIvfPqIndexParams: + cuvsIvfPqIndexParams_t base_params + cuvsMultiGpuDistributionMode mode + + cdef struct cuvsMultiGpuIvfPqSearchParams: + cuvsIvfPqSearchParams_t base_params + cuvsMultiGpuReplicatedSearchMode search_mode + cuvsMultiGpuShardedMergeMode merge_mode + int64_t n_rows_per_batch + + cdef struct cuvsMultiGpuIvfPqIndex: + uintptr_t addr + DLDataType dtype + + ctypedef cuvsMultiGpuIvfPqIndexParams* cuvsMultiGpuIvfPqIndexParams_t + ctypedef cuvsMultiGpuIvfPqSearchParams* cuvsMultiGpuIvfPqSearchParams_t + ctypedef cuvsMultiGpuIvfPqIndex* cuvsMultiGpuIvfPqIndex_t + + cuvsError_t cuvsMultiGpuIvfPqIndexParamsCreate( + cuvsMultiGpuIvfPqIndexParams_t* index_params) + + cuvsError_t cuvsMultiGpuIvfPqIndexParamsDestroy( + cuvsMultiGpuIvfPqIndexParams_t index_params) + + cuvsError_t cuvsMultiGpuIvfPqSearchParamsCreate( + cuvsMultiGpuIvfPqSearchParams_t* params) + + cuvsError_t cuvsMultiGpuIvfPqSearchParamsDestroy( + cuvsMultiGpuIvfPqSearchParams_t params) + + cuvsError_t cuvsMultiGpuIvfPqIndexCreate(cuvsMultiGpuIvfPqIndex_t* index) + + cuvsError_t cuvsMultiGpuIvfPqIndexDestroy(cuvsMultiGpuIvfPqIndex_t index) + + cuvsError_t cuvsMultiGpuIvfPqBuild(cuvsResources_t res, + cuvsMultiGpuIvfPqIndexParams_t params, + DLManagedTensor* dataset_tensor, + cuvsMultiGpuIvfPqIndex_t index) except + + + cuvsError_t cuvsMultiGpuIvfPqSearch( + cuvsResources_t res, + cuvsMultiGpuIvfPqSearchParams_t params, + cuvsMultiGpuIvfPqIndex_t index, + DLManagedTensor* queries_tensor, + DLManagedTensor* neighbors_tensor, + DLManagedTensor* distances_tensor) except + + + cuvsError_t cuvsMultiGpuIvfPqExtend( + cuvsResources_t res, + cuvsMultiGpuIvfPqIndex_t index, + DLManagedTensor* new_vectors_tensor, + DLManagedTensor* new_indices_tensor) except + + + cuvsError_t cuvsMultiGpuIvfPqSerialize( + cuvsResources_t res, + cuvsMultiGpuIvfPqIndex_t index, + const char* filename) except + + + cuvsError_t cuvsMultiGpuIvfPqDeserialize( + cuvsResources_t res, + const char* filename, + cuvsMultiGpuIvfPqIndex_t index) except + + + cuvsError_t cuvsMultiGpuIvfPqDistribute( + cuvsResources_t res, + const char* filename, + cuvsMultiGpuIvfPqIndex_t index) except + + + +cdef class IndexParams(SingleGpuIndexParams): + cdef cuvsMultiGpuIvfPqIndexParams_t mg_params + +cdef class SearchParams(SingleGpuSearchParams): + cdef cuvsMultiGpuIvfPqSearchParams_t mg_params + +cdef class Index: + cdef cuvsMultiGpuIvfPqIndex_t mg_index + cdef bool mg_trained diff --git a/python/cuvs/cuvs/neighbors/mg/ivf_pq/ivf_pq.pyx b/python/cuvs/cuvs/neighbors/mg/ivf_pq/ivf_pq.pyx new file mode 100644 index 0000000000..6e137ce492 --- /dev/null +++ b/python/cuvs/cuvs/neighbors/mg/ivf_pq/ivf_pq.pyx @@ -0,0 +1,572 @@ +# +# Copyright (c) 2025, NVIDIA CORPORATION. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +# cython: language_level=3 + +import numpy as np + +from libc.stdint cimport uint32_t +from libcpp.string cimport string + +from pylibraft.common import auto_convert_output +from pylibraft.common.cai_wrapper import wrap_array +from pylibraft.common.interruptible import cuda_interruptible + +from cuvs.common.exceptions import check_cuvs +from cuvs.common.mg_resources import auto_sync_multi_gpu_resources +from cuvs.neighbors.common import _check_input_array, _check_memory_location + +from cuvs.common cimport cydlpack +from cuvs.common.c_api cimport cuvsResources_t +from cuvs.neighbors.ivf_pq.ivf_pq cimport ( + IndexParams as SingleGpuIndexParams, + SearchParams as SingleGpuSearchParams, + cuvsIvfPqIndexParams_t, + cuvsIvfPqIndexParamsDestroy, + cuvsIvfPqSearchParams_t, + cuvsIvfPqSearchParamsDestroy, +) + +from .ivf_pq cimport ( + cuvsMultiGpuDistributionMode, + cuvsMultiGpuIvfPqBuild, + cuvsMultiGpuIvfPqDeserialize, + cuvsMultiGpuIvfPqDistribute, + cuvsMultiGpuIvfPqExtend, + cuvsMultiGpuIvfPqIndex, + cuvsMultiGpuIvfPqIndex_t, + cuvsMultiGpuIvfPqIndexCreate, + cuvsMultiGpuIvfPqIndexDestroy, + cuvsMultiGpuIvfPqIndexParams, + cuvsMultiGpuIvfPqIndexParams_t, + cuvsMultiGpuIvfPqIndexParamsCreate, + cuvsMultiGpuIvfPqIndexParamsDestroy, + cuvsMultiGpuIvfPqSearch, + cuvsMultiGpuIvfPqSearchParams, + cuvsMultiGpuIvfPqSearchParams_t, + cuvsMultiGpuIvfPqSearchParamsCreate, + cuvsMultiGpuIvfPqSearchParamsDestroy, + cuvsMultiGpuIvfPqSerialize, + cuvsMultiGpuReplicatedSearchMode, + cuvsMultiGpuShardedMergeMode, +) + + +cdef class IndexParams(SingleGpuIndexParams): + """ + Parameters to build multi-GPU IVF-PQ index for efficient search. + Extends single-GPU IndexParams with multi-GPU specific parameters. + + Parameters + ---------- + distribution_mode : str, default = "sharded" + Distribution mode for multi-GPU setup. + Valid values: ["replicated", "sharded"] + **kwargs : Additional parameters passed to single-GPU IndexParams + """ + + def __cinit__(self): + # Base class __cinit__ has already created self.params + # We need to destroy it and use our embedded params instead + if self.params != NULL: + check_cuvs(cuvsIvfPqIndexParamsDestroy(self.params)) + + # Create multi-GPU params which includes embedded base params + check_cuvs(cuvsMultiGpuIvfPqIndexParamsCreate(&self.mg_params)) + # Replace base pointer with embedded base params + self.params = self.mg_params.base_params + + def __dealloc__(self): + # Only destroy the mg_params, which will handle base_params cleanup + check_cuvs(cuvsMultiGpuIvfPqIndexParamsDestroy(self.mg_params)) + self.mg_params = NULL + self.params = NULL + + def __init__(self, *, distribution_mode="sharded", **kwargs): + super().__init__(**kwargs) + if distribution_mode == "replicated": + self.mg_params.mode = CUVS_NEIGHBORS_MG_REPLICATED + elif distribution_mode == "sharded": + self.mg_params.mode = CUVS_NEIGHBORS_MG_SHARDED + else: + raise ValueError( + "distribution_mode must be 'replicated' or 'sharded'") + + def get_handle(self): + return self.mg_params + + @property + def distribution_mode(self): + return ("replicated" if self.mg_params.mode == + CUVS_NEIGHBORS_MG_REPLICATED else "sharded") + + +cdef class Index: + """ + Multi-GPU IVF-PQ index object. Stores the trained multi-GPU IVF-PQ + index state which can be used to perform nearest neighbors searches + across multiple GPUs. + """ + + def __cinit__(self): + # Initialize multi-GPU index + check_cuvs(cuvsMultiGpuIvfPqIndexCreate(&self.mg_index)) + # Initialize multi-GPU trained state + self.mg_trained = False + + def __dealloc__(self): + check_cuvs(cuvsMultiGpuIvfPqIndexDestroy(self.mg_index)) + + def __repr__(self): + return "Index(type=MultiGpuIvfPq)" + + @property + def trained(self): + return self.mg_trained + + +@auto_sync_multi_gpu_resources +def build(IndexParams index_params, dataset, resources=None): + """ + Build the multi-GPU IVF-PQ index from the dataset for efficient search. + + Parameters + ---------- + index_params : :py:class:`cuvs.neighbors.ivf_pq.IndexParams` + dataset : Array interface compliant matrix shape (n_samples, dim) + Supported dtype [float32, float16, int8, uint8] + **IMPORTANT**: For multi-GPU IVF-PQ, the dataset MUST be + in host memory (CPU). + If using CuPy/device arrays, transfer to host with array.get() + or cp.asnumpy(array). + {resources_docstring} + + Returns + ------- + index: py:class:`cuvs.neighbors.ivf_pq.Index` + + Examples + -------- + + >>> import numpy as np + >>> from cuvs.neighbors.mg import ivf_pq + >>> n_samples = 50000 + >>> n_features = 50 + >>> n_queries = 1000 + >>> k = 10 + >>> # For multi-GPU IVF-PQ, use host (NumPy) arrays + >>> dataset = np.random.random_sample((n_samples, n_features)).astype( + ... np.float32) + >>> build_params = ivf_pq.IndexParams(metric="sqeuclidean") + >>> index = ivf_pq.build(build_params, dataset) + >>> distances, neighbors = ivf_pq.search( + ... ivf_pq.SearchParams(), + ... index, dataset, k) + >>> # Results are already in host memory (NumPy arrays) + """ + + dataset_ai = wrap_array(dataset) + _check_input_array(dataset_ai, [np.dtype('float32'), np.dtype('float16'), + np.dtype('byte'), np.dtype('ubyte')]) + + # Multi-GPU IVF-PQ requires dataset in host memory + _check_memory_location(dataset, expected_host=True, name="dataset") + + cdef Index idx = Index() + cdef cydlpack.DLManagedTensor* dataset_dlpack = \ + cydlpack.dlpack_c(dataset_ai) + cdef cuvsMultiGpuIvfPqIndexParams_t params = index_params.mg_params + + cdef cuvsResources_t res = resources.get_c_obj() + + # Build the multi-GPU index + with cuda_interruptible(): + check_cuvs(cuvsMultiGpuIvfPqBuild(res, params, + dataset_dlpack, + idx.mg_index)) + idx.mg_trained = True + + return idx + + +cdef class SearchParams(SingleGpuSearchParams): + """ + Parameters to search multi-GPU IVF-PQ index. + """ + + def __cinit__(self): + # Base class __cinit__ has already created self.params + # We need to destroy it and use our embedded params instead + if self.params != NULL: + check_cuvs(cuvsIvfPqSearchParamsDestroy(self.params)) + + # Create multi-GPU search params which includes embedded base params + check_cuvs(cuvsMultiGpuIvfPqSearchParamsCreate(&self.mg_params)) + # Replace base pointer with embedded base params + self.params = self.mg_params.base_params + + def __dealloc__(self): + # Only destroy the mg_params, which will handle base_params cleanup + check_cuvs(cuvsMultiGpuIvfPqSearchParamsDestroy(self.mg_params)) + self.mg_params = NULL + self.params = NULL + + def __init__(self, *, n_probes=20, search_mode="load_balancer", + merge_mode="merge_on_root_rank", n_rows_per_batch=1000, + **kwargs): + super().__init__(n_probes=n_probes, **kwargs) + # Use the property setters for consistent validation + self.search_mode = search_mode + self.merge_mode = merge_mode + self.n_rows_per_batch = n_rows_per_batch + + def get_handle(self): + return self.mg_params + + @property + def search_mode(self): + """Get the search mode for multi-GPU search.""" + return ("load_balancer" if self.mg_params.search_mode == + CUVS_NEIGHBORS_MG_LOAD_BALANCER else "round_robin") + + @search_mode.setter + def search_mode(self, value): + """Set the search mode for multi-GPU search.""" + if value == "load_balancer": + self.mg_params.search_mode = CUVS_NEIGHBORS_MG_LOAD_BALANCER + elif value == "round_robin": + self.mg_params.search_mode = CUVS_NEIGHBORS_MG_ROUND_ROBIN + else: + raise ValueError( + "search_mode must be 'load_balancer' or 'round_robin'") + + @property + def merge_mode(self): + """Get the merge mode for multi-GPU search.""" + return ("merge_on_root_rank" if self.mg_params.merge_mode == + CUVS_NEIGHBORS_MG_MERGE_ON_ROOT_RANK else "tree_merge") + + @merge_mode.setter + def merge_mode(self, value): + """Set the merge mode for multi-GPU search.""" + if value == "merge_on_root_rank": + self.mg_params.merge_mode = CUVS_NEIGHBORS_MG_MERGE_ON_ROOT_RANK + elif value == "tree_merge": + self.mg_params.merge_mode = CUVS_NEIGHBORS_MG_TREE_MERGE + else: + raise ValueError( + "merge_mode must be 'merge_on_root_rank' or 'tree_merge'") + + @property + def n_rows_per_batch(self): + """Get the number of rows per batch for multi-GPU search.""" + return self.mg_params.n_rows_per_batch + + @n_rows_per_batch.setter + def n_rows_per_batch(self, value): + """Set the number of rows per batch for multi-GPU search.""" + if not isinstance(value, int) or value <= 0: + raise ValueError("n_rows_per_batch must be a positive integer") + self.mg_params.n_rows_per_batch = value + + +@auto_sync_multi_gpu_resources +@auto_convert_output +def search(SearchParams search_params, Index index, queries, + k, neighbors=None, distances=None, resources=None): + """ + Search the multi-GPU IVF-PQ index for the k-nearest neighbors + of each query. + + Parameters + ---------- + search_params : :py:class:`cuvs.neighbors.ivf_pq.SearchParams` + index : :py:class:`cuvs.neighbors.ivf_pq.Index` + queries : Array interface compliant matrix shape (n_queries, dim) + Supported dtype [float32, float16, int8, uint8] + **IMPORTANT**: For multi-GPU IVF-PQ, queries MUST be + in host memory (CPU). + If using CuPy/device arrays, transfer to host with array.get() + or cp.asnumpy(array). + k : int + The number of neighbors to search for each query. + neighbors : Array interface compliant matrix shape (n_queries, k), optional + If provided, this array will be filled with the indices of + the k-nearest neighbors. + If not provided, a new host array will be allocated. + **IMPORTANT**: Must be in host memory (CPU) for multi-GPU IVF-PQ. + distances : Array interface compliant matrix shape (n_queries, k), optional + If provided, this array will be filled with the distances to + the k-nearest neighbors. + If not provided, a new host array will be allocated. + **IMPORTANT**: Must be in host memory (CPU) for multi-GPU IVF-PQ. + {resources_docstring} + + Returns + ------- + distances : numpy.ndarray + The distances to the k-nearest neighbors for each query + (in host memory). + neighbors : numpy.ndarray + The indices of the k-nearest neighbors for each query + (in host memory). + + Examples + -------- + + >>> import numpy as np + >>> from cuvs.neighbors.mg import ivf_pq + >>> n_samples = 50000 + >>> n_features = 50 + >>> n_queries = 1000 + >>> k = 10 + >>> # For multi-GPU IVF-PQ, use host (NumPy) arrays + >>> dataset = np.random.random_sample((n_samples, n_features)).astype( + ... np.float32) + >>> queries = np.random.random_sample((n_queries, n_features)).astype( + ... np.float32) + >>> build_params = ivf_pq.IndexParams(metric="sqeuclidean") + >>> index = ivf_pq.build(build_params, dataset) + >>> distances, neighbors = ivf_pq.search(ivf_pq.SearchParams(), + ... index, queries, k) + >>> # Results are already in host memory (NumPy arrays) + """ + + if not index.trained: + raise ValueError("Index needs to be built before searching") + + queries_ai = wrap_array(queries) + _check_input_array(queries_ai, [np.dtype('float32'), np.dtype('float16'), + np.dtype('byte'), np.dtype('ubyte')]) + + # Multi-GPU IVF-PQ requires queries in host memory + _check_memory_location(queries, expected_host=True, name="queries") + + # Get resources + cdef cuvsResources_t res = resources.get_c_obj() + + # Prepare output arrays + cdef uint32_t n_queries = queries.shape[0] + if neighbors is None: + # For multi-GPU, create host arrays instead of device arrays + neighbors = np.empty((n_queries, k), dtype='int64') + if distances is None: + # For multi-GPU, create host arrays instead of device arrays + distances = np.empty((n_queries, k), dtype='float32') + + neighbors_ai = wrap_array(neighbors) + _check_input_array(neighbors_ai, [np.dtype('int64')], + exp_rows=n_queries, exp_cols=k) + distances_ai = wrap_array(distances) + _check_input_array(distances_ai, [np.dtype('float32')], + exp_rows=n_queries, exp_cols=k) + + # Multi-GPU IVF-PQ requires output arrays in host memory + _check_memory_location(neighbors, expected_host=True, name="neighbors") + _check_memory_location(distances, expected_host=True, name="distances") + + cdef cydlpack.DLManagedTensor* queries_dlpack = \ + cydlpack.dlpack_c(queries_ai) + cdef cydlpack.DLManagedTensor* neighbors_dlpack = \ + cydlpack.dlpack_c(neighbors_ai) + cdef cydlpack.DLManagedTensor* distances_dlpack = \ + cydlpack.dlpack_c(distances_ai) + + # Perform search + with cuda_interruptible(): + check_cuvs(cuvsMultiGpuIvfPqSearch(res, search_params.mg_params, + index.mg_index, queries_dlpack, + neighbors_dlpack, distances_dlpack)) + + return (distances, neighbors) + + +@auto_sync_multi_gpu_resources +def extend(Index index, new_vectors, new_indices=None, + resources=None): + """ + Extend the multi-GPU IVF-PQ index with new vectors. + + Parameters + ---------- + index : :py:class:`cuvs.neighbors.ivf_pq.Index` + new_vectors : Array interface compliant matrix shape (n_new_vectors, dim) + Supported dtype [float32, float16, int8, uint8] + **IMPORTANT**: For multi-GPU IVF-PQ, new_vectors MUST be + in host memory (CPU). + If using CuPy/device arrays, transfer to host with array.get() + or cp.asnumpy(array). + new_indices : Array interface compliant matrix shape (n_new_vectors,) + , optional + If provided, these indices will be used for the new vectors. + If not provided, indices will be automatically assigned. + **IMPORTANT**: Must be in host memory (CPU) for multi-GPU IVF-PQ. + {resources_docstring} + + Examples + -------- + + >>> import numpy as np + >>> from cuvs.neighbors.mg import ivf_pq + >>> n_samples = 50000 + >>> n_features = 50 + >>> n_new_vectors = 1000 + >>> # For multi-GPU IVF-PQ, use host (NumPy) arrays + >>> dataset = np.random.random_sample((n_samples, n_features)).astype( + ... np.float32) + >>> new_vectors = np.random.random_sample( + ... (n_new_vectors, n_features)).astype(np.float32) + >>> new_indices = np.arange(n_samples, n_new_vectors, dtype=np.int64) + >>> build_params = ivf_pq.IndexParams(metric="sqeuclidean") + >>> index = ivf_pq.build(build_params, dataset) + >>> ivf_pq.extend(index, new_vectors, new_indices) + """ + + if not index.trained: + raise ValueError("Index needs to be built before extending") + + new_vectors_ai = wrap_array(new_vectors) + _check_input_array(new_vectors_ai, + [np.dtype('float32'), np.dtype('float16'), + np.dtype('byte'), np.dtype('ubyte')]) + + # Multi-GPU IVF-PQ requires new_vectors in host memory + _check_memory_location(new_vectors, expected_host=True, name="new_vectors") + + # Get resources + cdef cuvsResources_t res = resources.get_c_obj() + + cdef cydlpack.DLManagedTensor* new_vectors_dlpack = \ + cydlpack.dlpack_c(new_vectors_ai) + cdef cydlpack.DLManagedTensor* new_indices_dlpack = NULL + + if new_indices is not None: + new_indices_ai = wrap_array(new_indices) + _check_input_array(new_indices_ai, [np.dtype('int64')]) + # Multi-GPU IVF-PQ requires new_indices in host memory + _check_memory_location(new_indices, expected_host=True, + name="new_indices") + new_indices_dlpack = cydlpack.dlpack_c(new_indices_ai) + + with cuda_interruptible(): + check_cuvs(cuvsMultiGpuIvfPqExtend(res, index.mg_index, + new_vectors_dlpack, + new_indices_dlpack)) + + +@auto_sync_multi_gpu_resources +def save(Index index, filename, resources=None): + """ + Serialize the multi-GPU IVF-PQ index to a file. + + Parameters + ---------- + index : :py:class:`cuvs.neighbors.ivf_pq.Index` + filename : str + The filename to serialize the index to. + {resources_docstring} + + Examples + -------- + + >>> import numpy as np + >>> from cuvs.neighbors.mg import ivf_pq + >>> n_samples = 50000 + >>> n_features = 50 + >>> # For multi-GPU IVF-PQ, use host (NumPy) arrays + >>> dataset = np.random.random_sample((n_samples, n_features)).astype( + ... np.float32) + >>> build_params = ivf_pq.IndexParams(metric="sqeuclidean") + >>> index = ivf_pq.build(build_params, dataset) + >>> ivf_pq.save(index, "index.bin") + """ + + if not index.trained: + raise ValueError("Index needs to be built before serializing") + + # Get resources + cdef cuvsResources_t res = resources.get_c_obj() + + cdef string filename_str = filename.encode('utf-8') + check_cuvs(cuvsMultiGpuIvfPqSerialize(res, index.mg_index, + filename_str.c_str())) + + +@auto_sync_multi_gpu_resources +def load(filename, resources=None): + """ + Deserialize the multi-GPU IVF-PQ index from a file. + + Parameters + ---------- + filename : str + The filename to deserialize the index from. + {resources_docstring} + + Returns + ------- + index : Index + The deserialized index. + + Examples + -------- + + >>> from cuvs.neighbors.mg import ivf_pq + >>> index = ivf_pq.load("index.bin") # doctest: +SKIP + """ + + cdef Index index = Index() + cdef cuvsResources_t res = resources.get_c_obj() + + cdef string filename_str = filename.encode('utf-8') + check_cuvs(cuvsMultiGpuIvfPqDeserialize(res, filename_str.c_str(), + index.mg_index)) + index.mg_trained = True + + return index + + +@auto_sync_multi_gpu_resources +def distribute(filename, resources=None): + """ + Distribute a single-GPU IVF-PQ index across multiple GPUs from a file. + + Parameters + ---------- + filename : str + The filename to distribute the index from. + {resources_docstring} + + Returns + ------- + index : Index + The distributed index. + + Examples + -------- + + >>> from cuvs.neighbors.mg import ivf_pq + >>> index = ivf_pq.distribute("single_gpu_index.bin") # doctest: +SKIP + """ + + cdef Index index = Index() + cdef cuvsResources_t res = resources.get_c_obj() + + cdef string filename_str = filename.encode('utf-8') + check_cuvs(cuvsMultiGpuIvfPqDistribute(res, filename_str.c_str(), + index.mg_index)) + index.mg_trained = True + + return index diff --git a/python/cuvs/cuvs/tests/test_mg_cagra.py b/python/cuvs/cuvs/tests/test_mg_cagra.py new file mode 100644 index 0000000000..16d40f9c17 --- /dev/null +++ b/python/cuvs/cuvs/tests/test_mg_cagra.py @@ -0,0 +1,608 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +import os +import tempfile + +import numpy as np +import pytest +from sklearn.neighbors import NearestNeighbors +from sklearn.preprocessing import normalize + +from cuvs.common import MultiGpuResources +from cuvs.neighbors.mg import cagra as mg_cagra +from cuvs.tests.ann_utils import calc_recall, generate_data + + +# Check if multi-GPU functionality is available +def has_multiple_gpus(): + """Check if system has multiple GPUs available.""" + try: + import cupy as cp + + return cp.cuda.runtime.getDeviceCount() > 1 + except Exception: + return False + + +# Mark tests that require multiple GPUs +requires_multiple_gpus = pytest.mark.skipif( + not has_multiple_gpus(), reason="Multi-GPU tests require multiple GPUs" +) + + +def run_mg_cagra_build_search_test( + n_rows=10000, + n_cols=10, + n_queries=100, + k=10, + dtype=np.float32, + metric="sqeuclidean", # CAGRA only supports sqeuclidean and inner_product + distribution_mode="sharded", + search_mode="load_balancer", + merge_mode="tree_merge", + n_rows_per_batch=1000, + compare=True, + search_params=None, + graph_degree=64, + intermediate_graph_degree=128, +): + """ + Run a multi-GPU CAGRA build and search test. + + Note: Multi-GPU CAGRA requires host memory arrays (NumPy), not device + arrays. + """ + # Generate host memory arrays (NumPy) + dataset = generate_data((n_rows, n_cols), dtype) + if metric == "inner_product": + dataset = normalize(dataset, norm="l2", axis=1) + + queries = generate_data((n_queries, n_cols), dtype) + if metric == "inner_product": + queries = normalize(queries, norm="l2", axis=1) + + # Multi-GPU resources + resources = MultiGpuResources() + + # Build parameters + build_params = mg_cagra.IndexParams( + metric=metric, + distribution_mode=distribution_mode, + graph_degree=graph_degree, + intermediate_graph_degree=intermediate_graph_degree, + ) + + # Build index + index = mg_cagra.build(build_params, dataset, resources=resources) + assert index.trained + + # Search parameters + if search_params is None: + search_params = {} + search_params_obj = mg_cagra.SearchParams( + search_mode=search_mode, + merge_mode=merge_mode, + n_rows_per_batch=n_rows_per_batch, + **search_params, + ) + + # Perform search + distances, neighbors = mg_cagra.search( + search_params_obj, + index, + queries, + k, + resources=resources, + ) + + # Verify results are in host memory (NumPy arrays) + assert isinstance(distances, np.ndarray) + assert isinstance(neighbors, np.ndarray) + assert distances.shape == (n_queries, k) + assert neighbors.shape == (n_queries, k) + + if not compare: + return distances, neighbors + + # Calculate reference values with sklearn + skl_metric = { + "sqeuclidean": "sqeuclidean", + "inner_product": "cosine", + }[metric] + + nn_skl = NearestNeighbors( + n_neighbors=k, algorithm="brute", metric=skl_metric + ) + nn_skl.fit(dataset) + skl_idx = nn_skl.kneighbors(queries, return_distance=False) + + recall = calc_recall(neighbors, skl_idx) + # Multi-GPU implementation may have lower recall due to data + # distribution across GPUs + # This is acceptable as long as the functionality works correctly + assert recall > 0.3, f"Recall too low: {recall:.3f}" + + return distances, neighbors + + +@requires_multiple_gpus +@pytest.mark.parametrize("dtype", [np.float32]) +@pytest.mark.parametrize( + "metric", ["sqeuclidean"] +) # Start with just sqeuclidean +@pytest.mark.parametrize( + "distribution_mode", ["sharded"] +) # Start with just sharded +def test_mg_cagra_basic(dtype, metric, distribution_mode): + """Test basic multi-GPU CAGRA build and search functionality.""" + run_mg_cagra_build_search_test( + n_rows=2000, # Use smaller dataset for more reliable tests + n_cols=8, + n_queries=20, + k=5, + dtype=dtype, + metric=metric, + distribution_mode=distribution_mode, + graph_degree=32, # Smaller graph for faster tests + intermediate_graph_degree=64, + ) + + +@requires_multiple_gpus +@pytest.mark.parametrize( + "metric", ["inner_product"] +) # Only test supported metrics +@pytest.mark.parametrize("distribution_mode", ["replicated"]) +def test_mg_cagra_additional_metrics(metric, distribution_mode): + """Test additional metrics and distribution modes.""" + run_mg_cagra_build_search_test( + n_rows=2000, + n_cols=8, + n_queries=20, + k=5, + dtype=np.float32, + metric=metric, + distribution_mode=distribution_mode, + graph_degree=32, + intermediate_graph_degree=64, + ) + + +@requires_multiple_gpus +@pytest.mark.parametrize("dtype", [np.float32, np.float16, np.int8, np.uint8]) +def test_mg_cagra_dtypes(dtype): + """Test multi-GPU CAGRA with different data types.""" + run_mg_cagra_build_search_test( + n_rows=1500, + n_cols=8, + n_queries=15, + k=5, + dtype=dtype, + metric="sqeuclidean", + graph_degree=32, + intermediate_graph_degree=64, + ) + + +@requires_multiple_gpus +@pytest.mark.parametrize("distribution_mode", ["sharded", "replicated"]) +def test_mg_cagra_distribution_modes(distribution_mode): + """Test different distribution modes for multi-GPU CAGRA.""" + run_mg_cagra_build_search_test( + n_rows=1500, + n_cols=8, + n_queries=15, + k=5, + distribution_mode=distribution_mode, + graph_degree=32, + intermediate_graph_degree=64, + ) + + +@requires_multiple_gpus +@pytest.mark.parametrize("search_mode", ["load_balancer", "round_robin"]) +@pytest.mark.parametrize("merge_mode", ["merge_on_root_rank", "tree_merge"]) +def test_mg_cagra_search_params(search_mode, merge_mode): + """Test different multi-GPU search parameters.""" + run_mg_cagra_build_search_test( + n_rows=1500, + n_cols=8, + n_queries=15, + k=5, + search_mode=search_mode, + merge_mode=merge_mode, + n_rows_per_batch=500, + graph_degree=32, + intermediate_graph_degree=64, + ) + + +@requires_multiple_gpus +@pytest.mark.parametrize( + "metric", ["sqeuclidean"] +) # Only test supported metrics +def test_mg_cagra_metrics(metric): + """Test different distance metrics for multi-GPU CAGRA.""" + run_mg_cagra_build_search_test( + n_rows=1500, + n_cols=8, + n_queries=15, + k=5, + metric=metric, + graph_degree=32, + intermediate_graph_degree=64, + ) + + +@requires_multiple_gpus +def test_mg_cagra_serialize(): + """Test save/load functionality for multi-GPU CAGRA.""" + n_rows, n_cols = 2000, 8 + k = 5 + + # Generate data + dataset = generate_data((n_rows, n_cols), np.float32) + queries = generate_data((20, n_cols), np.float32) + + resources = MultiGpuResources() + + # Build original index + build_params = mg_cagra.IndexParams( + graph_degree=32, intermediate_graph_degree=64 + ) + original_index = mg_cagra.build(build_params, dataset, resources=resources) + + # Search with original index + search_params = mg_cagra.SearchParams(itopk_size=32) + orig_distances, orig_neighbors = mg_cagra.search( + search_params, original_index, queries, k, resources=resources + ) + + # Save index to temporary file + with tempfile.NamedTemporaryFile(suffix=".bin", delete=False) as f: + temp_filename = f.name + + try: + mg_cagra.save(original_index, temp_filename, resources=resources) + + # Load index from file + loaded_index = mg_cagra.load(temp_filename, resources=resources) + assert loaded_index.trained + + # Search with loaded index + loaded_distances, loaded_neighbors = mg_cagra.search( + search_params, loaded_index, queries, k, resources=resources + ) + + # Results should be identical + np.testing.assert_array_equal(orig_neighbors, loaded_neighbors) + np.testing.assert_allclose(orig_distances, loaded_distances, rtol=1e-6) + + finally: + if os.path.exists(temp_filename): + os.unlink(temp_filename) + + +@requires_multiple_gpus +def test_mg_cagra_distribute(): + """Test distribute functionality for multi-GPU CAGRA.""" + # Note: Distribute is for replicating a single-GPU index across + # multiple GPUs. + # This test builds a single-GPU index, serializes it, then distributes it. + + n_rows, n_cols = 2000, 8 + k = 5 + + # Generate data + dataset = generate_data((n_rows, n_cols), np.float32) + queries = generate_data((20, n_cols), np.float32) + + # Import single-GPU CAGRA to build and serialize a single-GPU index + from cuvs.common import Resources + from cuvs.neighbors import cagra + + # Build single-GPU index first + single_gpu_resources = Resources() + single_build_params = cagra.IndexParams( + metric="sqeuclidean", graph_degree=32, intermediate_graph_degree=64 + ) + + # Convert to device arrays for single-GPU build + try: + import cupy as cp + + device_dataset = cp.asarray(dataset) + single_index = cagra.build( + single_build_params, device_dataset, resources=single_gpu_resources + ) + except ImportError: + pytest.skip("CuPy not available for single-GPU index building") + + with tempfile.NamedTemporaryFile(suffix=".bin", delete=False) as f: + temp_filename = f.name + + try: + # Serialize single-GPU index + cagra.save(temp_filename, single_index, resources=single_gpu_resources) + + # Now distribute the single-GPU index across multiple GPUs + resources = MultiGpuResources() + distributed_index = mg_cagra.distribute( + temp_filename, resources=resources + ) + assert distributed_index.trained + + # Search should work with distributed index (using host memory arrays) + search_params = mg_cagra.SearchParams(itopk_size=32) + distances, neighbors = mg_cagra.search( + search_params, distributed_index, queries, k, resources=resources + ) + + assert distances.shape == (20, k) + assert neighbors.shape == (20, k) + + finally: + if os.path.exists(temp_filename): + os.unlink(temp_filename) + + +def test_memory_location_validation(): + """Test that multi-GPU CAGRA validates memory locations correctly.""" + try: + import cupy as cp + except ImportError: + pytest.skip("CuPy not available for memory location tests") + + n_rows, n_cols = 1500, 8 + + # Create host and device arrays + host_data = generate_data((n_rows, n_cols), np.float32) + device_data = cp.asarray(host_data) + + resources = MultiGpuResources() + build_params = mg_cagra.IndexParams( + graph_degree=32, intermediate_graph_degree=64 + ) + + # Test that device arrays are rejected for build + with pytest.raises(ValueError, match="host memory"): + mg_cagra.build(build_params, device_data, resources=resources) + + # Test that host arrays work for build + index = mg_cagra.build(build_params, host_data, resources=resources) + + # Test that device arrays are rejected for search + queries = generate_data((20, n_cols), np.float32) + device_queries = cp.asarray(queries) + search_params = mg_cagra.SearchParams(itopk_size=32) + + with pytest.raises(ValueError, match="host memory"): + mg_cagra.search( + search_params, index, device_queries, 5, resources=resources + ) + + # Test that host arrays work for search + distances, neighbors = mg_cagra.search( + search_params, index, queries, 5, resources=resources + ) + assert isinstance(distances, np.ndarray) + assert isinstance(neighbors, np.ndarray) + + +def test_parameter_validation(): + """Test parameter validation for multi-GPU CAGRA.""" + # Test invalid distribution mode + with pytest.raises(ValueError, match="distribution_mode must be"): + mg_cagra.IndexParams(distribution_mode="invalid") + + # Test invalid search mode + with pytest.raises(ValueError, match="search_mode must be"): + mg_cagra.SearchParams(search_mode="invalid") + + # Test invalid merge mode + with pytest.raises(ValueError, match="merge_mode must be"): + mg_cagra.SearchParams(merge_mode="invalid") + + +def test_parameter_properties(): + """Test that parameters can be accessed via properties.""" + # Test IndexParams properties + params = mg_cagra.IndexParams(distribution_mode="replicated") + assert params.distribution_mode == "replicated" + + params = mg_cagra.IndexParams(distribution_mode="sharded") + assert params.distribution_mode == "sharded" + + # Test SearchParams creation with different parameters + mg_cagra.SearchParams( + search_mode="round_robin", + merge_mode="tree_merge", + n_rows_per_batch=2000, + ) + # These don't have properties exposed, but creation should work + + +def test_untrained_index_error(): + """Test that using an untrained index raises appropriate errors.""" + resources = MultiGpuResources() + + # Create untrained index + index = mg_cagra.Index() + assert not index.trained + + queries = generate_data((100, 10), np.float32) + search_params = mg_cagra.SearchParams() + + # Test that search on untrained index fails + with pytest.raises(ValueError, match="Index needs to be built"): + mg_cagra.search(search_params, index, queries, 10, resources=resources) + + # Test that save on untrained index fails + with pytest.raises(ValueError, match="Index needs to be built"): + mg_cagra.save(index, "temp.bin", resources=resources) + + +@requires_multiple_gpus +def test_mg_cagra_with_prealloc_output(): + """Test multi-GPU CAGRA search with pre-allocated output arrays.""" + n_rows, n_cols = 1500, 8 + n_queries = 20 + k = 5 + + # Generate data in host memory + dataset = generate_data((n_rows, n_cols), np.float32) + queries = generate_data((n_queries, n_cols), np.float32) + + resources = MultiGpuResources() + + # Build index + build_params = mg_cagra.IndexParams( + graph_degree=32, intermediate_graph_degree=64 + ) + index = mg_cagra.build(build_params, dataset, resources=resources) + + # Pre-allocate output arrays in host memory + neighbors = np.empty((n_queries, k), dtype=np.int64) + distances = np.empty((n_queries, k), dtype=np.float32) + + # Search with pre-allocated arrays + search_params = mg_cagra.SearchParams(itopk_size=32) + ret_distances, ret_neighbors = mg_cagra.search( + search_params, + index, + queries, + k, + neighbors=neighbors, + distances=distances, + resources=resources, + ) + + # Should return the same arrays we passed in + assert ret_distances is distances + assert ret_neighbors is neighbors + assert distances.shape == (n_queries, k) + assert neighbors.shape == (n_queries, k) + + +def test_index_repr(): + """Test string representation of Index.""" + index = mg_cagra.Index() + assert repr(index) == "Index(type=MultiGpuCagra)" + + +def test_mg_cagra_simple(): + """Simple test to validate multi-GPU CAGRA works with very favorable + parameters. + """ + if not has_multiple_gpus(): + pytest.skip("Multi-GPU tests require multiple GPUs") + + # Use simple test case that should definitely work + n_rows, n_cols = 1000, 8 + n_queries, k = 20, 5 + + # Generate data + dataset = generate_data((n_rows, n_cols), np.float32) + queries = generate_data((n_queries, n_cols), np.float32) + + resources = MultiGpuResources() + + # Use small graph for reliable testing + build_params = mg_cagra.IndexParams( + metric="sqeuclidean", + graph_degree=16, + intermediate_graph_degree=32, + ) + + # Build index + index = mg_cagra.build(build_params, dataset, resources=resources) + + # Search with basic parameters + search_params = mg_cagra.SearchParams(itopk_size=16) + distances, neighbors = mg_cagra.search( + search_params, index, queries, k, resources=resources + ) + + # Basic sanity checks + assert distances.shape == (n_queries, k) + assert neighbors.shape == (n_queries, k) + assert isinstance(distances, np.ndarray) + assert isinstance(neighbors, np.ndarray) + + # Check that we get valid neighbors + assert np.all(neighbors >= 0) + assert np.all(neighbors < n_rows) + + # Distances should be non-negative and sorted + assert np.all(distances >= 0) + for i in range(n_queries): + assert np.all( + distances[i, :-1] <= distances[i, 1:] + ), f"Distances not sorted for query {i}" + + +# Integration test with multiple operations +@requires_multiple_gpus +def test_mg_cagra_integration(): + """Integration test covering build, search, and serialization.""" + n_rows, n_cols = 2000, 8 + k = 5 + + # Generate initial dataset + dataset = generate_data((n_rows, n_cols), np.float32) + queries = generate_data((20, n_cols), np.float32) + + resources = MultiGpuResources() + + # Build initial index + build_params = mg_cagra.IndexParams( + distribution_mode="sharded", + metric="sqeuclidean", + graph_degree=32, + intermediate_graph_degree=64, + ) + index = mg_cagra.build(build_params, dataset, resources=resources) + + # Initial search + search_params = mg_cagra.SearchParams( + itopk_size=32, + search_mode="load_balancer", + merge_mode="merge_on_root_rank", + ) + distances1, neighbors1 = mg_cagra.search( + search_params, index, queries, k, resources=resources + ) + + # Save and reload + with tempfile.NamedTemporaryFile(suffix=".bin", delete=False) as f: + temp_filename = f.name + + try: + mg_cagra.save(index, temp_filename, resources=resources) + reloaded_index = mg_cagra.load(temp_filename, resources=resources) + + # Search with reloaded index + distances2, neighbors2 = mg_cagra.search( + search_params, reloaded_index, queries, k, resources=resources + ) + + # Results from reloaded index should match + np.testing.assert_array_equal(neighbors1, neighbors2) + np.testing.assert_allclose(distances1, distances2, rtol=1e-6) + + finally: + if os.path.exists(temp_filename): + os.unlink(temp_filename) diff --git a/python/cuvs/cuvs/tests/test_mg_ivf_flat.py b/python/cuvs/cuvs/tests/test_mg_ivf_flat.py new file mode 100644 index 0000000000..8bec3663c0 --- /dev/null +++ b/python/cuvs/cuvs/tests/test_mg_ivf_flat.py @@ -0,0 +1,650 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +import os +import tempfile + +import numpy as np +import pytest +from sklearn.neighbors import NearestNeighbors +from sklearn.preprocessing import normalize + +from cuvs.common import MultiGpuResources +from cuvs.neighbors.mg import ivf_flat as mg_ivf_flat +from cuvs.tests.ann_utils import calc_recall, generate_data + + +# Check if multi-GPU functionality is available +def has_multiple_gpus(): + """Check if system has multiple GPUs available.""" + try: + import cupy as cp + + return cp.cuda.runtime.getDeviceCount() > 1 + except Exception: + return False + + +# Mark tests that require multiple GPUs +requires_multiple_gpus = pytest.mark.skipif( + not has_multiple_gpus(), reason="Multi-GPU tests require multiple GPUs" +) + + +def run_mg_ivf_flat_build_search_test( + n_rows=10000, + n_cols=10, + n_queries=100, + k=10, + dtype=np.float32, + metric="euclidean", + distribution_mode="sharded", + search_mode="load_balancer", + merge_mode="tree_merge", + n_rows_per_batch=1000, + compare=True, + add_data_on_build=True, + search_params=None, + n_lists=None, +): + """ + Run a multi-GPU IVF-Flat build and search test. + + Note: Multi-GPU IVF-Flat requires host memory arrays (NumPy), not + device arrays. + """ + # Generate host memory arrays (NumPy) + dataset = generate_data((n_rows, n_cols), dtype) + if metric == "inner_product": + dataset = normalize(dataset, norm="l2", axis=1) + + queries = generate_data((n_queries, n_cols), dtype) + if metric == "inner_product": + queries = normalize(queries, norm="l2", axis=1) + + # Multi-GPU resources + resources = MultiGpuResources() + + # Build parameters - use fewer clusters for better recall + # with smaller datasets + if n_lists is None: + # Use fewer clusters for smaller datasets to ensure enough points + # per cluster + n_lists = min(1024, max(64, n_rows // 50)) + + build_params = mg_ivf_flat.IndexParams( + metric=metric, + distribution_mode=distribution_mode, + add_data_on_build=add_data_on_build, + n_lists=n_lists, + ) + + # Build index + index = mg_ivf_flat.build(build_params, dataset, resources=resources) + assert index.trained + + # If not adding data on build, extend the index + if not add_data_on_build: + dataset_1 = dataset[: n_rows // 2, :] + dataset_2 = dataset[n_rows // 2 :, :] + indices_1 = np.arange(n_rows // 2, dtype=np.int64) + indices_2 = np.arange(n_rows // 2, n_rows, dtype=np.int64) + + mg_ivf_flat.extend(index, dataset_1, indices_1, resources=resources) + mg_ivf_flat.extend(index, dataset_2, indices_2, resources=resources) + + # Search parameters + if search_params is None: + search_params = {} + # Use higher n_probes for better recall in multi-GPU setting + if "n_probes" not in search_params: + # Use many clusters for good recall - search majority of clusters + search_params["n_probes"] = min(n_lists, max(20, (n_lists * 3) // 4)) + search_params_obj = mg_ivf_flat.SearchParams( + search_mode=search_mode, + merge_mode=merge_mode, + n_rows_per_batch=n_rows_per_batch, + **search_params, + ) + + # Perform search + distances, neighbors = mg_ivf_flat.search( + search_params_obj, + index, + queries, + k, + resources=resources, + ) + + # Verify results are in host memory (NumPy arrays) + assert isinstance(distances, np.ndarray) + assert isinstance(neighbors, np.ndarray) + assert distances.shape == (n_queries, k) + assert neighbors.shape == (n_queries, k) + + if not compare: + return distances, neighbors + + # Calculate reference values with sklearn + skl_metric = { + "sqeuclidean": "sqeuclidean", + "inner_product": "cosine", + "cosine": "cosine", + "euclidean": "euclidean", + }[metric] + + nn_skl = NearestNeighbors( + n_neighbors=k, algorithm="brute", metric=skl_metric + ) + nn_skl.fit(dataset) + skl_idx = nn_skl.kneighbors(queries, return_distance=False) + + recall = calc_recall(neighbors, skl_idx) + # Multi-GPU implementation may have lower recall due to data distribution + # across GPUs + # This is acceptable as long as the functionality works correctly + assert recall > 0.3, ( + f"Recall too low: {recall:.3f} (n_lists={n_lists}, " + f"n_probes={search_params.get('n_probes', 'default')})" + ) + + return distances, neighbors + + +@requires_multiple_gpus +@pytest.mark.parametrize("dtype", [np.float32]) +@pytest.mark.parametrize( + "metric", ["sqeuclidean"] +) # Start with just sqeuclidean +@pytest.mark.parametrize( + "distribution_mode", ["sharded"] +) # Start with just sharded +def test_mg_ivf_flat_basic(dtype, metric, distribution_mode): + """Test basic multi-GPU IVF-Flat build and search functionality.""" + run_mg_ivf_flat_build_search_test( + n_rows=2000, # Use smaller dataset for more reliable tests + n_cols=8, + n_queries=20, + k=5, + dtype=dtype, + metric=metric, + distribution_mode=distribution_mode, + n_lists=50, # Fixed small number of clusters + ) + + +@requires_multiple_gpus +@pytest.mark.parametrize("metric", ["inner_product", "euclidean", "cosine"]) +@pytest.mark.parametrize("distribution_mode", ["replicated"]) +def test_mg_ivf_flat_additional_metrics(metric, distribution_mode): + """Test additional metrics and distribution modes.""" + run_mg_ivf_flat_build_search_test( + n_rows=2000, + n_cols=8, + n_queries=20, + k=5, + dtype=np.float32, + metric=metric, + distribution_mode=distribution_mode, + n_lists=50, + ) + + +@requires_multiple_gpus +@pytest.mark.parametrize("dtype", [np.float32, np.float16, np.int8, np.uint8]) +def test_mg_ivf_flat_dtypes(dtype): + """Test multi-GPU IVF-Flat with different data types.""" + run_mg_ivf_flat_build_search_test( + n_rows=1500, + n_cols=8, + n_queries=15, + k=5, + dtype=dtype, + metric="sqeuclidean", + n_lists=30, + ) + + +@requires_multiple_gpus +@pytest.mark.parametrize("distribution_mode", ["sharded", "replicated"]) +def test_mg_ivf_flat_distribution_modes(distribution_mode): + """Test different distribution modes for multi-GPU IVF-Flat.""" + run_mg_ivf_flat_build_search_test( + n_rows=1500, + n_cols=8, + n_queries=15, + k=5, + distribution_mode=distribution_mode, + n_lists=30, + ) + + +@requires_multiple_gpus +@pytest.mark.parametrize("search_mode", ["load_balancer", "round_robin"]) +@pytest.mark.parametrize("merge_mode", ["merge_on_root_rank", "tree_merge"]) +def test_mg_ivf_flat_search_params(search_mode, merge_mode): + """Test different multi-GPU search parameters.""" + run_mg_ivf_flat_build_search_test( + n_rows=1500, + n_cols=8, + n_queries=15, + k=5, + search_mode=search_mode, + merge_mode=merge_mode, + n_rows_per_batch=500, + n_lists=30, + ) + + +@requires_multiple_gpus +@pytest.mark.parametrize("metric", ["euclidean", "sqeuclidean"]) +def test_mg_ivf_flat_metrics(metric): + """Test different distance metrics for multi-GPU IVF-Flat.""" + run_mg_ivf_flat_build_search_test( + n_rows=1500, + n_cols=8, + n_queries=15, + k=5, + metric=metric, + n_lists=30, + ) + + +@requires_multiple_gpus +def test_mg_ivf_flat_extend(): + """Test extending multi-GPU IVF-Flat index with new vectors.""" + run_mg_ivf_flat_build_search_test( + n_rows=1500, + n_cols=8, + n_queries=15, + k=5, + add_data_on_build=False, + n_lists=30, + ) + + +@requires_multiple_gpus +def test_mg_ivf_flat_serialize(): + """Test save/load functionality for multi-GPU IVF-Flat.""" + n_rows, n_cols = 2000, 8 + k = 5 + + # Generate data + dataset = generate_data((n_rows, n_cols), np.float32) + queries = generate_data((20, n_cols), np.float32) + + resources = MultiGpuResources() + + # Build original index + build_params = mg_ivf_flat.IndexParams(n_lists=50) + original_index = mg_ivf_flat.build( + build_params, dataset, resources=resources + ) + + # Search with original index + search_params = mg_ivf_flat.SearchParams(n_probes=37) + orig_distances, orig_neighbors = mg_ivf_flat.search( + search_params, original_index, queries, k, resources=resources + ) + + # Save index to temporary file + with tempfile.NamedTemporaryFile(suffix=".bin", delete=False) as f: + temp_filename = f.name + + try: + mg_ivf_flat.save(original_index, temp_filename, resources=resources) + + # Load index from file + loaded_index = mg_ivf_flat.load(temp_filename, resources=resources) + assert loaded_index.trained + + # Search with loaded index + loaded_distances, loaded_neighbors = mg_ivf_flat.search( + search_params, loaded_index, queries, k, resources=resources + ) + + # Results should be identical + np.testing.assert_array_equal(orig_neighbors, loaded_neighbors) + np.testing.assert_allclose(orig_distances, loaded_distances, rtol=1e-6) + + finally: + if os.path.exists(temp_filename): + os.unlink(temp_filename) + + +@requires_multiple_gpus +def test_mg_ivf_flat_distribute(): + """Test distribute functionality for multi-GPU IVF-Flat.""" + # Note: Distribute is for replicating a single-GPU index + # across multiple GPUs. + # This test builds a single-GPU index, serializes it, then distributes it. + + n_rows, n_cols = 2000, 8 + k = 5 + + # Generate data + dataset = generate_data((n_rows, n_cols), np.float32) + queries = generate_data((20, n_cols), np.float32) + + # Import single-GPU IVF-Flat to build and serialize a single-GPU index + from cuvs.common import Resources + from cuvs.neighbors import ivf_flat + + # Build single-GPU index first + single_gpu_resources = Resources() + single_build_params = ivf_flat.IndexParams( + metric="sqeuclidean", n_lists=50 + ) + + # Convert to device arrays for single-GPU build + try: + import cupy as cp + + device_dataset = cp.asarray(dataset) + single_index = ivf_flat.build( + single_build_params, device_dataset, resources=single_gpu_resources + ) + except ImportError: + pytest.skip("CuPy not available for single-GPU index building") + + with tempfile.NamedTemporaryFile(suffix=".bin", delete=False) as f: + temp_filename = f.name + + try: + # Serialize single-GPU index + ivf_flat.save( + temp_filename, single_index, resources=single_gpu_resources + ) + + # Now distribute the single-GPU index across multiple GPUs + resources = MultiGpuResources() + distributed_index = mg_ivf_flat.distribute( + temp_filename, resources=resources + ) + assert distributed_index.trained + + # Search should work with distributed index (using host memory arrays) + search_params = mg_ivf_flat.SearchParams(n_probes=37) + distances, neighbors = mg_ivf_flat.search( + search_params, distributed_index, queries, k, resources=resources + ) + + assert distances.shape == (20, k) + assert neighbors.shape == (20, k) + + finally: + if os.path.exists(temp_filename): + os.unlink(temp_filename) + + +def test_memory_location_validation(): + """Test that multi-GPU IVF-Flat validates memory locations correctly.""" + try: + import cupy as cp + except ImportError: + pytest.skip("CuPy not available for memory location tests") + + n_rows, n_cols = 1500, 8 + + # Create host and device arrays + host_data = generate_data((n_rows, n_cols), np.float32) + device_data = cp.asarray(host_data) + + resources = MultiGpuResources() + build_params = mg_ivf_flat.IndexParams(n_lists=30) + + # Test that device arrays are rejected for build + with pytest.raises(ValueError, match="host memory"): + mg_ivf_flat.build(build_params, device_data, resources=resources) + + # Test that host arrays work for build + index = mg_ivf_flat.build(build_params, host_data, resources=resources) + + # Test that device arrays are rejected for search + queries = generate_data((20, n_cols), np.float32) + device_queries = cp.asarray(queries) + search_params = mg_ivf_flat.SearchParams(n_probes=22) + + with pytest.raises(ValueError, match="host memory"): + mg_ivf_flat.search( + search_params, index, device_queries, 5, resources=resources + ) + + # Test that host arrays work for search + distances, neighbors = mg_ivf_flat.search( + search_params, index, queries, 5, resources=resources + ) + assert isinstance(distances, np.ndarray) + assert isinstance(neighbors, np.ndarray) + + +def test_parameter_validation(): + """Test parameter validation for multi-GPU IVF-Flat.""" + # Test invalid distribution mode + with pytest.raises(ValueError, match="distribution_mode must be"): + mg_ivf_flat.IndexParams(distribution_mode="invalid") + + # Test invalid search mode + with pytest.raises(ValueError, match="search_mode must be"): + mg_ivf_flat.SearchParams(search_mode="invalid") + + # Test invalid merge mode + with pytest.raises(ValueError, match="merge_mode must be"): + mg_ivf_flat.SearchParams(merge_mode="invalid") + + +def test_parameter_properties(): + """Test that parameters can be accessed via properties.""" + # Test IndexParams properties + params = mg_ivf_flat.IndexParams(distribution_mode="replicated") + assert params.distribution_mode == "replicated" + + params = mg_ivf_flat.IndexParams(distribution_mode="sharded") + assert params.distribution_mode == "sharded" + + # Test SearchParams creation with different parameters + mg_ivf_flat.SearchParams( + search_mode="round_robin", + merge_mode="tree_merge", + n_rows_per_batch=2000, + ) + # These don't have properties exposed, but creation should work + + +def test_untrained_index_error(): + """Test that using an untrained index raises appropriate errors.""" + resources = MultiGpuResources() + + # Create untrained index + index = mg_ivf_flat.Index() + assert not index.trained + + queries = generate_data((100, 10), np.float32) + search_params = mg_ivf_flat.SearchParams(n_probes=20) + + # Test that search on untrained index fails + with pytest.raises(ValueError, match="Index needs to be built"): + mg_ivf_flat.search( + search_params, index, queries, 10, resources=resources + ) + + # Test that extend on untrained index fails + new_vectors = generate_data((50, 10), np.float32) + with pytest.raises(ValueError, match="Index needs to be built"): + mg_ivf_flat.extend(index, new_vectors, resources=resources) + + # Test that save on untrained index fails + with pytest.raises(ValueError, match="Index needs to be built"): + mg_ivf_flat.save(index, "temp.bin", resources=resources) + + +@requires_multiple_gpus +def test_mg_ivf_flat_with_prealloc_output(): + """Test multi-GPU IVF-Flat search with pre-allocated output arrays.""" + n_rows, n_cols = 1500, 8 # Ensure n_rows > n_lists + n_queries = 20 + k = 5 + + # Generate data in host memory + dataset = generate_data((n_rows, n_cols), np.float32) + queries = generate_data((n_queries, n_cols), np.float32) + + resources = MultiGpuResources() + + # Build index with fewer clusters to avoid n_rows < n_lists error + build_params = mg_ivf_flat.IndexParams(n_lists=30) + index = mg_ivf_flat.build(build_params, dataset, resources=resources) + + # Pre-allocate output arrays in host memory + neighbors = np.empty((n_queries, k), dtype=np.int64) + distances = np.empty((n_queries, k), dtype=np.float32) + + # Search with pre-allocated arrays + search_params = mg_ivf_flat.SearchParams(n_probes=20) + ret_distances, ret_neighbors = mg_ivf_flat.search( + search_params, + index, + queries, + k, + neighbors=neighbors, + distances=distances, + resources=resources, + ) + + # Should return the same arrays we passed in + assert ret_distances is distances + assert ret_neighbors is neighbors + assert distances.shape == (n_queries, k) + assert neighbors.shape == (n_queries, k) + + +def test_index_repr(): + """Test string representation of Index.""" + index = mg_ivf_flat.Index() + assert repr(index) == "Index(type=MultiGpuIvfFlat)" + + +def test_mg_ivf_flat_simple(): + """Simple test to validate multi-GPU IVF-Flat works with very favorable + parameters. + """ + if not has_multiple_gpus(): + pytest.skip("Multi-GPU tests require multiple GPUs") + + # Use simple test case that should definitely work + n_rows, n_cols = 1000, 8 + n_queries, k = 20, 5 + + # Generate data + dataset = generate_data((n_rows, n_cols), np.float32) + queries = generate_data((n_queries, n_cols), np.float32) + + resources = MultiGpuResources() + + # Use very few clusters for high recall + build_params = mg_ivf_flat.IndexParams( + metric="sqeuclidean", + n_lists=32, # Very few clusters + ) + + # Build index + index = mg_ivf_flat.build(build_params, dataset, resources=resources) + + # Search with many probes for maximum recall + search_params = mg_ivf_flat.SearchParams( + n_probes=32 + ) # Search all clusters + distances, neighbors = mg_ivf_flat.search( + search_params, index, queries, k, resources=resources + ) + + # Basic sanity checks + assert distances.shape == (n_queries, k) + assert neighbors.shape == (n_queries, k) + assert isinstance(distances, np.ndarray) + assert isinstance(neighbors, np.ndarray) + + # Check that we get valid neighbors + assert np.all(neighbors >= 0) + assert np.all(neighbors < n_rows) + + # Distances should be non-negative and sorted + assert np.all(distances >= 0) + for i in range(n_queries): + assert np.all( + distances[i, :-1] <= distances[i, 1:] + ), f"Distances not sorted for query {i}" + + +# Integration test with multiple operations +@requires_multiple_gpus +def test_mg_ivf_flat_integration(): + """Integration test covering build, search, extend, and serialization.""" + n_rows, n_cols = 2000, 8 + k = 5 + + # Generate initial dataset + dataset = generate_data((n_rows, n_cols), np.float32) + queries = generate_data((20, n_cols), np.float32) + + resources = MultiGpuResources() + + # Build initial index + build_params = mg_ivf_flat.IndexParams( + distribution_mode="sharded", metric="sqeuclidean", n_lists=50 + ) + index = mg_ivf_flat.build(build_params, dataset, resources=resources) + + # Initial search + search_params = mg_ivf_flat.SearchParams( + n_probes=37, + search_mode="load_balancer", + merge_mode="merge_on_root_rank", + ) + distances1, neighbors1 = mg_ivf_flat.search( + search_params, index, queries, k, resources=resources + ) + + # Extend index with new vectors + new_vectors = generate_data((200, n_cols), np.float32) + # Provide indices for extend operation on non-empty index + new_indices = np.arange(n_rows, n_rows + 200, dtype=np.int64) + mg_ivf_flat.extend(index, new_vectors, new_indices, resources=resources) + + # Search after extend + distances2, neighbors2 = mg_ivf_flat.search( + search_params, index, queries, k, resources=resources + ) + + # Save and reload + with tempfile.NamedTemporaryFile(suffix=".bin", delete=False) as f: + temp_filename = f.name + + try: + mg_ivf_flat.save(index, temp_filename, resources=resources) + reloaded_index = mg_ivf_flat.load(temp_filename, resources=resources) + + # Search with reloaded index + distances3, neighbors3 = mg_ivf_flat.search( + search_params, reloaded_index, queries, k, resources=resources + ) + + # Results from extended and reloaded index should match + np.testing.assert_array_equal(neighbors2, neighbors3) + np.testing.assert_allclose(distances2, distances3, rtol=1e-6) + + finally: + if os.path.exists(temp_filename): + os.unlink(temp_filename) diff --git a/python/cuvs/cuvs/tests/test_mg_ivf_pq.py b/python/cuvs/cuvs/tests/test_mg_ivf_pq.py new file mode 100644 index 0000000000..382fb9eed7 --- /dev/null +++ b/python/cuvs/cuvs/tests/test_mg_ivf_pq.py @@ -0,0 +1,682 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +import os +import tempfile + +import numpy as np +import pytest +from sklearn.neighbors import NearestNeighbors +from sklearn.preprocessing import normalize + +from cuvs.common import MultiGpuResources +from cuvs.neighbors.mg import ivf_pq as mg_ivf_pq +from cuvs.tests.ann_utils import calc_recall, generate_data + + +# Check if multi-GPU functionality is available +def has_multiple_gpus(): + """Check if system has multiple GPUs available.""" + try: + import cupy as cp + + return cp.cuda.runtime.getDeviceCount() > 1 + except Exception: + return False + + +# Mark tests that require multiple GPUs +requires_multiple_gpus = pytest.mark.skipif( + not has_multiple_gpus(), reason="Multi-GPU tests require multiple GPUs" +) + + +def run_mg_ivf_pq_build_search_test( + n_rows=10000, + n_cols=10, + n_queries=100, + k=10, + dtype=np.float32, + metric="euclidean", + distribution_mode="sharded", + search_mode="load_balancer", + merge_mode="tree_merge", + n_rows_per_batch=1000, + compare=True, + add_data_on_build=True, + search_params=None, + n_lists=None, + pq_bits=8, + pq_dim=0, + codebook_kind="subspace", +): + """ + Run a multi-GPU IVF-PQ build and search test. + + Note: Multi-GPU IVF-PQ requires host memory arrays (NumPy), not device + arrays. + """ + # Generate host memory arrays (NumPy) + dataset = generate_data((n_rows, n_cols), dtype) + if metric == "inner_product": + dataset = normalize(dataset, norm="l2", axis=1) + + queries = generate_data((n_queries, n_cols), dtype) + if metric == "inner_product": + queries = normalize(queries, norm="l2", axis=1) + + # Multi-GPU resources + resources = MultiGpuResources() + + # Build parameters - use fewer clusters for better recall with smaller + # datasets + if n_lists is None: + # Use fewer clusters for smaller datasets to ensure enough points per + # cluster + n_lists = min(1024, max(64, n_rows // 50)) + + build_params = mg_ivf_pq.IndexParams( + metric=metric, + distribution_mode=distribution_mode, + add_data_on_build=add_data_on_build, + n_lists=n_lists, + pq_bits=pq_bits, + pq_dim=pq_dim, + codebook_kind=codebook_kind, + ) + + # Build index + index = mg_ivf_pq.build(build_params, dataset, resources=resources) + assert index.trained + + # If not adding data on build, extend the index + if not add_data_on_build: + dataset_1 = dataset[: n_rows // 2, :] + dataset_2 = dataset[n_rows // 2 :, :] + indices_1 = np.arange(n_rows // 2, dtype=np.int64) + indices_2 = np.arange(n_rows // 2, n_rows, dtype=np.int64) + + mg_ivf_pq.extend(index, dataset_1, indices_1, resources=resources) + mg_ivf_pq.extend(index, dataset_2, indices_2, resources=resources) + + # Search parameters + if search_params is None: + search_params = {} + # Use higher n_probes for better recall in multi-GPU setting + if "n_probes" not in search_params: + # Use many clusters for good recall - search majority of clusters + search_params["n_probes"] = min(n_lists, max(20, (n_lists * 3) // 4)) + search_params_obj = mg_ivf_pq.SearchParams( + search_mode=search_mode, + merge_mode=merge_mode, + n_rows_per_batch=n_rows_per_batch, + **search_params, + ) + + # Perform search + distances, neighbors = mg_ivf_pq.search( + search_params_obj, + index, + queries, + k, + resources=resources, + ) + + # Verify results are in host memory (NumPy arrays) + assert isinstance(distances, np.ndarray) + assert isinstance(neighbors, np.ndarray) + assert distances.shape == (n_queries, k) + assert neighbors.shape == (n_queries, k) + + if not compare: + return distances, neighbors + + # Calculate reference values with sklearn + skl_metric = { + "sqeuclidean": "sqeuclidean", + "inner_product": "cosine", + "cosine": "cosine", + "euclidean": "euclidean", + }[metric] + + nn_skl = NearestNeighbors( + n_neighbors=k, algorithm="brute", metric=skl_metric + ) + nn_skl.fit(dataset) + skl_idx = nn_skl.kneighbors(queries, return_distance=False) + + recall = calc_recall(neighbors, skl_idx) + + return distances, neighbors, recall + + +@requires_multiple_gpus +@pytest.mark.parametrize("dtype", [np.float32]) +@pytest.mark.parametrize( + "metric", ["sqeuclidean"] +) # Start with just sqeuclidean +@pytest.mark.parametrize( + "distribution_mode", ["sharded"] +) # Start with just sharded +def test_mg_ivf_pq_basic(dtype, metric, distribution_mode): + """Test basic multi-GPU IVF-PQ build and search functionality.""" + run_mg_ivf_pq_build_search_test( + n_rows=2000, # Use smaller dataset for more reliable tests + n_cols=32, + n_queries=20, + k=5, + dtype=dtype, + metric=metric, + distribution_mode=distribution_mode, + n_lists=50, # Fixed small number of clusters + compare=True, + ) + + +@requires_multiple_gpus +@pytest.mark.parametrize("metric", ["inner_product", "euclidean", "cosine"]) +@pytest.mark.parametrize("distribution_mode", ["replicated"]) +def test_mg_ivf_pq_additional_metrics(metric, distribution_mode): + """Test additional metrics and distribution modes for IVF-PQ.""" + run_mg_ivf_pq_build_search_test( + n_rows=2000, + n_cols=32, + n_queries=20, + k=5, + dtype=np.float32, + metric=metric, + distribution_mode=distribution_mode, + n_lists=50, + compare=False, # PQ may have lower recall, don't enforce strict recall + ) + + +@requires_multiple_gpus +@pytest.mark.parametrize("dtype", [np.float32, np.float16, np.int8, np.uint8]) +def test_mg_ivf_pq_dtypes(dtype): + """Test multi-GPU IVF-PQ with different data types.""" + run_mg_ivf_pq_build_search_test( + n_rows=1500, + n_cols=32, + n_queries=15, + k=5, + dtype=dtype, + metric="sqeuclidean", + n_lists=30, + compare=False, + ) + + +@requires_multiple_gpus +@pytest.mark.parametrize("distribution_mode", ["sharded", "replicated"]) +def test_mg_ivf_pq_distribution_modes(distribution_mode): + """Test different distribution modes for multi-GPU IVF-PQ.""" + run_mg_ivf_pq_build_search_test( + n_rows=1500, + n_cols=32, + n_queries=15, + k=5, + distribution_mode=distribution_mode, + n_lists=30, + compare=False, + ) + + +@requires_multiple_gpus +@pytest.mark.parametrize("search_mode", ["load_balancer", "round_robin"]) +@pytest.mark.parametrize("merge_mode", ["merge_on_root_rank", "tree_merge"]) +def test_mg_ivf_pq_search_params(search_mode, merge_mode): + """Test different multi-GPU search parameters for IVF-PQ.""" + run_mg_ivf_pq_build_search_test( + n_rows=1500, + n_cols=32, + n_queries=15, + k=5, + search_mode=search_mode, + merge_mode=merge_mode, + n_rows_per_batch=500, + n_lists=30, + compare=False, + ) + + +@requires_multiple_gpus +def test_mg_ivf_pq_pq_parameters(): + """Test different PQ-specific parameters.""" + for pq_bits in [4, 8]: + for pq_dim in [0, 8, 16]: # 0 means auto-select + for codebook_kind in ["subspace", "cluster"]: + run_mg_ivf_pq_build_search_test( + n_rows=1000, + n_cols=32, + n_queries=100, + k=10, + pq_bits=pq_bits, + pq_dim=pq_dim, + codebook_kind=codebook_kind, + compare=False, + ) + + +@requires_multiple_gpus +@pytest.mark.parametrize("metric", ["euclidean", "sqeuclidean"]) +def test_mg_ivf_pq_metrics(metric): + """Test different distance metrics for multi-GPU IVF-PQ.""" + run_mg_ivf_pq_build_search_test( + n_rows=1500, + n_cols=32, + n_queries=15, + k=5, + metric=metric, + n_lists=30, + compare=False, + ) + + +@requires_multiple_gpus +def test_mg_ivf_pq_extend(): + """Test extending index with new vectors.""" + run_mg_ivf_pq_build_search_test( + n_rows=1000, + n_cols=32, + n_queries=100, + k=10, + add_data_on_build=False, # This triggers extend functionality + compare=False, + ) + + +@requires_multiple_gpus +def test_mg_ivf_pq_serialize(): + """Test serialization and deserialization.""" + # Generate data + n_rows, n_cols = 1000, 32 + dataset = generate_data((n_rows, n_cols), np.float32) + queries = generate_data((100, n_cols), np.float32) + + resources = MultiGpuResources() + + # Build index + build_params = mg_ivf_pq.IndexParams( + metric="euclidean", + n_lists=100, + pq_bits=8, + pq_dim=16, + ) + index = mg_ivf_pq.build(build_params, dataset, resources=resources) + + # Search before serialization + search_params = mg_ivf_pq.SearchParams(n_probes=50) + distances_1, neighbors_1 = mg_ivf_pq.search( + search_params, index, queries, 10, resources=resources + ) + + # Serialize + with tempfile.NamedTemporaryFile(delete=False) as f: + filename = f.name + + try: + mg_ivf_pq.save(index, filename, resources=resources) + + # Load index + index_loaded = mg_ivf_pq.load(filename, resources=resources) + assert index_loaded.trained + + # Search after loading + distances_2, neighbors_2 = mg_ivf_pq.search( + search_params, index_loaded, queries, 10, resources=resources + ) + + # Results should be the same + assert np.array_equal(distances_1, distances_2) + assert np.array_equal(neighbors_1, neighbors_2) + + finally: + if os.path.exists(filename): + os.unlink(filename) + + +@requires_multiple_gpus +def test_mg_ivf_pq_distribute(): + """Test distribute functionality for multi-GPU IVF-PQ.""" + # Note: Distribute is for replicating a single-GPU index across multiple + # GPUs. + # This test builds a single-GPU index, serializes it, then distributes it. + # Multi-GPU distribute only supports float32 indexes. + + n_rows, n_cols = 2000, 32 + k = 5 + + # Generate data + dataset = generate_data((n_rows, n_cols), np.float32) + queries = generate_data((100, n_cols), np.float32) + + # Import single-GPU IVF-PQ to build and serialize a single-GPU index + from cuvs.common import Resources + from cuvs.neighbors import ivf_pq + + # Build single-GPU index first + single_gpu_resources = Resources() + single_build_params = ivf_pq.IndexParams( + metric="sqeuclidean", n_lists=50, pq_bits=8, pq_dim=16 + ) + + # Convert to device arrays for single-GPU build + try: + import cupy as cp + + device_dataset = cp.asarray(dataset, dtype=np.float32) + single_index = ivf_pq.build( + single_build_params, device_dataset, resources=single_gpu_resources + ) + except ImportError: + pytest.skip("CuPy not available for single-GPU index building") + + with tempfile.NamedTemporaryFile(suffix=".bin", delete=False) as f: + temp_filename = f.name + + try: + # Serialize single-GPU index + ivf_pq.save( + temp_filename, single_index, resources=single_gpu_resources + ) + + # Now distribute the single-GPU index across multiple GPUs + resources = MultiGpuResources() + distributed_index = mg_ivf_pq.distribute( + temp_filename, resources=resources + ) + assert distributed_index.trained + + # Search using the distributed index + search_params = mg_ivf_pq.SearchParams(n_probes=25) + distances, neighbors = mg_ivf_pq.search( + search_params, distributed_index, queries, k, resources=resources + ) + + # Verify results shape + assert distances.shape == (100, k) + assert neighbors.shape == (100, k) + + finally: + if os.path.exists(temp_filename): + os.unlink(temp_filename) + + +def test_memory_location_validation(): + """Test that multi-GPU IVF-PQ validates memory locations correctly.""" + try: + import cupy as cp + except ImportError: + pytest.skip("CuPy not available") + + # Generate device arrays (should fail) - use enough data points for n_lists + dataset_gpu = cp.random.random((1000, 32), dtype=cp.float32) + queries_gpu = cp.random.random((100, 32), dtype=cp.float32) + + # Create parameters with smaller n_lists for the small dataset + build_params = mg_ivf_pq.IndexParams( + n_lists=20 + ) # Smaller n_lists for 1000 points + search_params = mg_ivf_pq.SearchParams() + + # These should raise ValueError about memory location + with pytest.raises(ValueError, match="host memory"): + mg_ivf_pq.build(build_params, dataset_gpu) + + # For search test, we need a valid index first + dataset_cpu = cp.asnumpy(dataset_gpu) + resources = MultiGpuResources() if has_multiple_gpus() else None + if resources: + index = mg_ivf_pq.build(build_params, dataset_cpu, resources=resources) + + with pytest.raises(ValueError, match="host memory"): + mg_ivf_pq.search( + search_params, index, queries_gpu, 5, resources=resources + ) + + +def test_parameter_validation(): + """Test parameter validation for multi-GPU IVF-PQ.""" + # Test invalid distribution mode + with pytest.raises(ValueError, match="distribution_mode must be"): + mg_ivf_pq.IndexParams(distribution_mode="invalid") + + # Test invalid search mode + with pytest.raises(ValueError, match="search_mode must be"): + mg_ivf_pq.SearchParams(search_mode="invalid") + + # Test invalid merge mode + with pytest.raises(ValueError, match="merge_mode must be"): + mg_ivf_pq.SearchParams(merge_mode="invalid") + + # Test invalid codebook kind + with pytest.raises(ValueError, match="Incorrect codebook kind"): + mg_ivf_pq.IndexParams(codebook_kind="invalid") + + +def test_parameter_properties(): + """Test that parameters can be accessed via properties.""" + # Test IndexParams properties + params = mg_ivf_pq.IndexParams(distribution_mode="replicated") + assert params.distribution_mode == "replicated" + + params = mg_ivf_pq.IndexParams(distribution_mode="sharded") + assert params.distribution_mode == "sharded" + + # Test PQ-specific parameters + params = mg_ivf_pq.IndexParams( + pq_bits=4, pq_dim=16, codebook_kind="cluster" + ) + # These don't have properties exposed, but creation should work + + # Test SearchParams creation with different parameters + mg_ivf_pq.SearchParams( + search_mode="round_robin", + merge_mode="tree_merge", + n_rows_per_batch=2000, + ) + # These don't have properties exposed, but creation should work + + +def test_untrained_index_error(): + """Test that using an untrained index raises appropriate errors.""" + resources = MultiGpuResources() + + # Create untrained index + index = mg_ivf_pq.Index() + assert not index.trained + + queries = generate_data((100, 10), np.float32) + search_params = mg_ivf_pq.SearchParams(n_probes=20) + + # Test that search on untrained index fails + with pytest.raises(ValueError, match="Index needs to be built"): + mg_ivf_pq.search( + search_params, index, queries, 10, resources=resources + ) + + # Test that extend on untrained index fails + new_vectors = generate_data((50, 10), np.float32) + with pytest.raises(ValueError, match="Index needs to be built"): + mg_ivf_pq.extend(index, new_vectors, resources=resources) + + # Test that save on untrained index fails + with pytest.raises(ValueError, match="Index needs to be built"): + mg_ivf_pq.save(index, "temp.bin", resources=resources) + + +@requires_multiple_gpus +def test_mg_ivf_pq_with_prealloc_output(): + """Test multi-GPU IVF-PQ search with pre-allocated output arrays.""" + n_rows, n_cols = 1500, 32 # Ensure n_rows > n_lists + n_queries = 20 + k = 5 + + # Generate data in host memory + dataset = generate_data((n_rows, n_cols), np.float32) + queries = generate_data((n_queries, n_cols), np.float32) + + resources = MultiGpuResources() + + # Build index with fewer clusters to avoid n_rows < n_lists error + build_params = mg_ivf_pq.IndexParams(n_lists=30, pq_bits=8, pq_dim=16) + index = mg_ivf_pq.build(build_params, dataset, resources=resources) + + # Pre-allocate output arrays in host memory + neighbors = np.empty((n_queries, k), dtype=np.int64) + distances = np.empty((n_queries, k), dtype=np.float32) + + # Search with pre-allocated arrays + search_params = mg_ivf_pq.SearchParams(n_probes=20) + ret_distances, ret_neighbors = mg_ivf_pq.search( + search_params, + index, + queries, + k, + neighbors=neighbors, + distances=distances, + resources=resources, + ) + + # Should return the same arrays we passed in + assert ret_distances is distances + assert ret_neighbors is neighbors + assert distances.shape == (n_queries, k) + assert neighbors.shape == (n_queries, k) + + +def test_index_repr(): + """Test string representation of Index.""" + index = mg_ivf_pq.Index() + assert repr(index) == "Index(type=MultiGpuIvfPq)" + + +def test_mg_ivf_pq_simple(): + """Simple test to validate multi-GPU IVF-PQ works with very favorable + parameters. + """ + if not has_multiple_gpus(): + pytest.skip("Multi-GPU tests require multiple GPUs") + + # Use simple test case that should definitely work + n_rows, n_cols = 1000, 32 + n_queries, k = 20, 5 + + # Generate data + dataset = generate_data((n_rows, n_cols), np.float32) + queries = generate_data((n_queries, n_cols), np.float32) + + resources = MultiGpuResources() + + # Use very few clusters for high recall + build_params = mg_ivf_pq.IndexParams( + metric="sqeuclidean", + n_lists=32, # Very few clusters + pq_bits=8, + pq_dim=16, + ) + + # Build index + index = mg_ivf_pq.build(build_params, dataset, resources=resources) + + # Search with many probes for maximum recall + search_params = mg_ivf_pq.SearchParams(n_probes=32) # Search all clusters + distances, neighbors = mg_ivf_pq.search( + search_params, index, queries, k, resources=resources + ) + + # Basic sanity checks + assert distances.shape == (n_queries, k) + assert neighbors.shape == (n_queries, k) + assert isinstance(distances, np.ndarray) + assert isinstance(neighbors, np.ndarray) + + # Check that we get valid neighbors + assert np.all(neighbors >= 0) + assert np.all(neighbors < n_rows) + + # Distances should be non-negative and sorted + assert np.all(distances >= 0) + for i in range(n_queries): + assert np.all( + distances[i, :-1] <= distances[i, 1:] + ), f"Distances not sorted for query {i}" + + +# Integration test with multiple operations +@requires_multiple_gpus +def test_mg_ivf_pq_integration(): + """Integration test covering build, search, extend, and serialization.""" + n_rows, n_cols = 2000, 32 + k = 5 + + # Generate initial dataset + dataset = generate_data((n_rows, n_cols), np.float32) + queries = generate_data((20, n_cols), np.float32) + + resources = MultiGpuResources() + + # Build initial index + build_params = mg_ivf_pq.IndexParams( + distribution_mode="sharded", + metric="sqeuclidean", + n_lists=50, + pq_bits=8, + pq_dim=16, + ) + index = mg_ivf_pq.build(build_params, dataset, resources=resources) + + # Initial search + search_params = mg_ivf_pq.SearchParams( + n_probes=37, + search_mode="load_balancer", + merge_mode="merge_on_root_rank", + ) + distances1, neighbors1 = mg_ivf_pq.search( + search_params, index, queries, k, resources=resources + ) + + # Extend index with new vectors + new_vectors = generate_data((200, n_cols), np.float32) + # Provide indices for extend operation on non-empty index + new_indices = np.arange(n_rows, n_rows + 200, dtype=np.int64) + mg_ivf_pq.extend(index, new_vectors, new_indices, resources=resources) + + # Search after extend + distances2, neighbors2 = mg_ivf_pq.search( + search_params, index, queries, k, resources=resources + ) + + # Save and reload + with tempfile.NamedTemporaryFile(suffix=".bin", delete=False) as f: + temp_filename = f.name + + try: + mg_ivf_pq.save(index, temp_filename, resources=resources) + reloaded_index = mg_ivf_pq.load(temp_filename, resources=resources) + + # Search with reloaded index + distances3, neighbors3 = mg_ivf_pq.search( + search_params, reloaded_index, queries, k, resources=resources + ) + + # Results from extended and reloaded index should match + np.testing.assert_array_equal(neighbors2, neighbors3) + np.testing.assert_allclose(distances2, distances3, rtol=1e-6) + + finally: + if os.path.exists(temp_filename): + os.unlink(temp_filename) From 8fc4a2ad134f2d2d7763b026d0965b328ecc8db8 Mon Sep 17 00:00:00 2001 From: Jake Awe <50372925+AyodeAwe@users.noreply.github.com> Date: Mon, 22 Sep 2025 09:10:10 -0500 Subject: [PATCH 77/93] Configure repo for automatic release notes generation (#1334) This PR configures this repo (adds a `.github/release.yml` file) for automatic release notes generation. See https://docs.github.com/en/repositories/releasing-projects-on-github/automatically-generated-release-notes. This will fix such issues as hotfix releases including wrong content (e.g [cugraph v25.04.01](https://github.com/rapidsai/cugraph/releases/tag/v25.04.01) contains 50+ PRs from the full v25.04.00 release when it should only contain the single hotfix PR (#5017).) Example auto-generated release notes: https://github.com/rapidsai/literate-octo-potato/releases/tag/v25.10.00a --- .github/release.yml | 27 +++++++++++++++++++++++++++ 1 file changed, 27 insertions(+) create mode 100644 .github/release.yml diff --git a/.github/release.yml b/.github/release.yml new file mode 100644 index 0000000000..2c9a85805c --- /dev/null +++ b/.github/release.yml @@ -0,0 +1,27 @@ +# GitHub Auto-Generated Release Notes Configuration for RAPIDS +# This file configures how GitHub automatically generates release notes + +changelog: + exclude: + labels: + - ignore-for-release + - dependencies + authors: + - rapids-bot[bot] + - dependabot[bot] + categories: + - title: 🚨 Breaking Changes + labels: + - breaking + - title: 🐛 Bug Fixes + labels: + - bug + - title: 📖 Documentation + labels: + - doc + - title: 🚀 New Features + labels: + - feature request + - title: 🛠️ Improvements + labels: + - improvement From b431fda8f93d059567054d2c4ef71541fd4bd859 Mon Sep 17 00:00:00 2001 From: Mike Sarahan Date: Tue, 23 Sep 2025 16:19:45 -0500 Subject: [PATCH 78/93] Empty commit to trigger a build (#1363) This is an empty commit to trigger a build. It is used when builds get stuck with an old ABI. Rebuilding updates them to the new one. From b6a09dd8404190a64b385c6b9ce343877d2e6162 Mon Sep 17 00:00:00 2001 From: James Lamb Date: Wed, 24 Sep 2025 12:05:05 -0500 Subject: [PATCH 79/93] use CUDA 13.0.1 CI images (#1353) RAPIDS recently started building and testing against CUDA 13.0.1: * https://github.com/rapidsai/ci-imgs/pull/304 * https://github.com/rapidsai/shared-workflows/pull/423 This updates hard-coded `13.0.0` references in some CI jobs to `13.0.1`. Authors: - James Lamb (https://github.com/jameslamb) Approvers: - Bradley Dice (https://github.com/bdice) URL: https://github.com/rapidsai/cuvs/pull/1353 --- .github/workflows/build.yaml | 12 ++++++------ .github/workflows/pr.yaml | 12 ++++++------ .github/workflows/test.yaml | 2 +- 3 files changed, 13 insertions(+), 13 deletions(-) diff --git a/.github/workflows/build.yaml b/.github/workflows/build.yaml index 892764844b..9e914fbfcb 100644 --- a/.github/workflows/build.yaml +++ b/.github/workflows/build.yaml @@ -51,8 +51,8 @@ jobs: fail-fast: false matrix: cuda_version: - - '12.9.1' - - '13.0.0' + - &latest_cuda12 '12.9.1' + - &latest_cuda13 '13.0.1' with: build_type: ${{ inputs.build_type || 'branch' }} branch: ${{ inputs.branch }} @@ -72,8 +72,8 @@ jobs: fail-fast: false matrix: cuda_version: - - '12.9.1' - - '13.0.0' + - *latest_cuda12 + - *latest_cuda13 with: build_type: ${{ inputs.build_type || 'branch' }} branch: ${{ inputs.branch }} @@ -93,8 +93,8 @@ jobs: fail-fast: false matrix: cuda_version: - - '12.9.1' - - '13.0.0' + - *latest_cuda12 + - *latest_cuda13 with: build_type: ${{ inputs.build_type || 'branch' }} branch: ${{ inputs.branch }} diff --git a/.github/workflows/pr.yaml b/.github/workflows/pr.yaml index 15b3f35efb..80e2107af8 100644 --- a/.github/workflows/pr.yaml +++ b/.github/workflows/pr.yaml @@ -159,8 +159,8 @@ jobs: fail-fast: false matrix: cuda_version: - - '12.9.1' - - '13.0.0' + - &latest_cuda12 '12.9.1' + - &latest_cuda13 '13.0.1' with: build_type: pull-request node_type: "gpu-l4-latest-1" @@ -189,8 +189,8 @@ jobs: fail-fast: false matrix: cuda_version: - - '12.9.1' - - '13.0.0' + - *latest_cuda12 + - *latest_cuda13 with: build_type: pull-request node_type: "gpu-l4-latest-1" @@ -207,8 +207,8 @@ jobs: fail-fast: false matrix: cuda_version: - - '12.9.1' - - '13.0.0' + - *latest_cuda12 + - *latest_cuda13 with: build_type: pull-request node_type: "gpu-l4-latest-1" diff --git a/.github/workflows/test.yaml b/.github/workflows/test.yaml index c5892ca7ea..87f530a672 100644 --- a/.github/workflows/test.yaml +++ b/.github/workflows/test.yaml @@ -60,7 +60,7 @@ jobs: matrix: cuda_version: - '12.9.1' - - '13.0.0' + - '13.0.1' with: build_type: ${{ inputs.build_type }} branch: ${{ inputs.branch }} From 0317f1e96c5770f96f05363aeb9bf9dc58a66268 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Lorenzo=20Dematt=C3=A9?= Date: Thu, 25 Sep 2025 00:26:13 +0200 Subject: [PATCH 80/93] [Review][Java] Refactor: extract interface from CuVSMatrixBaseImpl (#1361) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit This PR extract an internal interface that is used as a base to implement all internal CuVSMatrix types; the interface introduces commonly used field accessors like `memorySegment()` and `toTensor()` that we do not want/could not appear on the public interface (e.g. because they expose or require Panama types or internal types). The new interface is implemented by all concrete matrix types, closing a gap that we had in https://github.com/rapidsai/cuvs/pull/1328 (which I realized while working on separate PRs like https://github.com/rapidsai/cuvs/pull/1283) Follow-up of https://github.com/rapidsai/cuvs/pull/1328 Authors: - Lorenzo Dematté (https://github.com/ldematte) Approvers: - Ben Frederickson (https://github.com/benfred) URL: https://github.com/rapidsai/cuvs/pull/1361 --- .../com/nvidia/cuvs/CuVSDeviceMatrix.java | 63 ------------ .../java/com/nvidia/cuvs/CuVSHostMatrix.java | 67 ------------- .../cuvs/internal/BruteForceIndexImpl.java | 6 +- .../nvidia/cuvs/internal/CagraIndexImpl.java | 17 ++-- .../cuvs/internal/CuVSDeviceMatrixImpl.java | 97 ++++++++++++++++++- .../cuvs/internal/CuVSHostMatrixImpl.java | 94 +++++++++++++++++- .../cuvs/internal/CuVSMatrixBaseImpl.java | 38 ++------ .../cuvs/internal/CuVSMatrixInternal.java | 61 ++++++++++++ .../nvidia/cuvs/internal/TieredIndexImpl.java | 4 +- 9 files changed, 267 insertions(+), 180 deletions(-) create mode 100644 java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CuVSMatrixInternal.java diff --git a/java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSDeviceMatrix.java b/java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSDeviceMatrix.java index ec34a3de1c..6b0782da52 100644 --- a/java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSDeviceMatrix.java +++ b/java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSDeviceMatrix.java @@ -30,67 +30,4 @@ default CuVSHostMatrix toHost() { toHost(hostMatrix); return hostMatrix; } - - default CuVSDeviceMatrix toDevice(CuVSResources resources) { - return new CuVSDeviceMatrixDelegate(this); - } - - class CuVSDeviceMatrixDelegate implements CuVSDeviceMatrix { - - private final CuVSDeviceMatrix deviceMatrix; - - private CuVSDeviceMatrixDelegate(CuVSDeviceMatrix deviceMatrix) { - this.deviceMatrix = deviceMatrix; - } - - @Override - public long size() { - return deviceMatrix.size(); - } - - @Override - public long columns() { - return deviceMatrix.columns(); - } - - @Override - public DataType dataType() { - return deviceMatrix.dataType(); - } - - @Override - public RowView getRow(long row) { - return deviceMatrix.getRow(row); - } - - @Override - public void toArray(int[][] array) { - deviceMatrix.toArray(array); - } - - @Override - public void toArray(float[][] array) { - deviceMatrix.toArray(array); - } - - @Override - public void toArray(byte[][] array) { - deviceMatrix.toArray(array); - } - - @Override - public void toHost(CuVSHostMatrix hostMatrix) { - deviceMatrix.toHost(hostMatrix); - } - - @Override - public void toDevice(CuVSDeviceMatrix deviceMatrix, CuVSResources cuVSResources) { - this.deviceMatrix.toDevice(deviceMatrix, cuVSResources); - } - - @Override - public void close() { - // Do nothing - } - } } diff --git a/java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSHostMatrix.java b/java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSHostMatrix.java index 91ff9012f3..7b98d82a60 100644 --- a/java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSHostMatrix.java +++ b/java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSHostMatrix.java @@ -21,76 +21,9 @@ public interface CuVSHostMatrix extends CuVSMatrix { int get(int row, int col); - default CuVSHostMatrix toHost() { - return new CuVSHostMatrixDelegate(this); - } - default CuVSDeviceMatrix toDevice(CuVSResources resources) { var deviceMatrix = CuVSMatrix.deviceBuilder(resources, size(), columns(), dataType()).build(); toDevice(deviceMatrix, resources); return deviceMatrix; } - - class CuVSHostMatrixDelegate implements CuVSHostMatrix { - private final CuVSHostMatrix hostMatrix; - - public CuVSHostMatrixDelegate(CuVSHostMatrix cuVSHostMatrix) { - this.hostMatrix = cuVSHostMatrix; - } - - @Override - public int get(int row, int col) { - return hostMatrix.get(row, col); - } - - @Override - public long size() { - return hostMatrix.size(); - } - - @Override - public long columns() { - return hostMatrix.columns(); - } - - @Override - public DataType dataType() { - return hostMatrix.dataType(); - } - - @Override - public RowView getRow(long row) { - return hostMatrix.getRow(row); - } - - @Override - public void toArray(int[][] array) { - hostMatrix.toArray(array); - } - - @Override - public void toArray(float[][] array) { - hostMatrix.toArray(array); - } - - @Override - public void toArray(byte[][] array) { - hostMatrix.toArray(array); - } - - @Override - public void toHost(CuVSHostMatrix hostMatrix) { - this.hostMatrix.toHost(hostMatrix); - } - - @Override - public void toDevice(CuVSDeviceMatrix deviceMatrix, CuVSResources cuVSResources) { - hostMatrix.toDevice(deviceMatrix, cuVSResources); - } - - @Override - public void close() { - // Do nothing - } - } } diff --git a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/BruteForceIndexImpl.java b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/BruteForceIndexImpl.java index bc7d6b0a8b..76e6003e19 100644 --- a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/BruteForceIndexImpl.java +++ b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/BruteForceIndexImpl.java @@ -78,8 +78,8 @@ private BruteForceIndexImpl( Objects.requireNonNull(dataset); try (dataset) { this.resources = resources; - assert dataset instanceof CuVSMatrixBaseImpl; - this.bruteForceIndexReference = build((CuVSMatrixBaseImpl) dataset, bruteForceIndexParams); + assert dataset instanceof CuVSMatrixInternal; + this.bruteForceIndexReference = build((CuVSMatrixInternal) dataset, bruteForceIndexParams); } } @@ -124,7 +124,7 @@ public void close() { * index */ private IndexReference build( - CuVSMatrixBaseImpl dataset, BruteForceIndexParams bruteForceIndexParams) { + CuVSMatrixInternal dataset, BruteForceIndexParams bruteForceIndexParams) { long rows = dataset.size(); long cols = dataset.columns(); diff --git a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CagraIndexImpl.java b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CagraIndexImpl.java index 99b0ecc71e..4045d8b83d 100644 --- a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CagraIndexImpl.java +++ b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CagraIndexImpl.java @@ -36,7 +36,6 @@ import com.nvidia.cuvs.internal.common.CloseableRMMAllocation; import com.nvidia.cuvs.internal.common.CompositeCloseableHandle; import com.nvidia.cuvs.internal.panama.*; - import java.io.FileInputStream; import java.io.InputStream; import java.io.OutputStream; @@ -78,8 +77,8 @@ private CagraIndexImpl( CagraIndexParams indexParameters, CuVSMatrix dataset, CuVSResources resources) { Objects.requireNonNull(dataset); this.resources = resources; - assert dataset instanceof CuVSMatrixBaseImpl; - this.cagraIndexReference = build(indexParameters, (CuVSMatrixBaseImpl) dataset); + assert dataset instanceof CuVSMatrixInternal; + this.cagraIndexReference = build(indexParameters, (CuVSMatrixInternal) dataset); } /** @@ -124,11 +123,11 @@ private CagraIndexImpl( this.resources = resources; - assert graph instanceof CuVSMatrixBaseImpl; - assert dataset instanceof CuVSMatrixBaseImpl; + assert graph instanceof CuVSMatrixInternal; + assert dataset instanceof CuVSMatrixInternal; this.cagraIndexReference = - fromGraph(metric, (CuVSMatrixBaseImpl) graph, (CuVSMatrixBaseImpl) dataset); + fromGraph(metric, (CuVSMatrixInternal) graph, (CuVSMatrixInternal) dataset); } private void checkNotDestroyed() { @@ -161,7 +160,7 @@ public void close() { * @return an instance of {@link IndexReference} that holds the pointer to the * index */ - private IndexReference build(CagraIndexParams indexParameters, CuVSMatrixBaseImpl dataset) { + private IndexReference build(CagraIndexParams indexParameters, CuVSMatrixInternal dataset) { long rows = dataset.size(); try (var indexParams = segmentFromIndexParams(indexParameters); @@ -410,8 +409,8 @@ public CuVSDeviceMatrix getGraph() { private IndexReference fromGraph( CagraIndexParams.CuvsDistanceType metric, - CuVSMatrixBaseImpl graph, - CuVSMatrixBaseImpl dataset) { + CuVSMatrixInternal graph, + CuVSMatrixInternal dataset) { try (var localArena = Arena.ofConfined()) { var index = createCagraIndex(); try (var resourcesAccess = resources.access()) { diff --git a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CuVSDeviceMatrixImpl.java b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CuVSDeviceMatrixImpl.java index 0f97608ae1..27f2adbd02 100644 --- a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CuVSDeviceMatrixImpl.java +++ b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CuVSDeviceMatrixImpl.java @@ -224,7 +224,7 @@ public void toHost(CuVSHostMatrix hostMatrix) { throw new IllegalArgumentException("[hostMatrix] must have the same dataType"); } try (var localArena = Arena.ofConfined()) { - var hostMatrixTensor = ((CuVSHostMatrixImpl) hostMatrix).toTensor(localArena); + var hostMatrixTensor = ((CuVSMatrixInternal) hostMatrix).toTensor(localArena); try (var resourceAccess = resources.access()) { var cuvsRes = resourceAccess.handle(); @@ -236,9 +236,14 @@ public void toHost(CuVSHostMatrix hostMatrix) { } } + @Override + public CuVSDeviceMatrix toDevice(CuVSResources resources) { + return new CuVSDeviceMatrixDelegate(this); + } + @Override public void toDevice(CuVSDeviceMatrix targetMatrix, CuVSResources cuVSResources) { - copyMatrix(this, (CuVSMatrixBaseImpl) targetMatrix, cuVSResources); + copyMatrix(this, (CuVSMatrixInternal) targetMatrix, cuVSResources); } @Override @@ -248,4 +253,92 @@ public void close() { hostBuffer = MemorySegment.NULL; } } + + private static class CuVSDeviceMatrixDelegate implements CuVSDeviceMatrix, CuVSMatrixInternal { + private final CuVSDeviceMatrixImpl deviceMatrix; + + private CuVSDeviceMatrixDelegate(CuVSDeviceMatrixImpl deviceMatrix) { + this.deviceMatrix = deviceMatrix; + } + + @Override + public long size() { + return deviceMatrix.size(); + } + + @Override + public long columns() { + return deviceMatrix.columns(); + } + + @Override + public DataType dataType() { + return deviceMatrix.dataType(); + } + + @Override + public RowView getRow(long row) { + return deviceMatrix.getRow(row); + } + + @Override + public void toArray(int[][] array) { + deviceMatrix.toArray(array); + } + + @Override + public void toArray(float[][] array) { + deviceMatrix.toArray(array); + } + + @Override + public void toArray(byte[][] array) { + deviceMatrix.toArray(array); + } + + @Override + public void toHost(CuVSHostMatrix hostMatrix) { + deviceMatrix.toHost(hostMatrix); + } + + @Override + public void toDevice(CuVSDeviceMatrix deviceMatrix, CuVSResources cuVSResources) { + deviceMatrix.toDevice(deviceMatrix, cuVSResources); + } + + @Override + public CuVSDeviceMatrix toDevice(CuVSResources cuVSResources) { + return this; + } + + @Override + public MemorySegment memorySegment() { + return deviceMatrix.memorySegment(); + } + + @Override + public ValueLayout valueLayout() { + return deviceMatrix.valueLayout(); + } + + @Override + public int bits() { + return deviceMatrix.bits(); + } + + @Override + public int code() { + return 0; + } + + @Override + public MemorySegment toTensor(Arena arena) { + return deviceMatrix.toTensor(arena); + } + + @Override + public void close() { + // Do nothing + } + } } diff --git a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CuVSHostMatrixImpl.java b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CuVSHostMatrixImpl.java index 408feda2bf..86e7e2a059 100644 --- a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CuVSHostMatrixImpl.java +++ b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CuVSHostMatrixImpl.java @@ -144,17 +144,22 @@ public void toArray(byte[][] array) { } } + @Override + public CuVSHostMatrix toHost() { + return new CuVSHostMatrixDelegate(this); + } + @Override public void toHost(CuVSHostMatrix hostMatrix) { - var targetMatrix = (CuVSHostMatrixImpl) hostMatrix; + var targetMatrix = (CuVSMatrixInternal) hostMatrix; var valueByteSize = valueLayout.byteSize(); MemorySegment.copy( - this.memorySegment, 0L, targetMatrix.memorySegment, 0L, size * columns * valueByteSize); + this.memorySegment, 0L, targetMatrix.memorySegment(), 0L, size * columns * valueByteSize); } @Override public void toDevice(CuVSDeviceMatrix deviceMatrix, CuVSResources cuVSResources) { - copyMatrix(this, (CuVSMatrixBaseImpl) deviceMatrix, cuVSResources); + copyMatrix(this, (CuVSMatrixInternal) deviceMatrix, cuVSResources); } @Override @@ -170,4 +175,87 @@ public MemorySegment toTensor(Arena arena) { return prepareTensor( arena, memorySegment, new long[] {size, columns}, code(), bits(), kDLCPU()); } + + private static class CuVSHostMatrixDelegate implements CuVSHostMatrix, CuVSMatrixInternal { + private final CuVSHostMatrixImpl hostMatrix; + + public CuVSHostMatrixDelegate(CuVSHostMatrixImpl cuVSHostMatrix) { + this.hostMatrix = cuVSHostMatrix; + } + + @Override + public int get(int row, int col) { + return hostMatrix.get(row, col); + } + + @Override + public long size() { + return hostMatrix.size(); + } + + @Override + public long columns() { + return hostMatrix.columns(); + } + + @Override + public DataType dataType() { + return hostMatrix.dataType(); + } + + @Override + public RowView getRow(long row) { + return hostMatrix.getRow(row); + } + + @Override + public void toArray(int[][] array) { + hostMatrix.toArray(array); + } + + @Override + public void toArray(float[][] array) { + hostMatrix.toArray(array); + } + + @Override + public void toArray(byte[][] array) { + hostMatrix.toArray(array); + } + + @Override + public void toHost(CuVSHostMatrix hostMatrix) { + this.hostMatrix.toHost(hostMatrix); + } + + @Override + public CuVSHostMatrix toHost() { + return this; + } + + @Override + public void toDevice(CuVSDeviceMatrix deviceMatrix, CuVSResources cuVSResources) { + hostMatrix.toDevice(deviceMatrix, cuVSResources); + } + + @Override + public MemorySegment memorySegment() { + return hostMatrix.memorySegment(); + } + + @Override + public ValueLayout valueLayout() { + return hostMatrix.valueLayout(); + } + + @Override + public MemorySegment toTensor(Arena arena) { + return hostMatrix.toTensor(arena); + } + + @Override + public void close() { + // Do nothing + } + } } diff --git a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CuVSMatrixBaseImpl.java b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CuVSMatrixBaseImpl.java index b46ab63c9f..2e510a99b2 100644 --- a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CuVSMatrixBaseImpl.java +++ b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CuVSMatrixBaseImpl.java @@ -30,7 +30,7 @@ import java.lang.foreign.*; import java.util.Locale; -public abstract class CuVSMatrixBaseImpl implements CuVSMatrix { +abstract class CuVSMatrixBaseImpl implements CuVSMatrixInternal { protected final MemorySegment memorySegment; protected final DataType dataType; protected final ValueLayout valueLayout; @@ -51,13 +51,13 @@ protected CuVSMatrixBaseImpl( } protected static void copyMatrix( - CuVSMatrixBaseImpl sourceMatrix, CuVSMatrixBaseImpl targetMatrix, CuVSResources resources) { - if (targetMatrix.columns() != sourceMatrix.columns - || targetMatrix.size() != sourceMatrix.size) { + CuVSMatrixInternal sourceMatrix, CuVSMatrixInternal targetMatrix, CuVSResources resources) { + if (targetMatrix.columns() != sourceMatrix.columns() + || targetMatrix.size() != sourceMatrix.size()) { throw new IllegalArgumentException( "Source and target matrices must have the same dimensions"); } - if (targetMatrix.dataType() != sourceMatrix.dataType) { + if (targetMatrix.dataType() != sourceMatrix.dataType()) { throw new IllegalArgumentException("Source and target matrices must have the same dataType"); } @@ -88,32 +88,16 @@ public DataType dataType() { return dataType; } + @Override public MemorySegment memorySegment() { return memorySegment; } + @Override public ValueLayout valueLayout() { return valueLayout; } - /** - * Size (in bits) for the element type of this matrix - */ - protected int bits() { - return (int) (valueLayout.byteSize() * 8); - } - - /** - * DLTensor data type {@code code} for the element type of this matrix - */ - protected int code() { - return switch (dataType) { - case FLOAT -> kDLFloat(); - case INT -> kDLInt(); - case UINT, BYTE -> kDLUInt(); - }; - } - protected static ValueLayout valueLayoutFromType(DataType dataType) { return switch (dataType) { case FLOAT -> C_FLOAT; @@ -128,14 +112,6 @@ protected static SequenceLayout sequenceLayoutFromType( .withByteAlignment(32); } - /** - * Creates a {@link DLManagedTensor} representing the matrix data and shape, to be - * passed to the CuVS C API. - * @param arena The Arena to use to allocate DL data structures - * @return a {@link MemorySegment} for the newly allocated DLManagedTensor - */ - public abstract MemorySegment toTensor(Arena arena); - /** * Creates a {@link CuVSMatrix} from data and infos from a {@link DLManagedTensor} * diff --git a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CuVSMatrixInternal.java b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CuVSMatrixInternal.java new file mode 100644 index 0000000000..ed53836719 --- /dev/null +++ b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CuVSMatrixInternal.java @@ -0,0 +1,61 @@ +/* + * Copyright (c) 2025, NVIDIA CORPORATION. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.nvidia.cuvs.internal; + +import static com.nvidia.cuvs.internal.panama.headers_h.*; + +import com.nvidia.cuvs.CuVSMatrix; +import com.nvidia.cuvs.internal.panama.DLManagedTensor; +import java.lang.foreign.Arena; +import java.lang.foreign.MemorySegment; +import java.lang.foreign.ValueLayout; + +/** + * Internal interface for {@link CuVSMatrix}, for shared functionality that + * require/expose Panama types or internal cuvs types. + */ +interface CuVSMatrixInternal extends CuVSMatrix { + + MemorySegment memorySegment(); + + ValueLayout valueLayout(); + + /** + * Size (in bits) for the element type of this matrix + */ + default int bits() { + return (int) (valueLayout().byteSize() * 8); + } + + /** + * DLTensor data type {@code code} for the element type of this matrix + */ + default int code() { + return switch (dataType()) { + case FLOAT -> kDLFloat(); + case INT -> kDLInt(); + case UINT, BYTE -> kDLUInt(); + }; + } + + /** + * Creates a {@link DLManagedTensor} representing the matrix data and shape, to be + * passed to the CuVS C API. + * @param arena The Arena to use to allocate DL data structures + * @return a {@link MemorySegment} for the newly allocated DLManagedTensor + */ + MemorySegment toTensor(Arena arena); +} diff --git a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/TieredIndexImpl.java b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/TieredIndexImpl.java index 5edbcabd71..31967eee13 100644 --- a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/TieredIndexImpl.java +++ b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/TieredIndexImpl.java @@ -122,7 +122,7 @@ private IndexReference build() { long cols = dataset.columns(); // Get host data - MemorySegment hostDataSeg = ((CuVSHostMatrixImpl) dataset).memorySegment(); + MemorySegment hostDataSeg = ((CuVSMatrixInternal) dataset).memorySegment(); try (var resourceAccess = resources.access(); var indexParamsHandle = @@ -330,7 +330,7 @@ private void performExtend(CuVSMatrix extendDataset) { long cols = extendDataset.columns(); // Get host data - MemorySegment hostDataSeg = ((CuVSMatrixBaseImpl) extendDataset).memorySegment(); + MemorySegment hostDataSeg = ((CuVSMatrixInternal) extendDataset).memorySegment(); try (var resourceAccess = resources.access()) { long cuvsRes = resourceAccess.handle(); From 12ebfa1ee38131cb5dbb1f46ca16d833b98f3bb4 Mon Sep 17 00:00:00 2001 From: Simon Adorf Date: Wed, 24 Sep 2025 20:53:01 -0500 Subject: [PATCH 81/93] Fix k-means++ root rank race condition (#1359) Fixes https://github.com/rapidsai/cuvs/issues/1358 Authors: - Simon Adorf (https://github.com/csadorf) Approvers: - Divye Gala (https://github.com/divyegala) - Victor Lafargue (https://github.com/viclafargue) URL: https://github.com/rapidsai/cuvs/pull/1359 --- cpp/src/cluster/detail/kmeans_mg.cuh | 17 ++++++++++++++--- 1 file changed, 14 insertions(+), 3 deletions(-) diff --git a/cpp/src/cluster/detail/kmeans_mg.cuh b/cpp/src/cluster/detail/kmeans_mg.cuh index 8a39f91900..797e531f58 100644 --- a/cpp/src/cluster/detail/kmeans_mg.cuh +++ b/cpp/src/cluster/detail/kmeans_mg.cuh @@ -175,9 +175,20 @@ void initKMeansPlusPlus(const raft::resources& handle, // X which will be used as the initial centroid for kmeans++ // 1.3 - Communicate the initial centroid chosen by rank-r' to all other // ranks - std::mt19937 gen(params.rng_state.seed); - std::uniform_int_distribution<> dis(0, n_rank - 1); - int rp = dis(gen); + // Choose rp on rank 0 and broadcast to all ranks to guarantee agreement + int rp = 0; + if (my_rank == KMEANS_COMM_ROOT) { + std::mt19937 gen(params.rng_state.seed); + std::uniform_int_distribution<> dis(0, n_rank - 1); + rp = dis(gen); + } + { + rmm::device_scalar rp_d(stream); + raft::copy(rp_d.data(), &rp, 1, stream); + comm.bcast(rp_d.data(), 1, /*root=*/KMEANS_COMM_ROOT, stream); + raft::copy(&rp, rp_d.data(), 1, stream); + raft::resource::sync_stream(handle); + } // buffer to flag the sample that is chosen as initial centroids std::vector h_isSampleCentroid(n_samples); From c55cfd7f4c8f0ea8a1696b9dc2b86cceb5d7c1b5 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Lorenzo=20Dematt=C3=A9?= Date: Thu, 25 Sep 2025 11:48:16 +0200 Subject: [PATCH 82/93] [Java] Adding tests to use CuVSDeviceMatrix (device memory) directly as a CagraIndex input dataset (#1340) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit When we introduced `CuVSDeviceMatrix` in #1232, we made it possible to use device-memory-backed dataset as an input for index build: since we accept a `CuVSMatrix`, and we have correct `toTensor` implementations for CPU and GPU, and the underlying functions in libcuvs support different memory types and sizes (through DLManagedTensor information), this became supported "naturally". However, we never tested this explicitly. This PR adds tests to check and show that using CuVSDeviceMatrix (device memory) directly as a CagraIndex input dataset works as intended. (similar tests for other index types will be added as follow-ups) Authors: - Lorenzo Dematté (https://github.com/ldematte) - MithunR (https://github.com/mythrocks) Approvers: - Chris Hegarty (https://github.com/ChrisHegarty) - MithunR (https://github.com/mythrocks) URL: https://github.com/rapidsai/cuvs/pull/1340 --- .../common/CloseableRMMAllocation.java | 1 + .../nvidia/cuvs/CagraBuildAndSearchIT.java | 149 +++++++++++------- .../com/nvidia/cuvs/CagraRandomizedIT.java | 36 ++++- .../java/com/nvidia/cuvs/CuVSMatrixIT.java | 73 ++++----- 4 files changed, 155 insertions(+), 104 deletions(-) diff --git a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/common/CloseableRMMAllocation.java b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/common/CloseableRMMAllocation.java index 6dcc4c892f..98f269b573 100644 --- a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/common/CloseableRMMAllocation.java +++ b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/common/CloseableRMMAllocation.java @@ -79,6 +79,7 @@ private boolean mustClose() { public void close() { if (mustClose()) { checkCuVSError(cuvsRMMFree(cuvsResourceHandle, pointer, numBytes), "cuvsRMMFree"); + pointer = MemorySegment.NULL; } } 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 c0c0c8a37a..45076e6482 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 @@ -16,15 +16,13 @@ package com.nvidia.cuvs; import static com.carrotsearch.randomizedtesting.RandomizedTest.assumeTrue; +import static com.nvidia.cuvs.CuVSMatrixIT.assertSame2dArray; import static org.junit.Assert.*; import com.carrotsearch.randomizedtesting.RandomizedRunner; import com.nvidia.cuvs.CagraIndexParams.CagraGraphBuildAlgo; import com.nvidia.cuvs.CagraIndexParams.CuvsDistanceType; import com.nvidia.cuvs.CagraMergeParams.MergeStrategy; -import java.io.File; -import java.io.FileInputStream; -import java.io.FileOutputStream; import java.io.InputStream; import java.lang.foreign.Arena; import java.lang.foreign.Linker; @@ -397,12 +395,12 @@ public void testPrefilteringReducesResults() throws Throwable { .withMetric(CuvsDistanceType.L2Expanded) .build(); - try (CuVSResources resources = CheckedCuVSResources.create()) { - CagraIndex index = - CagraIndex.newBuilder(resources) - .withDataset(dataset) - .withIndexParams(indexParams) - .build(); + try (CuVSResources resources = CheckedCuVSResources.create(); + CagraIndex index = + CagraIndex.newBuilder(resources) + .withDataset(dataset) + .withIndexParams(indexParams) + .build()) { // No prefilter (all points allowed) CagraSearchParams searchParams = new CagraSearchParams.Builder().build(); @@ -514,6 +512,10 @@ private void cleanup(CagraIndex index, CagraIndex loadedIndex) throws Throwable loadedIndex.close(); } + /** + * Tests that an index built starting from a native MemorySegment is identical to one built from + * Java heap arrays + */ @Test public void testNativeDatasetEquivalent() throws Throwable { float[][] sampleData = createSampleData(); @@ -555,6 +557,37 @@ public void testNativeDatasetEquivalent() throws Throwable { } } + /** + * Tests that an index built starting from device memory ({@link CuVSDeviceMatrix}) is identical to one + * built from Java heap arrays + */ + @Test + public void testDeviceDatasetEquivalent() throws Throwable { + float[][] sampleData = createSampleData(); + + try (var resources = CuVSResources.create(); + var javaDataset = CuVSMatrix.ofArray(sampleData); + var deviceDataset = javaDataset.toDevice(resources)) { + + // Indexing with an on-heap and native datasets produce the same results + var javaIndex = indexOnce(javaDataset, resources); + var deviceIndex = indexOnce(deviceDataset, resources); + + int size = (int) javaIndex.getGraph().size(); + assertEquals(size, (int) deviceIndex.getGraph().size()); + + int columns = (int) javaIndex.getGraph().columns(); + assertEquals(columns, (int) deviceIndex.getGraph().columns()); + + var javaIndexGraph = new int[size][columns]; + var deviceIndexGraph = new int[size][columns]; + javaIndex.getGraph().toArray(javaIndexGraph); + deviceIndex.getGraph().toArray(deviceIndexGraph); + + assertSame2dArray(size, columns, javaIndexGraph, deviceIndexGraph); + } + } + @Test public void testMergingIndexes() throws Throwable { float[][] vector1 = { @@ -626,22 +659,23 @@ public void testMergingIndexes() throws Throwable { // --- Serialization/deserialization check --- String indexFileName = UUID.randomUUID() + ".cag"; - mergedIndex.serialize(new FileOutputStream(indexFileName)); + var indexFile = Path.of(indexFileName); - File indexFile = new File(indexFileName); - InputStream inputStream = new FileInputStream(indexFile); - CagraIndex loadedMergedIndex = CagraIndex.newBuilder(resources).from(inputStream).build(); + try (var out = Files.newOutputStream(indexFile)) { + mergedIndex.serialize(out); + } - SearchResults resultsFromLoaded = loadedMergedIndex.search(query); - assertEquals(expectedResults, resultsFromLoaded.getResults()); + try (InputStream inputStream = Files.newInputStream(indexFile)) { + CagraIndex loadedMergedIndex = CagraIndex.newBuilder(resources).from(inputStream).build(); - if (indexFile.exists()) { - indexFile.delete(); + SearchResults resultsFromLoaded = loadedMergedIndex.search(query); + assertEquals(expectedResults, resultsFromLoaded.getResults()); + mergedIndex.close(); + loadedMergedIndex.close(); } + Files.deleteIfExists(indexFile); index1.close(); index2.close(); - mergedIndex.close(); - loadedMergedIndex.close(); } } @@ -710,50 +744,51 @@ public void testMergeStrategies() throws Throwable { .build(); log.trace("Merging indexes with PHYSICAL strategy..."); - CagraIndex physicalMergedIndex = - CagraIndex.merge(new CagraIndex[] {index1, index2}, physicalMergeParams); - log.trace("Physical merge completed successfully"); + try (CagraIndex physicalMergedIndex = + CagraIndex.merge(new CagraIndex[] {index1, index2}, physicalMergeParams)) { + log.trace("Physical merge completed successfully"); + + CagraSearchParams searchParams = new CagraSearchParams.Builder().build(); + + CagraQuery query = + new CagraQuery.Builder(resources) + .withTopK(3) + .withSearchParams(searchParams) + .withQueryVectors(queries) + .withMapping(SearchResults.IDENTITY_MAPPING) + .build(); + + log.trace("Searching physically merged index..."); + SearchResults physicalResults = physicalMergedIndex.search(query); + assertNotNull("Physical merge search results should not be null", physicalResults); + assertEquals( + "Physical merge search results should match expected", + expectedResults, + physicalResults.getResults()); - CagraSearchParams searchParams = new CagraSearchParams.Builder().build(); + // --- Serialization/deserialization check for both merged indexes --- + String physicalIndexFileName = UUID.randomUUID() + ".cag"; + var physicalIndexFile = Path.of(physicalIndexFileName); - CagraQuery query = - new CagraQuery.Builder(resources) - .withTopK(3) - .withSearchParams(searchParams) - .withQueryVectors(queries) - .withMapping(SearchResults.IDENTITY_MAPPING) - .build(); + try (var out = Files.newOutputStream(physicalIndexFile)) { + physicalMergedIndex.serialize(out); + } - log.trace("Searching physically merged index..."); - SearchResults physicalResults = physicalMergedIndex.search(query); - assertNotNull("Physical merge search results should not be null", physicalResults); - assertEquals( - "Physical merge search results should match expected", - expectedResults, - physicalResults.getResults()); - - // --- Serialization/deserialization check for both merged indexes --- - String physicalIndexFileName = UUID.randomUUID().toString() + ".cag"; - physicalMergedIndex.serialize(new FileOutputStream(physicalIndexFileName)); - - File physicalIndexFile = new File(physicalIndexFileName); - InputStream physicalInputStream = new FileInputStream(physicalIndexFile); - CagraIndex loadedPhysicalIndex = - CagraIndex.newBuilder(resources).from(physicalInputStream).build(); - - SearchResults resultsFromLoadedPhysical = loadedPhysicalIndex.search(query); - assertEquals( - "Loaded physical index search results should match expected", - expectedResults, - resultsFromLoadedPhysical.getResults()); - - if (physicalIndexFile.exists()) { - physicalIndexFile.delete(); + try (InputStream physicalInputStream = Files.newInputStream(physicalIndexFile); + CagraIndex loadedPhysicalIndex = + CagraIndex.newBuilder(resources).from(physicalInputStream).build()) { + + Files.deleteIfExists(physicalIndexFile); + + SearchResults resultsFromLoadedPhysical = loadedPhysicalIndex.search(query); + assertEquals( + "Loaded physical index search results should match expected", + expectedResults, + resultsFromLoadedPhysical.getResults()); + } } index1.close(); index2.close(); - physicalMergedIndex.close(); - loadedPhysicalIndex.close(); } } } 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 6e02793478..d7d41c8c2f 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 @@ -39,17 +39,26 @@ public void setup() { log.trace("Random context initialized for test."); } + enum TestDatasetMemoryKind { + HEAP, + NATIVE, + DEVICE + } + @Test public void testResultsTopKWithRandomValues() throws Throwable { - boolean[] useNativeMemoryDatasets = {true, false}; + TestDatasetMemoryKind[] testDatasetMemoryKinds = { + TestDatasetMemoryKind.HEAP, TestDatasetMemoryKind.NATIVE, TestDatasetMemoryKind.DEVICE + }; for (int i = 0; i < 100; i++) { - for (boolean use : useNativeMemoryDatasets) { - tmpResultsTopKWithRandomValues(use); + for (var datasetMemoryKind : testDatasetMemoryKinds) { + tmpResultsTopKWithRandomValues(datasetMemoryKind); } } } - private void tmpResultsTopKWithRandomValues(boolean useNativeMemoryDataset) throws Throwable { + private void tmpResultsTopKWithRandomValues(TestDatasetMemoryKind datasetMemoryKind) + throws Throwable { int DATASET_SIZE_LIMIT = 10_000; int DIMENSIONS_LIMIT = 2048; int NUM_QUERIES_LIMIT = 10; @@ -90,7 +99,7 @@ private void tmpResultsTopKWithRandomValues(boolean useNativeMemoryDataset) thro log.debug("Dataset size: {}x{}", datasetSize, dimensions); log.debug("Query size: {}x{}", numQueries, dimensions); log.debug("TopK: {}", topK); - log.debug("Use native memory dataset? " + useNativeMemoryDataset); + log.debug("Use memory dataset: " + datasetMemoryKind.name()); // Debugging: Log dataset and queries if (log.isDebugEnabled()) { @@ -119,8 +128,8 @@ private void tmpResultsTopKWithRandomValues(boolean useNativeMemoryDataset) thro .withCagraGraphBuildAlgo(CagraGraphBuildAlgo.NN_DESCENT) .build(); - CagraIndex index; - if (useNativeMemoryDataset) { + final CagraIndex index; + if (datasetMemoryKind == TestDatasetMemoryKind.NATIVE) { var datasetBuilder = CuVSMatrix.hostBuilder(vectors.length, vectors[0].length, CuVSMatrix.DataType.FLOAT); for (float[] v : vectors) { @@ -131,7 +140,20 @@ private void tmpResultsTopKWithRandomValues(boolean useNativeMemoryDataset) thro .withDataset(datasetBuilder.build()) .withIndexParams(indexParams) .build(); + } else if (datasetMemoryKind == TestDatasetMemoryKind.DEVICE) { + var datasetBuilder = + CuVSMatrix.deviceBuilder( + resources, vectors.length, vectors[0].length, CuVSMatrix.DataType.FLOAT); + for (float[] v : vectors) { + datasetBuilder.addVector(v); + } + index = + CagraIndex.newBuilder(resources) + .withDataset(datasetBuilder.build()) + .withIndexParams(indexParams) + .build(); } else { + assert datasetMemoryKind == TestDatasetMemoryKind.HEAP; index = CagraIndex.newBuilder(resources) .withDataset(vectors) 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 499c7d366f..917acd4fc4 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 @@ -211,11 +211,7 @@ public void testIntDatasetCopy() { try (var dataset = CuVSMatrix.ofArray(intData)) { var intDataCopy = new int[(int) dataset.size()][(int) dataset.columns()]; dataset.toArray(intDataCopy); - for (int n = 0; n < dataset.size(); ++n) { - for (int i = 0; i < dataset.columns(); ++i) { - assertEquals(intData[n][i], intDataCopy[n][i]); - } - } + assertSame2dArray(dataset.size(), dataset.columns(), intData, intDataCopy); } } @@ -254,10 +250,32 @@ public void testFloatDatasetCopy() { try (var dataset = CuVSMatrix.ofArray(floatData)) { var dataCopy = new float[(int) dataset.size()][(int) dataset.columns()]; dataset.toArray(dataCopy); - for (int n = 0; n < dataset.size(); ++n) { - for (int i = 0; i < dataset.columns(); ++i) { - assertEquals(floatData[n][i], dataCopy[n][i], DELTA); - } + assertSame2dArray(dataset.size(), dataset.columns(), floatData, dataCopy); + } + } + + static void assertSame2dArray(long rows, long cols, float[][] array1, float[][] array2) { + assertEquals(rows, array1.length); + assertEquals(cols, array1[0].length); + assertEquals(rows, array2.length); + assertEquals(cols, array2[0].length); + + for (int n = 0; n < rows; ++n) { + for (int i = 0; i < cols; ++i) { + assertEquals(array1[n][i], array2[n][i], DELTA); + } + } + } + + static void assertSame2dArray(long rows, long cols, int[][] array1, int[][] array2) { + assertEquals(rows, array1.length); + assertEquals(cols, array1[0].length); + assertEquals(rows, array2.length); + assertEquals(cols, array2[0].length); + + for (int n = 0; n < rows; ++n) { + for (int i = 0; i < cols; ++i) { + assertEquals(array1[n][i], array2[n][i]); } } } @@ -280,11 +298,7 @@ private void testFloatDatasetBuilder(int rows, int cols, CuVSMatrix.Builder b try (var dataset = builder.build()) { dataset.toArray(roundTripData); - for (int n = 0; n < dataset.size(); ++n) { - for (int i = 0; i < dataset.columns(); ++i) { - assertEquals(data[n][i], roundTripData[n][i], DELTA); - } - } + assertSame2dArray(dataset.size(), dataset.columns(), data, roundTripData); } } @@ -323,12 +337,7 @@ private void testIntDatasetBuilder(int rows, int cols, CuVSMatrix.Builder bui try (var dataset = builder.build()) { dataset.toArray(roundTripData); - - for (int n = 0; n < dataset.size(); ++n) { - for (int i = 0; i < dataset.columns(); ++i) { - assertEquals(data[n][i], roundTripData[n][i]); - } - } + assertSame2dArray(dataset.size(), dataset.columns(), data, roundTripData); } } @@ -421,11 +430,7 @@ public void testDeviceToHost() throws Throwable { hostMatrix.toArray(roundTripData); - for (int n = 0; n < hostMatrix.size(); ++n) { - for (int i = 0; i < hostMatrix.columns(); ++i) { - assertEquals(data[n][i], roundTripData[n][i], 1e-9); - } - } + assertSame2dArray(hostMatrix.size(), hostMatrix.columns(), data, roundTripData); } } } @@ -452,11 +457,7 @@ public void testHostToDevice() throws Throwable { deviceMatrix.toArray(roundTripData); - for (int n = 0; n < deviceMatrix.size(); ++n) { - for (int i = 0; i < deviceMatrix.columns(); ++i) { - assertEquals(data[n][i], roundTripData[n][i], 1e-9); - } - } + assertSame2dArray(deviceMatrix.size(), deviceMatrix.columns(), data, roundTripData); } } } @@ -481,11 +482,7 @@ public void testHostToHostReturnsWeakReferenceSameData() { hostMatrix2.toArray(roundTripData); - for (int n = 0; n < hostMatrix2.size(); ++n) { - for (int i = 0; i < hostMatrix2.columns(); ++i) { - assertEquals(data[n][i], roundTripData[n][i], 1e-9); - } - } + assertSame2dArray(hostMatrix2.size(), hostMatrix2.columns(), data, roundTripData); } } @@ -517,11 +514,7 @@ public void testDeviceToDeviceReturnsWeakReferenceSameData() throws Throwable { deviceMatrix2.toArray(roundTripData); - for (int n = 0; n < deviceMatrix2.size(); ++n) { - for (int i = 0; i < deviceMatrix2.columns(); ++i) { - assertEquals(data[n][i], roundTripData[n][i], 1e-9); - } - } + assertSame2dArray(deviceMatrix2.size(), deviceMatrix2.columns(), data, roundTripData); } } } From c07ab70786d5553d4bb726a9d25aa643352ef946 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Lorenzo=20Dematt=C3=A9?= Date: Thu, 25 Sep 2025 18:50:17 +0200 Subject: [PATCH 83/93] [Java] Buffered device matrix builder (#1332) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The current `CuVSDeviceMatrix` builder implementation uses a non-native `MemorySegment` created directly from on-heap java arrays to avoid an intermediate allocation and copy to a native (off-heap) segment. However, this means that copy from host (Java heap) memory to device memory happens row by row (vector by vector), which is not ideal. This PR introduces a new builder implementation; it still returns a `CuVSDeviceMatrix` backed by managed RMM device memory, but it uses an intermediate buffer (pinned, implemented by a new `PinnedMemoryBuffer` class extracted/refactored from the existing code) to batch data before copying it to the GPU. Even if this means an intermediate allocation and copy to a native (off-heap) segment (so, data is copied twice), using pinned memory is up to 2x faster; benchmarks reveal that copying data directly from on-heap memory and using this intermediate buffer have similar throughput: Before (heap to GPU, row-by-row) ``` Benchmark (dims) (size) Mode Cnt Score Error Units CuVSDeviceMatrixBenchmarks.matrixDeviceBuilder 2048 16384 thrpt 5 34.932 ± 0.574 ops/s ``` After (heap to native memory to GPU, batched) ``` Benchmark (dims) (size) Mode Cnt Score Error Units CuVSDeviceMatrixBenchmarks.matrixDeviceBuilder 2048 16384 thrpt 5 35.622 ± 0.304 ops/s ``` So essentially identical; however, I think that despite the added complexity the change is still worth it as this method significantly reduces the number of calls to the GPU, reducing the risk of device locking due to syncs. But I'm open to feedback, comments and suggestions! Authors: - Lorenzo Dematté (https://github.com/ldematte) - MithunR (https://github.com/mythrocks) - Ben Frederickson (https://github.com/benfred) Approvers: - MithunR (https://github.com/mythrocks) URL: https://github.com/rapidsai/cuvs/pull/1332 --- .../cuvs/internal/CuVSDeviceMatrixImpl.java | 54 +++---------- .../internal/common/PinnedMemoryBuffer.java | 79 +++++++++++++++++++ .../com/nvidia/cuvs/spi/JDKProvider.java | 71 +++++++++++++---- 3 files changed, 145 insertions(+), 59 deletions(-) create mode 100644 java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/common/PinnedMemoryBuffer.java diff --git a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CuVSDeviceMatrixImpl.java b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CuVSDeviceMatrixImpl.java index 27f2adbd02..2e72ea69f5 100644 --- a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CuVSDeviceMatrixImpl.java +++ b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CuVSDeviceMatrixImpl.java @@ -15,21 +15,17 @@ */ package com.nvidia.cuvs.internal; -import static com.nvidia.cuvs.internal.common.LinkerHelper.C_POINTER; import static com.nvidia.cuvs.internal.common.Util.*; import static com.nvidia.cuvs.internal.panama.headers_h.*; import com.nvidia.cuvs.*; +import com.nvidia.cuvs.internal.common.PinnedMemoryBuffer; import com.nvidia.cuvs.internal.panama.DLManagedTensor; import com.nvidia.cuvs.internal.panama.DLTensor; import java.lang.foreign.*; public class CuVSDeviceMatrixImpl extends CuVSMatrixBaseImpl implements CuVSDeviceMatrix { - private static final int CHUNK_BYTES = - 8 * 1024 * 1024; // Based on benchmarks, 8MB seems the minimum size to optimize PCIe bandwidth - private final long hostBufferBytes; - private long bufferedMatrixRowStart = 0; private long bufferedMatrixRowEnd = 0; @@ -38,7 +34,7 @@ public class CuVSDeviceMatrixImpl extends CuVSMatrixBaseImpl implements CuVSDevi private final long rowStride; private final long columnStride; - private MemorySegment hostBuffer = MemorySegment.NULL; + private final PinnedMemoryBuffer hostBuffer; protected CuVSDeviceMatrixImpl( CuVSResources resources, @@ -63,18 +59,7 @@ protected CuVSDeviceMatrixImpl( this.resources = resources; this.rowStride = rowStride; this.columnStride = columnStride; - - long rowBytes = columns * valueLayout.byteSize(); - long matrixBytes = size * rowBytes; - if (matrixBytes < CHUNK_BYTES) { - this.hostBufferBytes = matrixBytes; - } else if (rowBytes > CHUNK_BYTES) { - // We need to buffer at least one row at time - this.hostBufferBytes = rowBytes; - } else { - var rowCount = (CHUNK_BYTES / rowBytes); - this.hostBufferBytes = rowBytes * rowCount; - } + this.hostBuffer = new PinnedMemoryBuffer(size, columns, valueLayout); } @Override @@ -84,27 +69,10 @@ public MemorySegment toTensor(Arena arena) { arena, memorySegment, new long[] {size, columns}, strides, code(), bits(), kDLCUDA()); } - private static MemorySegment createPinnedBuffer(long bufferBytes) { - try (var localArena = Arena.ofConfined()) { - MemorySegment pointer = localArena.allocate(C_POINTER); - checkCudaError(cudaMallocHost(pointer, bufferBytes), "cudaMallocHost"); - return pointer.get(C_POINTER, 0); - } - } - - private static void destroyPinnedBuffer(MemorySegment bufferSegment) { - checkCudaError(cudaFreeHost(bufferSegment), "cudaFreeHost"); - } - private void populateBuffer(long startRow) { - if (hostBuffer == MemorySegment.NULL) { - // System.out.println("Creating a buffer of size " + hostBufferBytes); - hostBuffer = createPinnedBuffer(hostBufferBytes); - } - try (var localArena = Arena.ofConfined()) { long rowBytes = columns * valueLayout.byteSize(); - var endRow = Math.min(startRow + (hostBufferBytes / rowBytes), size); + var endRow = Math.min(startRow + (hostBuffer.size() / rowBytes), size); var rowCount = endRow - startRow; // System.out.printf( @@ -123,7 +91,12 @@ private void populateBuffer(long startRow) { MemorySegment bufferTensor = prepareTensor( - localArena, hostBuffer, new long[] {rowCount, columns}, code(), bits(), kDLCPU()); + localArena, + hostBuffer.address(), + new long[] {rowCount, columns}, + code(), + bits(), + kDLCPU()); try (var resourceAccess = resources.access()) { checkCuVSError( @@ -146,7 +119,7 @@ public RowView getRow(long row) { var startRow = row - bufferedMatrixRowStart; return new SliceRowView( - hostBuffer.asSlice(startRow * columns * valueByteSize, columns * valueByteSize), + hostBuffer.address().asSlice(startRow * columns * valueByteSize, columns * valueByteSize), columns, valueLayout, dataType, @@ -248,10 +221,7 @@ public void toDevice(CuVSDeviceMatrix targetMatrix, CuVSResources cuVSResources) @Override public void close() { - if (hostBuffer != MemorySegment.NULL) { - destroyPinnedBuffer(hostBuffer); - hostBuffer = MemorySegment.NULL; - } + hostBuffer.close(); } private static class CuVSDeviceMatrixDelegate implements CuVSDeviceMatrix, CuVSMatrixInternal { diff --git a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/common/PinnedMemoryBuffer.java b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/common/PinnedMemoryBuffer.java new file mode 100644 index 0000000000..23f4068bad --- /dev/null +++ b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/common/PinnedMemoryBuffer.java @@ -0,0 +1,79 @@ +/* + * Copyright (c) 2025, NVIDIA CORPORATION. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.nvidia.cuvs.internal.common; + +import static com.nvidia.cuvs.internal.common.LinkerHelper.C_POINTER; +import static com.nvidia.cuvs.internal.common.Util.checkCudaError; +import static com.nvidia.cuvs.internal.panama.headers_h.*; + +import java.lang.foreign.Arena; +import java.lang.foreign.MemorySegment; +import java.lang.foreign.ValueLayout; + +public class PinnedMemoryBuffer implements AutoCloseable { + + private static final int CHUNK_BYTES = + 8 * 1024 * 1024; // Based on benchmarks, 8MB seems the minimum size to optimize PCIe bandwidth + private final long hostBufferBytes; + + private MemorySegment hostBuffer = MemorySegment.NULL; + + public PinnedMemoryBuffer(long rows, long columns, ValueLayout valueLayout) { + + long rowBytes = columns * valueLayout.byteSize(); + long matrixBytes = rows * rowBytes; + if (matrixBytes < CHUNK_BYTES) { + this.hostBufferBytes = matrixBytes; + } else if (rowBytes > CHUNK_BYTES) { + // We need to buffer at least one row at time + this.hostBufferBytes = rowBytes; + } else { + var rowCount = (CHUNK_BYTES / rowBytes); + this.hostBufferBytes = rowBytes * rowCount; + } + } + + private static MemorySegment createPinnedBuffer(long bufferBytes) { + try (var localArena = Arena.ofConfined()) { + MemorySegment pointer = localArena.allocate(C_POINTER); + checkCudaError(cudaMallocHost(pointer, bufferBytes), "cudaMallocHost"); + return pointer.get(C_POINTER, 0); + } + } + + private static void destroyPinnedBuffer(MemorySegment bufferSegment) { + checkCudaError(cudaFreeHost(bufferSegment), "cudaFreeHost"); + } + + public MemorySegment address() { + if (hostBuffer == MemorySegment.NULL) { + hostBuffer = createPinnedBuffer(hostBufferBytes); + } + return hostBuffer; + } + + public long size() { + return hostBufferBytes; + } + + @Override + public void close() { + if (hostBuffer != MemorySegment.NULL) { + destroyPinnedBuffer(hostBuffer); + hostBuffer = MemorySegment.NULL; + } + } +} diff --git a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/spi/JDKProvider.java b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/spi/JDKProvider.java index 6f3d3fa6e3..ed8af603e0 100644 --- a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/spi/JDKProvider.java +++ b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/spi/JDKProvider.java @@ -18,9 +18,11 @@ import static com.nvidia.cuvs.internal.common.Util.*; import static com.nvidia.cuvs.internal.panama.headers_h.cuvsVersionGet; import static com.nvidia.cuvs.internal.panama.headers_h.uint16_t; +import static com.nvidia.cuvs.internal.panama.headers_h_1.cudaStreamSynchronize; import com.nvidia.cuvs.*; import com.nvidia.cuvs.internal.*; +import com.nvidia.cuvs.internal.common.PinnedMemoryBuffer; import com.nvidia.cuvs.internal.common.Util; import java.io.IOException; import java.lang.foreign.Arena; @@ -216,7 +218,7 @@ public CuVSHostMatrix build() { public CuVSMatrix.Builder newDeviceMatrixBuilder( CuVSResources resources, long size, long columns, CuVSMatrix.DataType dataType) throws UnsupportedOperationException { - return new HeapSegmentBuilder(resources, size, columns, dataType); + return new BufferedSegmentBuilder(resources, size, columns, dataType); } @Override @@ -227,7 +229,7 @@ public CuVSMatrix.Builder newDeviceMatrixBuilder( int rowStride, int columnStride, CuVSMatrix.DataType dataType) { - return new HeapSegmentBuilder(resources, size, columns, rowStride, columnStride, dataType); + return new BufferedSegmentBuilder(resources, size, columns, rowStride, columnStride, dataType); } @Override @@ -279,28 +281,38 @@ public CuVSMatrix newMatrixFromArray(byte[][] vectors) { /** * This {@link CuVSDeviceMatrix} builder implementation returns a {@link CuVSDeviceMatrix} backed by managed RMM - * device memory. It uses a non-native {@link MemorySegment} created directly from on-heap java arrays to avoid - * an intermediate allocation and copy to a native (off-heap) segment. - * It requires the copy function ({@code cudaMemcpyAsync}) to have the {@code Critical} linker option in order - * to allow the access to on-heap memory (see {@link Util#cudaMemcpyAsync}). + * device memory. It uses a {@link PinnedMemoryBuffer} to batch data before copying it to the GPU. */ - private static class HeapSegmentBuilder implements CuVSMatrix.Builder { + private static class BufferedSegmentBuilder implements CuVSMatrix.Builder { + private final long columns; private final long size; private final CuVSDeviceMatrixImpl matrix; private final MemorySegment stream; - private int current; - private HeapSegmentBuilder( + private final long rowBytes; + private int currentRow; + + private final PinnedMemoryBuffer hostBuffer; + private final long bufferRowCount; + private int currentBufferRow; + + private BufferedSegmentBuilder( CuVSResources resources, long size, long columns, CuVSMatrix.DataType dataType) { this.columns = columns; this.size = size; this.matrix = CuVSDeviceMatrixRMMImpl.create(resources, size, columns, dataType); this.stream = Util.getStream(resources); - this.current = 0; + this.currentRow = 0; + + this.hostBuffer = new PinnedMemoryBuffer(size, columns, matrix.valueLayout()); + + this.rowBytes = columns * matrix.valueLayout().byteSize(); + this.bufferRowCount = Math.min((hostBuffer.size() / rowBytes), size); + this.currentBufferRow = 0; } - private HeapSegmentBuilder( + private BufferedSegmentBuilder( CuVSResources resources, long size, long columns, @@ -313,7 +325,13 @@ private HeapSegmentBuilder( CuVSDeviceMatrixRMMImpl.create( resources, size, columns, rowStride, columnStride, dataType); this.stream = Util.getStream(resources); - this.current = 0; + this.currentRow = 0; + + this.hostBuffer = new PinnedMemoryBuffer(size, columns, matrix.valueLayout()); + + this.rowBytes = columns * matrix.valueLayout().byteSize(); + this.bufferRowCount = Math.min((hostBuffer.size() / rowBytes), size); + this.currentBufferRow = 0; } @Override @@ -347,19 +365,38 @@ public void addVector(int[] vector) { } private void internalAddVector(MemorySegment vector) { - if (current >= size) { + if (currentRow >= size) { throw new ArrayIndexOutOfBoundsException(); } + var hostBufferOffset = currentBufferRow * rowBytes; + MemorySegment.copy(vector, 0, hostBuffer.address(), hostBufferOffset, rowBytes); - long rowBytes = columns * matrix.valueLayout().byteSize(); + currentRow++; + currentBufferRow++; + if (currentBufferRow == bufferRowCount) { + flushBuffer(); + } + } - var dstOffset = ((current++) * rowBytes); - var dst = matrix.memorySegment().asSlice(dstOffset); - cudaMemcpyAsync(dst, vector, rowBytes, CudaMemcpyKind.HOST_TO_DEVICE, stream); + private void flushBuffer() { + if (currentBufferRow > 0) { + var deviceMemoryOffset = (currentRow - currentBufferRow) * rowBytes; + var dst = matrix.memorySegment().asSlice(deviceMemoryOffset); + cudaMemcpyAsync( + dst, + hostBuffer.address(), + currentBufferRow * rowBytes, + CudaMemcpyKind.HOST_TO_DEVICE, + stream); + currentBufferRow = 0; + checkCudaError(cudaStreamSynchronize(stream), "cudaStreamSynchronize"); + } } @Override public CuVSDeviceMatrix build() { + flushBuffer(); + hostBuffer.close(); return matrix; } } From e32b14dbb9ff3f414998dfcbd56c5aff02df72ed Mon Sep 17 00:00:00 2001 From: Tarang Jain <40517122+tarang-jain@users.noreply.github.com> Date: Thu, 25 Sep 2025 16:42:12 -0700 Subject: [PATCH 84/93] [DOCS] Update IVF Docs with Supported Distance Metrics (#1365) Fixes #1355 Authors: - Tarang Jain (https://github.com/tarang-jain) Approvers: - Corey J. Nolet (https://github.com/cjnolet) URL: https://github.com/rapidsai/cuvs/pull/1365 --- cpp/include/cuvs/neighbors/ivf_pq.hpp | 50 +++++++++++++++++++- python/cuvs/cuvs/neighbors/ivf_pq/ivf_pq.pyx | 5 +- 2 files changed, 53 insertions(+), 2 deletions(-) diff --git a/cpp/include/cuvs/neighbors/ivf_pq.hpp b/cpp/include/cuvs/neighbors/ivf_pq.hpp index 2db1f33eb2..593403f2aa 100644 --- a/cpp/include/cuvs/neighbors/ivf_pq.hpp +++ b/cpp/include/cuvs/neighbors/ivf_pq.hpp @@ -1,5 +1,5 @@ /* - * Copyright (c) 2024, NVIDIA CORPORATION. + * Copyright (c) 2024-2025, NVIDIA CORPORATION. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -570,6 +570,12 @@ auto build(raft::resources const& handle, /** * @brief Build the index from the dataset for efficient search. * + * NB: Currently, the following distance metrics are supported: + * - L2Expanded + * - L2Unexpanded + * - InnerProduct + * - CosineExpanded + * * Usage example: * @code{.cpp} * using namespace cuvs::neighbors; @@ -617,6 +623,12 @@ auto build(raft::resources const& handle, /** * @brief Build the index from the dataset for efficient search. * + * NB: Currently, the following distance metrics are supported: + * - L2Expanded + * - L2Unexpanded + * - InnerProduct + * - CosineExpanded + * * Usage example: * @code{.cpp} * using namespace cuvs::neighbors; @@ -663,6 +675,12 @@ auto build(raft::resources const& handle, /** * @brief Build the index from the dataset for efficient search. * + * NB: Currently, the following distance metrics are supported: + * - L2Expanded + * - L2Unexpanded + * - InnerProduct + * - CosineExpanded + * * Usage example: * @code{.cpp} * using namespace cuvs::neighbors; @@ -710,6 +728,12 @@ auto build(raft::resources const& handle, /** * @brief Build the index from the dataset for efficient search. * + * NB: Currently, the following distance metrics are supported: + * - L2Expanded + * - L2Unexpanded + * - InnerProduct + * - CosineExpanded + * * Usage example: * @code{.cpp} * using namespace cuvs::neighbors; @@ -763,6 +787,12 @@ auto build(raft::resources const& handle, /** * @brief Build the index from the dataset for efficient search. * + * NB: Currently, the following distance metrics are supported: + * - L2Expanded + * - L2Unexpanded + * - InnerProduct + * - CosineExpanded + * * Note, if index_params.add_data_on_build is set to true, the user can set a * stream pool in the input raft::resource with at least one stream to enable kernel and copy * overlapping. @@ -824,6 +854,12 @@ auto build(raft::resources const& handle, /** * @brief Build the index from the dataset for efficient search. * + * NB: Currently, the following distance metrics are supported: + * - L2Expanded + * - L2Unexpanded + * - InnerProduct + * - CosineExpanded + * * Usage example: * @code{.cpp} * using namespace cuvs::neighbors; @@ -871,6 +907,12 @@ auto build(raft::resources const& handle, /** * @brief Build the index from the dataset for efficient search. * + * NB: Currently, the following distance metrics are supported: + * - L2Expanded + * - L2Unexpanded + * - InnerProduct + * - CosineExpanded + * * Note, if index_params.add_data_on_build is set to true, the user can set a * stream pool in the input raft::resource with at least one stream to enable kernel and copy * overlapping. @@ -932,6 +974,12 @@ auto build(raft::resources const& handle, /** * @brief Build the index from the dataset for efficient search. * + * NB: Currently, the following distance metrics are supported: + * - L2Expanded + * - L2Unexpanded + * - InnerProduct + * - CosineExpanded + * * Note, if index_params.add_data_on_build is set to true, the user can set a * stream pool in the input raft::resource with at least one stream to enable kernel and copy * overlapping. diff --git a/python/cuvs/cuvs/neighbors/ivf_pq/ivf_pq.pyx b/python/cuvs/cuvs/neighbors/ivf_pq/ivf_pq.pyx index dd3b17f949..5e13b3bf3b 100644 --- a/python/cuvs/cuvs/neighbors/ivf_pq/ivf_pq.pyx +++ b/python/cuvs/cuvs/neighbors/ivf_pq/ivf_pq.pyx @@ -58,7 +58,8 @@ cdef class IndexParams: The number of clusters used in the coarse quantizer. metric : str, default="sqeuclidean" String denoting the metric type. - Valid values for metric: ["sqeuclidean", "inner_product", "euclidean"], + Valid values for metric: ["sqeuclidean", "inner_product", + "euclidean", "cosine"], where: - sqeuclidean is the euclidean distance without the square root @@ -66,6 +67,8 @@ cdef class IndexParams: - euclidean is the euclidean distance - inner product distance is defined as distance(a, b) = \\sum_i a_i * b_i. + - cosine distance is defined as + distance(a, b) = 1 - \\sum_i a_i * b_i / ( ||a||_2 * ||b||_2). kmeans_n_iters : int, default = 20 The number of iterations searching for kmeans centers during index From ee1ad84baf748c8e39fffc2a8f7a582d91df07ad Mon Sep 17 00:00:00 2001 From: Micka Date: Fri, 26 Sep 2025 10:50:23 +0200 Subject: [PATCH 85/93] Regroup overloads in docs (#1377) This PR introduce a custom js script to regroup overloads into a single link. Another alternative would be to have the brief of every function in there instead of grouping them, so that it could look like ``` - fit(float, int) - fit(double, int) - fit(float, int64) - ... ``` Before: grafik After: grafik Authors: - Micka (https://github.com/lowener) Approvers: - Ben Frederickson (https://github.com/benfred) URL: https://github.com/rapidsai/cuvs/pull/1377 --- .../cuvs/preprocessing/quantize/binary.h | 2 +- docs/source/_static/collapse_overloads.js | 20 +++++++++++++++++++ docs/source/conf.py | 1 + docs/source/cpp_api/cluster_agglomerative.rst | 4 ++-- docs/source/cpp_api/cluster_kmeans.rst | 2 +- 5 files changed, 25 insertions(+), 4 deletions(-) create mode 100644 docs/source/_static/collapse_overloads.js diff --git a/cpp/include/cuvs/preprocessing/quantize/binary.h b/cpp/include/cuvs/preprocessing/quantize/binary.h index 9a5834ef3e..7c6e522c85 100644 --- a/cpp/include/cuvs/preprocessing/quantize/binary.h +++ b/cpp/include/cuvs/preprocessing/quantize/binary.h @@ -95,7 +95,7 @@ cuvsError_t cuvsBinaryQuantizerCreate(cuvsBinaryQuantizer_t* quantizer); cuvsError_t cuvsBinaryQuantizerDestroy(cuvsBinaryQuantizer_t quantizer); /** - * @brief Trains a scalar quantizer to be used later for quantizing the dataset. + * @brief Trains a binary quantizer to be used later for quantizing the dataset. * * @param[in] res raft resource * @param[in] params configure binary quantizer, e.g. threshold diff --git a/docs/source/_static/collapse_overloads.js b/docs/source/_static/collapse_overloads.js new file mode 100644 index 0000000000..ea8135f407 --- /dev/null +++ b/docs/source/_static/collapse_overloads.js @@ -0,0 +1,20 @@ +document.addEventListener("DOMContentLoaded", () => { + const toc = document.querySelector(".bd-toc-nav"); + if (!toc) return; + + // Get all TOC links + const links = toc.querySelectorAll("a"); + + const seen = new Set(); + links.forEach(link => { + let text = link.textContent.trim(); + let norm = text.replace(/\(\)$/, ""); // strip trailing () + + if (seen.has(norm)) { + // hide duplicate + link.parentElement.style.display = "none"; + } else { + seen.add(norm); + } + }); + }); diff --git a/docs/source/conf.py b/docs/source/conf.py index a7247ed719..b5404edd22 100644 --- a/docs/source/conf.py +++ b/docs/source/conf.py @@ -192,6 +192,7 @@ def setup(app): app.add_js_file( "https://docs.rapids.ai/assets/js/custom.js", loading_method="defer" ) + app.add_js_file("collapse_overloads.js") # The following is used by sphinx.ext.linkcode to provide links to github diff --git a/docs/source/cpp_api/cluster_agglomerative.rst b/docs/source/cpp_api/cluster_agglomerative.rst index 08a811780a..e4bdf0c020 100644 --- a/docs/source/cpp_api/cluster_agglomerative.rst +++ b/docs/source/cpp_api/cluster_agglomerative.rst @@ -1,5 +1,5 @@ -Cluster -======= +Agglomerative +============= .. role:: py(code) :language: c++ diff --git a/docs/source/cpp_api/cluster_kmeans.rst b/docs/source/cpp_api/cluster_kmeans.rst index 6c342baabd..bae14dcb9d 100644 --- a/docs/source/cpp_api/cluster_kmeans.rst +++ b/docs/source/cpp_api/cluster_kmeans.rst @@ -1,4 +1,4 @@ -Cluster +K-Means ======= .. role:: py(code) From f328ff1e27b2ad65f53c34bf1bff8db6e38b95df Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Lorenzo=20Dematt=C3=A9?= Date: Fri, 26 Sep 2025 16:15:37 +0200 Subject: [PATCH 86/93] [Review][Java] Expand `CagraIndex#search` to accept more types (int8 and int32) (#1283) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit With the introduction of `CuVSMatrix` as the type to provide an input dataset, we expanded the types that can be used in indexing. With this PR we aim to do the same query/search: - Use a `CuVSMatrix` for `CagraQuery` too - Use the correct types for sizes in `CagraIndex#search` There is more we can do there, e.g. use `CuVSDeviceMatrix` for storing the results so we don't have to explicitly copy them (all duplicated code). That would be follow-up work. Authors: - Lorenzo Dematté (https://github.com/ldematte) - Ben Frederickson (https://github.com/benfred) Approvers: - Ben Frederickson (https://github.com/benfred) - Chris Hegarty (https://github.com/ChrisHegarty) URL: https://github.com/rapidsai/cuvs/pull/1283 --- .../main/java/com/nvidia/cuvs/CagraQuery.java | 21 +- .../nvidia/cuvs/internal/CagraIndexImpl.java | 89 +++---- .../cuvs/internal/CuVSDeviceMatrixImpl.java | 10 - .../cuvs/internal/CuVSMatrixInternal.java | 6 +- .../com/nvidia/cuvs/internal/common/Util.java | 14 +- .../nvidia/cuvs/CagraBuildAndSearchIT.java | 218 +++++++++--------- .../cuvs/CagraMultiThreadStabilityIT.java | 31 +-- .../com/nvidia/cuvs/CagraRandomizedIT.java | 4 +- .../com/nvidia/cuvs/CheckedCuVSResources.java | 13 +- .../java/com/nvidia/cuvs/CuVSResourcesIT.java | 2 +- 10 files changed, 216 insertions(+), 192 deletions(-) diff --git a/java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraQuery.java b/java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraQuery.java index f81ec3191c..78f4de759e 100644 --- a/java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraQuery.java +++ b/java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraQuery.java @@ -15,7 +15,6 @@ */ package com.nvidia.cuvs; -import java.util.Arrays; import java.util.BitSet; import java.util.Objects; import java.util.function.LongToIntFunction; @@ -34,7 +33,7 @@ public class CagraQuery { private final CagraSearchParams cagraSearchParameters; private final LongToIntFunction mapping; - private final float[][] queryVectors; + private final CuVSMatrix queryVectors; private final int topK; private final BitSet prefilter; private final int numDocs; @@ -53,9 +52,9 @@ public class CagraQuery { * @param numDocs Total number of dataset vectors; used to align the prefilter correctly * @param resources CuVSResources instance to use for this query */ - public CagraQuery( + private CagraQuery( CagraSearchParams cagraSearchParameters, - float[][] queryVectors, + CuVSMatrix queryVectors, LongToIntFunction mapping, int topK, BitSet prefilter, @@ -81,11 +80,9 @@ public CagraSearchParams getCagraSearchParameters() { } /** - * Gets the query vector 2D float array. - * - * @return 2D float array + * Gets the query vector matrix. */ - public float[][] getQueryVectors() { + public CuVSMatrix getQueryVectors() { return queryVectors; } @@ -137,7 +134,7 @@ public String toString() { return "CuVSQuery [cagraSearchParameters=" + cagraSearchParameters + ", queryVectors=" - + Arrays.toString(queryVectors) + + queryVectors.toString() + ", mapping=" + mapping + ", topK=" @@ -151,7 +148,7 @@ public String toString() { public static class Builder { private CagraSearchParams cagraSearchParams; - private float[][] queryVectors; + private CuVSMatrix queryVectors; private LongToIntFunction mapping = SearchResults.IDENTITY_MAPPING; private int topK = 2; private BitSet prefilter; @@ -186,10 +183,10 @@ public Builder withSearchParams(CagraSearchParams cagraSearchParams) { /** * Registers the query vectors to be passed in the search call. * - * @param queryVectors 2D float query vector array + * @param queryVectors 2D query vector array * @return an instance of this Builder */ - public Builder withQueryVectors(float[][] queryVectors) { + public Builder withQueryVectors(CuVSMatrix queryVectors) { this.queryVectors = queryVectors; return this; } diff --git a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CagraIndexImpl.java b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CagraIndexImpl.java index 4045d8b83d..a9ac722e17 100644 --- a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CagraIndexImpl.java +++ b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CagraIndexImpl.java @@ -17,16 +17,13 @@ import static com.nvidia.cuvs.internal.CuVSParamsHelper.*; import static com.nvidia.cuvs.internal.common.CloseableRMMAllocation.allocateRMMSegment; -import static com.nvidia.cuvs.internal.common.LinkerHelper.C_FLOAT; -import static com.nvidia.cuvs.internal.common.LinkerHelper.C_FLOAT_BYTE_SIZE; import static com.nvidia.cuvs.internal.common.LinkerHelper.C_INT; import static com.nvidia.cuvs.internal.common.LinkerHelper.C_INT_BYTE_SIZE; +import static com.nvidia.cuvs.internal.common.Util.CudaMemcpyKind.DEVICE_TO_HOST; import static com.nvidia.cuvs.internal.common.Util.CudaMemcpyKind.HOST_TO_DEVICE; -import static com.nvidia.cuvs.internal.common.Util.CudaMemcpyKind.INFER_DIRECTION; import static com.nvidia.cuvs.internal.common.Util.buildMemorySegment; import static com.nvidia.cuvs.internal.common.Util.checkCuVSError; import static com.nvidia.cuvs.internal.common.Util.concatenate; -import static com.nvidia.cuvs.internal.common.Util.cudaMemcpy; import static com.nvidia.cuvs.internal.common.Util.prepareTensor; import static com.nvidia.cuvs.internal.panama.headers_h.*; @@ -35,6 +32,7 @@ import com.nvidia.cuvs.internal.common.CloseableHandle; import com.nvidia.cuvs.internal.common.CloseableRMMAllocation; import com.nvidia.cuvs.internal.common.CompositeCloseableHandle; +import com.nvidia.cuvs.internal.common.Util; import com.nvidia.cuvs.internal.panama.*; import java.io.FileInputStream; import java.io.InputStream; @@ -232,19 +230,18 @@ public SearchResults search(CagraQuery query) throws Throwable { try (var localArena = Arena.ofConfined()) { checkNotDestroyed(); int topK = query.getTopK(); - long numQueries = query.getQueryVectors().length; + var queryVectors = (CuVSMatrixInternal) query.getQueryVectors(); + long numQueries = queryVectors.size(); long numBlocks = topK * numQueries; - int vectorDimension = numQueries > 0 ? query.getQueryVectors()[0].length : 0; SequenceLayout neighborsSequenceLayout = MemoryLayout.sequenceLayout(numBlocks, C_INT); - SequenceLayout distancesSequenceLayout = MemoryLayout.sequenceLayout(numBlocks, C_FLOAT); + SequenceLayout distancesSequenceLayout = + MemoryLayout.sequenceLayout(numBlocks, queryVectors.valueLayout()); MemorySegment neighborsMemorySegment = localArena.allocate(neighborsSequenceLayout); MemorySegment distancesMemorySegment = localArena.allocate(distancesSequenceLayout); - MemorySegment floatsSeg = buildMemorySegment(localArena, query.getQueryVectors()); - final long queriesBytes = C_FLOAT_BYTE_SIZE * numQueries * vectorDimension; final long neighborsBytes = C_INT_BYTE_SIZE * numQueries * topK; - final long distancesBytes = C_FLOAT_BYTE_SIZE * numQueries * topK; + final long distancesBytes = queryVectors.valueLayout().byteSize() * numQueries * topK; final boolean hasPreFilter = query.getPrefilter() != null; final BitSet[] prefilters = hasPreFilter ? new BitSet[] {query.getPrefilter()} : EMPTY_PREFILTER_BITSET; @@ -254,8 +251,10 @@ public SearchResults search(CagraQuery query) throws Throwable { try (var resourcesAccessor = query.getResources().access()) { var cuvsRes = resourcesAccessor.handle(); + var cuvsStream = Util.getStream(cuvsRes); - try (var queriesDP = allocateRMMSegment(cuvsRes, queriesBytes); + try (var deviceQueryVectors = + (CuVSMatrixInternal) queryVectors.toDevice(query.getResources()); var neighborsDP = allocateRMMSegment(cuvsRes, neighborsBytes); var distancesDP = allocateRMMSegment(cuvsRes, distancesBytes); var prefilterDP = @@ -263,12 +262,7 @@ public SearchResults search(CagraQuery query) throws Throwable { ? allocateRMMSegment(cuvsRes, prefilterBytes) : CloseableRMMAllocation.EMPTY) { - cudaMemcpy(queriesDP.handle(), floatsSeg, queriesBytes, INFER_DIRECTION); - - long[] queriesShape = {numQueries, vectorDimension}; - MemorySegment queriesTensor = - prepareTensor( - localArena, queriesDP.handle(), queriesShape, kDLFloat(), 32, kDLCUDA()); + var queryTensor = deviceQueryVectors.toTensor(localArena); long[] neighborsShape = {numQueries, topK}; MemorySegment neighborsTensor = prepareTensor( @@ -276,32 +270,34 @@ public SearchResults search(CagraQuery query) throws Throwable { long[] distancesShape = {numQueries, topK}; MemorySegment distancesTensor = prepareTensor( - localArena, distancesDP.handle(), distancesShape, kDLFloat(), 32, kDLCUDA()); - - var returnValue = cuvsStreamSync(cuvsRes); - checkCuVSError(returnValue, "cuvsStreamSync"); + localArena, + distancesDP.handle(), + distancesShape, + deviceQueryVectors.code(), + deviceQueryVectors.bits(), + kDLCUDA()); // prepare the prefiltering data - MemorySegment prefilterDataMemorySegment = MemorySegment.NULL; - if (hasPreFilter) { - BitSet concatenatedFilters = concatenate(prefilters, query.getNumDocs()); - long[] filters = concatenatedFilters.toLongArray(); - prefilterDataMemorySegment = buildMemorySegment(localArena, filters); - } - MemorySegment prefilter = cuvsFilter.allocate(localArena); - MemorySegment prefilterTensor; if (!hasPreFilter) { cuvsFilter.type(prefilter, 0); // NO_FILTER cuvsFilter.addr(prefilter, 0); } else { + BitSet concatenatedFilters = concatenate(prefilters, query.getNumDocs()); + long[] filters = concatenatedFilters.toLongArray(); + var prefilterDataMemorySegment = buildMemorySegment(localArena, filters); + long[] prefilterShape = {prefilterLen}; - cudaMemcpy( - prefilterDP.handle(), prefilterDataMemorySegment, prefilterBytes, HOST_TO_DEVICE); + Util.cudaMemcpyAsync( + prefilterDP.handle(), + prefilterDataMemorySegment, + prefilterBytes, + HOST_TO_DEVICE, + cuvsStream); - prefilterTensor = + MemorySegment prefilterTensor = prepareTensor( localArena, prefilterDP.handle(), prefilterShape, kDLUInt(), 32, kDLCUDA()); @@ -309,25 +305,36 @@ public SearchResults search(CagraQuery query) throws Throwable { cuvsFilter.addr(prefilter, prefilterTensor.address()); } - returnValue = cuvsStreamSync(cuvsRes); - checkCuVSError(returnValue, "cuvsStreamSync"); + // TODO: do we need this stream sync here? + checkCuVSError(cuvsStreamSync(cuvsRes), "cuvsStreamSync"); - returnValue = + var returnValue = cuvsCagraSearch( cuvsRes, segmentFromSearchParams(localArena, query.getCagraSearchParameters()), cagraIndexReference.getMemorySegment(), - queriesTensor, + queryTensor, neighborsTensor, distancesTensor, prefilter); checkCuVSError(returnValue, "cuvsCagraSearch"); - returnValue = cuvsStreamSync(cuvsRes); - checkCuVSError(returnValue, "cuvsStreamSync"); - - cudaMemcpy(neighborsMemorySegment, neighborsDP.handle(), neighborsBytes, INFER_DIRECTION); - cudaMemcpy(distancesMemorySegment, distancesDP.handle(), distancesBytes, INFER_DIRECTION); + // TODO: we can avoid/defer this using CuVSDeviceMatrix for neighborsDP and distancesDP + // TODO: also, should we use cuvsMatrixCopy instead? + Util.cudaMemcpyAsync( + neighborsMemorySegment, + neighborsDP.handle(), + neighborsBytes, + DEVICE_TO_HOST, + cuvsStream); + Util.cudaMemcpyAsync( + distancesMemorySegment, + distancesDP.handle(), + distancesBytes, + DEVICE_TO_HOST, + cuvsStream); + + checkCuVSError(cuvsStreamSync(cuvsRes), "cuvsStreamSync"); } } diff --git a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CuVSDeviceMatrixImpl.java b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CuVSDeviceMatrixImpl.java index 2e72ea69f5..ef1b10941e 100644 --- a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CuVSDeviceMatrixImpl.java +++ b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CuVSDeviceMatrixImpl.java @@ -291,16 +291,6 @@ public ValueLayout valueLayout() { return deviceMatrix.valueLayout(); } - @Override - public int bits() { - return deviceMatrix.bits(); - } - - @Override - public int code() { - return 0; - } - @Override public MemorySegment toTensor(Arena arena) { return deviceMatrix.toTensor(arena); diff --git a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CuVSMatrixInternal.java b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CuVSMatrixInternal.java index ed53836719..f700811d4d 100644 --- a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CuVSMatrixInternal.java +++ b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CuVSMatrixInternal.java @@ -44,7 +44,11 @@ default int bits() { * DLTensor data type {@code code} for the element type of this matrix */ default int code() { - return switch (dataType()) { + return code(dataType()); + } + + static int code(DataType dataType) { + return switch (dataType) { case FLOAT -> kDLFloat(); case INT -> kDLInt(); case UINT, BYTE -> kDLUInt(); diff --git a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/common/Util.java b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/common/Util.java index 1ffb5ad559..7510a846f6 100644 --- a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/common/Util.java +++ b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/common/Util.java @@ -177,10 +177,18 @@ public static void cudaMemcpyAsync( * Helper to get the CUDA stream associated with a {@link CuVSResources} */ public static MemorySegment getStream(CuVSResources resources) { - try (var resourcesAccess = resources.access(); - var localArena = Arena.ofConfined()) { + try (var resourcesAccess = resources.access()) { + return getStream(resourcesAccess.handle()); + } + } + + /** + * Helper to get the CUDA stream associated with a {@link CuVSResources} handle + */ + public static MemorySegment getStream(long resourcesHandle) { + try (var localArena = Arena.ofConfined()) { var streamPointer = localArena.allocate(cudaStream_t); - checkCuVSError(cuvsStreamGet(resourcesAccess.handle(), streamPointer), "cuvsStreamGet"); + checkCuVSError(cuvsStreamGet(resourcesHandle, streamPointer), "cuvsStreamGet"); return streamPointer.get(cudaStream_t, 0); } } 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 45076e6482..c9de5f6eeb 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 @@ -23,7 +23,6 @@ import com.nvidia.cuvs.CagraIndexParams.CagraGraphBuildAlgo; import com.nvidia.cuvs.CagraIndexParams.CuvsDistanceType; import com.nvidia.cuvs.CagraMergeParams.MergeStrategy; -import java.io.InputStream; import java.lang.foreign.Arena; import java.lang.foreign.Linker; import java.lang.foreign.MemoryLayout; @@ -404,36 +403,40 @@ public void testPrefilteringReducesResults() throws Throwable { // No prefilter (all points allowed) CagraSearchParams searchParams = new CagraSearchParams.Builder().build(); - CagraQuery fullQuery = - new CagraQuery.Builder(resources) - .withTopK(2) - .withSearchParams(searchParams) - .withQueryVectors(queries) - .build(); - SearchResults fullSearchResults = index.search(fullQuery); - List> fullResults = fullSearchResults.getResults(); - log.debug("Full results: {}", fullResults); + // No prefilter (all points allowed) + try (var queryVectors = CuVSMatrix.ofArray(queries)) { + CagraQuery fullQuery = + new CagraQuery.Builder(resources) + .withTopK(2) + .withSearchParams(searchParams) + .withQueryVectors(queryVectors) + .build(); + + SearchResults fullSearchResults = index.search(fullQuery); + List> fullResults = fullSearchResults.getResults(); + log.debug("Full results: {}", fullResults); - // Apply prefilter: only allow ids 0 and 2 (bitset: 1100) - BitSet prefilter = new BitSet(4); - prefilter.set(0); - prefilter.set(2); + // Apply prefilter: only allow ids 0 and 2 (bitset: 1100) + BitSet prefilter = new BitSet(4); + prefilter.set(0); + prefilter.set(2); - CagraQuery filteredQuery = - new CagraQuery.Builder(resources) - .withTopK(2) - .withSearchParams(searchParams) - .withQueryVectors(queries) - .withPrefilter(prefilter, 4) - .build(); + CagraQuery filteredQuery = + new CagraQuery.Builder(resources) + .withTopK(2) + .withSearchParams(searchParams) + .withQueryVectors(queryVectors) + .withPrefilter(prefilter, 4) + .build(); - SearchResults filteredSearchResults = index.search(filteredQuery); - List> filteredResults = filteredSearchResults.getResults(); - log.debug("Filtered results: {}", filteredResults); + SearchResults filteredSearchResults = index.search(filteredQuery); + List> filteredResults = filteredSearchResults.getResults(); + log.debug("Filtered results: {}", filteredResults); - assertEquals(expectedResults, fullResults); - assertEquals(expectedFilteredResults, filteredResults); + assertEquals(expectedResults, fullResults); + assertEquals(expectedFilteredResults, filteredResults); + } } } @@ -483,27 +486,29 @@ private void queryAndCompare( CagraSearchParams searchParams = new CagraSearchParams.Builder().build(); // Create a query object with the query vectors - CagraQuery cuvsQuery = - new CagraQuery.Builder(resources) - .withTopK(3) - .withSearchParams(searchParams) - .withQueryVectors(queries) - .withMapping(mapping) - .build(); + try (var queryVectors = CuVSMatrix.ofArray(queries)) { + CagraQuery cuvsQuery = + new CagraQuery.Builder(resources) + .withTopK(3) + .withSearchParams(searchParams) + .withQueryVectors(queryVectors) + .withMapping(mapping) + .build(); - // Perform the search - SearchResults results = index1.search(cuvsQuery); + // Perform the search + SearchResults results = index1.search(cuvsQuery); - // Check results - log.debug(results.getResults().toString()); - checkResults(expectedResults, results.getResults()); + // Check results + log.debug(results.getResults().toString()); + checkResults(expectedResults, results.getResults()); - // Search from the second index - results = index2.search(cuvsQuery); + // Search from the second index + results = index2.search(cuvsQuery); - // Check results - log.debug(results.getResults().toString()); - checkResults(expectedResults, results.getResults()); + // Check results + log.debug(results.getResults().toString()); + checkResults(expectedResults, results.getResults()); + } } private void cleanup(CagraIndex index, CagraIndex loadedIndex) throws Throwable { @@ -643,39 +648,42 @@ public void testMergingIndexes() throws Throwable { CagraSearchParams searchParams = new CagraSearchParams.Builder().build(); - CagraQuery query = - new CagraQuery.Builder(resources) - .withTopK(3) - .withSearchParams(searchParams) - .withQueryVectors(queries) - .withMapping(SearchResults.IDENTITY_MAPPING) - .build(); + try (var queryVectors = CuVSMatrix.ofArray(queries)) { + CagraQuery query = + new CagraQuery.Builder(resources) + .withTopK(3) + .withSearchParams(searchParams) + .withQueryVectors(queryVectors) + .withMapping(SearchResults.IDENTITY_MAPPING) + .build(); - log.trace("Searching merged index..."); - SearchResults results = mergedIndex.search(query); - log.debug("Search results: " + results.getResults().toString()); + log.trace("Searching merged index..."); + SearchResults results = mergedIndex.search(query); + log.debug("Search results: " + results.getResults().toString()); - assertEquals(expectedResults, results.getResults()); + assertEquals(expectedResults, results.getResults()); - // --- Serialization/deserialization check --- - String indexFileName = UUID.randomUUID() + ".cag"; - var indexFile = Path.of(indexFileName); + // --- Serialization/deserialization check --- + String indexFileName = UUID.randomUUID() + ".cag"; + var indexFile = Path.of(indexFileName); - try (var out = Files.newOutputStream(indexFile)) { - mergedIndex.serialize(out); - } + try (var out = Files.newOutputStream(indexFile)) { + mergedIndex.serialize(out); + } - try (InputStream inputStream = Files.newInputStream(indexFile)) { - CagraIndex loadedMergedIndex = CagraIndex.newBuilder(resources).from(inputStream).build(); + try (var inputStream = Files.newInputStream(indexFile)) { + CagraIndex loadedMergedIndex = CagraIndex.newBuilder(resources).from(inputStream).build(); - SearchResults resultsFromLoaded = loadedMergedIndex.search(query); - assertEquals(expectedResults, resultsFromLoaded.getResults()); - mergedIndex.close(); - loadedMergedIndex.close(); + SearchResults resultsFromLoaded = loadedMergedIndex.search(query); + assertEquals(expectedResults, resultsFromLoaded.getResults()); + mergedIndex.close(); + loadedMergedIndex.close(); + } finally { + Files.deleteIfExists(indexFile); + } + index1.close(); + index2.close(); } - Files.deleteIfExists(indexFile); - index1.close(); - index2.close(); } } @@ -750,45 +758,47 @@ public void testMergeStrategies() throws Throwable { CagraSearchParams searchParams = new CagraSearchParams.Builder().build(); - CagraQuery query = - new CagraQuery.Builder(resources) - .withTopK(3) - .withSearchParams(searchParams) - .withQueryVectors(queries) - .withMapping(SearchResults.IDENTITY_MAPPING) - .build(); - - log.trace("Searching physically merged index..."); - SearchResults physicalResults = physicalMergedIndex.search(query); - assertNotNull("Physical merge search results should not be null", physicalResults); - assertEquals( - "Physical merge search results should match expected", - expectedResults, - physicalResults.getResults()); - - // --- Serialization/deserialization check for both merged indexes --- - String physicalIndexFileName = UUID.randomUUID() + ".cag"; - var physicalIndexFile = Path.of(physicalIndexFileName); - - try (var out = Files.newOutputStream(physicalIndexFile)) { - physicalMergedIndex.serialize(out); - } - - try (InputStream physicalInputStream = Files.newInputStream(physicalIndexFile); - CagraIndex loadedPhysicalIndex = - CagraIndex.newBuilder(resources).from(physicalInputStream).build()) { - - Files.deleteIfExists(physicalIndexFile); - - SearchResults resultsFromLoadedPhysical = loadedPhysicalIndex.search(query); + try (var queryVectors = CuVSMatrix.ofArray(queries)) { + CagraQuery query = + new CagraQuery.Builder(resources) + .withTopK(3) + .withSearchParams(searchParams) + .withQueryVectors(queryVectors) + .withMapping(SearchResults.IDENTITY_MAPPING) + .build(); + + log.trace("Searching physically merged index..."); + SearchResults physicalResults = physicalMergedIndex.search(query); + assertNotNull("Physical merge search results should not be null", physicalResults); assertEquals( - "Loaded physical index search results should match expected", + "Physical merge search results should match expected", expectedResults, - resultsFromLoadedPhysical.getResults()); + physicalResults.getResults()); + + // --- Serialization/deserialization check for both merged indexes --- + String physicalIndexFileName = UUID.randomUUID() + ".cag"; + var physicalIndexFile = Path.of(physicalIndexFileName); + + try (var out = Files.newOutputStream(physicalIndexFile)) { + physicalMergedIndex.serialize(out); + } + + try (var physicalInputStream = Files.newInputStream(physicalIndexFile); + CagraIndex loadedPhysicalIndex = + CagraIndex.newBuilder(resources).from(physicalInputStream).build()) { + + SearchResults resultsFromLoadedPhysical = loadedPhysicalIndex.search(query); + assertEquals( + "Loaded physical index search results should match expected", + expectedResults, + resultsFromLoadedPhysical.getResults()); + } finally { + Files.deleteIfExists(physicalIndexFile); + } } + index1.close(); + index2.close(); } - index1.close(); - index2.close(); } } } 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 32f3d86fe0..1452411b72 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 @@ -186,13 +186,14 @@ private void testQueryingUsingMultipleThreads(QueryAction queryAction) throws Th private void performQueryWithPrivateResource(CagraIndex index) throws Throwable { float[][] queries = generateRandomDataset(queryBatchSize); - try (CuVSResources threadResources = CheckedCuVSResources.create()) { - CagraSearchParams searchParams = new CagraSearchParams.Builder().build(); + CagraSearchParams searchParams = new CagraSearchParams.Builder().build(); + try (CuVSResources threadResources = CheckedCuVSResources.create(); + var queryVectors = CuVSMatrix.ofArray(queries)) { CagraQuery query = new CagraQuery.Builder(threadResources) .withTopK(topK) .withSearchParams(searchParams) - .withQueryVectors(queries) + .withQueryVectors(queryVectors) .build(); // This call should now work with per-thread resources @@ -207,17 +208,19 @@ private void performQueryWithSharedSynchronizedResource( float[][] queries = generateRandomDataset(queryBatchSize); CagraSearchParams searchParams = new CagraSearchParams.Builder().build(); - CagraQuery query = - new CagraQuery.Builder(threadResources) - .withTopK(topK) - .withSearchParams(searchParams) - .withQueryVectors(queries) - .build(); - - // This call should now work with per-thread resources - SearchResults results = index.search(query); - assertNotNull("Query should return results", results); - assertFalse("Query should return some results", results.getResults().isEmpty()); + try (var queryVectors = CuVSMatrix.ofArray(queries)) { + CagraQuery query = + new CagraQuery.Builder(threadResources) + .withTopK(topK) + .withSearchParams(searchParams) + .withQueryVectors(queryVectors) + .build(); + + // This call should now work with per-thread resources + SearchResults results = index.search(query); + assertNotNull("Query should return results", results); + assertFalse("Query should return some results", results.getResults().isEmpty()); + } } private float[][] generateRandomDataset(int size) { 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 d7d41c8c2f..b18b03adbc 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 @@ -162,11 +162,11 @@ private void tmpResultsTopKWithRandomValues(TestDatasetMemoryKind datasetMemoryK } log.trace("Index built successfully."); - try { + try (var queryVectors = CuVSMatrix.ofArray(queries)) { // Execute search and retrieve results CagraQuery.Builder queryBuilder = new CagraQuery.Builder(resources) - .withQueryVectors(queries) + .withQueryVectors(queryVectors) .withTopK(topK) .withSearchParams(new CagraSearchParams.Builder().build()); diff --git a/java/cuvs-java/src/test/java/com/nvidia/cuvs/CheckedCuVSResources.java b/java/cuvs-java/src/test/java/com/nvidia/cuvs/CheckedCuVSResources.java index 581e29973c..8b10c7f9b0 100644 --- a/java/cuvs-java/src/test/java/com/nvidia/cuvs/CheckedCuVSResources.java +++ b/java/cuvs-java/src/test/java/com/nvidia/cuvs/CheckedCuVSResources.java @@ -42,10 +42,15 @@ private void checkNotDestroyed() { @Override public ScopedAccess access() { checkNotDestroyed(); - var previousThreadId = currentThreadId.compareAndExchange(0, Thread.currentThread().threadId()); - if (previousThreadId != 0) { + var currentThreadId = Thread.currentThread().threadId(); + var previousThreadId = this.currentThreadId.compareAndExchange(0, currentThreadId); + if (previousThreadId != 0 && previousThreadId != currentThreadId) { throw new IllegalStateException( - "This resource is already accessed by thread [" + previousThreadId + "]"); + "This resource is already accessed by thread [" + + previousThreadId + + "]. Current thread id: [" + + currentThreadId + + "]"); } return new ScopedAccess() { @Override @@ -56,7 +61,7 @@ public long handle() { @Override public void close() { - currentThreadId.set(0); + CheckedCuVSResources.this.currentThreadId.set(0); } }; } 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 bc012f6a0e..633ecd52b1 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 @@ -65,7 +65,7 @@ public void testConcurrentAccessViaCheckedCuVSResourcesIsForbidden() throws Thro future.get(); }); assertEquals(IllegalStateException.class, exception.getCause().getClass()); - assertEquals(expectedError, exception.getCause().getMessage()); + assertTrue(exception.getCause().getMessage().startsWith(expectedError)); log.debug("Outer access finished"); } } From 90c812d00b6aa316ce45486ac6ea160c676304b4 Mon Sep 17 00:00:00 2001 From: Kyle Edwards Date: Fri, 26 Sep 2025 16:58:07 -0400 Subject: [PATCH 87/93] Update `cpp/cmake/config.json` for `ConfigureTest()` (#1385) Authors: - Kyle Edwards (https://github.com/KyleFromNVIDIA) Approvers: - Robert Maynard (https://github.com/robertmaynard) URL: https://github.com/rapidsai/cuvs/pull/1385 --- cpp/cmake/config.json | 19 +-- cpp/tests/CMakeLists.txt | 289 +++++++++++++++++++++++---------------- 2 files changed, 177 insertions(+), 131 deletions(-) diff --git a/cpp/cmake/config.json b/cpp/cmake/config.json index 3c568d9766..a10c274cb1 100644 --- a/cpp/cmake/config.json +++ b/cpp/cmake/config.json @@ -1,22 +1,15 @@ { "parse": { "additional_commands": { - "CPMFindPackage": { + "ConfigureTest": { + "flags": ["NOCUDA", "C_LIB", "FETCH_CODEBOOKS"], "kwargs": { "NAME": 1, - "GITHUB_REPOSITORY": "?", - "GIT_TAG": "?", - "VERSION": "?", - "GIT_SHALLOW": "?", - "OPTIONS": "*", - "FIND_PACKAGE_ARGUMENTS": "*" + "GPUS": "?", + "PERCENT": "?", + "ADDITIONAL_DEP": "?", + "PATH": "*" } - }, - "ConfigureTest": { - "flags": ["TEST_NAME", "TEST_SRC"] - }, - "ConfigureBench": { - "flags": ["BENCH_NAME", "BENCH_SRC"] } } }, diff --git a/cpp/tests/CMakeLists.txt b/cpp/tests/CMakeLists.txt index 33f9a12193..3921bae864 100644 --- a/cpp/tests/CMakeLists.txt +++ b/cpp/tests/CMakeLists.txt @@ -21,9 +21,9 @@ rapids_test_init() function(ConfigureTest) - set(options OPTIONAL NOCUDA C_LIB FETCH_CODEBOOKS) + set(options NOCUDA C_LIB FETCH_CODEBOOKS) set(oneValueArgs NAME GPUS PERCENT ADDITIONAL_DEP) - set(multiValueArgs PATH TARGETS CONFIGURATIONS) + set(multiValueArgs PATH) cmake_parse_arguments(_CUVS_TEST "${options}" "${oneValueArgs}" "${multiValueArgs}" ${ARGN}) if(NOT DEFINED _CUVS_TEST_GPUS AND NOT DEFINED _CUVS_TEST_PERCENT) @@ -109,221 +109,274 @@ endfunction() if(BUILD_TESTS) ConfigureTest( - NAME NEIGHBORS_TEST PATH neighbors/brute_force.cu neighbors/brute_force_prefiltered.cu - neighbors/sparse_brute_force.cu neighbors/refine.cu GPUS 1 PERCENT 100 + NAME NEIGHBORS_TEST + PATH neighbors/brute_force.cu neighbors/brute_force_prefiltered.cu + neighbors/sparse_brute_force.cu neighbors/refine.cu + GPUS 1 + PERCENT 100 ) - ConfigureTest(NAME NEIGHBORS_TIERED_INDEX_TEST PATH neighbors/tiered_index.cu GPUS 1 PERCENT 100) + ConfigureTest( + NAME NEIGHBORS_TIERED_INDEX_TEST + PATH neighbors/tiered_index.cu + GPUS 1 + PERCENT 100 + ) ConfigureTest( - NAME CLUSTER_TEST PATH cluster/kmeans.cu cluster/kmeans_balanced.cu cluster/kmeans_find_k.cu - cluster/linkage.cu cluster/connect_knn.cu GPUS 1 PERCENT 100 + NAME CLUSTER_TEST + PATH cluster/kmeans.cu cluster/kmeans_balanced.cu cluster/kmeans_find_k.cu cluster/linkage.cu + cluster/connect_knn.cu + GPUS 1 + PERCENT 100 ) ConfigureTest( - NAME - NEIGHBORS_ANN_IVF_FLAT_TEST - PATH - neighbors/ann_ivf_flat/test_float_int64_t.cu - neighbors/ann_ivf_flat/test_half_int64_t.cu - neighbors/ann_ivf_flat/test_int8_t_int64_t.cu - neighbors/ann_ivf_flat/test_uint8_t_int64_t.cu - GPUS - 1 - PERCENT - 100 + NAME NEIGHBORS_ANN_IVF_FLAT_TEST + PATH neighbors/ann_ivf_flat/test_float_int64_t.cu + neighbors/ann_ivf_flat/test_half_int64_t.cu + neighbors/ann_ivf_flat/test_int8_t_int64_t.cu + neighbors/ann_ivf_flat/test_uint8_t_int64_t.cu + GPUS 1 + PERCENT 100 ) ConfigureTest( - NAME - NEIGHBORS_ANN_IVF_PQ_TEST - PATH - neighbors/ann_ivf_pq/test_float_int64_t.cu - neighbors/ann_ivf_pq/test_int8_t_int64_t.cu - neighbors/ann_ivf_pq/test_uint8_t_int64_t.cu - GPUS - 1 - PERCENT - 100 + NAME NEIGHBORS_ANN_IVF_PQ_TEST + PATH neighbors/ann_ivf_pq/test_float_int64_t.cu neighbors/ann_ivf_pq/test_int8_t_int64_t.cu + neighbors/ann_ivf_pq/test_uint8_t_int64_t.cu + GPUS 1 + PERCENT 100 ) ConfigureTest( - NAME NEIGHBORS_ANN_SCANN_TEST PATH neighbors/ann_scann/test_float_int64_t.cu GPUS 1 PERCENT 1 + NAME NEIGHBORS_ANN_SCANN_TEST + PATH neighbors/ann_scann/test_float_int64_t.cu + GPUS 1 + PERCENT 1 ) ConfigureTest( - NAME NEIGHBORS_ANN_BRUTE_FORCE_TEST PATH neighbors/ann_brute_force/test_float.cu - neighbors/ann_brute_force/test_half.cu GPUS 1 PERCENT 100 + NAME NEIGHBORS_ANN_BRUTE_FORCE_TEST + PATH neighbors/ann_brute_force/test_float.cu neighbors/ann_brute_force/test_half.cu + GPUS 1 + PERCENT 100 ) ConfigureTest( - NAME NEIGHBORS_ANN_CAGRA_TEST_BUGS PATH neighbors/ann_cagra/bug_extreme_inputs_oob.cu - neighbors/ann_cagra/bug_multi_cta_crash.cu GPUS 1 PERCENT 100 + NAME NEIGHBORS_ANN_CAGRA_TEST_BUGS + PATH neighbors/ann_cagra/bug_extreme_inputs_oob.cu neighbors/ann_cagra/bug_multi_cta_crash.cu + GPUS 1 + PERCENT 100 ) ConfigureTest( - NAME NEIGHBORS_ANN_CAGRA_FLOAT_UINT32_TEST PATH neighbors/ann_cagra/test_float_uint32_t.cu GPUS - 1 PERCENT 100 + NAME NEIGHBORS_ANN_CAGRA_FLOAT_UINT32_TEST + PATH neighbors/ann_cagra/test_float_uint32_t.cu + GPUS 1 + PERCENT 100 ) ConfigureTest( - NAME NEIGHBORS_ANN_CAGRA_HELPERS_TEST PATH neighbors/ann_cagra/test_optimize_uint32_t.cu GPUS 1 + NAME NEIGHBORS_ANN_CAGRA_HELPERS_TEST + PATH neighbors/ann_cagra/test_optimize_uint32_t.cu + GPUS 1 PERCENT 100 ) ConfigureTest( - NAME NEIGHBORS_ANN_CAGRA_HALF_UINT32_TEST PATH neighbors/ann_cagra/test_half_uint32_t.cu GPUS 1 + NAME NEIGHBORS_ANN_CAGRA_HALF_UINT32_TEST + PATH neighbors/ann_cagra/test_half_uint32_t.cu + GPUS 1 PERCENT 100 ) ConfigureTest( - NAME NEIGHBORS_ANN_CAGRA_INT8_UINT32_TEST PATH neighbors/ann_cagra/test_int8_t_uint32_t.cu GPUS - 1 PERCENT 100 + NAME NEIGHBORS_ANN_CAGRA_INT8_UINT32_TEST + PATH neighbors/ann_cagra/test_int8_t_uint32_t.cu + GPUS 1 + PERCENT 100 ) ConfigureTest( - NAME NEIGHBORS_ANN_CAGRA_UINT8_UINT32_TEST PATH neighbors/ann_cagra/test_uint8_t_uint32_t.cu - GPUS 1 PERCENT 100 + NAME NEIGHBORS_ANN_CAGRA_UINT8_UINT32_TEST + PATH neighbors/ann_cagra/test_uint8_t_uint32_t.cu + GPUS 1 + PERCENT 100 ) ConfigureTest( - NAME - NEIGHBORS_ANN_NN_DESCENT_TEST - PATH - neighbors/ann_nn_descent/test_float_uint32_t.cu - neighbors/ann_nn_descent/test_int8_t_uint32_t.cu - neighbors/ann_nn_descent/test_uint8_t_uint32_t.cu - GPUS - 1 - PERCENT - 100 + NAME NEIGHBORS_ANN_NN_DESCENT_TEST + PATH neighbors/ann_nn_descent/test_float_uint32_t.cu + neighbors/ann_nn_descent/test_int8_t_uint32_t.cu + neighbors/ann_nn_descent/test_uint8_t_uint32_t.cu + GPUS 1 + PERCENT 100 ) ConfigureTest( - NAME - NEIGHBORS_ANN_VAMANA_TEST - PATH - neighbors/ann_vamana/test_float_uint32_t.cu - neighbors/ann_vamana/test_int8_t_uint32_t.cu - neighbors/ann_vamana/test_uint8_t_uint32_t.cu - GPUS - 1 - PERCENT - 100 + NAME NEIGHBORS_ANN_VAMANA_TEST + PATH neighbors/ann_vamana/test_float_uint32_t.cu neighbors/ann_vamana/test_int8_t_uint32_t.cu + neighbors/ann_vamana/test_uint8_t_uint32_t.cu + GPUS 1 + PERCENT 100 FETCH_CODEBOOKS ) - ConfigureTest(NAME NEIGHBORS_BALL_COVER_TEST PATH neighbors/ball_cover.cu GPUS 1 PERCENT 100) + ConfigureTest( + NAME NEIGHBORS_BALL_COVER_TEST + PATH neighbors/ball_cover.cu + GPUS 1 + PERCENT 100 + ) ConfigureTest( - NAME NEIGHBORS_EPSILON_NEIGHBORHOOD_TEST PATH neighbors/epsilon_neighborhood.cu GPUS 1 PERCENT - 100 + NAME NEIGHBORS_EPSILON_NEIGHBORHOOD_TEST + PATH neighbors/epsilon_neighborhood.cu + GPUS 1 + PERCENT 100 ) ConfigureTest( - NAME NEIGHBORS_ALL_NEIGHBORS_TEST PATH neighbors/all_neighbors/test_float.cu GPUS 1 PERCENT 100 + NAME NEIGHBORS_ALL_NEIGHBORS_TEST + PATH neighbors/all_neighbors/test_float.cu + GPUS 1 + PERCENT 100 ) ConfigureTest( - NAME - NEIGHBORS_DYNAMIC_BATCHING_TEST - PATH - neighbors/dynamic_batching/test_brute_force.cu - neighbors/dynamic_batching/test_cagra.cu - neighbors/dynamic_batching/test_ivf_flat.cu - neighbors/dynamic_batching/test_ivf_pq.cu - GPUS - 1 - PERCENT - 100 + NAME NEIGHBORS_DYNAMIC_BATCHING_TEST + PATH neighbors/dynamic_batching/test_brute_force.cu neighbors/dynamic_batching/test_cagra.cu + neighbors/dynamic_batching/test_ivf_flat.cu neighbors/dynamic_batching/test_ivf_pq.cu + GPUS 1 + PERCENT 100 ) if(BUILD_CAGRA_HNSWLIB) - ConfigureTest(NAME NEIGHBORS_HNSW_TEST PATH neighbors/hnsw.cu GPUS 1 PERCENT 100) + ConfigureTest( + NAME NEIGHBORS_HNSW_TEST + PATH neighbors/hnsw.cu + GPUS 1 + PERCENT 100 + ) target_link_libraries(NEIGHBORS_HNSW_TEST PRIVATE hnswlib::hnswlib) target_compile_definitions(NEIGHBORS_HNSW_TEST PUBLIC CUVS_BUILD_CAGRA_HNSWLIB) endif() if(BUILD_MG_ALGOS) ConfigureTest( - NAME NEIGHBORS_MG_TEST PATH neighbors/mg/test_float.cu GPUS 1 PERCENT 100 ADDITIONAL_DEP - NCCL::NCCL + NAME NEIGHBORS_MG_TEST + PATH neighbors/mg/test_float.cu + GPUS 1 + PERCENT 100 + ADDITIONAL_DEP NCCL::NCCL ) endif() ConfigureTest( - NAME - DISTANCE_TEST - PATH - distance/dist_canberra.cu - distance/dist_correlation.cu - distance/dist_cos.cu - distance/dist_hamming.cu - distance/dist_hellinger.cu - distance/dist_inner_product.cu - distance/dist_jensen_shannon.cu - distance/dist_kl_divergence.cu - distance/dist_l1.cu - distance/dist_l2_exp.cu - distance/dist_l2_sqrt_exp.cu - distance/dist_l_inf.cu - distance/dist_lp_unexp.cu - distance/dist_russell_rao.cu - distance/gram.cu - distance/masked_nn.cu - distance/sparse_distance.cu - sparse/neighbors/cross_component_nn.cu - GPUS - 1 - PERCENT - 100 + NAME DISTANCE_TEST + PATH distance/dist_canberra.cu + distance/dist_correlation.cu + distance/dist_cos.cu + distance/dist_hamming.cu + distance/dist_hellinger.cu + distance/dist_inner_product.cu + distance/dist_jensen_shannon.cu + distance/dist_kl_divergence.cu + distance/dist_l1.cu + distance/dist_l2_exp.cu + distance/dist_l2_sqrt_exp.cu + distance/dist_l_inf.cu + distance/dist_lp_unexp.cu + distance/dist_russell_rao.cu + distance/gram.cu + distance/masked_nn.cu + distance/sparse_distance.cu + sparse/neighbors/cross_component_nn.cu + GPUS 1 + PERCENT 100 ) ConfigureTest( - NAME SPARSE_TEST PATH sparse/cluster/cluster_solvers.cu sparse/cluster/eigen_solvers.cu - sparse/cluster/spectral.cu GPUS 1 PERCENT 100 + NAME SPARSE_TEST + PATH sparse/cluster/cluster_solvers.cu sparse/cluster/eigen_solvers.cu + sparse/cluster/spectral.cu + GPUS 1 + PERCENT 100 ) ConfigureTest( - NAME PREPROCESSING_TEST PATH preprocessing/scalar_quantization.cu - preprocessing/binary_quantization.cu preprocessing/spectral_embedding.cu GPUS 1 PERCENT 100 + NAME PREPROCESSING_TEST + PATH preprocessing/scalar_quantization.cu preprocessing/binary_quantization.cu + preprocessing/spectral_embedding.cu + GPUS 1 + PERCENT 100 ) ConfigureTest( - NAME STATS_TEST PATH stats/trustworthiness.cu stats/silhouette_score.cu GPUS 1 PERCENT 100 + NAME STATS_TEST + PATH stats/trustworthiness.cu stats/silhouette_score.cu + GPUS 1 + PERCENT 100 ) endif() if(TARGET cuvs::c_api) enable_language(C) - ConfigureTest(NAME INTEROP_TEST PATH core/interop.cu C_LIB) ConfigureTest( - NAME DISTANCE_C_TEST PATH distance/run_pairwise_distance_c.c distance/pairwise_distance_c.cu + NAME INTEROP_TEST + PATH core/interop.cu + C_LIB + ) + ConfigureTest( + NAME DISTANCE_C_TEST + PATH distance/run_pairwise_distance_c.c distance/pairwise_distance_c.cu C_LIB ) ConfigureTest( - NAME BRUTEFORCE_C_TEST PATH neighbors/run_brute_force_c.c neighbors/brute_force_c.cu C_LIB + NAME BRUTEFORCE_C_TEST + PATH neighbors/run_brute_force_c.c neighbors/brute_force_c.cu + C_LIB ) ConfigureTest( - NAME IVF_FLAT_C_TEST PATH neighbors/run_ivf_flat_c.c neighbors/ann_ivf_flat_c.cu C_LIB + NAME IVF_FLAT_C_TEST + PATH neighbors/run_ivf_flat_c.c neighbors/ann_ivf_flat_c.cu + C_LIB ) - ConfigureTest(NAME IVF_PQ_C_TEST PATH neighbors/run_ivf_pq_c.c neighbors/ann_ivf_pq_c.cu C_LIB) + ConfigureTest( + NAME IVF_PQ_C_TEST + PATH neighbors/run_ivf_pq_c.c neighbors/ann_ivf_pq_c.cu + C_LIB + ) - ConfigureTest(NAME CAGRA_C_TEST PATH neighbors/ann_cagra_c.cu C_LIB) + ConfigureTest( + NAME CAGRA_C_TEST + PATH neighbors/ann_cagra_c.cu + C_LIB + ) - ConfigureTest(NAME MG_C_TEST PATH neighbors/run_mg_c.c neighbors/ann_mg_c.cu C_LIB) + ConfigureTest( + NAME MG_C_TEST + PATH neighbors/run_mg_c.c neighbors/ann_mg_c.cu + C_LIB + ) ConfigureTest( - NAME ALL_NEIGHBORS_C_TEST PATH neighbors/run_all_neighbors_c.c neighbors/all_neighbors_c.cu + NAME ALL_NEIGHBORS_C_TEST + PATH neighbors/run_all_neighbors_c.c neighbors/all_neighbors_c.cu C_LIB ) if(BUILD_CAGRA_HNSWLIB) - ConfigureTest(NAME HNSW_C_TEST PATH neighbors/ann_hnsw_c.cu C_LIB) + ConfigureTest( + NAME HNSW_C_TEST + PATH neighbors/ann_hnsw_c.cu + C_LIB + ) target_link_libraries(NEIGHBORS_HNSW_TEST PRIVATE hnswlib::hnswlib) target_compile_definitions(NEIGHBORS_HNSW_TEST PUBLIC CUVS_BUILD_CAGRA_HNSWLIB) endif() From 172c023df8fabf73a0ecabbc0d1edb39f2344dfc Mon Sep 17 00:00:00 2001 From: Tarang Jain <40517122+tarang-jain@users.noreply.github.com> Date: Fri, 26 Sep 2025 19:02:41 -0700 Subject: [PATCH 88/93] Vamana C / Python API (#1112) Fixes https://github.com/rapidsai/cuvs/issues/853 Authors: - Tarang Jain (https://github.com/tarang-jain) Approvers: - Ben Frederickson (https://github.com/benfred) - Ben Karsin (https://github.com/bkarsin) URL: https://github.com/rapidsai/cuvs/pull/1112 --- cpp/CMakeLists.txt | 1 + cpp/include/cuvs/neighbors/vamana.h | 227 ++++++++++++++ cpp/src/neighbors/vamana_c.cpp | 186 ++++++++++++ docs/source/c_api/neighbors_vamana_c.rst | 43 +++ python/cuvs/cuvs/neighbors/CMakeLists.txt | 1 + python/cuvs/cuvs/neighbors/__init__.py | 2 + .../cuvs/cuvs/neighbors/vamana/CMakeLists.txt | 24 ++ .../cuvs/cuvs/neighbors/vamana/__init__.pxd | 0 python/cuvs/cuvs/neighbors/vamana/__init__.py | 23 ++ python/cuvs/cuvs/neighbors/vamana/vamana.pxd | 80 +++++ python/cuvs/cuvs/neighbors/vamana/vamana.pyx | 282 ++++++++++++++++++ python/cuvs/cuvs/tests/test_vamana.py | 92 ++++++ 12 files changed, 961 insertions(+) create mode 100644 cpp/include/cuvs/neighbors/vamana.h create mode 100644 cpp/src/neighbors/vamana_c.cpp create mode 100644 docs/source/c_api/neighbors_vamana_c.rst create mode 100644 python/cuvs/cuvs/neighbors/vamana/CMakeLists.txt create mode 100644 python/cuvs/cuvs/neighbors/vamana/__init__.pxd create mode 100644 python/cuvs/cuvs/neighbors/vamana/__init__.py create mode 100644 python/cuvs/cuvs/neighbors/vamana/vamana.pxd create mode 100644 python/cuvs/cuvs/neighbors/vamana/vamana.pyx create mode 100644 python/cuvs/cuvs/tests/test_vamana.py diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt index 85e1c171ab..d68c8df37a 100644 --- a/cpp/CMakeLists.txt +++ b/cpp/CMakeLists.txt @@ -735,6 +735,7 @@ target_compile_definitions(cuvs::cuvs INTERFACE $<$:NVTX_ENAB src/neighbors/cagra_c.cpp $<$:src/neighbors/hnsw_c.cpp> src/neighbors/nn_descent_c.cpp + src/neighbors/vamana_c.cpp src/neighbors/refine/refine_c.cpp src/neighbors/tiered_index_c.cpp src/neighbors/all_neighbors_c.cpp diff --git a/cpp/include/cuvs/neighbors/vamana.h b/cpp/include/cuvs/neighbors/vamana.h new file mode 100644 index 0000000000..3847aa6d1f --- /dev/null +++ b/cpp/include/cuvs/neighbors/vamana.h @@ -0,0 +1,227 @@ +/* + * Copyright (c) 2025, NVIDIA CORPORATION. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include +#include +#include +#include +#include +#include + +#ifdef __cplusplus +extern "C" { +#endif + +/** + * @defgroup vamana_c_index_params C API for Vamana index build + * @{ + */ + +/** + * @brief Supplemental parameters to build Vamana Index + * + * `graph_degree`: Maximum degree of graph; corresponds to the R parameter of + * Vamana algorithm in the literature. + * `visited_size`: Maximum number of visited nodes per search during Vamana algorithm. + * Loosely corresponds to the L parameter in the literature. + * `vamana_iters`: The number of times all vectors are inserted into the graph. If > 1, + * all vectors are re-inserted to improve graph quality. + * `max_fraction`: The maximum batch size is this fraction of the total dataset size. Larger + * gives faster build but lower graph quality. + * `alpha`: Used to determine how aggressive the pruning will be. + */ +struct cuvsVamanaIndexParams { + /** Distance type. */ + cuvsDistanceType metric; + /** Maximum degree of output graph corresponds to the R parameter in the original Vamana + * literature. */ + uint32_t graph_degree; + /** Maximum number of visited nodes per search corresponds to the L parameter in the Vamana + * literature **/ + uint32_t visited_size; + /** Number of Vamana vector insertion iterations (each iteration inserts all vectors). */ + float vamana_iters; + /** Alpha for pruning parameter */ + float alpha; + /** Maximum fraction of dataset inserted per batch. * + * Larger max batch decreases graph quality, but improves speed */ + float max_fraction; + /** Base of growth rate of batch sizes **/ + float batch_base; + /** Size of candidate queue structure - should be (2^x)-1 */ + uint32_t queue_size; + /** Max batchsize of reverse edge processing (reduces memory footprint) */ + uint32_t reverse_batchsize; +}; + +typedef struct cuvsVamanaIndexParams* cuvsVamanaIndexParams_t; + +/** + * @brief Allocate Vamana Index params, and populate with default values + * + * @param[in] params cuvsVamanaIndexParams_t to allocate + * @return cuvsError_t + */ +cuvsError_t cuvsVamanaIndexParamsCreate(cuvsVamanaIndexParams_t* params); + +/** + * @brief De-allocate Vamana Index params + * + * @param[in] params cuvsVamanaIndexParams_t to de-allocate + * @return cuvsError_t + */ +cuvsError_t cuvsVamanaIndexParamsDestroy(cuvsVamanaIndexParams_t params); + +/** + * @} + */ + +/** + * @defgroup vamana_c_index Vamana index + * @{ + */ + +/** + * @brief Struct to hold address of cuvs::neighbors::vamana::index and its active trained dtype + * + */ +typedef struct { + uintptr_t addr; + DLDataType dtype; + +} cuvsVamanaIndex; + +typedef cuvsVamanaIndex* cuvsVamanaIndex_t; + +/** + * @brief Allocate Vamana index + * + * @param[in] index cuvsVamanaIndex_t to allocate + * @return cuvsError_t + */ +cuvsError_t cuvsVamanaIndexCreate(cuvsVamanaIndex_t* index); + +/** + * @brief De-allocate Vamana index + * + * @param[in] index cuvsVamanaIndex_t to de-allocate + * @return cuvsError_t + */ +cuvsError_t cuvsVamanaIndexDestroy(cuvsVamanaIndex_t index); + +/** + * @brief Get the dimension of the index + * + * @param[in] index cuvsVamanaIndex_t to get dimension of + * @param[out] dim pointer to dimension to set + * @return cuvsError_t + */ +cuvsError_t cuvsVamanaIndexGetDims(cuvsVamanaIndex_t index, int* dim); + +/** + * @} + */ + +/** + * @defgroup vamana_c_index_build Vamana index build + * @{ + */ + +/** + * @brief Build Vamana index + * + * Build the index from the dataset for efficient DiskANN search. + * + * The build utilities the Vamana insertion-based algorithm to create the graph. The algorithm + * starts with an empty graph and iteratively inserts batches of nodes. Each batch involves + * performing a greedy search for each vector to be inserted, and inserting it with edges to + * all nodes traversed during the search. Reverse edges are also inserted and robustPrune is applied + * to improve graph quality. The index_params struct controls the degree of the final graph. + * + * The following distance metrics are supported: + * - L2 + * + * Usage example: + * @code{.c} + * // Create cuvsResources_t + * cuvsResources_t res; + * cuvsResourcesCreate(&res); + * + * // Assume a row-major dataset [n_rows, n_cols] is defined as `float* dataset` + * cuvsVamanaIndexParams_t index_params; + * cuvsVamanaIndexParamsCreate(&index_params); + * index_params->metric = L2Expanded; // set distance metric + * cuvsVamanaIndex_t index; + * cuvsVamanaIndexCreate(&index); + * cuvsVamanaBuild(res, index_params, dataset, index); + * @endcode + * + * @param[in] res cuvsResources_t opaque C handle + * @param[in] params cuvsVamanaIndexParams_t used to build Vamana index + * @param[in] dataset DLManagedTensor* training dataset + * @param[out] index cuvsVamanaIndex_t Vamana index + * @return cuvsError_t + */ +cuvsError_t cuvsVamanaBuild(cuvsResources_t res, + cuvsVamanaIndexParams_t params, + DLManagedTensor* dataset, + cuvsVamanaIndex_t index); + +/** + * @} + */ + +/** + * @defgroup vamana_c_index_serialize Vamana index serialize + * @{ + */ + +/** + * @brief Save Vamana index to file + * + * Matches the file format used by the DiskANN open-source repository, allowing cross-compatibility. + * + * Serialized Index is to be used by the DiskANN open-source repository for graph search. + * + * @code{.c} + * // Create cuvsResources_t + * cuvsResources_t res; + * cuvsResourcesCreate(&res); + * + * // create an index with `cuvsVamanaBuild` + * cuvsVamanaSerialize(res, "/path/to/index", index, true); + * @endcode + * + * @param[in] res cuvsResources_t opaque C handle + * @param[in] filename the file prefix for where the index is saved + * @param[in] index cuvsVamanaIndex_t to serialize + * @param[in] include_dataset whether to include the dataset in the serialized index + * @return cuvsError_t + */ +cuvsError_t cuvsVamanaSerialize(cuvsResources_t res, + const char* filename, + cuvsVamanaIndex_t index, + bool include_dataset); + +/** + * @} + */ + +#ifdef __cplusplus +} +#endif diff --git a/cpp/src/neighbors/vamana_c.cpp b/cpp/src/neighbors/vamana_c.cpp new file mode 100644 index 0000000000..4ad497d17e --- /dev/null +++ b/cpp/src/neighbors/vamana_c.cpp @@ -0,0 +1,186 @@ +/* + * Copyright (c) 2025, NVIDIA CORPORATION. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include +#include + +#include +#include +#include + +#include +#include +#include +#include +#include +#include + +namespace { + +template +void* _build(cuvsResources_t res, cuvsVamanaIndexParams* params, DLManagedTensor* dataset_tensor) +{ + auto res_ptr = reinterpret_cast(res); + + auto dataset = dataset_tensor->dl_tensor; + cuvs::neighbors::vamana::index_params index_params; + index_params.metric = static_cast((int)params->metric); + index_params.graph_degree = params->graph_degree; + index_params.visited_size = params->visited_size; + index_params.vamana_iters = params->vamana_iters; + index_params.alpha = params->alpha; + index_params.max_fraction = params->max_fraction; + index_params.batch_base = params->batch_base; + index_params.queue_size = params->queue_size; + index_params.reverse_batchsize = params->reverse_batchsize; + auto index = new cuvs::neighbors::vamana::index(*res_ptr); + + if (cuvs::core::is_dlpack_device_compatible(dataset)) { + using mdspan_type = raft::device_matrix_view; + auto mds = cuvs::core::from_dlpack(dataset_tensor); + *index = cuvs::neighbors::vamana::build(*res_ptr, index_params, mds); + } else if (cuvs::core::is_dlpack_host_compatible(dataset)) { + using mdspan_type = raft::host_matrix_view; + auto mds = cuvs::core::from_dlpack(dataset_tensor); + *index = cuvs::neighbors::vamana::build(*res_ptr, index_params, mds); + } + + return index; +} + +template +void _serialize(cuvsResources_t res, + const char* filename, + cuvsVamanaIndex_t index, + bool include_dataset) +{ + auto res_ptr = reinterpret_cast(res); + auto index_ptr = reinterpret_cast*>(index->addr); + + cuvs::neighbors::vamana::serialize(*res_ptr, std::string(filename), *index_ptr, include_dataset); +} + +} // namespace + +extern "C" cuvsError_t cuvsVamanaIndexCreate(cuvsVamanaIndex_t* index) +{ + return cuvs::core::translate_exceptions([=] { + *index = new cuvsVamanaIndex{}; + (*index)->addr = 0; + (*index)->dtype = {}; + }); +} + +extern "C" cuvsError_t cuvsVamanaIndexDestroy(cuvsVamanaIndex_t index_c_ptr) +{ + return cuvs::core::translate_exceptions([=] { + auto index = *index_c_ptr; + + if (index.addr != 0) { + if (index.dtype.code == kDLFloat && index.dtype.bits == 32) { + delete reinterpret_cast*>(index.addr); + } else if (index.dtype.code == kDLInt && index.dtype.bits == 8) { + delete reinterpret_cast*>(index.addr); + } else if (index.dtype.code == kDLUInt && index.dtype.bits == 8) { + delete reinterpret_cast*>(index.addr); + } + } + + delete index_c_ptr; + }); +} + +extern "C" cuvsError_t cuvsVamanaIndexGetDims(cuvsVamanaIndex_t index, int* dim) +{ + return cuvs::core::translate_exceptions([=] { + if (index->dtype.code == kDLFloat && index->dtype.bits == 32) { + auto index_ptr = + reinterpret_cast*>(index->addr); + *dim = index_ptr->dim(); + } else if (index->dtype.code == kDLInt && index->dtype.bits == 8) { + auto index_ptr = + reinterpret_cast*>(index->addr); + *dim = index_ptr->dim(); + } else if (index->dtype.code == kDLUInt && index->dtype.bits == 8) { + auto index_ptr = + reinterpret_cast*>(index->addr); + *dim = index_ptr->dim(); + } + }); +} + +extern "C" cuvsError_t cuvsVamanaBuild(cuvsResources_t res, + cuvsVamanaIndexParams_t params, + DLManagedTensor* dataset_tensor, + cuvsVamanaIndex_t index) +{ + return cuvs::core::translate_exceptions([=] { + auto dataset = dataset_tensor->dl_tensor; + index->dtype = dataset.dtype; + + if (dataset.dtype.code == kDLFloat && dataset.dtype.bits == 32) { + index->addr = reinterpret_cast(_build(res, params, dataset_tensor)); + } else if (dataset.dtype.code == kDLInt && dataset.dtype.bits == 8) { + index->addr = reinterpret_cast(_build(res, params, dataset_tensor)); + } else if (dataset.dtype.code == kDLUInt && dataset.dtype.bits == 8) { + index->addr = reinterpret_cast(_build(res, params, dataset_tensor)); + } else { + RAFT_FAIL("Unsupported dataset DLtensor dtype: %d and bits: %d", + dataset.dtype.code, + dataset.dtype.bits); + } + }); +} + +extern "C" cuvsError_t cuvsVamanaSerialize(cuvsResources_t res, + const char* filename, + cuvsVamanaIndex_t index, + bool include_dataset) +{ + return cuvs::core::translate_exceptions([=] { + if (index->dtype.code == kDLFloat && index->dtype.bits == 32) { + _serialize(res, filename, index, include_dataset); + } else if (index->dtype.code == kDLInt && index->dtype.bits == 8) { + _serialize(res, filename, index, include_dataset); + } else if (index->dtype.code == kDLUInt && index->dtype.bits == 8) { + _serialize(res, filename, index, include_dataset); + } else { + RAFT_FAIL( + "Unsupported index DLtensor dtype: %d and bits: %d", index->dtype.code, index->dtype.bits); + } + }); +} + +extern "C" cuvsError_t cuvsVamanaIndexParamsCreate(cuvsVamanaIndexParams_t* params) +{ + return cuvs::core::translate_exceptions([=] { + *params = new cuvsVamanaIndexParams{}; + (*params)->metric = L2Expanded; + (*params)->graph_degree = 32; + (*params)->visited_size = 64; + (*params)->vamana_iters = 1; + (*params)->alpha = 1.2f; + (*params)->max_fraction = 0.06f; + (*params)->batch_base = 2.0f; + (*params)->queue_size = 127; + (*params)->reverse_batchsize = 1000000; + }); +} + +extern "C" cuvsError_t cuvsVamanaIndexParamsDestroy(cuvsVamanaIndexParams_t params) +{ + return cuvs::core::translate_exceptions([=] { delete params; }); +} diff --git a/docs/source/c_api/neighbors_vamana_c.rst b/docs/source/c_api/neighbors_vamana_c.rst new file mode 100644 index 0000000000..8e489d315a --- /dev/null +++ b/docs/source/c_api/neighbors_vamana_c.rst @@ -0,0 +1,43 @@ +Vamana +====== + +Vamana is the graph construction algorithm behind the well-known DiskANN vector search solution. The cuVS implementation of Vamana/DiskANN is a custom GPU-acceleration version of the algorithm that aims to reduce index construction time using NVIDIA GPUs. + + +.. role:: py(code) + :language: c + :class: highlight + +``#include `` + +Index build parameters +---------------------- + +.. doxygengroup:: vamana_c_index_params + :project: cuvs + :members: + :content-only: + +Index +----- + +.. doxygengroup:: vamana_c_index + :project: cuvs + :members: + :content-only: + +Index build +----------- + +.. doxygengroup:: vamana_c_index_build + :project: cuvs + :members: + :content-only: + +Index serialize +--------------- + +.. doxygengroup:: vamana_c_index_serialize + :project: cuvs + :members: + :content-only: diff --git a/python/cuvs/cuvs/neighbors/CMakeLists.txt b/python/cuvs/cuvs/neighbors/CMakeLists.txt index ee48687d69..a6ef180aeb 100644 --- a/python/cuvs/cuvs/neighbors/CMakeLists.txt +++ b/python/cuvs/cuvs/neighbors/CMakeLists.txt @@ -21,6 +21,7 @@ add_subdirectory(nn_descent) add_subdirectory(tiered_index) add_subdirectory(mg) add_subdirectory(all_neighbors) +add_subdirectory(vamana) # Set the list of Cython files to build set(cython_sources refine.pyx) diff --git a/python/cuvs/cuvs/neighbors/__init__.py b/python/cuvs/cuvs/neighbors/__init__.py index b34a3b3598..1e58c21d9c 100644 --- a/python/cuvs/cuvs/neighbors/__init__.py +++ b/python/cuvs/cuvs/neighbors/__init__.py @@ -22,6 +22,7 @@ ivf_pq, mg, nn_descent, + vamana, ) from .refine import refine @@ -36,4 +37,5 @@ "nn_descent", "all_neighbors", "refine", + "vamana", ] diff --git a/python/cuvs/cuvs/neighbors/vamana/CMakeLists.txt b/python/cuvs/cuvs/neighbors/vamana/CMakeLists.txt new file mode 100644 index 0000000000..e1bc25af6c --- /dev/null +++ b/python/cuvs/cuvs/neighbors/vamana/CMakeLists.txt @@ -0,0 +1,24 @@ +# ============================================================================= +# Copyright (c) 2025, NVIDIA CORPORATION. +# +# Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except +# in compliance with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software distributed under the License +# is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express +# or implied. See the License for the specific language governing permissions and limitations under +# the License. +# ============================================================================= + +# Set the list of Cython files to build +set(cython_sources vamana.pyx) +set(linked_libraries cuvs::cuvs cuvs::c_api) + +# Build all of the Cython targets +rapids_cython_create_modules( + CXX + SOURCE_FILES "${cython_sources}" + LINKED_LIBRARIES "${linked_libraries}" MODULE_PREFIX neighbors_vamana_ +) diff --git a/python/cuvs/cuvs/neighbors/vamana/__init__.pxd b/python/cuvs/cuvs/neighbors/vamana/__init__.pxd new file mode 100644 index 0000000000..e69de29bb2 diff --git a/python/cuvs/cuvs/neighbors/vamana/__init__.py b/python/cuvs/cuvs/neighbors/vamana/__init__.py new file mode 100644 index 0000000000..d12956e21b --- /dev/null +++ b/python/cuvs/cuvs/neighbors/vamana/__init__.py @@ -0,0 +1,23 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. + +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at + +# http://www.apache.org/licenses/LICENSE-2.0 + +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +from .vamana import Index, IndexParams, build, save + +__all__ = [ + "Index", + "IndexParams", + "build", + "save", +] diff --git a/python/cuvs/cuvs/neighbors/vamana/vamana.pxd b/python/cuvs/cuvs/neighbors/vamana/vamana.pxd new file mode 100644 index 0000000000..0e0c776e89 --- /dev/null +++ b/python/cuvs/cuvs/neighbors/vamana/vamana.pxd @@ -0,0 +1,80 @@ +# +# Copyright (c) 2025, NVIDIA CORPORATION. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +# cython: language_level=3 + +from libc.stdint cimport int32_t, uint32_t, uintptr_t +from libcpp cimport bool + +from cuvs.common.c_api cimport cuvsError_t, cuvsResources_t +from cuvs.common.cydlpack cimport DLDataType, DLManagedTensor +from cuvs.distance_type cimport cuvsDistanceType + + +cdef extern from "cuvs/neighbors/vamana.h" nogil: + + ctypedef struct cuvsVamanaIndexParams: + cuvsDistanceType metric + uint32_t graph_degree + uint32_t visited_size + float vamana_iters + float alpha + float max_fraction + float batch_base + uint32_t queue_size + uint32_t reverse_batchsize + + ctypedef cuvsVamanaIndexParams* cuvsVamanaIndexParams_t + + ctypedef struct cuvsVamanaIndex: + uintptr_t addr + DLDataType dtype + + ctypedef cuvsVamanaIndex* cuvsVamanaIndex_t + + cuvsError_t cuvsVamanaIndexParamsCreate(cuvsVamanaIndexParams_t* params) + + cuvsError_t cuvsVamanaIndexParamsDestroy(cuvsVamanaIndexParams_t params) + + cuvsError_t cuvsVamanaIndexCreate(cuvsVamanaIndex_t* index) + + cuvsError_t cuvsVamanaIndexDestroy(cuvsVamanaIndex_t index) + + cuvsError_t cuvsVamanaIndexGetDims(cuvsVamanaIndex_t index, int* dim) + + cuvsError_t cuvsVamanaBuild(cuvsResources_t res, + cuvsVamanaIndexParams_t params, + DLManagedTensor* dataset, + cuvsVamanaIndex_t index) + + cuvsError_t cuvsVamanaSerialize(cuvsResources_t res, + const char* filename, + cuvsVamanaIndex_t index, + bool include_dataset) + + +cdef class Index: + """ + Vamana index object. This object stores the trained Vamana index state + which can be used to perform nearest neighbors searches. + """ + + cdef cuvsVamanaIndex_t index + cdef bool trained + cdef str active_index_type + + +cdef class IndexParams: + cdef cuvsVamanaIndexParams* params diff --git a/python/cuvs/cuvs/neighbors/vamana/vamana.pyx b/python/cuvs/cuvs/neighbors/vamana/vamana.pyx new file mode 100644 index 0000000000..fc65c228d0 --- /dev/null +++ b/python/cuvs/cuvs/neighbors/vamana/vamana.pyx @@ -0,0 +1,282 @@ +# +# Copyright (c) 2025, NVIDIA CORPORATION. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +# cython: language_level=3 + +import warnings + +import numpy as np + +cimport cuvs.common.cydlpack + +from cuvs.common.resources import auto_sync_resources + +from cython.operator cimport dereference as deref +from libcpp cimport bool +from libcpp.string cimport string + +from cuvs.common cimport cydlpack +from cuvs.distance_type cimport cuvsDistanceType + +from pylibraft.common import auto_convert_output, cai_wrapper, device_ndarray +from pylibraft.common.cai_wrapper import wrap_array +from pylibraft.common.interruptible import cuda_interruptible + +from cuvs.distance import DISTANCE_NAMES, DISTANCE_TYPES +from cuvs.neighbors.common import _check_input_array + +from libc.stdint cimport uint32_t, uintptr_t + +from cuvs.common.exceptions import check_cuvs + + +cdef class Index: + """ + Vamana index object. This object stores the trained Vamana index state + which can be used to perform nearest neighbors searches. + """ + + def __cinit__(self): + self.trained = False + check_cuvs(cuvsVamanaIndexCreate(&self.index)) + + def __dealloc__(self): + if self.index is not NULL: + check_cuvs(cuvsVamanaIndexDestroy(self.index)) + + @property + def trained(self): + return self.trained + + def __repr__(self): + attr_str = [attr + "=" + str(getattr(self, attr)) + for attr in ["trained"]] + return "Index(type=Vamana, " + (", ".join(attr_str)) + ")" + + +cdef class IndexParams: + """ + Parameters for building a Vamana index + + Parameters + ---------- + metric : str, default="sqeuclidean" + String denoting the metric type. Supported metrics include: + - "sqeuclidean" + - "l2" + graph_degree : int, default=32 + Maximum degree of graph; corresponds to the R parameter of + Vamana algorithm in the literature. + visited_size : int, default=64 + Maximum number of visited nodes per search during Vamana algorithm. + Loosely corresponds to the L parameter in the literature. + vamana_iters : float, default=1 + Number of Vamana vector insertion iterations (each iteration inserts + all vectors). + alpha : float, default=1.2 + Alpha for pruning parameter. Used to determine how aggressive the + pruning will be. + max_fraction : float, default=0.06 + Maximum fraction of dataset inserted per batch. Larger max batch + decreases graph quality, but improves speed. + batch_base : float, default=2.0 + Base of growth rate of batch sizes. + queue_size : int, default=127 + Size of candidate queue structure - should be (2^x)-1. + reverse_batchsize : int, default=1000000 + Max batchsize of reverse edge processing (reduces memory footprint). + """ + + def __cinit__(self): + check_cuvs(cuvsVamanaIndexParamsCreate(&self.params)) + + def __dealloc__(self): + check_cuvs(cuvsVamanaIndexParamsDestroy(self.params)) + + def __init__(self, *, + metric="sqeuclidean", + graph_degree=32, + visited_size=64, + vamana_iters=1, + alpha=1.2, + max_fraction=0.06, + batch_base=2.0, + queue_size=127, + reverse_batchsize=1000000): + if metric in DISTANCE_TYPES: + self.params.metric = DISTANCE_TYPES[metric] + else: + raise ValueError("metric %s not supported" % metric) + + self.params.graph_degree = graph_degree + self.params.visited_size = visited_size + self.params.vamana_iters = vamana_iters + self.params.alpha = alpha + self.params.max_fraction = max_fraction + self.params.batch_base = batch_base + self.params.queue_size = queue_size + self.params.reverse_batchsize = reverse_batchsize + + @property + def metric(self): + return DISTANCE_NAMES[self.params.metric] + + @property + def graph_degree(self): + return self.params.graph_degree + + @property + def visited_size(self): + return self.params.visited_size + + @property + def vamana_iters(self): + return self.params.vamana_iters + + @property + def alpha(self): + return self.params.alpha + + @property + def max_fraction(self): + return self.params.max_fraction + + @property + def batch_base(self): + return self.params.batch_base + + @property + def queue_size(self): + return self.params.queue_size + + @property + def reverse_batchsize(self): + return self.params.reverse_batchsize + + +@auto_sync_resources +def build(IndexParams index_params, dataset, resources=None): + """ + Build the Vamana index from the dataset for efficient search. + + The build utilities the Vamana insertion-based algorithm to create + the graph. The algorithm starts with an empty graph and iteratively + inserts batches of nodes. Each batch involves performing a greedy + search for each vector to be inserted, and inserting it with edges to + all nodes traversed during the search. Reverse edges are also inserted + and robustPrune is applied to improve graph quality. The index_params + struct controls the degree of the final graph. + + The following distance metrics are supported: + - L2Expanded + + Parameters + ---------- + index_params : IndexParams object + dataset : CUDA array interface compliant matrix shape (n_samples, dim) + Supported dtype [float, int8, uint8] + {resources_docstring} + + Returns + ------- + index: cuvs.vamana.Index + + Examples + -------- + + >>> import cupy as cp + >>> from cuvs.neighbors import vamana + >>> n_samples = 50000 + >>> n_features = 50 + >>> dataset = cp.random.random_sample((n_samples, n_features), + ... dtype=cp.float32) + >>> build_params = vamana.IndexParams(metric="sqeuclidean") + >>> index = vamana.build(build_params, dataset) + >>> # Serialize index to file for later use with CPU DiskANN + >>> vamana.save("my_index.bin", index) + """ + + # todo(dgd): we can make the check of dtype a parameter of wrap_array + # in RAFT to make this a single call + dataset_ai = wrap_array(dataset) + _check_input_array(dataset_ai, [np.dtype('float32'), + np.dtype('int8'), + np.dtype('uint8')]) + + cdef Index idx = Index() + cdef cydlpack.DLManagedTensor* dataset_dlpack = \ + cydlpack.dlpack_c(dataset_ai) + cdef cuvsVamanaIndexParams* params = index_params.params + + cdef cuvsResources_t res = resources.get_c_obj() + + with cuda_interruptible(): + check_cuvs(cuvsVamanaBuild( + res, + params, + dataset_dlpack, + idx.index + )) + idx.trained = True + idx.active_index_type = dataset_ai.dtype.name + + return idx + + +@auto_sync_resources +def save(filename, Index index, bool include_dataset=True, resources=None): + """ + Saves the index to a file. + + Matches the file format used by the DiskANN open-source repository, + allowing cross-compatibility. + + Parameters + ---------- + filename : string + Name of the file. + index : Index + Trained Vamana index. + include_dataset : bool + Whether or not to write out the dataset along with the index. Including + the dataset in the serialized index will use extra disk space, and + might not be desired if you already have a copy of the dataset on + disk. + {resources_docstring} + + Examples + -------- + >>> import cupy as cp + >>> from cuvs.neighbors import vamana + >>> n_samples = 50000 + >>> n_features = 50 + >>> dataset = cp.random.random_sample((n_samples, n_features), + ... dtype=cp.float32) + >>> # Build index + >>> index = vamana.build(vamana.IndexParams(), dataset) + >>> # Serialize and save the vamana index + >>> vamana.save("my_index.bin", index) + """ + cdef string c_filename = filename.encode('utf-8') + cdef cuvsResources_t res = resources.get_c_obj() + check_cuvs(cuvsVamanaSerialize(res, + c_filename.c_str(), + index.index, + include_dataset)) + + +# Note: Vamana index currently only supports build and serialize operations. +# Search functionality is not yet implemented in cuVS and should be performed +# using the DiskANN library with the serialized index. diff --git a/python/cuvs/cuvs/tests/test_vamana.py b/python/cuvs/cuvs/tests/test_vamana.py new file mode 100644 index 0000000000..236688b543 --- /dev/null +++ b/python/cuvs/cuvs/tests/test_vamana.py @@ -0,0 +1,92 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import numpy as np +import pytest +from pylibraft.common import device_ndarray + +from cuvs.neighbors import vamana + + +def _gen_data(shape, dtype): + rng = np.random.default_rng(12345) + if dtype == np.float32: + return rng.random(shape, dtype=np.float32) + if dtype == np.int8: + # keep small magnitude to avoid overflow if used elsewhere + return rng.integers(low=-10, high=10, size=shape, dtype=np.int8) + if dtype == np.uint8: + return rng.integers(low=0, high=20, size=shape, dtype=np.uint8) + raise AssertionError("unexpected dtype in test helper") + + +@pytest.mark.parametrize("dtype", [np.float32, np.int8, np.uint8]) +def test_vamana_build_basic(dtype): + n_rows, n_cols = 1000, 16 + data = _gen_data((n_rows, n_cols), dtype) + data_dev = device_ndarray(data) + + params = vamana.IndexParams(metric="sqeuclidean") + idx = vamana.build(params, data_dev) + + # Basic sanity: object type and flags + assert isinstance(idx, vamana.Index) + assert idx.trained is True + # repr should include trained flag + assert "Index(type=Vamana" in repr(idx) + + +@pytest.mark.parametrize("dtype", [np.float32, np.int8, np.uint8]) +@pytest.mark.skip( + reason="Skipping host build test because of CUDA error " + "in C++ API. Reference issue: " + "https://github.com/rapidsai/cuvs/issues/1380" +) +def test_vamana_build_basic_host(dtype): + n_rows, n_cols = 512, 12 + data = _gen_data((n_rows, n_cols), dtype) # host array + + params = vamana.IndexParams(metric="sqeuclidean") + idx = vamana.build(params, data) + + assert isinstance(idx, vamana.Index) + assert idx.trained is True + + +@pytest.mark.parametrize("include_dataset", [True, False]) +def test_vamana_serialize(tmp_path, include_dataset): + n_rows, n_cols = 256, 8 + data = _gen_data((n_rows, n_cols), np.float32) + data_dev = device_ndarray(data) + + params = vamana.IndexParams() + idx = vamana.build(params, data_dev) + + out_path = tmp_path / "vamana_index.bin" + vamana.save(str(out_path), idx, include_dataset=include_dataset) + + assert out_path.exists() and out_path.stat().st_size > 0 + + +def test_vamana_build_rejects_unsupported_dtype(): + # float16 is not in the accepted list in the build wrapper; expect failure + n_rows, n_cols = 64, 8 + data = _gen_data((n_rows, n_cols), np.float32).astype(np.float16) + from pylibraft.common import device_ndarray + + data_dev = device_ndarray(data) + params = vamana.IndexParams() + + with pytest.raises(Exception): + _ = vamana.build(params, data_dev) From 91d72b170548a59e74572c447439fa94f74b90a1 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Lorenzo=20Dematt=C3=A9?= Date: Sat, 27 Sep 2025 19:45:22 +0200 Subject: [PATCH 89/93] [Java] Support row strides in CuVSMatrix (#1345) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit This PR adds full support for row strides to `CuVSMatrix`. There was already some support for them in `CuVSDeviceMatrix`, as we needed to at least to save them to handle matrices built from tensors (e.g. resulting from a cuvs computation) and/or to create tensors for that kind of data (again, to pass them to the C API). With this PR we introduce ways to: - allocate host and device matrices with row stride != from the number of columns (e.g. padding/additional info at the end of row vectors) - map a pre-allocated memory region to a `CuVSMatrix` taking row stride into account - copy between different matrices with different strides - populate data taking the target matrix stride into account - read data from matrices (e.g. read rows or copy to heap arrays) taking the source matrix stride into account The API changes include also the introduction of column strides ("gaps" between elements withing the same vector); however this is not implemented yet. To the best of my knowledge, cuvs functions do not support these _yet_. I wanted to introduce them as we need them in some scenarios (e.g. round-tripping from `DLTensor`s), and I did not want to change the API once again if/when we are going to support them. However, I am open to drop `columnStride` if we do not want it. Let me know. Authors: - Lorenzo Dematté (https://github.com/ldematte) - Ben Frederickson (https://github.com/benfred) - MithunR (https://github.com/mythrocks) Approvers: - Chris Hegarty (https://github.com/ChrisHegarty) - MithunR (https://github.com/mythrocks) URL: https://github.com/rapidsai/cuvs/pull/1345 --- .../main/java/com/nvidia/cuvs/CuVSMatrix.java | 48 ++- .../com/nvidia/cuvs/spi/CuVSProvider.java | 23 ++ .../nvidia/cuvs/spi/UnsupportedProvider.java | 11 + .../cuvs/internal/CuVSDeviceMatrixImpl.java | 30 +- .../internal/CuVSDeviceMatrixRMMImpl.java | 22 +- .../internal/CuVSHostMatrixArenaImpl.java | 35 +- .../cuvs/internal/CuVSHostMatrixImpl.java | 101 +++++- .../cuvs/internal/CuVSMatrixBaseImpl.java | 5 +- .../cuvs/internal/CuVSMatrixInternal.java | 4 +- .../com/nvidia/cuvs/spi/JDKProvider.java | 295 ++++++++++------- .../java/com/nvidia/cuvs/CuVSMatrixIT.java | 307 +++++++++++++++--- .../java/com/nvidia/cuvs/DatasetHelper.java | 24 ++ 12 files changed, 691 insertions(+), 214 deletions(-) diff --git a/java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSMatrix.java b/java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSMatrix.java index 3ccdb9ae28..e37b696308 100644 --- a/java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSMatrix.java +++ b/java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSMatrix.java @@ -99,10 +99,10 @@ interface Builder { } /** - * Returns a builder to create a new instance of a dataset + * Returns a builder to create a new instance of a host-memory matrix * - * @param size Number of vectors in the dataset - * @param columns Size of each vector in the dataset + * @param size Number of rows (e.g. vectors in a dataset) + * @param columns Number of columns (e.g. dimension of each vector in the dataset) * @param dataType The data type of the dataset elements * @return a builder for creating a {@link CuVSHostMatrix} */ @@ -110,12 +110,28 @@ static Builder hostBuilder(long size, long columns, DataType dat return CuVSProvider.provider().newHostMatrixBuilder(size, columns, dataType); } + /** + * Returns a builder to create a new instance of a host-memory matrix + * + * @param size Number of rows (e.g. vectors in a dataset) + * @param columns Number of columns (e.g. dimension of each vector in the dataset) + * @param rowStride The stride (in number of elements) for each row. Must be -1 or > than {@code columns} + * @param columnStride The stride for each column. Currently, it is not supported (must be -1) + * @param dataType The data type of the dataset elements + * @return a builder for creating a {@link CuVSDeviceMatrix} + */ + static Builder hostBuilder( + long size, long columns, int rowStride, int columnStride, DataType dataType) { + return CuVSProvider.provider() + .newHostMatrixBuilder(size, columns, rowStride, columnStride, dataType); + } + /** * Returns a builder to create a new instance of a dataset * * @param resources CuVS resources used to allocate the device memory needed - * @param size Number of vectors in the dataset - * @param columns Size of each vector in the dataset + * @param size Number of rows (e.g. vectors in a dataset) + * @param columns Number of columns (e.g. dimension of each vector in the dataset) * @param dataType The data type of the dataset elements * @return a builder for creating a {@link CuVSDeviceMatrix} */ @@ -124,6 +140,28 @@ static Builder deviceBuilder( return CuVSProvider.provider().newDeviceMatrixBuilder(resources, size, columns, dataType); } + /** + * Returns a builder to create a new instance of a dataset + * + * @param resources CuVS resources used to allocate the device memory needed + * @param size Number of rows (e.g. vectors in a dataset) + * @param columns Number of columns (e.g. dimension of each vector in the dataset) + * @param rowStride The stride (in number of elements) for each row. Must be -1 or > than {@code columns} + * @param columnStride The stride for each column. Currently, it is not supported (must be -1) + * @param dataType The data type of the dataset elements + * @return a builder for creating a {@link CuVSDeviceMatrix} + */ + static Builder deviceBuilder( + CuVSResources resources, + long size, + long columns, + int rowStride, + int columnStride, + DataType dataType) { + return CuVSProvider.provider() + .newDeviceMatrixBuilder(resources, size, columns, rowStride, columnStride, dataType); + } + /** * Gets the size of the dataset * 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 f12389d226..034aaf0d7d 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 @@ -50,6 +50,10 @@ default Path nativeLibraryPath() { CuVSMatrix.Builder newHostMatrixBuilder( long size, long dimensions, CuVSMatrix.DataType dataType); + /** Create a {@link CuVSMatrix.Builder} instance for a host memory matrix **/ + CuVSMatrix.Builder newHostMatrixBuilder( + long size, long columns, int rowStride, int columnStride, CuVSMatrix.DataType dataType); + /** Create a {@link CuVSMatrix.Builder} instance for a device memory matrix **/ CuVSMatrix.Builder newDeviceMatrixBuilder( CuVSResources cuVSResources, long size, long dimensions, CuVSMatrix.DataType dataType); @@ -80,6 +84,25 @@ CuVSMatrix.Builder newDeviceMatrixBuilder( */ MethodHandle newNativeMatrixBuilder(); + /** + * Returns the factory method used to build a CuVSMatrix from native memory, with strides. + * The factory method will have this signature: + * {@code CuVSMatrix createNativeMatrix(memorySegment, size, dimensions, rowStride, columnStride, dataType)}, + * where {@code memorySegment} is a {@code java.lang.foreign.MemorySegment} containing {@code int size} vectors of + * {@code int dimensions} length of type {@link CuVSMatrix.DataType}. Rows have a stride of {@code rowStride}, + * where 0 indicates "no stride" (a stride equal to the number of columns), and columns have a stride of + * {@code columnStride} + *

+ * In order to expose this factory in a way that is compatible with Java 21, the factory method is returned as a + * {@link MethodHandle} with {@link MethodType} equal to + * {@code (CuVSMatrix.class, MemorySegment.class, int.class, int.class, int.class, int.class, DataType.class)}. + * The caller will need to invoke the factory via the {@link MethodHandle#invokeExact} method: + * {@code var matrix = (CuVSMatrix)newNativeMatrixBuilder().invokeExact(memorySegment, size, dimensions, rowStride, columnStride, dataType)} + *

+ * @return a MethodHandle which can be invoked to build a CuVSMatrix from an external {@code MemorySegment} + */ + MethodHandle newNativeMatrixBuilderWithStrides(); + /** Create a {@link CuVSMatrix} from an on-heap array **/ CuVSMatrix newMatrixFromArray(float[][] vectors); diff --git a/java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/UnsupportedProvider.java b/java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/UnsupportedProvider.java index 0f1edda99c..8ba651ce6b 100644 --- a/java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/UnsupportedProvider.java +++ b/java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/UnsupportedProvider.java @@ -66,6 +66,12 @@ public CuVSMatrix.Builder newHostMatrixBuilder( throw new UnsupportedOperationException(reasons); } + @Override + public CuVSMatrix.Builder newHostMatrixBuilder( + long size, long columns, int rowStride, int columnStride, CuVSMatrix.DataType dataType) { + throw new UnsupportedOperationException(reasons); + } + @Override public CuVSMatrix.Builder newDeviceMatrixBuilder( CuVSResources cuVSResources, long size, long dimensions, CuVSMatrix.DataType dataType) { @@ -93,6 +99,11 @@ public MethodHandle newNativeMatrixBuilder() { throw new UnsupportedOperationException(reasons); } + @Override + public MethodHandle newNativeMatrixBuilderWithStrides() { + throw new UnsupportedOperationException(reasons); + } + @Override public CuVSMatrix newMatrixFromArray(float[][] vectors) { throw new UnsupportedOperationException(reasons); diff --git a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CuVSDeviceMatrixImpl.java b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CuVSDeviceMatrixImpl.java index ef1b10941e..de9639c188 100644 --- a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CuVSDeviceMatrixImpl.java +++ b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CuVSDeviceMatrixImpl.java @@ -62,6 +62,11 @@ protected CuVSDeviceMatrixImpl( this.hostBuffer = new PinnedMemoryBuffer(size, columns, valueLayout); } + @Override + public long rowStride() { + return rowStride; + } + @Override public MemorySegment toTensor(Arena arena) { var strides = rowStride >= 0 ? new long[] {rowStride, columnStride} : null; @@ -189,24 +194,8 @@ private void copyRow(Object array, Arena localArena, int r, MemorySegment tmpRow } @Override - public void toHost(CuVSHostMatrix hostMatrix) { - if (hostMatrix.columns() != columns || hostMatrix.size() != size) { - throw new IllegalArgumentException("[hostMatrix] must have the same dimensions"); - } - if (hostMatrix.dataType() != dataType) { - throw new IllegalArgumentException("[hostMatrix] must have the same dataType"); - } - try (var localArena = Arena.ofConfined()) { - var hostMatrixTensor = ((CuVSMatrixInternal) hostMatrix).toTensor(localArena); - - try (var resourceAccess = resources.access()) { - var cuvsRes = resourceAccess.handle(); - var deviceMatrixTensor = toTensor(localArena); - checkCuVSError( - cuvsMatrixCopy(cuvsRes, deviceMatrixTensor, hostMatrixTensor), "cuvsMatrixCopy"); - checkCuVSError(cuvsStreamSync(cuvsRes), "cuvsStreamSync"); - } - } + public void toHost(CuVSHostMatrix targetMatrix) { + copyMatrix(this, (CuVSMatrixInternal) targetMatrix, resources); } @Override @@ -291,6 +280,11 @@ public ValueLayout valueLayout() { return deviceMatrix.valueLayout(); } + @Override + public long rowStride() { + return deviceMatrix.rowStride(); + } + @Override public MemorySegment toTensor(Arena arena) { return deviceMatrix.toTensor(arena); diff --git a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CuVSDeviceMatrixRMMImpl.java b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CuVSDeviceMatrixRMMImpl.java index 7ad84a0e1f..97558776f6 100644 --- a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CuVSDeviceMatrixRMMImpl.java +++ b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CuVSDeviceMatrixRMMImpl.java @@ -79,11 +79,27 @@ public static CuVSDeviceMatrixImpl create( long rowStride, long columnStride, DataType dataType) { + + var valueLayout = valueLayoutFromType(dataType); + var elementSize = valueLayout.byteSize(); + + final long rowSize; + if (rowStride <= 0) { + rowSize = columns * elementSize; + } else if (rowStride >= columns) { + rowSize = rowStride * elementSize; + } else { + throw new IllegalArgumentException("Row stride cannot be less than the number of columns"); + } + if (columnStride != -1) { + throw new UnsupportedOperationException( + "Stridden columns are currently not supported; columnStride must be equal to -1"); + } + try (var resourcesAccess = resources.access()) { - var valueLayout = valueLayoutFromType(dataType); + var rmmAllocation = - CloseableRMMAllocation.allocateRMMSegment( - resourcesAccess.handle(), size * columns * valueLayout.byteSize()); + CloseableRMMAllocation.allocateRMMSegment(resourcesAccess.handle(), size * rowSize); return new CuVSDeviceMatrixRMMImpl( resources, rmmAllocation, size, columns, rowStride, columnStride, dataType, valueLayout); } diff --git a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CuVSHostMatrixArenaImpl.java b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CuVSHostMatrixArenaImpl.java index 1dc2350638..37708d59a1 100644 --- a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CuVSHostMatrixArenaImpl.java +++ b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CuVSHostMatrixArenaImpl.java @@ -31,20 +31,51 @@ public CuVSHostMatrixArenaImpl(long size, long columns, DataType dataType) { this( size, columns, + -1, + -1, dataType, valueLayoutFromType(dataType), - sequenceLayoutFromType(size, columns, dataType), + sequenceLayoutFromType(size, columns, -1, dataType), + Arena.ofShared()); + } + + public CuVSHostMatrixArenaImpl( + long size, long columns, int rowStride, int columnStride, DataType dataType) { + this( + size, + columns, + rowStride, + columnStride, + dataType, + valueLayoutFromType(dataType), + sequenceLayoutFromType(size, columns, rowStride, dataType), Arena.ofShared()); } private CuVSHostMatrixArenaImpl( long size, long columns, + int rowStride, + int columnStride, DataType dataType, ValueLayout valueLayout, SequenceLayout layout, Arena arena) { - super(arena.allocate(layout), size, columns, dataType, valueLayout, layout); + super( + arena.allocate(layout), + size, + columns, + rowStride, + columnStride, + dataType, + valueLayout, + layout); + + if (columnStride != -1) { + throw new UnsupportedOperationException( + "Stridden columns are currently not supported; columnStride must be equal to -1"); + } + this.arenaReference = new AtomicReference<>(arena); } diff --git a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CuVSHostMatrixImpl.java b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CuVSHostMatrixImpl.java index 86e7e2a059..0f5945aae5 100644 --- a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CuVSHostMatrixImpl.java +++ b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CuVSHostMatrixImpl.java @@ -32,26 +32,65 @@ public class CuVSHostMatrixImpl extends CuVSMatrixBaseImpl implements CuVSHostMatrix { protected final VarHandle accessor$vh; + private final int rowStride; + private final int columnStride; + private final long rowBytes; + private final long rowSize; + private final long elementSize; + public CuVSHostMatrixImpl( MemorySegment memorySegment, long size, long columns, DataType dataType) { this( memorySegment, size, columns, + -1, + -1, dataType, valueLayoutFromType(dataType), - MemoryLayout.sequenceLayout(size * columns, valueLayoutFromType(dataType)) - .withByteAlignment(32)); + sequenceLayoutFromType(size, columns, -1, dataType)); + } + + public CuVSHostMatrixImpl( + MemorySegment memorySegment, + long size, + long columns, + int rowStride, + int columnStride, + DataType dataType) { + this( + memorySegment, + size, + columns, + rowStride, + columnStride, + dataType, + valueLayoutFromType(dataType), + sequenceLayoutFromType(size, columns, rowStride, dataType)); } protected CuVSHostMatrixImpl( MemorySegment memorySegment, long size, long columns, + int rowStride, + int columnStride, DataType dataType, ValueLayout valueLayout, MemoryLayout sequenceLayout) { super(memorySegment, dataType, valueLayout, size, columns); + + if (rowStride > 0 && rowStride < columns) { + throw new IllegalArgumentException("Row stride cannot be less than the number of columns"); + } + + this.rowStride = rowStride; + this.columnStride = columnStride; + + this.elementSize = valueLayout.byteSize(); + this.rowSize = rowStride > 0 ? rowStride * elementSize : columns * elementSize; + this.rowBytes = columns * elementSize; + this.accessor$vh = sequenceLayout.varHandle(MemoryLayout.PathElement.sequenceElement()); } @@ -87,10 +126,8 @@ public void toArray(int[][] array) { : String.format( Locale.ROOT, "Input array is of the wrong type for dataType [%s]", dataType.toString()); - var valueByteSize = valueLayout.byteSize(); for (int r = 0; r < size; ++r) { - MemorySegment.copy( - memorySegment, valueLayout, r * columns * valueByteSize, array[r], 0, (int) columns); + MemorySegment.copy(memorySegment, valueLayout, r * rowSize, array[r], 0, (int) columns); } } @@ -112,10 +149,8 @@ public void toArray(float[][] array) { : String.format( Locale.ROOT, "Input array is of the wrong type for dataType [%s]", dataType.toString()); - var valueByteSize = valueLayout.byteSize(); for (int r = 0; r < size; ++r) { - MemorySegment.copy( - memorySegment, valueLayout, r * columns * valueByteSize, array[r], 0, (int) columns); + MemorySegment.copy(memorySegment, valueLayout, r * rowSize, array[r], 0, (int) columns); } } @@ -137,10 +172,8 @@ public void toArray(byte[][] array) { : String.format( Locale.ROOT, "Input array is of the wrong type for dataType [%s]", dataType.toString()); - var valueByteSize = valueLayout.byteSize(); for (int r = 0; r < size; ++r) { - MemorySegment.copy( - memorySegment, valueLayout, r * columns * valueByteSize, array[r], 0, (int) columns); + MemorySegment.copy(memorySegment, valueLayout, r * rowSize, array[r], 0, (int) columns); } } @@ -152,9 +185,33 @@ public CuVSHostMatrix toHost() { @Override public void toHost(CuVSHostMatrix hostMatrix) { var targetMatrix = (CuVSMatrixInternal) hostMatrix; - var valueByteSize = valueLayout.byteSize(); - MemorySegment.copy( - this.memorySegment, 0L, targetMatrix.memorySegment(), 0L, size * columns * valueByteSize); + + if (targetMatrix.columns() != this.columns || targetMatrix.size() != this.size) { + throw new IllegalArgumentException( + "Source and target matrices must have the same dimensions"); + } + if (targetMatrix.dataType() != this.dataType) { + throw new IllegalArgumentException("Source and target matrices must have the same dataType"); + } + + if (this.rowStride <= 0 && targetMatrix.rowStride() <= 0) { + // copy whole matrix + MemorySegment.copy(this.memorySegment, 0L, targetMatrix.memorySegment(), 0L, size * rowSize); + } else { + // copy row-by-row + long targetRowSize = + targetMatrix.rowStride() > 0 + ? targetMatrix.rowStride() * elementSize + : columns * elementSize; + for (int r = 0; r < size; ++r) { + MemorySegment.copy( + this.memorySegment, + r * rowSize, + targetMatrix.memorySegment(), + r * targetRowSize, + rowBytes); + } + } } @Override @@ -167,13 +224,20 @@ public void close() {} @Override public int get(int row, int col) { - return (int) accessor$vh.get(memorySegment, 0L, (long) row * columns + col); + var rowPitch = rowStride > 0 ? rowStride : columns; + return (int) accessor$vh.get(memorySegment, 0L, (long) row * rowPitch + col); + } + + @Override + public long rowStride() { + return rowStride; } @Override public MemorySegment toTensor(Arena arena) { + var strides = rowStride >= 0 ? new long[] {rowStride, columnStride} : null; return prepareTensor( - arena, memorySegment, new long[] {size, columns}, code(), bits(), kDLCPU()); + arena, memorySegment, new long[] {size, columns}, strides, code(), bits(), kDLCPU()); } private static class CuVSHostMatrixDelegate implements CuVSHostMatrix, CuVSMatrixInternal { @@ -248,6 +312,11 @@ public ValueLayout valueLayout() { return hostMatrix.valueLayout(); } + @Override + public long rowStride() { + return hostMatrix.rowStride(); + } + @Override public MemorySegment toTensor(Arena arena) { return hostMatrix.toTensor(arena); diff --git a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CuVSMatrixBaseImpl.java b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CuVSMatrixBaseImpl.java index 2e510a99b2..934603e3c4 100644 --- a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CuVSMatrixBaseImpl.java +++ b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CuVSMatrixBaseImpl.java @@ -107,8 +107,9 @@ protected static ValueLayout valueLayoutFromType(DataType dataType) { } protected static SequenceLayout sequenceLayoutFromType( - long size, long columns, DataType dataType) { - return MemoryLayout.sequenceLayout(size * columns, valueLayoutFromType(dataType)) + long size, long columns, int rowStride, DataType dataType) { + var elements = rowStride > 0 ? rowStride * size : columns * size; + return MemoryLayout.sequenceLayout(elements, valueLayoutFromType(dataType)) .withByteAlignment(32); } diff --git a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CuVSMatrixInternal.java b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CuVSMatrixInternal.java index f700811d4d..472caf41d8 100644 --- a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CuVSMatrixInternal.java +++ b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CuVSMatrixInternal.java @@ -27,7 +27,7 @@ * Internal interface for {@link CuVSMatrix}, for shared functionality that * require/expose Panama types or internal cuvs types. */ -interface CuVSMatrixInternal extends CuVSMatrix { +public interface CuVSMatrixInternal extends CuVSMatrix { MemorySegment memorySegment(); @@ -55,6 +55,8 @@ static int code(DataType dataType) { }; } + long rowStride(); + /** * Creates a {@link DLManagedTensor} representing the matrix data and shape, to be * passed to the CuVS C API. diff --git a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/spi/JDKProvider.java b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/spi/JDKProvider.java index ed8af603e0..8a036a62ef 100644 --- a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/spi/JDKProvider.java +++ b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/spi/JDKProvider.java @@ -16,6 +16,7 @@ package com.nvidia.cuvs.spi; import static com.nvidia.cuvs.internal.common.Util.*; +import static com.nvidia.cuvs.internal.panama.headers_h.cudaMemcpy2DAsync; import static com.nvidia.cuvs.internal.panama.headers_h.cuvsVersionGet; import static com.nvidia.cuvs.internal.panama.headers_h.uint16_t; import static com.nvidia.cuvs.internal.panama.headers_h_1.cudaStreamSynchronize; @@ -39,7 +40,39 @@ final class JDKProvider implements CuVSProvider { - private static final MethodHandle createNativeDataset$mh = createNativeDatasetBuilder(); + private static final MethodHandle createNativeDataset$mh; + private static final MethodHandle createNativeDatasetWithStrides$mh; + + static { + try { + var lookup = MethodHandles.lookup(); + createNativeDataset$mh = + lookup.findStatic( + JDKProvider.class, + "createNativeDataset", + MethodType.methodType( + CuVSMatrix.class, + MemorySegment.class, + int.class, + int.class, + CuVSMatrix.DataType.class)); + + createNativeDatasetWithStrides$mh = + lookup.findStatic( + JDKProvider.class, + "createNativeDatasetWithStrides", + MethodType.methodType( + CuVSMatrix.class, + MemorySegment.class, + int.class, + int.class, + int.class, + int.class, + CuVSMatrix.DataType.class)); + } catch (NoSuchMethodException | IllegalAccessException e) { + throw new RuntimeException(e); + } + } private JDKProvider() {} @@ -84,27 +117,22 @@ private static String readCuVSVersionFromManifest() { } } - static MethodHandle createNativeDatasetBuilder() { - try { - var lookup = MethodHandles.lookup(); - var mt = - MethodType.methodType( - CuVSMatrix.class, - MemorySegment.class, - int.class, - int.class, - CuVSMatrix.DataType.class); - return lookup.findStatic(JDKProvider.class, "createNativeDataset", mt); - } catch (NoSuchMethodException | IllegalAccessException e) { - throw new RuntimeException(e); - } - } - private static CuVSMatrix createNativeDataset( MemorySegment memorySegment, int size, int dimensions, CuVSMatrix.DataType dataType) { return new CuVSHostMatrixImpl(memorySegment, size, dimensions, dataType); } + private static CuVSMatrix createNativeDatasetWithStrides( + MemorySegment memorySegment, + int size, + int dimensions, + int rowStride, + int columnStride, + CuVSMatrix.DataType dataType) { + return new CuVSHostMatrixImpl( + memorySegment, size, dimensions, rowStride, columnStride, dataType); + } + @Override public CuVSResources newCuVSResources(Path tempDirectory) { Objects.requireNonNull(tempDirectory); @@ -160,65 +188,22 @@ public GPUInfoProvider gpuInfoProvider() { @Override public CuVSMatrix.Builder newHostMatrixBuilder( - long size, long columns, CuVSMatrix.DataType dataType) throws UnsupportedOperationException { - - return new CuVSMatrix.Builder<>() { - final CuVSHostMatrixArenaImpl matrix = new CuVSHostMatrixArenaImpl(size, columns, dataType); - int current = 0; - - @Override - public void addVector(float[] vector) { - if (vector.length != columns) { - throw new IllegalArgumentException( - String.format( - Locale.ROOT, "Expected a vector of size [%d], got [%d]", columns, vector.length)); - } - internalAddVector(vector); - } - - @Override - public void addVector(byte[] vector) { - if (vector.length != columns) { - throw new IllegalArgumentException( - String.format( - Locale.ROOT, "Expected a vector of size [%d], got [%d]", columns, vector.length)); - } - internalAddVector(vector); - } + long size, long columns, CuVSMatrix.DataType dataType) { - @Override - public void addVector(int[] vector) { - if (vector.length != columns) { - throw new IllegalArgumentException( - String.format( - Locale.ROOT, "Expected a vector of size [%d], got [%d]", columns, vector.length)); - } - internalAddVector(vector); - } + return new HostMatrixBuilder(size, columns, dataType); + } - private void internalAddVector(Object vector) { - if (current >= size) throw new ArrayIndexOutOfBoundsException(); - MemorySegment.copy( - vector, - 0, - matrix.memorySegment(), - matrix.valueLayout(), - ((current++) * columns * matrix.valueLayout().byteSize()), - (int) columns); - } + @Override + public CuVSMatrix.Builder newHostMatrixBuilder( + long size, long columns, int rowStride, int columnStride, CuVSMatrix.DataType dataType) { - @Override - public CuVSHostMatrix build() { - return matrix; - } - }; + return new HostMatrixBuilder(size, columns, rowStride, columnStride, dataType); } @Override public CuVSMatrix.Builder newDeviceMatrixBuilder( - CuVSResources resources, long size, long columns, CuVSMatrix.DataType dataType) - throws UnsupportedOperationException { - return new BufferedSegmentBuilder(resources, size, columns, dataType); + CuVSResources resources, long size, long columns, CuVSMatrix.DataType dataType) { + return new DeviceMatrixBuilder(resources, size, columns, dataType); } @Override @@ -229,7 +214,7 @@ public CuVSMatrix.Builder newDeviceMatrixBuilder( int rowStride, int columnStride, CuVSMatrix.DataType dataType) { - return new BufferedSegmentBuilder(resources, size, columns, rowStride, columnStride, dataType); + return new DeviceMatrixBuilder(resources, size, columns, rowStride, columnStride, dataType); } @Override @@ -237,6 +222,11 @@ public MethodHandle newNativeMatrixBuilder() { return createNativeDataset$mh; } + @Override + public MethodHandle newNativeMatrixBuilderWithStrides() { + return createNativeDatasetWithStrides$mh; + } + @Override public CuVSMatrix newMatrixFromArray(float[][] vectors) { Objects.requireNonNull(vectors); @@ -279,62 +269,37 @@ public CuVSMatrix newMatrixFromArray(byte[][] vectors) { return dataset; } - /** - * This {@link CuVSDeviceMatrix} builder implementation returns a {@link CuVSDeviceMatrix} backed by managed RMM - * device memory. It uses a {@link PinnedMemoryBuffer} to batch data before copying it to the GPU. - */ - private static class BufferedSegmentBuilder implements CuVSMatrix.Builder { - - private final long columns; - private final long size; - private final CuVSDeviceMatrixImpl matrix; - private final MemorySegment stream; + private abstract static class MatrixBuilder { - private final long rowBytes; - private int currentRow; + protected final long columns; + protected final long size; + protected final T matrix; + protected final long elementSize; + protected final long rowSize; + protected final long rowBytes; + protected int currentRow; - private final PinnedMemoryBuffer hostBuffer; - private final long bufferRowCount; - private int currentBufferRow; - - private BufferedSegmentBuilder( - CuVSResources resources, long size, long columns, CuVSMatrix.DataType dataType) { + protected MatrixBuilder(T matrix, long size, long columns) { this.columns = columns; this.size = size; - this.matrix = CuVSDeviceMatrixRMMImpl.create(resources, size, columns, dataType); - this.stream = Util.getStream(resources); + this.matrix = matrix; + this.elementSize = matrix.valueLayout().byteSize(); + this.rowSize = columns * elementSize; + this.rowBytes = rowSize; this.currentRow = 0; - - this.hostBuffer = new PinnedMemoryBuffer(size, columns, matrix.valueLayout()); - - this.rowBytes = columns * matrix.valueLayout().byteSize(); - this.bufferRowCount = Math.min((hostBuffer.size() / rowBytes), size); - this.currentBufferRow = 0; } - private BufferedSegmentBuilder( - CuVSResources resources, - long size, - long columns, - int rowStride, - int columnStride, - CuVSMatrix.DataType dataType) { + protected MatrixBuilder(T matrix, long size, long columns, int rowStride) { this.columns = columns; this.size = size; - this.matrix = - CuVSDeviceMatrixRMMImpl.create( - resources, size, columns, rowStride, columnStride, dataType); - this.stream = Util.getStream(resources); - this.currentRow = 0; + this.matrix = matrix; + this.elementSize = matrix.valueLayout().byteSize(); + this.rowSize = rowStride > 0 ? rowStride * elementSize : columns * elementSize; + this.rowBytes = columns * elementSize; - this.hostBuffer = new PinnedMemoryBuffer(size, columns, matrix.valueLayout()); - - this.rowBytes = columns * matrix.valueLayout().byteSize(); - this.bufferRowCount = Math.min((hostBuffer.size() / rowBytes), size); - this.currentBufferRow = 0; + this.currentRow = 0; } - @Override public void addVector(float[] vector) { if (vector.length != columns) { throw new IllegalArgumentException( @@ -344,7 +309,6 @@ public void addVector(float[] vector) { internalAddVector(MemorySegment.ofArray(vector)); } - @Override public void addVector(byte[] vector) { if (vector.length != columns) { throw new IllegalArgumentException( @@ -354,7 +318,6 @@ public void addVector(byte[] vector) { internalAddVector(MemorySegment.ofArray(vector)); } - @Override public void addVector(int[] vector) { if (vector.length != columns) { throw new IllegalArgumentException( @@ -364,7 +327,55 @@ public void addVector(int[] vector) { internalAddVector(MemorySegment.ofArray(vector)); } - private void internalAddVector(MemorySegment vector) { + protected abstract void internalAddVector(MemorySegment vector); + } + + /** + * This {@link CuVSDeviceMatrix} builder implementation returns a {@link CuVSDeviceMatrix} backed by managed RMM + * device memory. It uses a {@link PinnedMemoryBuffer} to batch data before copying it to the GPU. + */ + private static final class DeviceMatrixBuilder extends MatrixBuilder + implements CuVSMatrix.Builder { + + private final MemorySegment stream; + + private final PinnedMemoryBuffer hostBuffer; + private final long bufferRowCount; + private int currentBufferRow; + + private DeviceMatrixBuilder( + CuVSResources resources, long size, long columns, CuVSMatrix.DataType dataType) { + super(CuVSDeviceMatrixRMMImpl.create(resources, size, columns, dataType), size, columns); + this.stream = Util.getStream(resources); + + this.hostBuffer = new PinnedMemoryBuffer(size, columns, matrix.valueLayout()); + this.bufferRowCount = Math.min((hostBuffer.size() / rowBytes), size); + this.currentBufferRow = 0; + } + + private DeviceMatrixBuilder( + CuVSResources resources, + long size, + long columns, + int rowStride, + int columnStride, + CuVSMatrix.DataType dataType) { + super( + CuVSDeviceMatrixRMMImpl.create( + resources, size, columns, rowStride, columnStride, dataType), + size, + columns, + rowStride); + + this.stream = Util.getStream(resources); + + this.hostBuffer = new PinnedMemoryBuffer(size, columns, matrix.valueLayout()); + this.bufferRowCount = Math.min((hostBuffer.size() / rowBytes), size); + this.currentBufferRow = 0; + } + + @Override + protected void internalAddVector(MemorySegment vector) { if (currentRow >= size) { throw new ArrayIndexOutOfBoundsException(); } @@ -380,14 +391,20 @@ private void internalAddVector(MemorySegment vector) { private void flushBuffer() { if (currentBufferRow > 0) { - var deviceMemoryOffset = (currentRow - currentBufferRow) * rowBytes; + var deviceMemoryOffset = (currentRow - currentBufferRow) * rowSize; var dst = matrix.memorySegment().asSlice(deviceMemoryOffset); - cudaMemcpyAsync( - dst, - hostBuffer.address(), - currentBufferRow * rowBytes, - CudaMemcpyKind.HOST_TO_DEVICE, - stream); + checkCudaError( + cudaMemcpy2DAsync( + dst, + rowSize, + hostBuffer.address(), + rowBytes, + rowBytes, + currentBufferRow, + CudaMemcpyKind.HOST_TO_DEVICE.kind, + stream), + "cudaMemcpy2DAsync"); + currentBufferRow = 0; checkCudaError(cudaStreamSynchronize(stream), "cudaStreamSynchronize"); } @@ -400,4 +417,34 @@ public CuVSDeviceMatrix build() { return matrix; } } + + private static class HostMatrixBuilder extends MatrixBuilder + implements CuVSMatrix.Builder { + + private HostMatrixBuilder(long size, long columns, CuVSMatrix.DataType dataType) { + super(new CuVSHostMatrixArenaImpl(size, columns, dataType), size, columns); + } + + private HostMatrixBuilder( + long size, long columns, int rowStride, int columnStride, CuVSMatrix.DataType dataType) { + super( + new CuVSHostMatrixArenaImpl(size, columns, rowStride, columnStride, dataType), + size, + columns, + rowStride); + } + + @Override + protected void internalAddVector(MemorySegment vector) { + if (currentRow >= size) { + throw new ArrayIndexOutOfBoundsException(); + } + MemorySegment.copy(vector, 0, matrix.memorySegment(), ((currentRow++) * rowSize), rowBytes); + } + + @Override + public CuVSHostMatrix build() { + return matrix; + } + } } 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 917acd4fc4..9e46845110 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 @@ -20,6 +20,8 @@ import static org.junit.Assert.assertEquals; import com.carrotsearch.randomizedtesting.RandomizedRunner; +import java.lang.foreign.*; +import java.util.Locale; import org.junit.Before; import org.junit.Test; import org.junit.runner.RunWith; @@ -71,9 +73,11 @@ private float[][] createFloatMatrix() { private float[][] createFloatMatrix(int rows, int cols) { float[][] result = new float[rows][cols]; + float value = 1; for (int r = 0; r < rows; ++r) { for (int c = 0; c < cols; ++c) { - result[r][c] = randomFloat(); + result[r][c] = value; + value += 1; } } return result; @@ -419,18 +423,7 @@ public void testDeviceToHost() throws Throwable { try (var deviceMatrix = builder.build(); var hostMatrix = deviceMatrix.toHost()) { - - assertEquals(data.length, deviceMatrix.size()); - assertEquals(data[0].length, deviceMatrix.columns()); - - assertEquals(deviceMatrix.size(), hostMatrix.size()); - assertEquals(deviceMatrix.columns(), hostMatrix.columns()); - - var roundTripData = new float[size][columns]; - - hostMatrix.toArray(roundTripData); - - assertSame2dArray(hostMatrix.size(), hostMatrix.columns(), data, roundTripData); + checkSameData(deviceMatrix, hostMatrix, data, size, columns); } } } @@ -446,18 +439,7 @@ public void testHostToDevice() throws Throwable { try (var hostMatrix = CuVSMatrix.ofArray(data); var deviceMatrix = hostMatrix.toDevice(resources)) { - - assertEquals(data.length, deviceMatrix.size()); - assertEquals(data[0].length, deviceMatrix.columns()); - - assertEquals(deviceMatrix.size(), hostMatrix.size()); - assertEquals(deviceMatrix.columns(), hostMatrix.columns()); - - var roundTripData = new float[size][columns]; - - deviceMatrix.toArray(roundTripData); - - assertSame2dArray(deviceMatrix.size(), deviceMatrix.columns(), data, roundTripData); + checkSameData(hostMatrix, deviceMatrix, data, size, columns); } } } @@ -471,18 +453,7 @@ public void testHostToHostReturnsWeakReferenceSameData() { try (var hostMatrix = CuVSMatrix.ofArray(data); var hostMatrix2 = hostMatrix.toHost()) { - - assertEquals(data.length, hostMatrix2.size()); - assertEquals(data[0].length, hostMatrix2.columns()); - - assertEquals(hostMatrix2.size(), hostMatrix.size()); - assertEquals(hostMatrix2.columns(), hostMatrix.columns()); - - var roundTripData = new float[size][columns]; - - hostMatrix2.toArray(roundTripData); - - assertSame2dArray(hostMatrix2.size(), hostMatrix2.columns(), data, roundTripData); + checkSameData(hostMatrix, hostMatrix2, data, size, columns); } } @@ -503,18 +474,268 @@ public void testDeviceToDeviceReturnsWeakReferenceSameData() throws Throwable { try (var deviceMatrix = builder.build(); var deviceMatrix2 = deviceMatrix.toDevice(resources)) { + checkSameData(deviceMatrix, deviceMatrix2, data, size, columns); + } + } + } + + @Test + public void testHostToHostWithDifferentStrides() { + + int size = randomIntBetween(1, 32 * 1024); + int columns = randomIntBetween(16, 2048); + int rowStride1 = randomIntBetween(columns, columns * 2); + int rowStride2 = randomIntBetween(columns, columns * 2); + final float[][] data = createFloatMatrix(size, columns); - assertEquals(data.length, deviceMatrix.size()); - assertEquals(data[0].length, deviceMatrix.columns()); + var builder1 = CuVSMatrix.hostBuilder(size, columns, rowStride1, -1, CuVSMatrix.DataType.FLOAT); + for (int i = 0; i < size; ++i) { + var array = data[i]; + builder1.addVector(array); + } + + var builder2 = CuVSMatrix.hostBuilder(size, columns, rowStride2, -1, CuVSMatrix.DataType.FLOAT); + + try (var matrix1 = builder1.build(); + var matrix2 = builder2.build()) { + + matrix1.toHost(matrix2); + checkSameData(matrix1, matrix2, data, size, columns); + } + } - assertEquals(deviceMatrix.size(), deviceMatrix2.size()); - assertEquals(deviceMatrix.columns(), deviceMatrix2.columns()); + @Test + public void testHostBuilderWithDifferentStrides() throws Throwable { + + int size = randomIntBetween(1, 32 * 1024); + int columns = randomIntBetween(16, 2048); + int rowStride1 = randomIntBetween(columns, columns * 2); + int rowStride2 = randomIntBetween(columns, columns * 2); + final float[][] data = createFloatMatrix(size, columns); + + var builder1 = CuVSMatrix.hostBuilder(size, columns, rowStride1, -1, CuVSMatrix.DataType.FLOAT); + for (int i = 0; i < size; ++i) { + var array = data[i]; + builder1.addVector(array); + } + + var builder2 = CuVSMatrix.hostBuilder(size, columns, rowStride2, -1, CuVSMatrix.DataType.FLOAT); + for (int i = 0; i < size; ++i) { + var array = data[i]; + builder2.addVector(array); + } + + try (var matrix1 = builder1.build(); + var matrix2 = builder2.build()) { + checkSameData(matrix1, matrix2, data, size, columns); + } + } + + @Test + public void testDeviceToHostWithDifferentStrides() throws Throwable { + + int size = randomIntBetween(1, 32 * 1024); + int columns = randomIntBetween(16, 2048); + int rowStride1 = randomIntBetween(columns, columns * 2); + int rowStride2 = randomIntBetween(columns, columns * 2); + final float[][] data = createFloatMatrix(size, columns); + + try (var resources = CuVSResources.create()) { + var builder1 = + CuVSMatrix.deviceBuilder( + resources, size, columns, rowStride1, -1, CuVSMatrix.DataType.FLOAT); + for (int i = 0; i < size; ++i) { + var array = data[i]; + builder1.addVector(array); + } + + var builder2 = + CuVSMatrix.hostBuilder(size, columns, rowStride2, -1, CuVSMatrix.DataType.FLOAT); + + try (var matrix1 = builder1.build(); + var matrix2 = builder2.build()) { + + matrix1.toHost(matrix2); + checkSameData(matrix1, matrix2, data, size, columns); + } + } + } + + @Test + public void testDeviceToDeviceWithDifferentStrides() throws Throwable { + + int size = randomIntBetween(1, 32 * 1024); + int columns = randomIntBetween(16, 2048); + int rowStride1 = randomIntBetween(columns, columns * 2); + int rowStride2 = randomIntBetween(columns, columns * 2); + final float[][] data = createFloatMatrix(size, columns); + + try (var resources = CuVSResources.create()) { + var builder1 = + CuVSMatrix.deviceBuilder( + resources, size, columns, rowStride1, -1, CuVSMatrix.DataType.FLOAT); + for (int i = 0; i < size; ++i) { + var array = data[i]; + builder1.addVector(array); + } + + var builder2 = + CuVSMatrix.deviceBuilder( + resources, size, columns, rowStride2, -1, CuVSMatrix.DataType.FLOAT); + + try (var matrix1 = builder1.build(); + var matrix2 = builder2.build()) { + + matrix1.toDevice(matrix2, resources); + checkSameData(matrix1, matrix2, data, size, columns); + } + } + } + + private static void checkSameData( + CuVSMatrix matrix1, CuVSMatrix matrix2, float[][] data, int size, int columns) { + assertEquals(data.length, matrix1.size()); + assertEquals(data[0].length, matrix1.columns()); + + assertEquals(matrix1.size(), matrix2.size()); + assertEquals(matrix1.columns(), matrix2.columns()); + + var roundTripData1 = new float[size][columns]; + var roundTripData2 = new float[size][columns]; + + matrix1.toArray(roundTripData1); + matrix2.toArray(roundTripData2); + + for (int n = 0; n < matrix2.size(); ++n) { + for (int i = 0; i < matrix2.columns(); ++i) { + var diff1 = Math.abs(data[n][i] - roundTripData1[n][i]); + if (diff1 > DELTA) { + throw new AssertionError( + String.format( + Locale.ROOT, + "Matrix1 mismatch. Expected:<%f> but was:<%f> at row [%d]of[%d], col [%d]of[%d]", + data[n][i], + roundTripData2[n][i], + n, + size, + i, + columns)); + } + var diff2 = Math.abs(data[n][i] - roundTripData2[n][i]); + if (diff2 > DELTA) { + throw new AssertionError( + String.format( + Locale.ROOT, + "Matrix2 mismatch. Expected:<%f> but was:<%f> at row [%d]of[%d], col [%d]of[%d]", + data[n][i], + roundTripData2[n][i], + n, + size, + i, + columns)); + } + } + } + } + + @Test + public void testDeviceBuilderWithDifferentStrides() throws Throwable { + + int size = randomIntBetween(1, 32 * 1024); + int columns = randomIntBetween(16, 2048); + int rowStride1 = randomIntBetween(columns, columns * 2); + int rowStride2 = randomIntBetween(columns, columns * 2); + final float[][] data = createFloatMatrix(size, columns); + + try (var resources = CuVSResources.create()) { + var builder1 = + CuVSMatrix.deviceBuilder( + resources, size, columns, rowStride1, -1, CuVSMatrix.DataType.FLOAT); + for (int i = 0; i < size; ++i) { + var array = data[i]; + builder1.addVector(array); + } + + var builder2 = + CuVSMatrix.deviceBuilder( + resources, size, columns, rowStride2, -1, CuVSMatrix.DataType.FLOAT); + for (int i = 0; i < size; ++i) { + var array = data[i]; + builder2.addVector(array); + } + + try (var matrix1 = builder1.build(); + var matrix2 = builder2.build()) { + checkSameData(matrix1, matrix2, data, size, columns); + } + } + } + + @Test + public void testHostMatrixFromNativeDataset() { + int size = randomIntBetween(1, 32 * 1024); + int columns = randomIntBetween(16, 2048); + final float[][] data = createFloatMatrix(size, columns); + + ValueLayout.OfFloat C_FLOAT = + (ValueLayout.OfFloat) Linker.nativeLinker().canonicalLayouts().get("float"); + + MemoryLayout dataMemoryLayout = MemoryLayout.sequenceLayout((long) size * columns, C_FLOAT); + + try (Arena arena = Arena.ofShared()) { + MemorySegment dataMemorySegment = arena.allocate(dataMemoryLayout); + for (int r = 0; r < size; r++) { + MemorySegment.copy( + data[r], 0, dataMemorySegment, C_FLOAT, (r * columns * C_FLOAT.byteSize()), columns); + } + + try (var nativeDataset = + DatasetHelper.fromMemorySegment( + dataMemorySegment, size, columns, CuVSMatrix.DataType.FLOAT)) { var roundTripData = new float[size][columns]; + nativeDataset.toArray(roundTripData); - deviceMatrix2.toArray(roundTripData); + for (int n = 0; n < nativeDataset.size(); ++n) { + for (int i = 0; i < nativeDataset.columns(); ++i) { + assertEquals(data[n][i], roundTripData[n][i], DELTA); + } + } + } + } + } + + @Test + public void testHostMatrixFromNativeDatasetWithStride() { + int size = randomIntBetween(1, 32 * 1024); + int columns = randomIntBetween(16, 2048); + int rowStride = randomIntBetween(columns, columns * 2); + final float[][] data = createFloatMatrix(size, columns); + + ValueLayout.OfFloat C_FLOAT = + (ValueLayout.OfFloat) Linker.nativeLinker().canonicalLayouts().get("float"); - assertSame2dArray(deviceMatrix2.size(), deviceMatrix2.columns(), data, roundTripData); + MemoryLayout dataMemoryLayout = MemoryLayout.sequenceLayout((long) size * rowStride, C_FLOAT); + + try (Arena arena = Arena.ofShared()) { + MemorySegment dataMemorySegment = arena.allocate(dataMemoryLayout); + for (int r = 0; r < size; r++) { + MemorySegment.copy( + data[r], 0, dataMemorySegment, C_FLOAT, (r * rowStride * C_FLOAT.byteSize()), columns); + } + + try (var nativeDataset = + DatasetHelper.fromMemorySegment( + dataMemorySegment, size, columns, rowStride, -1, CuVSMatrix.DataType.FLOAT)) { + + var roundTripData = new float[size][columns]; + nativeDataset.toArray(roundTripData); + + for (int n = 0; n < nativeDataset.size(); ++n) { + for (int i = 0; i < nativeDataset.columns(); ++i) { + assertEquals(data[n][i], roundTripData[n][i], DELTA); + } + } } } } diff --git a/java/cuvs-java/src/test/java/com/nvidia/cuvs/DatasetHelper.java b/java/cuvs-java/src/test/java/com/nvidia/cuvs/DatasetHelper.java index 9138f9952a..fd25e2df55 100644 --- a/java/cuvs-java/src/test/java/com/nvidia/cuvs/DatasetHelper.java +++ b/java/cuvs-java/src/test/java/com/nvidia/cuvs/DatasetHelper.java @@ -23,6 +23,8 @@ public class DatasetHelper { private static final MethodHandle createDataset$mh = CuVSProvider.provider().newNativeMatrixBuilder(); + private static final MethodHandle createDatasetWithStrides$mh = + CuVSProvider.provider().newNativeMatrixBuilderWithStrides(); public static CuVSMatrix fromMemorySegment( MemorySegment memorySegment, int size, int dimensions, CuVSMatrix.DataType dataType) { @@ -38,4 +40,26 @@ public static CuVSMatrix fromMemorySegment( } } } + + public static CuVSMatrix fromMemorySegment( + MemorySegment memorySegment, + int size, + int dimensions, + int rowStride, + int columnStride, + CuVSMatrix.DataType dataType) { + try { + return (CuVSMatrix) + createDatasetWithStrides$mh.invokeExact( + memorySegment, size, dimensions, rowStride, columnStride, dataType); + } catch (Throwable e) { + if (e instanceof Error err) { + throw err; + } else if (e instanceof RuntimeException re) { + throw re; + } else { + throw new RuntimeException(e); + } + } + } } From fdb8bfb8541f680f03fcd2ecd1e76657aff31251 Mon Sep 17 00:00:00 2001 From: Victor Lafargue Date: Mon, 29 Sep 2025 19:06:20 +0200 Subject: [PATCH 90/93] Treat warnings as errors in doc builds + MG ANNs Python API doc update (#1350) This PR : - updates doc build to treat warnings as errors - fixes an issue in the MG ANNs Python API doc - fixes several doc issues that resulted in warnings Closes https://github.com/rapidsai/cuvs/issues/1382 Authors: - Victor Lafargue (https://github.com/viclafargue) Approvers: - Ben Frederickson (https://github.com/benfred) - James Lamb (https://github.com/jameslamb) - Corey J. Nolet (https://github.com/cjnolet) URL: https://github.com/rapidsai/cuvs/pull/1350 --- build.sh | 4 ++-- ci/build_docs.sh | 8 ++++---- cpp/include/cuvs/neighbors/cagra.hpp | 4 ++++ cpp/include/cuvs/neighbors/hnsw.hpp | 4 ++++ cpp/include/cuvs/neighbors/vamana.hpp | 4 ++++ .../cuvs/preprocessing/spectral_embedding.hpp | 10 +++++----- docs/Makefile | 2 +- docs/source/neighbors/all_neighbors.rst | 2 +- docs/source/python_api/neighbors_mg_cagra.rst | 16 ++++++++-------- .../python_api/neighbors_mg_ivf_flat.rst | 18 +++++++++--------- docs/source/python_api/neighbors_mg_ivf_pq.rst | 18 +++++++++--------- 11 files changed, 51 insertions(+), 39 deletions(-) diff --git a/build.sh b/build.sh index 328c64b43b..35e6fc84fd 100755 --- a/build.sh +++ b/build.sh @@ -517,10 +517,10 @@ if hasArg docs; then cd "${DOXYGEN_BUILD_DIR}" doxygen Doxyfile cd "${SPHINX_BUILD_DIR}" - sphinx-build -W -b html source _html + make html cd "${REPODIR}"/rust cargo doc -p cuvs --no-deps - rsync -av "${RUST_BUILD_DIR}"/doc/ "${SPHINX_BUILD_DIR}"/_html/_static/rust + rsync -av "${RUST_BUILD_DIR}"/doc/ "${SPHINX_BUILD_DIR}"/build/html/_static/rust fi ################################################################################ diff --git a/ci/build_docs.sh b/ci/build_docs.sh index c53e6a0e9d..9fe65d1b04 100755 --- a/ci/build_docs.sh +++ b/ci/build_docs.sh @@ -1,5 +1,5 @@ #!/bin/bash -# Copyright (c) 2023-2024, NVIDIA CORPORATION. +# Copyright (c) 2023-2025, NVIDIA CORPORATION. set -euo pipefail @@ -48,10 +48,10 @@ popd rapids-logger "Build Python docs" pushd docs -sphinx-build -b dirhtml source _html -mv ../rust/target/doc ./_html/_static/rust +make dirhtml +mv ../rust/target/doc ./build/dirhtml/_static/rust mkdir -p "${RAPIDS_DOCS_DIR}/cuvs/"html -mv _html/* "${RAPIDS_DOCS_DIR}/cuvs/html" +mv build/dirhtml/* "${RAPIDS_DOCS_DIR}/cuvs/html" popd RAPIDS_VERSION_NUMBER="${RAPIDS_VERSION_MAJOR_MINOR}" rapids-upload-docs diff --git a/cpp/include/cuvs/neighbors/cagra.hpp b/cpp/include/cuvs/neighbors/cagra.hpp index ac86492ac6..1c0890b393 100644 --- a/cpp/include/cuvs/neighbors/cagra.hpp +++ b/cpp/include/cuvs/neighbors/cagra.hpp @@ -1964,6 +1964,10 @@ void serialize_to_hnswlib( std::optional> dataset = std::nullopt); +/** + * @} + */ + /** * @defgroup cagra_cpp_index_merge CAGRA index build functions * @{ diff --git a/cpp/include/cuvs/neighbors/hnsw.hpp b/cpp/include/cuvs/neighbors/hnsw.hpp index 7be7ccff54..b89a434156 100644 --- a/cpp/include/cuvs/neighbors/hnsw.hpp +++ b/cpp/include/cuvs/neighbors/hnsw.hpp @@ -163,6 +163,10 @@ struct extend_params { int num_threads = 0; }; +/** + * @} + */ + /** * @defgroup hnsw_cpp_index_load Load CAGRA index as hnswlib index * @{ diff --git a/cpp/include/cuvs/neighbors/vamana.hpp b/cpp/include/cuvs/neighbors/vamana.hpp index 25b3d4f7d6..625ceff988 100644 --- a/cpp/include/cuvs/neighbors/vamana.hpp +++ b/cpp/include/cuvs/neighbors/vamana.hpp @@ -488,6 +488,10 @@ auto build(raft::resources const& res, raft::host_matrix_view dataset) -> cuvs::neighbors::vamana::index; +/** + * @} + */ + /** * @defgroup vamana_cpp_serialize Vamana serialize functions * @{ diff --git a/cpp/include/cuvs/preprocessing/spectral_embedding.hpp b/cpp/include/cuvs/preprocessing/spectral_embedding.hpp index 3162f74f02..58f0880079 100644 --- a/cpp/include/cuvs/preprocessing/spectral_embedding.hpp +++ b/cpp/include/cuvs/preprocessing/spectral_embedding.hpp @@ -23,11 +23,6 @@ namespace cuvs::preprocessing::spectral_embedding { -/** - * @defgroup spectral_embedding Spectral Embedding - * @{ - */ - /** * @brief Parameters for spectral embedding algorithm * @@ -71,6 +66,11 @@ struct params { uint64_t seed; }; +/** + * @defgroup spectral_embedding Spectral Embedding + * @{ + */ + /** * @brief Perform spectral embedding on input dataset * diff --git a/docs/Makefile b/docs/Makefile index ac40850f1f..86a635a54d 100644 --- a/docs/Makefile +++ b/docs/Makefile @@ -2,7 +2,7 @@ # # You can set these variables from the command line. -SPHINXOPTS = +SPHINXOPTS = -W SPHINXBUILD = sphinx-build SPHINXPROJ = cuvs SOURCEDIR = source diff --git a/docs/source/neighbors/all_neighbors.rst b/docs/source/neighbors/all_neighbors.rst index c95d07488e..a70414fe06 100644 --- a/docs/source/neighbors/all_neighbors.rst +++ b/docs/source/neighbors/all_neighbors.rst @@ -16,7 +16,7 @@ All-neighbors supports multiple underlying algorithms: The algorithm partitions the dataset into clusters and distributes the work across multiple GPUs when possible, making it suitable for large-scale graph construction tasks. -[ :doc:`C API <../c_api/neighbors_all_neighbors_c>` |:doc:`C++ API <../cpp_api/neighbors_all_neighbors>` | :doc:`Python API <../python_api/neighbors_all_neighbors>` ] +[ :doc:`C API <../c_api/neighbors_all_neighbors_c>` | :doc:`C++ API <../cpp_api/neighbors_all_neighbors>` | :doc:`Python API <../python_api/neighbors_all_neighbors>` ] Algorithm Overview ------------------ diff --git a/docs/source/python_api/neighbors_mg_cagra.rst b/docs/source/python_api/neighbors_mg_cagra.rst index 2f03fa389b..763e0e2157 100644 --- a/docs/source/python_api/neighbors_mg_cagra.rst +++ b/docs/source/python_api/neighbors_mg_cagra.rst @@ -14,42 +14,42 @@ Multi-GPU CAGRA extends the graph-based CAGRA algorithm to work across multiple Index build parameters ###################### -.. autoclass:: cuvs.neighbors.mg_cagra.IndexParams +.. autoclass:: cuvs.neighbors.mg.cagra.IndexParams :members: Index search parameters ####################### -.. autoclass:: cuvs.neighbors.mg_cagra.SearchParams +.. autoclass:: cuvs.neighbors.mg.cagra.SearchParams :members: Index ##### -.. autoclass:: cuvs.neighbors.mg_cagra.Index +.. autoclass:: cuvs.neighbors.mg.cagra.Index :members: Index build ########### -.. autofunction:: cuvs.neighbors.mg_cagra.build +.. autofunction:: cuvs.neighbors.mg.cagra.build Index search ############ -.. autofunction:: cuvs.neighbors.mg_cagra.search +.. autofunction:: cuvs.neighbors.mg.cagra.search Index save ########## -.. autofunction:: cuvs.neighbors.mg_cagra.save +.. autofunction:: cuvs.neighbors.mg.cagra.save Index load ########## -.. autofunction:: cuvs.neighbors.mg_cagra.load +.. autofunction:: cuvs.neighbors.mg.cagra.load Index distribute ################ -.. autofunction:: cuvs.neighbors.mg_cagra.distribute +.. autofunction:: cuvs.neighbors.mg.cagra.distribute diff --git a/docs/source/python_api/neighbors_mg_ivf_flat.rst b/docs/source/python_api/neighbors_mg_ivf_flat.rst index 37515ce546..68eea86fec 100644 --- a/docs/source/python_api/neighbors_mg_ivf_flat.rst +++ b/docs/source/python_api/neighbors_mg_ivf_flat.rst @@ -14,47 +14,47 @@ Multi-GPU IVF-Flat extends the IVF-Flat algorithm to work across multiple GPUs, Index build parameters ###################### -.. autoclass:: cuvs.neighbors.mg_ivf_flat.IndexParams +.. autoclass:: cuvs.neighbors.mg.ivf_flat.IndexParams :members: Index search parameters ####################### -.. autoclass:: cuvs.neighbors.mg_ivf_flat.SearchParams +.. autoclass:: cuvs.neighbors.mg.ivf_flat.SearchParams :members: Index ##### -.. autoclass:: cuvs.neighbors.mg_ivf_flat.Index +.. autoclass:: cuvs.neighbors.mg.ivf_flat.Index :members: Index build ########### -.. autofunction:: cuvs.neighbors.mg_ivf_flat.build +.. autofunction:: cuvs.neighbors.mg.ivf_flat.build Index search ############ -.. autofunction:: cuvs.neighbors.mg_ivf_flat.search +.. autofunction:: cuvs.neighbors.mg.ivf_flat.search Index extend ############ -.. autofunction:: cuvs.neighbors.mg_ivf_flat.extend +.. autofunction:: cuvs.neighbors.mg.ivf_flat.extend Index save ########## -.. autofunction:: cuvs.neighbors.mg_ivf_flat.save +.. autofunction:: cuvs.neighbors.mg.ivf_flat.save Index load ########## -.. autofunction:: cuvs.neighbors.mg_ivf_flat.load +.. autofunction:: cuvs.neighbors.mg.ivf_flat.load Index distribute ################ -.. autofunction:: cuvs.neighbors.mg_ivf_flat.distribute +.. autofunction:: cuvs.neighbors.mg.ivf_flat.distribute diff --git a/docs/source/python_api/neighbors_mg_ivf_pq.rst b/docs/source/python_api/neighbors_mg_ivf_pq.rst index d7d13b4734..8343d59753 100644 --- a/docs/source/python_api/neighbors_mg_ivf_pq.rst +++ b/docs/source/python_api/neighbors_mg_ivf_pq.rst @@ -14,47 +14,47 @@ Multi-GPU IVF-PQ extends the IVF-PQ (Inverted File with Product Quantization) al Index build parameters ###################### -.. autoclass:: cuvs.neighbors.mg_ivf_pq.IndexParams +.. autoclass:: cuvs.neighbors.mg.ivf_pq.IndexParams :members: Index search parameters ####################### -.. autoclass:: cuvs.neighbors.mg_ivf_pq.SearchParams +.. autoclass:: cuvs.neighbors.mg.ivf_pq.SearchParams :members: Index ##### -.. autoclass:: cuvs.neighbors.mg_ivf_pq.Index +.. autoclass:: cuvs.neighbors.mg.ivf_pq.Index :members: Index build ########### -.. autofunction:: cuvs.neighbors.mg_ivf_pq.build +.. autofunction:: cuvs.neighbors.mg.ivf_pq.build Index search ############ -.. autofunction:: cuvs.neighbors.mg_ivf_pq.search +.. autofunction:: cuvs.neighbors.mg.ivf_pq.search Index extend ############ -.. autofunction:: cuvs.neighbors.mg_ivf_pq.extend +.. autofunction:: cuvs.neighbors.mg.ivf_pq.extend Index save ########## -.. autofunction:: cuvs.neighbors.mg_ivf_pq.save +.. autofunction:: cuvs.neighbors.mg.ivf_pq.save Index load ########## -.. autofunction:: cuvs.neighbors.mg_ivf_pq.load +.. autofunction:: cuvs.neighbors.mg.ivf_pq.load Index distribute ################ -.. autofunction:: cuvs.neighbors.mg_ivf_pq.distribute +.. autofunction:: cuvs.neighbors.mg.ivf_pq.distribute From 668e1b5c07a9c077b1dd6ab947745b1ae4b658cc Mon Sep 17 00:00:00 2001 From: Tarang Jain <40517122+tarang-jain@users.noreply.github.com> Date: Tue, 30 Sep 2025 13:22:18 -0700 Subject: [PATCH 91/93] Add Vamana C rst Docs to toctree (#1399) Authors: - Tarang Jain (https://github.com/tarang-jain) Approvers: - Ben Frederickson (https://github.com/benfred) URL: https://github.com/rapidsai/cuvs/pull/1399 --- docs/source/c_api/neighbors.rst | 1 + 1 file changed, 1 insertion(+) diff --git a/docs/source/c_api/neighbors.rst b/docs/source/c_api/neighbors.rst index bb6dbec5c2..09bd47e2c8 100644 --- a/docs/source/c_api/neighbors.rst +++ b/docs/source/c_api/neighbors.rst @@ -16,3 +16,4 @@ Nearest Neighbors neighbors_hnsw_c.rst neighbors_mg.rst neighbors_all_neighbors_c.rst + neighbors_vamana_c.rst From 815d86dd276ee9f7abf55e5f1bce9176d206a969 Mon Sep 17 00:00:00 2001 From: Bradley Dice Date: Tue, 30 Sep 2025 18:58:03 -0500 Subject: [PATCH 92/93] Patch FAISS for missing Thrust includes (#1398) This PR patches FAISS for some missing Thrust includes. This is needed to support building with CCCL 3.1.x. RAPIDS 25.10 will still use CCCL 3.0.x but we want to make it possible to build with 3.1.x in 25.10. xref: #1329 (This backports a fix from that testing PR, which is targeting 25.12) I am upstreaming the patch here: https://github.com/facebookresearch/faiss/pull/4597 Authors: - Bradley Dice (https://github.com/bdice) - Corey J. Nolet (https://github.com/cjnolet) Approvers: - Tarang Jain (https://github.com/tarang-jain) - Divye Gala (https://github.com/divyegala) URL: https://github.com/rapidsai/cuvs/pull/1398 --- ...fix-missing-includes-for-thrust-copy.patch | 40 +++++++++++++++++++ cpp/cmake/patches/faiss_override.json | 19 ++++++--- 2 files changed, 53 insertions(+), 6 deletions(-) create mode 100644 cpp/cmake/patches/faiss/fix-missing-includes-for-thrust-copy.patch diff --git a/cpp/cmake/patches/faiss/fix-missing-includes-for-thrust-copy.patch b/cpp/cmake/patches/faiss/fix-missing-includes-for-thrust-copy.patch new file mode 100644 index 0000000000..9f250b88be --- /dev/null +++ b/cpp/cmake/patches/faiss/fix-missing-includes-for-thrust-copy.patch @@ -0,0 +1,40 @@ +From 87c5331dfa6ec30fd9839ecb89af519e65f010e7 Mon Sep 17 00:00:00 2001 +From: Bradley Dice +Date: Mon, 29 Sep 2025 15:39:17 -0500 +Subject: [PATCH] Fix missing includes for thrust::copy + +--- + faiss/gpu/impl/BinaryCuvsCagra.cu | 3 +++ + faiss/gpu/impl/CuvsCagra.cu | 3 +++ + 2 files changed, 6 insertions(+) + +diff --git a/faiss/gpu/impl/BinaryCuvsCagra.cu b/faiss/gpu/impl/BinaryCuvsCagra.cu +index 0ca21dc5f..b331fdc8f 100644 +--- a/faiss/gpu/impl/BinaryCuvsCagra.cu ++++ b/faiss/gpu/impl/BinaryCuvsCagra.cu +@@ -32,6 +32,9 @@ + #include + #include + ++#include ++#include ++ + namespace faiss { + namespace gpu { + +diff --git a/faiss/gpu/impl/CuvsCagra.cu b/faiss/gpu/impl/CuvsCagra.cu +index 9ac4e1c5a..755817f43 100644 +--- a/faiss/gpu/impl/CuvsCagra.cu ++++ b/faiss/gpu/impl/CuvsCagra.cu +@@ -31,6 +31,9 @@ + #include + #include + ++#include ++#include ++ + namespace faiss { + namespace gpu { + +-- +2.49.0 diff --git a/cpp/cmake/patches/faiss_override.json b/cpp/cmake/patches/faiss_override.json index f3a1aaf795..5af8bd8543 100644 --- a/cpp/cmake/patches/faiss_override.json +++ b/cpp/cmake/patches/faiss_override.json @@ -1,9 +1,16 @@ { - "packages" : { - "faiss" : { - "version": "1.12.0", - "git_url": "https://github.com/facebookresearch/faiss.git", - "git_tag": "v1.12.0" - } + "packages" : { + "faiss" : { + "version": "1.12.0", + "git_url": "https://github.com/facebookresearch/faiss.git", + "git_tag": "v1.12.0", + "patches" : [ + { + "file" : "${current_json_dir}/faiss/fix-missing-includes-for-thrust-copy.patch", + "issue" : "Fix missing Thrust includes. https://github.com/facebookresearch/faiss/pull/4597", + "fixed_in" : "" + } + ] } } +} From b9a8f8062d50d9c0ec5ade8752d3027be7ebd1cc Mon Sep 17 00:00:00 2001 From: Jake Awe Date: Wed, 8 Oct 2025 15:23:48 +0000 Subject: [PATCH 93/93] Update Changelog [skip ci] --- CHANGELOG.md | 112 +++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 112 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index 2cf7bc0f62..a047a3f68a 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,3 +1,115 @@ +# cuvs 25.10.00 (8 Oct 2025) + +## 🚨 Breaking Changes + +- [Java] Uniform toHost/toDevice to work across all CuVSMatrix classes ([#1328](https://github.com/rapidsai/cuvs/pull/1328)) [@ldematte](https://github.com/ldematte) +- Use int64_t for getters in the cagra/ivf_flat c-api ([#1272](https://github.com/rapidsai/cuvs/pull/1272)) [@benfred](https://github.com/benfred) +- [REVIEW][Java] Rename destroyIndex() to close(), extend AutoCloseable ([#1252](https://github.com/rapidsai/cuvs/pull/1252)) [@ldematte](https://github.com/ldematte) +- Removing deprecated batching code specific to NN Descent ([#1249](https://github.com/rapidsai/cuvs/pull/1249)) [@jinsolp](https://github.com/jinsolp) +- MG C API ([#1160](https://github.com/rapidsai/cuvs/pull/1160)) [@viclafargue](https://github.com/viclafargue) + +## 🐛 Bug Fixes + +- Update `cpp/cmake/config.json` for `ConfigureTest()` ([#1385](https://github.com/rapidsai/cuvs/pull/1385)) [@KyleFromNVIDIA](https://github.com/KyleFromNVIDIA) +- Fix k-means++ root rank race condition ([#1359](https://github.com/rapidsai/cuvs/pull/1359)) [@csadorf](https://github.com/csadorf) +- Fix inadvertent uses of copy constructor in mdarrays across cuVS ([#1330](https://github.com/rapidsai/cuvs/pull/1330)) [@achirkin](https://github.com/achirkin) +- Refactor CMakeLists to support static only builds of libcuvs ([#1317](https://github.com/rapidsai/cuvs/pull/1317)) [@robertmaynard](https://github.com/robertmaynard) +- Fix pq_compressed file created by SSD Vamana build ([#1313](https://github.com/rapidsai/cuvs/pull/1313)) [@bkarsin](https://github.com/bkarsin) +- Revert "[Java]Binary and scalar quantization ([#1104)" (#1274](https://github.com/rapidsai/cuvs/pull/1104)" (#1274)) [@mythrocks](https://github.com/mythrocks) +- Fix balanced kmeans API ([#1268](https://github.com/rapidsai/cuvs/pull/1268)) [@enp1s0](https://github.com/enp1s0) +- [BUG] NEIGHBORS_ALL_NEIGHBORS_TEST build ignores --no-mg ([#1230](https://github.com/rapidsai/cuvs/pull/1230)) [@enp1s0](https://github.com/enp1s0) +- Add error handling for CPU-only bench ([#1203](https://github.com/rapidsai/cuvs/pull/1203)) [@lowener](https://github.com/lowener) +- Make duplicate removal in all neighbors robust to distance drift across batches ([#1185](https://github.com/rapidsai/cuvs/pull/1185)) [@jinsolp](https://github.com/jinsolp) +- Fix single GPU sharded search merge ([#1094](https://github.com/rapidsai/cuvs/pull/1094)) [@viclafargue](https://github.com/viclafargue) +- Processing mutual reachability dist in connect_knn_graph for host data ([#1093](https://github.com/rapidsai/cuvs/pull/1093)) [@jinsolp](https://github.com/jinsolp) +- [Java][Fix] Multithreaded querying fails without synchronization ([#1082](https://github.com/rapidsai/cuvs/pull/1082)) [@chatman](https://github.com/chatman) + +## 📖 Documentation + +- Add Vamana C rst Docs to toctree ([#1399](https://github.com/rapidsai/cuvs/pull/1399)) [@tarang-jain](https://github.com/tarang-jain) +- Regroup overloads in docs ([#1377](https://github.com/rapidsai/cuvs/pull/1377)) [@lowener](https://github.com/lowener) +- [DOCS] Update IVF Docs with Supported Distance Metrics ([#1365](https://github.com/rapidsai/cuvs/pull/1365)) [@tarang-jain](https://github.com/tarang-jain) +- Add note about Linux-only support for cuVS pre-compiled packages ([#1306](https://github.com/rapidsai/cuvs/pull/1306)) [@stic](https://github.com/stic) +- Documentation Fixes ([#1301](https://github.com/rapidsai/cuvs/pull/1301)) [@benfred](https://github.com/benfred) +- Docs for spectral embedding ([#1299](https://github.com/rapidsai/cuvs/pull/1299)) [@aamijar](https://github.com/aamijar) +- Improve NN Descent documentation ([#1246](https://github.com/rapidsai/cuvs/pull/1246)) [@jinsolp](https://github.com/jinsolp) +- Add NN.extend and kmeans to Python doc ([#1217](https://github.com/rapidsai/cuvs/pull/1217)) [@lowener](https://github.com/lowener) + +## 🚀 New Features + +- [REVIEW] [Java] Option to build fat-jars with native dependencies included ([#1296](https://github.com/rapidsai/cuvs/pull/1296)) [@mythrocks](https://github.com/mythrocks) +- Updates needed when building with CUDA 13 ([#1219](https://github.com/rapidsai/cuvs/pull/1219)) [@robertmaynard](https://github.com/robertmaynard) +- [Java] Add CAGRA index graph accessor/build from graph (host memory) ([#1216](https://github.com/rapidsai/cuvs/pull/1216)) [@ldematte](https://github.com/ldematte) +- [Feat] Add Dockerfile for reproducible installation ([#1195](https://github.com/rapidsai/cuvs/pull/1195)) [@rhdong](https://github.com/rhdong) +- Vamana C / Python API ([#1112](https://github.com/rapidsai/cuvs/pull/1112)) [@tarang-jain](https://github.com/tarang-jain) +- Mutual Reachability in all-neighbors API ([#1016](https://github.com/rapidsai/cuvs/pull/1016)) [@jinsolp](https://github.com/jinsolp) +- CAGRA Build + DiskANN Search cuvs-bench Wrapper ([#899](https://github.com/rapidsai/cuvs/pull/899)) [@tarang-jain](https://github.com/tarang-jain) + +## 🛠️ Improvements + +- Patch FAISS for missing Thrust includes ([#1398](https://github.com/rapidsai/cuvs/pull/1398)) [@bdice](https://github.com/bdice) +- Empty commit to trigger a build ([#1363](https://github.com/rapidsai/cuvs/pull/1363)) [@msarahan](https://github.com/msarahan) +- [Review][Java] Refactor: extract interface from CuVSMatrixBaseImpl ([#1361](https://github.com/rapidsai/cuvs/pull/1361)) [@ldematte](https://github.com/ldematte) +- use CUDA 13.0.1 CI images ([#1353](https://github.com/rapidsai/cuvs/pull/1353)) [@jameslamb](https://github.com/jameslamb) +- Treat warnings as errors in doc builds + MG ANNs Python API doc update ([#1350](https://github.com/rapidsai/cuvs/pull/1350)) [@viclafargue](https://github.com/viclafargue) +- [Java] Support row strides in CuVSMatrix ([#1345](https://github.com/rapidsai/cuvs/pull/1345)) [@ldematte](https://github.com/ldematte) +- [Java] Adding tests to use CuVSDeviceMatrix (device memory) directly as a CagraIndex input dataset ([#1340](https://github.com/rapidsai/cuvs/pull/1340)) [@ldematte](https://github.com/ldematte) +- Improve performance of assigning clusters in batched all_neighbors ([#1336](https://github.com/rapidsai/cuvs/pull/1336)) [@jinsolp](https://github.com/jinsolp) +- Configure repo for automatic release notes generation ([#1334](https://github.com/rapidsai/cuvs/pull/1334)) [@AyodeAwe](https://github.com/AyodeAwe) +- [Java] Buffered device matrix builder ([#1332](https://github.com/rapidsai/cuvs/pull/1332)) [@ldematte](https://github.com/ldematte) +- Remove UCX-Py/UCXX from release script ([#1331](https://github.com/rapidsai/cuvs/pull/1331)) [@pentschev](https://github.com/pentschev) +- [Java] Uniform toHost/toDevice to work across all CuVSMatrix classes ([#1328](https://github.com/rapidsai/cuvs/pull/1328)) [@ldematte](https://github.com/ldematte) +- [Java] Add `libcuvs` <-> cuvs-java version check ([#1327](https://github.com/rapidsai/cuvs/pull/1327)) [@ldematte](https://github.com/ldematte) +- update dependencies: use cuda-toolkit wheels ([#1326](https://github.com/rapidsai/cuvs/pull/1326)) [@jameslamb](https://github.com/jameslamb) +- Add cagra.extend bindings for python ([#1324](https://github.com/rapidsai/cuvs/pull/1324)) [@benfred](https://github.com/benfred) +- [Java] Make `cudaGetDeviceProperties` compatible with CUDA 12 and 13 based on symbol presence ([#1323](https://github.com/rapidsai/cuvs/pull/1323)) [@ldematte](https://github.com/ldematte) +- Use branch-25.10 again ([#1319](https://github.com/rapidsai/cuvs/pull/1319)) [@jameslamb](https://github.com/jameslamb) +- [Review][Java] Add detailed error message for `libcuvs` load failure to UnsupportedProvider/UnsupportedOperationExceptions ([#1316](https://github.com/rapidsai/cuvs/pull/1316)) [@ldematte](https://github.com/ldematte) +- [Java] Add reason(s) to UnsupportedProvider/UnsupportedOperationExceptions ([#1314](https://github.com/rapidsai/cuvs/pull/1314)) [@ldematte](https://github.com/ldematte) +- MG Python API ([#1307](https://github.com/rapidsai/cuvs/pull/1307)) [@viclafargue](https://github.com/viclafargue) +- Fix redundant memset ([#1305](https://github.com/rapidsai/cuvs/pull/1305)) [@vinaydes](https://github.com/vinaydes) +- Expose pq-centers to C and Python ([#1303](https://github.com/rapidsai/cuvs/pull/1303)) [@benfred](https://github.com/benfred) +- Port `raft::neighbors::epsilon_neighborhood` to cuvs ([#1294](https://github.com/rapidsai/cuvs/pull/1294)) [@aamijar](https://github.com/aamijar) +- ScaNN: Overlapped gather for AVQ ([#1286](https://github.com/rapidsai/cuvs/pull/1286)) [@rmaschal](https://github.com/rmaschal) +- Update rapids-dependency-file-generator ([#1285](https://github.com/rapidsai/cuvs/pull/1285)) [@KyleFromNVIDIA](https://github.com/KyleFromNVIDIA) +- [Review][Java] Expand `CagraIndex#search` to accept more types (int8 and int32) ([#1283](https://github.com/rapidsai/cuvs/pull/1283)) [@ldematte](https://github.com/ldematte) +- All neighbors C and Python bindings ([#1282](https://github.com/rapidsai/cuvs/pull/1282)) [@viclafargue](https://github.com/viclafargue) +- Build and test with CUDA 13.0.0 ([#1273](https://github.com/rapidsai/cuvs/pull/1273)) [@jameslamb](https://github.com/jameslamb) +- Use int64_t for getters in the cagra/ivf_flat c-api ([#1272](https://github.com/rapidsai/cuvs/pull/1272)) [@benfred](https://github.com/benfred) +- [Java][C] Expose GPUInfo ([#1267](https://github.com/rapidsai/cuvs/pull/1267)) [@ldematte](https://github.com/ldematte) +- Add rust bindings for kmeans ([#1266](https://github.com/rapidsai/cuvs/pull/1266)) [@benfred](https://github.com/benfred) +- [REVIEW] Add a public API for CAGRA graph optimize ([#1260](https://github.com/rapidsai/cuvs/pull/1260)) [@abc99lr](https://github.com/abc99lr) +- Fix debug build ([#1258](https://github.com/rapidsai/cuvs/pull/1258)) [@lowener](https://github.com/lowener) +- Run `cargo fmt` in the pre-commit hooks ([#1255](https://github.com/rapidsai/cuvs/pull/1255)) [@benfred](https://github.com/benfred) +- [REVIEW][Java] Rename destroyIndex() to close(), extend AutoCloseable ([#1252](https://github.com/rapidsai/cuvs/pull/1252)) [@ldematte](https://github.com/ldematte) +- Removing deprecated batching code specific to NN Descent ([#1249](https://github.com/rapidsai/cuvs/pull/1249)) [@jinsolp](https://github.com/jinsolp) +- Pin Latest Faiss Version ([#1247](https://github.com/rapidsai/cuvs/pull/1247)) [@tarang-jain](https://github.com/tarang-jain) +- Use build cluster in devcontainers ([#1240](https://github.com/rapidsai/cuvs/pull/1240)) [@trxcllnt](https://github.com/trxcllnt) +- Use rapids_cuda_enable_fatbin_compression ([#1239](https://github.com/rapidsai/cuvs/pull/1239)) [@robertmaynard](https://github.com/robertmaynard) +- [Improvement] Replace Calls to thrust::transform with raft::linalg ([#1238](https://github.com/rapidsai/cuvs/pull/1238)) [@tarang-jain](https://github.com/tarang-jain) +- [Java] CuVSMatrix for device memory ([#1232](https://github.com/rapidsai/cuvs/pull/1232)) [@ldematte](https://github.com/ldematte) +- Drop log level for CAGRA trace messages ([#1229](https://github.com/rapidsai/cuvs/pull/1229)) [@mythrocks](https://github.com/mythrocks) +- Update rapids_config to handle user defined branch name ([#1227](https://github.com/rapidsai/cuvs/pull/1227)) [@robertmaynard](https://github.com/robertmaynard) +- Reduce noise in Java tests ([#1226](https://github.com/rapidsai/cuvs/pull/1226)) [@mythrocks](https://github.com/mythrocks) +- Update another instance of rapids-build-backend ([#1225](https://github.com/rapidsai/cuvs/pull/1225)) [@KyleFromNVIDIA](https://github.com/KyleFromNVIDIA) +- Fix missing sync_stream in ScaNN build ([#1224](https://github.com/rapidsai/cuvs/pull/1224)) [@rmaschal](https://github.com/rmaschal) +- [Java] Exception-safe RMM Allocations ([#1215](https://github.com/rapidsai/cuvs/pull/1215)) [@mythrocks](https://github.com/mythrocks) +- Update rapids-build-backend to 0.4.1 ([#1214](https://github.com/rapidsai/cuvs/pull/1214)) [@KyleFromNVIDIA](https://github.com/KyleFromNVIDIA) +- Removing filter templates in ivfpq for binary size reduction ([#1211](https://github.com/rapidsai/cuvs/pull/1211)) [@jinsolp](https://github.com/jinsolp) +- ci(labeler): update labeler action to [@v5 ([#1210](https://github.com/rapidsai/cuvs/pull/1210)) @gforsyth](https://github.com/v5 ([#1210](https://github.com/rapidsai/cuvs/pull/1210)) @gforsyth) +- [Java] Add decorator for synchronized CuVSResource access ([#1209](https://github.com/rapidsai/cuvs/pull/1209)) [@ldematte](https://github.com/ldematte) +- Remove unused templated instances ([#1204](https://github.com/rapidsai/cuvs/pull/1204)) [@vinaydes](https://github.com/vinaydes) +- Spectral Embedding precomputed graph api ([#1197](https://github.com/rapidsai/cuvs/pull/1197)) [@aamijar](https://github.com/aamijar) +- cuvs RAPIDS dependency now use the new rapids-cmake branch info ([#1189](https://github.com/rapidsai/cuvs/pull/1189)) [@robertmaynard](https://github.com/robertmaynard) +- Allow latest OS in devcontainers ([#1169](https://github.com/rapidsai/cuvs/pull/1169)) [@bdice](https://github.com/bdice) +- MG C API ([#1160](https://github.com/rapidsai/cuvs/pull/1160)) [@viclafargue](https://github.com/viclafargue) +- Merge support for tiered index ([#1155](https://github.com/rapidsai/cuvs/pull/1155)) [@benfred](https://github.com/benfred) +- Update build infra to support new branching strategy ([#1143](https://github.com/rapidsai/cuvs/pull/1143)) [@robertmaynard](https://github.com/robertmaynard) +- [Java]Binary and scalar quantization ([#1104](https://github.com/rapidsai/cuvs/pull/1104)) [@punAhuja](https://github.com/punAhuja) +- Improve Vamana index build performance and recall ([#1032](https://github.com/rapidsai/cuvs/pull/1032)) [@bkarsin](https://github.com/bkarsin) +- Use GCC 14 in conda builds. ([#1030](https://github.com/rapidsai/cuvs/pull/1030)) [@vyasr](https://github.com/vyasr) +- Update SNMG ANN API testing ([#1023](https://github.com/rapidsai/cuvs/pull/1023)) [@viclafargue](https://github.com/viclafargue) + # cuvs 25.08.00 (6 Aug 2025) ## 🚨 Breaking Changes