diff --git a/.github/workflows/pr.yaml b/.github/workflows/pr.yaml index a6438c1a8d..40c2cb142e 100644 --- a/.github/workflows/pr.yaml +++ b/.github/workflows/pr.yaml @@ -11,450 +11,18 @@ jobs: pr-builder: needs: - build-details - - check-nightly-ci - - changed-files - - checks - - check-c-abi - - security-suite - conda-cpp-build - conda-cpp-tests - - conda-cpp-checks - - conda-python-build - - conda-python-tests - - rocky8-clib-standalone-build - - rocky8-clib-tests - - conda-java-build - - conda-java-tests - - conda-lucene-build - - conda-lucene-tests - - rust-build - - go-build - - docs-build - - docs-build-matrix - - wheel-build-libcuvs - - wheel-build-cuvs - - wheel-tests-cuvs - - devcontainer - - telemetry-setup uses: rapidsai/shared-workflows/.github/workflows/pr-builder.yaml@main permissions: contents: read if: always() with: needs: ${{ toJSON(needs) }} - telemetry-setup: - continue-on-error: true - runs-on: ubuntu-latest - permissions: - contents: read - env: - OTEL_SERVICE_NAME: 'pr-cuvs' - steps: - - name: Telemetry setup - if: ${{ vars.TELEMETRY_ENABLED == 'true' }} - uses: rapidsai/shared-actions/telemetry-dispatch-stash-base-env-vars@main build-details: uses: rapidsai/shared-workflows/.github/workflows/compute-build-details.yaml@main - check-nightly-ci: - runs-on: ubuntu-latest - permissions: - actions: read - contents: read - id-token: write - pull-requests: read - env: - GH_TOKEN: ${{ github.token }} - steps: - - name: Get PR Info - id: get-pr-info - uses: nv-gha-runners/get-pr-info@main - - name: Check if nightly CI is passing - uses: rapidsai/shared-actions/check_nightly_success/dispatch@main - with: - repo: ${{ github.repository }} - target-branch: ${{ fromJSON(steps.get-pr-info.outputs.pr-info).base.ref }} - max-days-without-success: 30 - changed-files: - needs: telemetry-setup - permissions: - actions: read - contents: read - packages: read - pull-requests: read - uses: rapidsai/shared-workflows/.github/workflows/changed-files.yaml@main - with: - files_yaml: | - build_docs: - - '**' - - '!**/*/REVIEW_GUIDELINES.md' - - '!.coderabbit.yaml' - - '!.devcontainer/**' - - '!.github/CODEOWNERS' - - '!.github/ISSUE_TEMPLATE/**' - - '!.github/PULL_REQUEST_TEMPLATE.md' - - '!.github/copy-pr-bot.yaml' - - '!.github/labeler.yml' - - '!.github/ops-bot.yaml' - - '!.github/release.yml' - - '!.github/workflows/build.yaml' - - '!.github/workflows/check-c-abi.yaml' - - '!.github/workflows/labeler.yml' - - '!.github/workflows/publish-rust.yaml' - - '!.github/workflows/security-suite.yml' - - '!.github/workflows/store-c-abi-baseline.yaml' - - '!.github/workflows/test.yaml' - - '!.github/workflows/trigger-breaking-change-alert.yaml' - - '!.github/workflows/update-c-abi-baseline.yaml' - - '!.pre-commit-config.yaml' - - '!.shellcheckrc' - - '!.yamllint.yaml' - - '!Dockerfile' - - '!SECURITY.md' - - '!ci/build_wheel*.sh' - - '!ci/check_style.sh' - - '!ci/release/update-version.sh' - - '!ci/run*.sh' - - '!ci/test*.sh' - - '!ci/validate_wheel.sh' - - '!thirdparty/LICENSES/**' - test_cpp: - - '.github/workflows/pr.yaml' - - 'RAPIDS_BRANCH' - - 'VERSION' - - 'build.sh' - - 'c/**' - - 'ci/build_cpp.sh' - - 'ci/get_test_data.sh' - - 'ci/test_cpp.sh' - - 'cmake/**' - - 'conda/recipes/libcuvs/**' - - 'cpp/**' - - 'dependencies.yaml' - - '!**/*/REVIEW_GUIDELINES.md' - - '!cpp/doxygen/**' - test_standalone_c: - - '.github/workflows/pr.yaml' - - 'RAPIDS_BRANCH' - - 'VERSION' - - 'build.sh' - - 'c/**' - - 'ci/standalone_c/**' - - 'ci/test_standalone_c.sh' - - 'cmake/**' - - 'cpp/**' - - 'dependencies.yaml' - - '!**/*/REVIEW_GUIDELINES.md' - - '!cpp/doxygen/**' - test_java: - - '**' - - '!**/*/REVIEW_GUIDELINES.md' - - '!.coderabbit.yaml' - - '!.devcontainer/**' - - '!.github/CODEOWNERS' - - '!.github/ISSUE_TEMPLATE/**' - - '!.github/PULL_REQUEST_TEMPLATE.md' - - '!.github/copy-pr-bot.yaml' - - '!.github/labeler.yml' - - '!.github/ops-bot.yaml' - - '!.github/release.yml' - - '!.github/workflows/build.yaml' - - '!.github/workflows/check-c-abi.yaml' - - '!.github/workflows/labeler.yml' - - '!.github/workflows/publish-rust.yaml' - - '!.github/workflows/security-suite.yml' - - '!.github/workflows/store-c-abi-baseline.yaml' - - '!.github/workflows/test.yaml' - - '!.github/workflows/trigger-breaking-change-alert.yaml' - - '!.github/workflows/update-c-abi-baseline.yaml' - - '!.pre-commit-config.yaml' - - '!.shellcheckrc' - - '!.yamllint.yaml' - - '!CODEOWNERS' - - '!Dockerfile' - - '!README.md' - - '!SECURITY.md' - - '!ci/build_docs.sh' - - '!ci/build_go.sh' - - '!ci/build_python.sh' - - '!ci/build_rust.sh' - - '!ci/build_wheel*.sh' - - '!ci/checkstyle.sh' - - '!ci/run_cuvs_pytests.sh' - - '!ci/test_cpp.sh' - - '!ci/test_python.sh' - - '!ci/test_wheel_cuvs.sh' - - '!ci/validate_wheel.sh' - - '!conda/environments/**' - - '!conda/recipes/cuvs/**' - - '!conda/recipes/cuvs-bench/**' - - '!conda/recipes/cuvs-bench-cpu/**' - - '!docs/**' - - '!examples/.clang-format' - - '!examples/.clang-tidy' - - '!examples/.clangd' - - '!examples/README.md' - - '!examples/build.sh' - - '!examples/c/**' - - '!examples/cmake/**' - - '!examples/cpp/**' - - '!examples/go/**' - - '!fern/**' - - '!go/**' - - '!img/**' - - '!notebooks/**' - - '!pyproject.toml' - - '!python/**' - - '!rust/**' - - '!thirdparty/LICENSES/**' - test_python_conda: - - '**' - - '!**/*/REVIEW_GUIDELINES.md' - - '!.coderabbit.yaml' - - '!.devcontainer/**' - - '!.github/CODEOWNERS' - - '!.github/ISSUE_TEMPLATE/**' - - '!.github/PULL_REQUEST_TEMPLATE.md' - - '!.github/copy-pr-bot.yaml' - - '!.github/labeler.yml' - - '!.github/ops-bot.yaml' - - '!.github/release.yml' - - '!.github/workflows/build.yaml' - - '!.github/workflows/check-c-abi.yaml' - - '!.github/workflows/labeler.yml' - - '!.github/workflows/publish-rust.yaml' - - '!.github/workflows/security-suite.yml' - - '!.github/workflows/store-c-abi-baseline.yaml' - - '!.github/workflows/test.yaml' - - '!.github/workflows/trigger-breaking-change-alert.yaml' - - '!.github/workflows/update-c-abi-baseline.yaml' - - '!.pre-commit-config.yaml' - - '!.shellcheckrc' - - '!.yamllint.yaml' - - '!ci/build_go.sh' - - '!ci/build_java.sh' - - '!ci/build_lucene.sh' - - '!ci/build_rust.sh' - - '!ci/check_style.sh' - - '!ci/release/update-version.sh' - - '!ci/run_ctests.sh' - - '!ci/test_cpp.sh' - - '!ci/test_java.sh' - - '!ci/test_lucene.sh' - - '!ci/test_standalone_c.sh' - - '!ci/test_wheel_cuvs.sh' - - '!ci/validate_wheel.sh' - - '!cpp/.clang-format' - - '!cpp/.clang-tidy' - - '!CODEOWNERS' - - '!Dockerfile' - - '!README.md' - - '!SECURITY.md' - - '!docs/**' - - '!examples/**' - - '!fern/**' - - '!go/**' - - '!img/**' - - '!java/**' - - '!notebooks/**' - - '!rust/**' - - '!thirdparty/LICENSES/**' - test_python_wheels: - - '**' - - '!**/*/REVIEW_GUIDELINES.md' - - '!.coderabbit.yaml' - - '!.devcontainer/**' - - '!.github/CODEOWNERS' - - '!.github/ISSUE_TEMPLATE/**' - - '!.github/PULL_REQUEST_TEMPLATE.md' - - '!.github/copy-pr-bot.yaml' - - '!.github/labeler.yml' - - '!.github/ops-bot.yaml' - - '!.github/release.yml' - - '!.github/workflows/build.yaml' - - '!.github/workflows/check-c-abi.yaml' - - '!.github/workflows/labeler.yml' - - '!.github/workflows/publish-rust.yaml' - - '!.github/workflows/security-suite.yml' - - '!.github/workflows/store-c-abi-baseline.yaml' - - '!.github/workflows/test.yaml' - - '!.github/workflows/trigger-breaking-change-alert.yaml' - - '!.github/workflows/update-c-abi-baseline.yaml' - - '!.pre-commit-config.yaml' - - '!.shellcheckrc' - - '!.yamllint.yaml' - - '!ci/build_cpp.sh' - - '!ci/build_docs.sh' - - '!ci/build_go.sh' - - '!ci/build_java.sh' - - '!ci/build_lucene.sh' - - '!ci/build_python.sh' - - '!ci/build_rust.sh' - - '!ci/build_standalone_c.sh' - - '!ci/check_style.sh' - - '!ci/release/update-version.sh' - - '!ci/run_ctests.sh' - - '!ci/test_cpp.sh' - - '!ci/test_java.sh' - - '!ci/test_lucene.sh' - - '!ci/test_python.sh' - - '!ci/test_standalone_c.sh' - - '!ci/validate_wheel.sh' - - '!CODEOWNERS' - - '!Dockerfile' - - '!README.md' - - '!SECURITY.md' - - '!conda/**' - - '!cpp/.clang-format' - - '!cpp/.clang-tidy' - - '!docs/**' - - '!examples/**' - - '!fern/**' - - '!java/**' - - '!go/**' - - '!img/**' - - '!notebooks/**' - - '!rust/**' - - '!thirdparty/LICENSES/**' - test_rust: - - '**' - - '!**/*/REVIEW_GUIDELINES.md' - - '!.coderabbit.yaml' - - '!.devcontainer/**' - - '!.github/CODEOWNERS' - - '!.github/ISSUE_TEMPLATE/**' - - '!.github/PULL_REQUEST_TEMPLATE.md' - - '!.github/copy-pr-bot.yaml' - - '!.github/labeler.yml' - - '!.github/ops-bot.yaml' - - '!.github/release.yml' - - '!.github/workflows/build.yaml' - - '!.github/workflows/check-c-abi.yaml' - - '!.github/workflows/labeler.yml' - - '!.github/workflows/publish-rust.yaml' - - '!.github/workflows/security-suite.yml' - - '!.github/workflows/store-c-abi-baseline.yaml' - - '!.github/workflows/test.yaml' - - '!.github/workflows/trigger-breaking-change-alert.yaml' - - '!.github/workflows/update-c-abi-baseline.yaml' - - '!.pre-commit-config.yaml' - - '!.shellcheckrc' - - '!.yamllint.yaml' - - '!CODEOWNERS' - - '!Dockerfile' - - '!README.md' - - '!SECURITY.md' - - '!ci/build_docs.sh' - - '!ci/build_go.sh' - - '!ci/build_java.sh' - - '!ci/build_lucene.sh' - - '!ci/build_python.sh' - - '!ci/build_wheel*.sh' - - '!ci/check_style.sh' - - '!ci/release/update-version.sh' - - '!ci/run_ctests.sh' - - '!ci/run_cuvs_pytests.sh' - - '!ci/test_cpp.sh' - - '!ci/test_java.sh' - - '!ci/test_lucene.sh' - - '!ci/test_python.sh' - - '!ci/test_wheel_cuvs.sh' - - '!ci/validate_wheel.sh' - - '!docs/**' - - '!examples/**' - - '!fern/**' - - '!go/**' - - '!img/**' - - '!java/**' - - '!notebooks/**' - - '!python/**' - - '!thirdparty/LICENSES/**' - test_go: - - '**' - - '!**/*/REVIEW_GUIDELINES.md' - - '!.coderabbit.yaml' - - '!.devcontainer/**' - - '!.github/CODEOWNERS' - - '!.github/ISSUE_TEMPLATE/**' - - '!.github/PULL_REQUEST_TEMPLATE.md' - - '!.github/copy-pr-bot.yaml' - - '!.github/labeler.yml' - - '!.github/ops-bot.yaml' - - '!.github/release.yml' - - '!.github/workflows/build.yaml' - - '!.github/workflows/check-c-abi.yaml' - - '!.github/workflows/labeler.yml' - - '!.github/workflows/publish-rust.yaml' - - '!.github/workflows/security-suite.yml' - - '!.github/workflows/store-c-abi-baseline.yaml' - - '!.github/workflows/test.yaml' - - '!.github/workflows/trigger-breaking-change-alert.yaml' - - '!.github/workflows/update-c-abi-baseline.yaml' - - '!.pre-commit-config.yaml' - - '!.shellcheckrc' - - '!.yamllint.yaml' - - '!CODEOWNERS' - - '!Dockerfile' - - '!README.md' - - '!SECURITY.md' - - '!ci/build_docs.sh' - - '!ci/build_java.sh' - - '!ci/build_lucene.sh' - - '!ci/build_python.sh' - - '!ci/build_rust.sh' - - '!ci/build_wheel*.sh' - - '!ci/check_style.sh' - - '!ci/release/update-version.sh' - - '!ci/run_ctests.sh' - - '!ci/run_cuvs_pytests.sh' - - '!ci/test_cpp.sh' - - '!ci/test_java.sh' - - '!ci/test_lucene.sh' - - '!ci/test_python.sh' - - '!ci/test_standalone_c.sh' - - '!ci/test_wheel_cuvs.sh' - - '!ci/validate_wheel.sh' - - '!docs/**' - - '!examples/.clang-format' - - '!examples/.clang-tidy' - - '!examples/.clangd' - - '!examples/README.md' - - '!examples/build.sh' - - '!examples/c/**' - - '!examples/cmake/**' - - '!examples/cpp/**' - - '!examples/java/**' - - '!fern/**' - - '!img/**' - - '!java/**' - - '!notebooks/**' - - '!python/**' - - '!rust/**' - - '!thirdparty/LICENSES/**' - checks: - needs: telemetry-setup - uses: rapidsai/shared-workflows/.github/workflows/checks.yaml@main - permissions: - contents: read - with: - enable_check_generated_files: false - ignored_pr_jobs: "telemetry-summarize" - check-c-abi: - needs: telemetry-setup - permissions: - actions: read - contents: read - issues: write - pull-requests: read - uses: ./.github/workflows/check-c-abi.yaml - security-suite: - permissions: - actions: read - contents: read - id-token: write # OIDC -> Vault -> nvcr.io image pull - security-events: write # publish redacted SARIF to code scanning - uses: ./.github/workflows/security-suite.yml conda-cpp-build: - needs: [build-details, checks] + needs: [build-details] permissions: actions: read contents: read @@ -468,8 +36,9 @@ jobs: build-datetime: ${{ needs.build-details.outputs.build-datetime }} node_type: cpu16 script: ci/build_cpp.sh + matrix_filter: map(select(.CUDA_VER == "13.3.0")) conda-cpp-tests: - needs: [conda-cpp-build, changed-files] + needs: [conda-cpp-build] permissions: actions: read contents: read @@ -478,16 +47,11 @@ jobs: pull-requests: read secrets: inherit # zizmor: ignore[secrets-inherit] uses: rapidsai/shared-workflows/.github/workflows/conda-cpp-tests.yaml@main - if: fromJSON(needs.changed-files.outputs.changed_file_groups).test_cpp - # Split the C++ tests across several jobs for each CUDA/GPU configuration - strategy: - fail-fast: false - matrix: - shard: [1, 2, 3, 4] - num_shards: [4] with: build_type: pull-request - script: "ci/test_cpp.sh ${{ matrix.shard }} ${{ matrix.num_shards }}" + script: "ci/test_cpp.sh" + matrix_type: nightly + matrix_filter: 'map(select(.CUDA_VER == "13.3.0" and (.GPU == "h100" or .GPU == "a100" or .GPU == "rtxpro6000"))) | [.[] | if .GPU == "h100" then (., . + {COPY: 1}) else . end]' cache-paths: .cache/cuda-jit cache-key-prefix: libcuvs-cudajit-v1 # Caches are immutable. Re-key when the JIT-relevant sources change so a completed run can @@ -518,279 +82,3 @@ jobs: cache-read-only: true cache-environment: | CUDA_CACHE_PATH=.cache/cuda-jit - conda-cpp-checks: - needs: conda-cpp-build - permissions: - actions: read - contents: read - id-token: write - packages: read - pull-requests: read - uses: rapidsai/shared-workflows/.github/workflows/conda-cpp-post-build-checks.yaml@main - with: - build_type: pull-request - symbol_exclusions: (void (thrust::|cub::)) - package_name: libcuvs - conda-python-build: - needs: [build-details, conda-cpp-build] - permissions: - actions: read - contents: read - id-token: write - packages: read - pull-requests: read - secrets: inherit # zizmor: ignore[secrets-inherit] - uses: rapidsai/shared-workflows/.github/workflows/conda-python-build.yaml@main - with: - build_type: pull-request - build-datetime: ${{ needs.build-details.outputs.build-datetime }} - script: ci/build_python.sh - # Build a conda package for each CUDA x ARCH x minimum supported Python version - matrix_filter: group_by({CUDA_VER, ARCH}) | map(min_by(.PY_VER | split(".") | map(tonumber))) - conda-python-tests: - needs: [conda-python-build, changed-files] - permissions: - actions: read - contents: read - id-token: write - packages: read - pull-requests: read - secrets: inherit # zizmor: ignore[secrets-inherit] - uses: rapidsai/shared-workflows/.github/workflows/conda-python-tests.yaml@main - if: fromJSON(needs.changed-files.outputs.changed_file_groups).test_python_conda - with: - build_type: pull-request - script: ci/test_python.sh - rocky8-clib-standalone-build: - needs: checks - permissions: - actions: read - contents: read - id-token: write - packages: read - pull-requests: read - uses: ./.github/workflows/clib-build.yaml - with: - build_type: pull-request - rocky8-clib-tests: - needs: [rocky8-clib-standalone-build, changed-files] - permissions: - actions: read - contents: read - id-token: write - packages: read - pull-requests: read - uses: ./.github/workflows/matrix-job.yaml - if: fromJSON(needs.changed-files.outputs.changed_file_groups).test_standalone_c - with: - build_type: pull-request - container-image-repository: "rapidsai/citestwheel" - script: "ci/test_standalone_c.sh" - conda-java-build: - needs: [conda-cpp-build, changed-files] - permissions: - actions: read - contents: read - id-token: write - packages: read - pull-requests: read - uses: ./.github/workflows/matrix-job.yaml - if: fromJSON(needs.changed-files.outputs.changed_file_groups).test_java || fromJSON(needs.changed-files.outputs.changed_file_groups).test_cpp - # 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. - with: - build_type: pull-request - node_type: cpu8 - matrix_filter: map(select(.ARCH == "amd64")) - script: "ci/build_java.sh" - artifact-name-prefix: "cuvs-java" - file_to_upload: "java/cuvs-java/target/" - conda-java-tests: - needs: conda-java-build - permissions: - actions: read - contents: read - id-token: write - packages: read - pull-requests: read - uses: ./.github/workflows/matrix-job.yaml - with: - build_type: pull-request - script: "ci/test_java.sh" - conda-lucene-build: - # Depends on the Java job for the cuvs-java artifact that cuvs-lucene builds against. - needs: conda-java-build - permissions: - actions: read - contents: read - id-token: write - packages: read - pull-requests: read - uses: ./.github/workflows/matrix-job.yaml - # 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. - with: - build_type: pull-request - node_type: cpu8 - matrix_filter: map(select(.ARCH == "amd64")) - script: "ci/build_lucene.sh" - artifact-name-prefix: "cuvs-lucene" - file_to_upload: "java/cuvs-lucene/target/" - conda-lucene-tests: - needs: conda-lucene-build - permissions: - actions: read - contents: read - id-token: write - packages: read - pull-requests: read - uses: ./.github/workflows/matrix-job.yaml - with: - build_type: pull-request - script: "ci/test_lucene.sh" - rust-build: - needs: [conda-cpp-build, changed-files] - permissions: - actions: read - contents: read - id-token: write - packages: read - pull-requests: read - uses: ./.github/workflows/matrix-job.yaml - if: fromJSON(needs.changed-files.outputs.changed_file_groups).test_rust || fromJSON(needs.changed-files.outputs.changed_file_groups).test_cpp - # 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. - with: - build_type: pull-request - matrix_filter: map(select(.ARCH == "amd64")) - script: "ci/build_rust.sh" - go-build: - needs: [conda-cpp-build, changed-files] - permissions: - actions: read - contents: read - id-token: write - packages: read - pull-requests: read - uses: ./.github/workflows/matrix-job.yaml - if: fromJSON(needs.changed-files.outputs.changed_file_groups).test_go || fromJSON(needs.changed-files.outputs.changed_file_groups).test_cpp - # 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. - with: - build_type: pull-request - matrix_filter: map(select(.ARCH == "amd64")) - script: "ci/build_go.sh" - docs-build-matrix: - permissions: - contents: read - uses: rapidsai/shared-workflows/.github/workflows/compute-matrix.yaml@main - with: - build_type: pull-request - matrix_name: docs-build - docs-build: - needs: [docs-build-matrix] - permissions: - actions: read - contents: read - id-token: write - packages: read - pull-requests: read - secrets: inherit # zizmor: ignore[secrets-inherit] - uses: rapidsai/shared-workflows/.github/workflows/custom-job.yaml@main - if: fromJSON(needs.changed-files.outputs.changed_file_groups).build_docs - strategy: - fail-fast: false - matrix: ${{ fromJSON(needs.docs-build-matrix.outputs.matrix) }} - with: - build_type: pull-request - node_type: "gpu-${{ matrix.GPU }}-${{ matrix.DRIVER }}-1" - arch: "${{ matrix.arch }}" - container_image: "rapidsai/ci-conda:26.12-cuda${{ matrix.CUDA_VER }}-${{ matrix.LINUX_VER }}-py${{ matrix.PY_VER }}" - script: "ci/build_docs.sh" - wheel-build-libcuvs: - needs: [build-details, checks] - permissions: - actions: read - contents: read - id-token: write - packages: read - pull-requests: read - secrets: inherit # zizmor: ignore[secrets-inherit] - uses: rapidsai/shared-workflows/.github/workflows/wheels-build.yaml@main - with: - build_type: pull-request - build-datetime: ${{ needs.build-details.outputs.build-datetime }} - node_type: cpu16 - script: ci/build_wheel_libcuvs.sh - # build for every combination of arch and CUDA version, but only for the latest Python - matrix_filter: group_by([.ARCH, (.CUDA_VER|split(".")|map(tonumber)|.[0])]) | map(max_by(.PY_VER|split(".")|map(tonumber))) - package-name: libcuvs - package-type: cpp - wheel-build-cuvs: - needs: [build-details, wheel-build-libcuvs] - permissions: - actions: read - contents: read - id-token: write - packages: read - pull-requests: read - secrets: inherit # zizmor: ignore[secrets-inherit] - uses: rapidsai/shared-workflows/.github/workflows/wheels-build.yaml@main - with: - build_type: pull-request - build-datetime: ${{ needs.build-details.outputs.build-datetime }} - node_type: cpu8 - script: ci/build_wheel_cuvs.sh - package-name: cuvs - package-type: python - # Build a wheel for each CUDA x ARCH x minimum supported Python version - matrix_filter: group_by({CUDA_VER, ARCH}) | map(min_by(.PY_VER | split(".") | map(tonumber))) - wheel-tests-cuvs: - needs: [wheel-build-cuvs, changed-files] - permissions: - actions: read - contents: read - id-token: write - packages: read - pull-requests: read - secrets: inherit # zizmor: ignore[secrets-inherit] - uses: rapidsai/shared-workflows/.github/workflows/wheels-test.yaml@main - if: fromJSON(needs.changed-files.outputs.changed_file_groups).test_python_wheels - with: - build_type: pull-request - script: ci/test_wheel_cuvs.sh - devcontainer: - permissions: - actions: read - contents: read - id-token: write - packages: read - pull-requests: read - secrets: inherit # zizmor: ignore[secrets-inherit] - needs: telemetry-setup - uses: rapidsai/shared-workflows/.github/workflows/build-in-devcontainer.yaml@main - with: - arch: '["amd64", "arm64"]' - cuda: '["13.3"]' - node_type: "cpu8" - env: | - SCCACHE_DIST_MAX_RETRIES=inf - SCCACHE_SERVER_LOG=sccache=debug - SCCACHE_DIST_FALLBACK_TO_LOCAL_COMPILE=false - build_command: | - 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 - needs: pr-builder - permissions: - actions: write - if: ${{ vars.TELEMETRY_ENABLED == 'true' && !cancelled() }} - continue-on-error: true - steps: - - name: Telemetry summarize - uses: rapidsai/shared-actions/telemetry-dispatch-summarize@main - env: - GH_TOKEN: ${{ github.token }} diff --git a/ci/test_cpp.sh b/ci/test_cpp.sh index 0e97a0ddef..ac4ccf6214 100755 --- a/ci/test_cpp.sh +++ b/ci/test_cpp.sh @@ -4,11 +4,6 @@ set -euo pipefail -# Usage: ci/test_cpp.sh [SHARD NUM_SHARDS] -# With SHARD and NUM_SHARDS (1 <= SHARD <= NUM_SHARDS), only every NUM_SHARDS-th libcuvs test is -# run, starting from test number SHARD, so the tests can be split across several CI jobs. -SHARD=${1:-1} -NUM_SHARDS=${2:-1} . /opt/conda/etc/profile.d/conda.sh @@ -52,23 +47,34 @@ nvidia-smi # RAPIDS_DATASET_ROOT_DIR is used by test scripts RAPIDS_DATASET_ROOT_DIR=${RAPIDS_TESTS_DIR}/dataset export RAPIDS_DATASET_ROOT_DIR -./ci/get_test_data.sh --NEIGHBORS_ANN_VAMANA_TEST EXITCODE=0 trap "EXITCODE=1" ERR set +e -# Run Python build utilities tests (once, in the first shard) -if [[ "${SHARD}" == "1" ]]; then - rapids-logger "Run libcuvs Python build utilities tests" - pytest cpp/tests/python -fi - -# Run libcuvs gtests from libcuvs-tests package -rapids-logger "Run libcuvs tests (shard ${SHARD} of ${NUM_SHARDS})" +# Run the k-means and IVF test binaries once in full. +rapids-logger "Run k-means and IVF test binaries" pushd "$CONDA_PREFIX"/bin/gtests/libcuvs -timeout -v --signal=SIGINT --kill-after=60s 100m ctest -j8 --output-on-failure -I "${SHARD},,${NUM_SHARDS}" +timeout -v --signal=SIGINT --kill-after=60s 120m ctest -j4 --output-on-failure -R "^(CLUSTER_TEST|CLUSTER_KMEANS_MG_TEST|CLUSTER_KMEANS_MNMG_TEST|NEIGHBORS_ANN_IVF_FLAT_TEST|NEIGHBORS_ANN_IVF_PQ_TEST|NEIGHBORS_ANN_IVF_SQ_TEST|NEIGHBORS_ANN_IVF_RABITQ_TEST)$" popd +# Repeat the reproducibility tests. +repeat_test() { + local binary=$1 filter=$2 repeat=$3 + local log="$CONDA_PREFIX/bin/gtests/libcuvs/${binary}_repeat.log" + rapids-logger "Repeat ${binary} --gtest_filter=${filter} ${repeat} times" + "$CONDA_PREFIX"/bin/gtests/libcuvs/"${binary}" --gtest_filter="${filter}" --gtest_repeat="${repeat}" > "${log}" 2>&1 + echo "Exit status: $?" + tail -n 5 "${log}" + grep -E "Failure|differ|Expected" "${log}" | head -n 50 || true + echo "Failures:" + grep -E "^\[ FAILED \] .* \([0-9]+ ms\)$" "${log}" | sed -E 's/, where GetParam.*//; s/ \([0-9]+ ms\)//' | sort | uniq -c || true + echo "Passes:" + grep -E "^\[ OK \]" "${log}" | sed -E 's/ \([0-9]+ ms\)//' | sort | uniq -c || true +} +repeat_test CLUSTER_TEST "KmeansFitBatchedTests/*:KmeansBalancedReproducible.*" 50 +repeat_test NEIGHBORS_ANN_IVF_FLAT_TEST "IVFFlatBuild.Reproducible" 10 +repeat_test NEIGHBORS_ANN_IVF_PQ_TEST "IvfPqBuild.Reproducible*" 10 + rapids-logger "Test script exiting with value: $EXITCODE" exit ${EXITCODE} diff --git a/cpp/src/cluster/detail/kmeans.cuh b/cpp/src/cluster/detail/kmeans.cuh index 332fed16b5..4478cec2a6 100644 --- a/cpp/src/cluster/detail/kmeans.cuh +++ b/cpp/src/cluster/detail/kmeans.cuh @@ -1082,7 +1082,8 @@ void kmeans_fit( compute_centroid_shift(handle, raft::make_const_mdspan(centroids_const), raft::make_const_mdspan(new_centroids_view), - sqrdNormError.view()); + sqrdNormError.view(), + ws); std::swap(cur_centroids_ptr, new_centroids_ptr); diff --git a/cpp/src/cluster/detail/kmeans_balanced.cuh b/cpp/src/cluster/detail/kmeans_balanced.cuh index bb96c9115b..c80ea47dbb 100644 --- a/cpp/src/cluster/detail/kmeans_balanced.cuh +++ b/cpp/src/cluster/detail/kmeans_balanced.cuh @@ -6,6 +6,7 @@ #pragma once #include "kmeans_common.cuh" +#include "kmeans_fixed_point_sums.cuh" #include #include "../../core/nvtx.hpp" @@ -43,6 +44,7 @@ #include #include +#include #include #include #include @@ -211,6 +213,17 @@ auto calc_minibatch_size(const raft::resources& handle, return std::make_tuple(minibatch_size, mem_per_row); } +/** The dataset as a random-access iterator over MathT values (mapping only if T != MathT). */ +template +auto math_iterator(const T* dataset, MappingOpT mapping_op) +{ + if constexpr (std::is_same_v) { + return dataset; + } else { + return thrust::transform_iterator(dataset, mapping_op); + } +} + /** * @brief Given the data and labels, calculate cluster centers and sizes in one sweep. * @@ -236,6 +249,9 @@ auto calc_minibatch_size(const raft::resources& handle, * the weighted average principle. * @param[in] mapping_op Mapping operation from T to MathT * @param[inout] mr (optional) Memory resource to use for temporary allocations on the device + * @param[inout] workspace Scratch buffer, resized as needed (reuse it across calls to avoid + * allocating every time) + * @param[inout] accumulator Fixed-point accumulator whose scales were set for `dataset` */ template & workspace, + fixed_point_sums::accumulator& accumulator) { auto stream = raft::resource::get_cuda_stream(handle); @@ -265,8 +283,6 @@ void calc_centers_and_sizes(const raft::resources& handle, handle, raft::make_const_mdspan(centersView), clusterSizesView, centersView, raft::mul_op{}); } - rmm::device_uvector workspace(0, stream, mr); - // If we reset the counters, we can compute directly the new sizes in cluster_sizes. // If we don't reset, we compute in a temporary buffer and add in a separate step. rmm::device_uvector temp_cluster_sizes(0, stream, mr); @@ -276,32 +292,10 @@ void calc_centers_and_sizes(const raft::resources& handle, temp_sizes = temp_cluster_sizes.data(); } - // Apply mapping only when the data and math types are different. - if constexpr (std::is_same_v) { - raft::linalg::reduce_rows_by_key(dataset, - dim, - labels, - nullptr, - n_rows, - dim, - n_clusters, - centers, - stream.get(), - reset_counters); - } else { - // todo(lsugy): use iterator from KV output of fusedL2NN - thrust::transform_iterator mapping_itr(dataset, mapping_op); - raft::linalg::reduce_rows_by_key(mapping_itr, - dim, - labels, - nullptr, - n_rows, - dim, - n_clusters, - centers, - stream.get(), - reset_counters); - } + // Sum the rows of each cluster exactly (fixed-point), so that identical builds give + // bitwise-identical centers. Apply mapping only when the data and math types are different. + accumulator.sums_by_label( + handle, math_iterator(dataset, mapping_op), labels, n_rows, centers, reset_counters); // Compute weight of each cluster cuvs::cluster::kmeans::detail::countLabels( @@ -323,6 +317,43 @@ void calc_centers_and_sizes(const raft::resources& handle, raft::div_checkzero_op{}); } +/** Same as above, with a temporary workspace and accumulator. */ +template +void calc_centers_and_sizes(const raft::resources& handle, + MathT* centers, + CounterT* cluster_sizes, + IdxT n_clusters, + IdxT dim, + const T* dataset, + IdxT n_rows, + const LabelT* labels, + bool reset_counters, + MappingOpT mapping_op, + rmm::device_async_resource_ref mr) +{ + rmm::device_uvector workspace(0, raft::resource::get_cuda_stream(handle), mr); + fixed_point_sums::accumulator accumulator(handle, dim, n_clusters, mr); + accumulator.set_scale(handle, math_iterator(dataset, mapping_op), n_rows); + calc_centers_and_sizes(handle, + centers, + cluster_sizes, + n_clusters, + dim, + dataset, + n_rows, + labels, + reset_counters, + mapping_op, + mr, + workspace, + accumulator); +} + /** Computes the L2 norm of the dataset, converting to MathT if necessary */ template void compute_norm(const raft::resources& handle, @@ -640,7 +671,7 @@ __launch_bounds__((raft::WarpSize * BlockDimY)) RAFT_KERNEL IdxT average, MathT centroid_offset, IdxT seed, - IdxT* search_count, + IdxT stride, IdxT* update_count, MappingOpT mapping_op) { @@ -652,12 +683,17 @@ __launch_bounds__((raft::WarpSize * BlockDimY)) RAFT_KERNEL IdxT i = n_rows; IdxT j = raft::laneId(); if (j == 0) { + // The candidate sequence of each receiver is derived from its index (not from a shared atomic + // counter), so the donor points do not depend on the order in which the warps run. + // `stride >= n_clusters`, so different receivers use different counters (and thus, before + // wrapping around, different candidates); `stride` is coprime with `n_rows`, so each receiver + // visits every row within `n_rows` attempts and finds a donor whenever one exists. IdxT attempt = 0; do { - auto old = atomicAdd(search_count, IdxT{1}); - auto candidate = - static_cast((static_cast(seed) * static_cast(old + 1)) % - static_cast(n_rows)); + auto counter = static_cast(attempt) * static_cast(stride) + + static_cast(receiver_cluster); + auto candidate = static_cast((static_cast(seed) * (counter + 1)) % + static_cast(n_rows)); if (static_cast(cluster_sizes[labels[candidate]]) >= average) { i = candidate; } ++attempt; } while (i >= n_rows && attempt < n_rows); @@ -791,6 +827,7 @@ auto adjust_centers(const raft::resources& handle, MathT balance_upper_tolerance, MathT centroid_offset, cuvs::cluster::kmeans::balanced_donor_selection donor_selection, + uint32_t balancing_step, MappingOpT mapping_op, rmm::device_async_resource_ref device_memory) -> bool { @@ -802,7 +839,6 @@ auto adjust_centers(const raft::resources& handle, 601, 659, 733, 809, 863, 941, 1013, 1069, 1151, 1223, 1291, 1373, 1451, 1511, 1583, 1657, 1733, 1811, 1889, 1987, 2053, 2129, 2213, 2287, 2357, 2423, 2531, 2617, 2687, 2741}; - static IdxT i_primes = 0; auto average = static_cast(n_rows) / static_cast(n_clusters); auto lower_threshold = average * balance_lower_tolerance; @@ -837,10 +873,12 @@ auto adjust_centers(const raft::resources& handle, auto n_pairs = static_cast(host_receiver_clusters.size()); if (n_pairs == 0) { return false; } + // Rotate the pseudo-random offset with the balancing step (rather than with a process-wide + // counter), so that identical builds pick the same donor points. IdxT ofst; + size_t i_primes = balancing_step; do { - i_primes = (i_primes + 1) % kPrimes.size(); - ofst = kPrimes[i_primes]; + ofst = kPrimes[i_primes++ % kPrimes.size()]; } while (n_rows % ofst == 0); rmm::device_uvector receiver_clusters(n_pairs, stream, device_memory); @@ -851,8 +889,11 @@ auto adjust_centers(const raft::resources& handle, update_count.set_value_to_zero_async(stream); if (donor_selection == cuvs::cluster::kmeans::balanced_donor_selection::Random) { - rmm::device_scalar search_count(stream, device_memory); - search_count.set_value_to_zero_async(stream); + // The smallest stride >= n_clusters that is coprime with n_rows (see the kernel). + auto stride = std::max(n_clusters, 1); + while (std::gcd(static_cast(stride), static_cast(n_rows)) != 1) { + ++stride; + } const dim3 grid_dim(raft::ceildiv(n_clusters, static_cast(kBlockDimY)), 1, 1); adjust_centers_random_donor_kernel <<>>(centers, @@ -866,7 +907,7 @@ auto adjust_centers(const raft::resources& handle, static_cast(n_rows / n_clusters), centroid_offset, ofst, - search_count.data(), + stride, update_count.data(), mapping_op); return update_count.value(stream) > 0; // NB: rmm scalar performs the sync @@ -963,6 +1004,11 @@ void balancing_em_iters(const raft::resources& handle, RAFT_EXPECTS(params.centroid_offset > 0.0f && params.centroid_offset <= 1.0f, "Balanced k-means centroid offset must be in the range (0, 1]"); + // Reused by every M step, so that they are not reallocated on each iteration. The dataset does + // not change, so the fixed-point scales are computed once. + rmm::device_uvector workspace(0, raft::resource::get_cuda_stream(handle), device_memory); + fixed_point_sums::accumulator accumulator(handle, dim, n_clusters, device_memory); + accumulator.set_scale(handle, math_iterator(dataset, mapping_op), n_rows); uint32_t balancing_counter = balancing_pullback; for (uint32_t iter = 0; iter < n_iters; iter++) { // Balancing step - move the centers around to equalize cluster sizes @@ -979,6 +1025,7 @@ void balancing_em_iters(const raft::resources& handle, balance_upper_tolerance, static_cast(params.centroid_offset), params.donor_selection, + iter, mapping_op, device_memory)) { if (balancing_counter++ >= balancing_pullback) { @@ -1025,7 +1072,9 @@ void balancing_em_iters(const raft::resources& handle, cluster_labels, true, mapping_op, - device_memory); + device_memory, + workspace, + accumulator); } } diff --git a/cpp/src/cluster/detail/kmeans_centroid_sums.cuh b/cpp/src/cluster/detail/kmeans_centroid_sums.cuh new file mode 100644 index 0000000000..6248e00ae1 --- /dev/null +++ b/cpp/src/cluster/detail/kmeans_centroid_sums.cuh @@ -0,0 +1,416 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +#pragma once + +#include +#include +#include +#include +#include +#include + +#include + +#include +#include + +#include +#include +#include + +/** + * Deterministic per-cluster (weighted) row sums for the k-means centroid updates. + * + * Accumulating the rows of each cluster with atomics (as raft::linalg::reduce_rows_by_key does) + * makes the floating-point summation order, and therefore the centroids, vary from run to run. + * Through the convergence check this can change the number of Lloyd iterations, so two identical + * fits can return noticeably different centroids. Here the summation order is a function of the + * input only: + * + * 1. Row indices are stably sorted by label (cub::DeviceRadixSort), so each cluster becomes a + * contiguous segment. + * 2. Each segment is cut into pieces of at most `kRowsPerPiece` rows. One thread block reduces a + * piece with a fixed assignment of rows to threads and a fixed-order combine. + * 3. One thread block per cluster adds up that cluster's piece sums in the same fixed way, and + * adds the result to the output. + * + * The order depends only on the labels and the sizes, not on scheduling or on the launch + * configuration. + * + * Why not a CCCL reduction: cub::DeviceSegmentedReduce (run_to_run) and + * cub::DeviceReduce::ReduceByKey (run_to_run) reduce one scalar per segment or run. A + * [n_rows x n_cols] matrix would need n_clusters * n_cols segments that read the gathered rows + * column by column (strided, uncoalesced), and DeviceSegmentedReduce gives each segment to one + * thread block, which serializes large clusters when n_clusters is small. The per-cluster weight + * sum alone could use them, but it comes for free here in the same pass over the rows. + */ +namespace cuvs::cluster::kmeans::detail::sums_by_label { + +constexpr int kBlockThreads = 256; +constexpr int kRowsPerPiece = 256; + +/** Number of threads along the columns: smallest power of two >= n_cols (at most kBlockThreads). */ +inline int column_threads(int64_t n_cols) +{ + int tc = 1; + while (tc < n_cols && tc < kBlockThreads) { + tc <<= 1; + } + return tc; +} + +/** + * Sums `values(i, j)` over i in [begin, end) for each column j in [0, n_cols), using a fixed + * assignment of items to threads and a fixed-order combine, so the result does not depend on + * scheduling. The result for column j is passed to `write(j, sum)` by a single thread. + */ +template +__device__ void fixed_order_column_sums( + IndexT begin, IndexT end, IndexT n_cols, int col_threads, ValueOp values, WriteOp write) +{ + __shared__ DataT smem[kBlockThreads]; + const int lanes = kBlockThreads / col_threads; + const int tc = threadIdx.x % col_threads; + const int lane = threadIdx.x / col_threads; + for (IndexT col0 = 0; col0 < n_cols; col0 += col_threads) { + const IndexT j = col0 + tc; + DataT acc = DataT{0}; + if (j < n_cols) { +#pragma unroll 4 + for (IndexT i = begin + lane; i < end; i += lanes) { + acc += values(i, j); + } + } + smem[threadIdx.x] = acc; + __syncthreads(); + if (lane == 0 && j < n_cols) { + DataT sum = smem[tc]; + for (int l = 1; l < lanes; ++l) { + sum += smem[l * col_threads + tc]; + } + write(j, sum); + } + __syncthreads(); + } +} + +template +RAFT_KERNEL init_sort_kernel(LabelsIterator labels, IndexT n_rows, uint32_t* keys, IndexT* rows) +{ + for (IndexT i = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; i < n_rows; + i += static_cast(blockDim.x) * gridDim.x) { + keys[i] = static_cast(labels[i]); + rows[i] = i; + } +} + +/** Finds each cluster's segment of the sorted labels and the number of pieces it is cut into. */ +template +RAFT_KERNEL segment_kernel(const uint32_t* sorted_keys, + IndexT n_rows, + IndexT n_clusters, + IndexT* cluster_offsets, + IndexT* piece_counts) +{ + auto lower_bound = [=](IndexT c) { + IndexT lo = 0, hi = n_rows; + while (lo < hi) { + IndexT mid = lo + (hi - lo) / 2; + if (static_cast(sorted_keys[mid]) < c) { + lo = mid + 1; + } else { + hi = mid; + } + } + return lo; + }; + for (IndexT c = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; c <= n_clusters; + c += static_cast(blockDim.x) * gridDim.x) { + IndexT begin = lower_bound(c); + cluster_offsets[c] = begin; + if (c < n_clusters) { + IndexT end = lower_bound(c + 1); + piece_counts[c] = raft::ceildiv(end - begin, kRowsPerPiece); + } else { + piece_counts[c] = 0; + } + } +} + +template +RAFT_KERNEL piece_sums_kernel(DataIterator X, + IndexT n_cols, + const DataT* weights, + bool sum_weights, + const IndexT* sorted_rows, + const IndexT* cluster_offsets, + const IndexT* piece_offsets, + IndexT n_clusters, + int col_threads, + DataT* piece_sums) +{ + const IndexT n_pieces = piece_offsets[n_clusters]; + for (IndexT piece = blockIdx.x; piece < n_pieces; piece += gridDim.x) { + // The cluster of this piece is the last c with piece_offsets[c] <= piece. + IndexT lo = 0, hi = n_clusters; + while (hi - lo > 1) { + IndexT mid = lo + (hi - lo) / 2; + if (piece_offsets[mid] <= piece) { + lo = mid; + } else { + hi = mid; + } + } + const IndexT c = lo; + const IndexT begin = cluster_offsets[c] + (piece - piece_offsets[c]) * kRowsPerPiece; + const IndexT end = std::min(begin + kRowsPerPiece, cluster_offsets[c + 1]); + DataT* out = piece_sums + static_cast(piece) * (n_cols + 1); + fixed_order_column_sums( + begin, + end, + n_cols, + col_threads, + [=](IndexT i, IndexT j) { + const IndexT row = sorted_rows[i]; + const auto x = static_cast(X[static_cast(row) * n_cols + j]); + return weights != nullptr ? x * weights[row] : x; + }, + [=](IndexT j, DataT sum) { out[j] = sum; }); + // The sum of weights is stored after the column sums. + if (sum_weights) { + fixed_order_column_sums( + begin, + end, + IndexT{1}, + 1, + [=](IndexT i, IndexT) { return weights != nullptr ? weights[sorted_rows[i]] : DataT{1}; }, + [=](IndexT, DataT sum) { out[n_cols] = sum; }); + } + } +} + +template +RAFT_KERNEL cluster_sums_kernel(const DataT* piece_sums, + IndexT n_cols, + const IndexT* piece_offsets, + IndexT n_clusters, + int col_threads, + bool reset_sums, + DataT* centroid_sums, + DataT* weight_per_cluster) +{ + for (IndexT c = blockIdx.x; c < n_clusters; c += gridDim.x) { + const IndexT p_begin = piece_offsets[c]; + const IndexT p_end = piece_offsets[c + 1]; + const IndexT stride = n_cols + 1; + DataT* out = centroid_sums + static_cast(c) * n_cols; + fixed_order_column_sums( + p_begin, + p_end, + n_cols, + col_threads, + [=](IndexT p, IndexT j) { return piece_sums[static_cast(p) * stride + j]; }, + [=](IndexT j, DataT sum) { out[j] = reset_sums ? sum : out[j] + sum; }); + if (weight_per_cluster != nullptr) { + fixed_order_column_sums( + p_begin, + p_end, + IndexT{1}, + 1, + [=](IndexT p, IndexT) { return piece_sums[static_cast(p) * stride + n_cols]; }, + [=](IndexT, DataT sum) { + weight_per_cluster[c] = reset_sums ? sum : weight_per_cluster[c] + sum; + }); + } + } +} + +inline size_t align_bytes(size_t bytes) { return raft::alignTo(bytes, 256); } + +/** Rows per pass of `weighted_sums_by_label`; bounds the workspace and fits 32-bit indices. */ +constexpr int64_t kMaxRowsPerPass = int64_t{1} << 22; + +/** One pass of `weighted_sums_by_label` over at most kMaxRowsPerPass rows. */ +template +void weighted_sums_by_label_pass(raft::resources const& handle, + DataIterator X, + const DataT* weights, + LabelsIterator labels, + uint32_t n_rows, + uint32_t n_cols, + uint32_t n_clusters, + DataT* centroid_sums, + DataT* weight_per_cluster, + rmm::device_uvector& workspace, + bool reset_sums) +{ + using IdxT = uint32_t; + cudaStream_t stream = raft::resource::get_cuda_stream(handle).get(); + + int end_bit = 1; + while (end_bit < 32 && (uint64_t{1} << end_bit) < static_cast(n_clusters)) { + ++end_bit; + } + const IdxT max_pieces = raft::ceildiv(n_rows, kRowsPerPiece) + n_clusters; + + // Workspace layout + size_t sort_bytes = 0; + RAFT_CUDA_TRY(cub::DeviceRadixSort::SortPairs(nullptr, + sort_bytes, + static_cast(nullptr), + static_cast(nullptr), + static_cast(nullptr), + static_cast(nullptr), + n_rows, + 0, + end_bit, + stream)); + size_t scan_bytes = 0; + RAFT_CUDA_TRY(cub::DeviceScan::ExclusiveSum(nullptr, + scan_bytes, + static_cast(nullptr), + static_cast(nullptr), + n_clusters + 1, + stream)); + const size_t keys_bytes = align_bytes(sizeof(uint32_t) * n_rows); + const size_t rows_bytes = align_bytes(sizeof(IdxT) * n_rows); + const size_t offsets_bytes = align_bytes(sizeof(IdxT) * (n_clusters + 1)); + const size_t pieces_bytes = + align_bytes(sizeof(DataT) * static_cast(max_pieces) * (n_cols + 1)); + const size_t temp_bytes = align_bytes(std::max(sort_bytes, scan_bytes)); + const size_t total_bytes = + 2 * keys_bytes + 2 * rows_bytes + 3 * offsets_bytes + pieces_bytes + temp_bytes; + if (workspace.size() < total_bytes) { workspace.resize(total_bytes, stream); } + + char* ptr = workspace.data(); + auto take = [&ptr](size_t bytes) { + char* p = ptr; + ptr += bytes; + return p; + }; + auto* keys_in = reinterpret_cast(take(keys_bytes)); + auto* keys_out = reinterpret_cast(take(keys_bytes)); + auto* rows_in = reinterpret_cast(take(rows_bytes)); + auto* rows_out = reinterpret_cast(take(rows_bytes)); + auto* cluster_offsets = reinterpret_cast(take(offsets_bytes)); + auto* piece_counts = reinterpret_cast(take(offsets_bytes)); + auto* piece_offsets = reinterpret_cast(take(offsets_bytes)); + auto* piece_sums = reinterpret_cast(take(pieces_bytes)); + void* temp = take(temp_bytes); + + if (n_rows > 0) { + const int grid = static_cast( + std::min(raft::ceildiv(n_rows, kBlockThreads), int64_t{65536})); + init_sort_kernel<<>>(labels, n_rows, keys_in, rows_in); + RAFT_CUDA_TRY(cudaPeekAtLastError()); + RAFT_CUDA_TRY(cub::DeviceRadixSort::SortPairs( + temp, sort_bytes, keys_in, keys_out, rows_in, rows_out, n_rows, 0, end_bit, stream)); + } + { + const int grid = static_cast( + std::min(raft::ceildiv(n_clusters + 1, kBlockThreads), int64_t{65536})); + segment_kernel<<>>( + keys_out, n_rows, n_clusters, cluster_offsets, piece_counts); + RAFT_CUDA_TRY(cudaPeekAtLastError()); + } + RAFT_CUDA_TRY(cub::DeviceScan::ExclusiveSum( + temp, scan_bytes, piece_counts, piece_offsets, n_clusters + 1, stream)); + + const int col_threads = column_threads(n_cols); + { + const int grid = static_cast(std::min(max_pieces, int64_t{65536})); + piece_sums_kernel<<>>(X, + n_cols, + weights, + weight_per_cluster != nullptr, + rows_out, + cluster_offsets, + piece_offsets, + n_clusters, + col_threads, + piece_sums); + RAFT_CUDA_TRY(cudaPeekAtLastError()); + } + { + const int grid = static_cast(std::min(n_clusters, int64_t{65536})); + cluster_sums_kernel<<>>(piece_sums, + n_cols, + piece_offsets, + n_clusters, + col_threads, + reset_sums, + centroid_sums, + weight_per_cluster); + RAFT_CUDA_TRY(cudaPeekAtLastError()); + } +} + +/** + * @brief Deterministic (weighted) sum of the rows of X per label, and of the weights per label. + * + * centroid_sums[c, :] (+)= sum_{i : labels[i] == c} weights[i] * X[i, :] + * weight_per_cluster[c] (+)= sum_{i : labels[i] == c} weights[i] + * + * The summation order depends only on the labels and the sizes, not on scheduling, + * so the result is bitwise reproducible. Rows are processed in passes of at most kMaxRowsPerPass + * rows, which bounds the workspace (16 bytes per row of a pass, plus the piece sums). + * + * @tparam DataT Accumulation and output type + * @tparam IndexT Index type + * @tparam DataIterator Random-access iterator over the row-major input; its values are converted + * to DataT (e.g. a transform iterator that maps int8 data to float) + * @tparam LabelsIterator Random-access iterator over the labels + * + * @param[in] handle RAFT resources handle + * @param[in] X Row-major input [n_rows x n_cols] + * @param[in] weights Per-row weights [n_rows], or nullptr for unit weights + * @param[in] labels Label of each row, in [0, n_clusters) + * @param[in] n_rows Number of rows of X + * @param[in] n_cols Number of columns of X + * @param[in] n_clusters Number of labels + * @param[inout] centroid_sums Row-major output [n_clusters x n_cols] + * @param[inout] weight_per_cluster Output [n_clusters], or nullptr to skip the weight sums + * @param[inout] workspace Scratch buffer, resized as needed + * @param[in] reset_sums Overwrite the outputs if true, add to them otherwise + */ +template +void weighted_sums_by_label(raft::resources const& handle, + DataIterator X, + const DataT* weights, + LabelsIterator labels, + IndexT n_rows, + IndexT n_cols, + IndexT n_clusters, + DataT* centroid_sums, + DataT* weight_per_cluster, + rmm::device_uvector& workspace, + bool reset_sums) +{ + if (n_clusters <= 0) { return; } + RAFT_EXPECTS(static_cast(n_clusters) < (uint64_t{1} << 31) && + static_cast(n_cols) < (uint64_t{1} << 31), + "weighted_sums_by_label: too many clusters or columns"); + // Passes are accumulated in a fixed order, so the result stays deterministic. There is always at + // least one pass, so that the outputs are reset even when there are no rows. + const auto total_rows = static_cast(n_rows); + int64_t start = 0; + do { + const int64_t pass_rows = std::min(kMaxRowsPerPass, total_rows - start); + weighted_sums_by_label_pass(handle, + X + start * static_cast(n_cols), + weights != nullptr ? weights + start : nullptr, + labels + start, + static_cast(pass_rows), + static_cast(n_cols), + static_cast(n_clusters), + centroid_sums, + weight_per_cluster, + workspace, + reset_sums && start == 0); + start += pass_rows; + } while (start < total_rows); +} + +} // namespace cuvs::cluster::kmeans::detail::sums_by_label diff --git a/cpp/src/cluster/detail/kmeans_common.cuh b/cpp/src/cluster/detail/kmeans_common.cuh index 78edf8759e..efedd05b1b 100644 --- a/cpp/src/cluster/detail/kmeans_common.cuh +++ b/cpp/src/cluster/detail/kmeans_common.cuh @@ -6,6 +6,7 @@ #include "../../distance/distance.cuh" #include "../../distance/top_1_nn.cuh" +#include "kmeans_centroid_sums.cuh" #include #include #include @@ -20,6 +21,7 @@ #include #include #include +#include #include #include #include @@ -42,7 +44,13 @@ #include #include #include +#include +#include #include +#include +#include +#include +#include #include #include @@ -121,6 +129,85 @@ struct WeightedDistanceOp { } }; +/** + * @brief A `cuda::mr::resource` that hands out temporary storage from one caller-owned buffer. + * + * The buffer only grows, and deallocation is a no-op, so repeated calls reuse the same memory + * instead of allocating (and possibly synchronizing) every time. Only one allocation may be live + * at a time. + */ +class workspace_buffer_resource { + public: + explicit workspace_buffer_resource(rmm::device_uvector& buffer) : buffer_{&buffer} {} + + void* allocate(cuda::stream_ref stream, size_t bytes, size_t alignment) + { + const size_t needed = bytes + alignment; + if (buffer_->size() < needed) { buffer_->resize(needed, stream); } + const auto ptr = reinterpret_cast(buffer_->data()); + return reinterpret_cast(raft::alignTo(ptr, alignment)); + } + void deallocate(cuda::stream_ref, void*, size_t, size_t) noexcept {} + void* allocate_sync(size_t bytes, size_t alignment) + { + return allocate(buffer_->stream(), bytes, alignment); + } + void deallocate_sync(void*, size_t, size_t) noexcept {} + bool operator==(const workspace_buffer_resource& other) const noexcept + { + return buffer_ == other.buffer_; + } + bool operator!=(const workspace_buffer_resource& other) const noexcept + { + return !(*this == other); + } + + private: + rmm::device_uvector* buffer_; +}; + +/** + * @brief Device-wide reduction that is bitwise reproducible from run to run. + * + * Requests `cuda::execution::determinism::run_to_run` explicitly (also cub's current default): + * cub combines the partial results in a fixed order on a given GPU, with no atomics. This is + * stronger than raft's mapThenSumReduce, which combines block results with atomics. + * + * `gpu_to_gpu` (an RFA-based sum that is identical across architectures) is not requested: the + * inputs of these reductions (distances from GEMM / cuTile TF32 kernels and the centroids derived + * from them) already differ between architectures, so a fit cannot be reproducible across GPUs. + * + * Temporary storage comes from `workspace`, which is resized as needed. + */ +template +void run_to_run_reduce(raft::resources const& handle, + InputIterator input, + IndexT size, + OutputT* output, + ReductionOpT reduction_op, + OutputT init, + rmm::device_uvector& workspace) +{ + auto env = + cuda::std::execution::env{cuda::stream_ref{raft::resource::get_cuda_stream(handle)}, + cuda::execution::require(cuda::execution::determinism::run_to_run), + cuda::std::execution::prop{cuda::mr::get_memory_resource, + workspace_buffer_resource{workspace}}}; + RAFT_CUDA_TRY(cub::DeviceReduce::Reduce(input, output, size, reduction_op, init, env)); +} + +template +struct SquaredDiffOp { + const DataT* a; + const DataT* b; + + __device__ DataT operator()(IndexT offset) const + { + DataT diff = a[offset] - b[offset]; + return diff * diff; + } +}; + template void copyClusterLabels(raft::resources const& handle, const MinClusterAndDistanceResult& result, @@ -194,8 +281,14 @@ void weightSum( DataT wt_sum_h = DataT{0}; if constexpr (raft::is_device_mdspan_v) { - raft::linalg::mapThenSumReduce( - d_wt_sum.data_handle(), n_samples, raft::identity_op{}, stream.get(), weight.data_handle()); + rmm::device_uvector workspace(0, stream); + run_to_run_reduce(handle, + weight.data_handle(), + n_samples, + d_wt_sum.data_handle(), + cuda::std::plus<>{}, + DataT{0}, + workspace); if (check_positive) { raft::copy(&wt_sum_h, d_wt_sum.data_handle(), 1, stream); raft::resource::sync_stream(handle); @@ -270,29 +363,9 @@ void computeClusterCost(raft::resources const& handle, MainOpT main_op, ReductionOpT reduction_op) { - cudaStream_t stream = raft::resource::get_cuda_stream(handle).get(); cuda::transform_iterator itr(input, main_op); - - size_t temp_storage_bytes = 0; - RAFT_CUDA_TRY(cub::DeviceReduce::Reduce(nullptr, - temp_storage_bytes, - itr, - clusterCost.data_handle(), - size, - reduction_op, - OutputT(), - stream)); - - workspace.resize(temp_storage_bytes, stream); - - RAFT_CUDA_TRY(cub::DeviceReduce::Reduce(workspace.data(), - temp_storage_bytes, - itr, - clusterCost.data_handle(), - size, - reduction_op, - OutputT(), - stream)); + run_to_run_reduce( + handle, itr, size, clusterCost.data_handle(), reduction_op, OutputT(), workspace); } template & workspace, bool reset_sums = true) { - cudaStream_t stream = raft::resource::get_cuda_stream(handle).get(); - auto n_samples = X.extent(0); - - workspace.resize(n_samples, stream); - - raft::linalg::reduce_rows_by_key(X.data_handle(), - X.extent(1), - cluster_labels, - sample_weights.data_handle(), - workspace.data(), - X.extent(0), - X.extent(1), - n_clusters, - centroid_sums.data_handle(), - stream, - reset_sums); - - raft::linalg::reduce_cols_by_key(sample_weights.data_handle(), - cluster_labels, - weight_per_cluster.data_handle(), - static_cast(1), - static_cast(n_samples), - n_clusters, - stream, - reset_sums); + sums_by_label::weighted_sums_by_label(handle, + X.data_handle(), + sample_weights.data_handle(), + cluster_labels, + X.extent(0), + X.extent(1), + n_clusters, + centroid_sums.data_handle(), + weight_per_cluster.data_handle(), + workspace, + reset_sums); } /** * @brief Finalize centroids by dividing accumulated sums by counts. @@ -707,21 +769,26 @@ void finalize_centroids(raft::resources const& handle, * @brief Compute the squared norm difference between two centroid sets. * * Writes sum((old_centroids - new_centroids)^2) into @p sqrd_norm_out. - * Used for convergence checking. Fully asynchronous — no stream sync. + * Used for convergence checking. Fully asynchronous — no stream sync. Temporary storage comes + * from @p workspace, which is resized as needed. */ template void compute_centroid_shift(raft::resources const& handle, raft::device_matrix_view old_centroids, raft::device_matrix_view new_centroids, - raft::device_scalar_view sqrd_norm_out) + raft::device_scalar_view sqrd_norm_out, + rmm::device_uvector& workspace) { - cudaStream_t stream = raft::resource::get_cuda_stream(handle).get(); - raft::linalg::mapThenSumReduce(sqrd_norm_out.data_handle(), - old_centroids.size(), - raft::sqdiff_op{}, - stream, - old_centroids.data_handle(), - new_centroids.data_handle()); + run_to_run_reduce( + handle, + cuda::transform_iterator( + cuda::counting_iterator(0), + SquaredDiffOp{old_centroids.data_handle(), new_centroids.data_handle()}), + static_cast(old_centroids.size()), + sqrd_norm_out.data_handle(), + cuda::std::plus<>{}, + DataT{0}, + workspace); } /** diff --git a/cpp/src/cluster/detail/kmeans_fixed_point_sums.cuh b/cpp/src/cluster/detail/kmeans_fixed_point_sums.cuh new file mode 100644 index 0000000000..5d9589b7ac --- /dev/null +++ b/cpp/src/cluster/detail/kmeans_fixed_point_sums.cuh @@ -0,0 +1,249 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +#pragma once + +#include +#include +#include +#include +#include +#include + +#include +#include + +#include +#include +#include + +/** + * Deterministic per-cluster column sums by exact fixed-point accumulation. + * + * Each value is converted to a 64-bit integer with a per-column power-of-two scale, and the + * integers are added with atomics. Integer addition is associative, so the result does not depend + * on the order of the atomics (or on the GPU). The scale is chosen from the largest absolute value + * of each column so that n_rows values cannot overflow: each value is rounded to a multiple of + * about max|column| * n_rows * 2^-62, which is far below the rounding error of a float sum. + * + * This needs no sort, and with the scales computed once for a dataset that is summed repeatedly + * (as in the balanced k-means EM loop), each sum takes two kernels: the same launch count as the + * atomic floating-point sums it replaces. That matters for balanced k-means, which runs many small + * problems where launch overhead dominates. + */ +namespace cuvs::cluster::kmeans::detail::fixed_point_sums { + +/** Exponent e such that |x| * 2^e < 2^(61 - log2_rows) whenever |x| <= max_abs. */ +__host__ __device__ inline int column_exponent(double max_abs, int log2_rows) +{ + int exp2 = 0; + // max_abs < 2^exp2 (frexp gives exp2 = 0 for 0) + frexp(max_abs, &exp2); + return 61 - log2_rows - exp2; +} + +template +__device__ inline int64_t to_fixed(MathT x, int exponent) +{ + if constexpr (std::is_same_v) { + return __float2ll_rn(ldexpf(x, exponent)); + } else { + return __double2ll_rn(ldexp(static_cast(x), exponent)); + } +} + +/** Per-column maximum of |X| (exact; the order of the atomics does not matter for a maximum). */ +template +RAFT_KERNEL column_max_abs_kernel(DataIterator X, int64_t n_rows, int64_t n_cols, MathT* col_max) +{ + using BitsT = std::conditional_t; + const int64_t j = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + const int64_t chunk = raft::ceildiv(n_rows, gridDim.y); + const int64_t begin = static_cast(blockIdx.y) * chunk; + const int64_t end = std::min(begin + chunk, n_rows); + if (j >= n_cols) { return; } + MathT m = MathT{0}; + for (int64_t row = begin; row < end; ++row) { + const auto x = static_cast(X[row * n_cols + j]); + m = std::max(m, x < MathT{0} ? -x : x); + } + // Non-negative floating-point values compare like their bit patterns. + if constexpr (sizeof(MathT) == 4) { + atomicMax(reinterpret_cast(col_max + j), __float_as_uint(m)); + } else { + atomicMax(reinterpret_cast(col_max + j), static_cast(__double_as_longlong(m))); + } +} + +template +RAFT_KERNEL accumulate_kernel(DataIterator X, + LabelsIterator labels, + int64_t n_rows, + int64_t n_cols, + int64_t n_clusters, + const MathT* col_max, + int log2_rows, + bool use_shared, + unsigned long long* acc) +{ + extern __shared__ unsigned long long fixed_point_smem[]; + unsigned long long* shared_acc = fixed_point_smem; + int* exponents = + reinterpret_cast(fixed_point_smem + (use_shared ? n_clusters * n_cols : 0)); + const int64_t cells = n_clusters * n_cols; + for (int64_t j = threadIdx.x; j < n_cols; j += blockDim.x) { + exponents[j] = column_exponent(static_cast(col_max[j]), log2_rows); + } + if (use_shared) { + for (int64_t i = threadIdx.x; i < cells; i += blockDim.x) { + shared_acc[i] = 0; + } + } + __syncthreads(); + unsigned long long* dst = use_shared ? shared_acc : acc; + const int64_t total = n_rows * n_cols; + for (int64_t idx = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; idx < total; + idx += static_cast(gridDim.x) * blockDim.x) { + const int64_t row = idx / n_cols; + const int64_t j = idx - row * n_cols; + const auto c = static_cast(labels[row]); + const auto t = to_fixed(static_cast(X[idx]), exponents[j]); + // Two's complement wrap-around makes the unsigned atomic add a signed one. + atomicAdd(dst + c * n_cols + j, static_cast(t)); + } + if (use_shared) { + __syncthreads(); + for (int64_t i = threadIdx.x; i < cells; i += blockDim.x) { + if (shared_acc[i] != 0) { atomicAdd(acc + i, shared_acc[i]); } + } + } +} + +/** Converts the accumulators to MathT, writes or adds them to `sums`, and zeroes them again. */ +template +RAFT_KERNEL convert_kernel(unsigned long long* acc, + int64_t n_clusters, + int64_t n_cols, + const MathT* col_max, + int log2_rows, + bool reset_sums, + MathT* sums) +{ + const int64_t cells = n_clusters * n_cols; + for (int64_t i = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; i < cells; + i += static_cast(gridDim.x) * blockDim.x) { + const int64_t j = i % n_cols; + const auto fixed = static_cast(acc[i]); + acc[i] = 0; + const int exp = column_exponent(static_cast(col_max[j]), log2_rows); + const auto value = static_cast(ldexp(static_cast(fixed), -exp)); + sums[i] = reset_sums ? value : sums[i] + value; + } +} + +/** + * @brief Repeated deterministic per-cluster column sums over one dataset. + * + * `set_scale` computes the per-column scales of a dataset; `sums_by_label` then sums the rows of + * that dataset (or of any subset of at most the same number of rows, with values within the same + * column ranges) per label. The accumulators are zeroed after every sum, so consecutive sums need + * no memset. + * + * @tparam MathT Accumulation and output type (float or double) + */ +template +class accumulator { + public: + accumulator(raft::resources const& handle, + int64_t n_cols, + int64_t n_clusters, + rmm::device_async_resource_ref mr) + : n_cols_{n_cols}, + n_clusters_{n_clusters}, + col_max_(n_cols, raft::resource::get_cuda_stream(handle), mr), + acc_(n_cols * n_clusters, raft::resource::get_cuda_stream(handle), mr) + { + auto stream = raft::resource::get_cuda_stream(handle); + RAFT_CUDA_TRY( + cudaMemsetAsync(acc_.data(), 0, acc_.size() * sizeof(unsigned long long), stream.get())); + } + + /** Computes the per-column scales of the [n_rows x n_cols] row-major dataset X. */ + template + void set_scale(raft::resources const& handle, DataIterator X, int64_t n_rows) + { + auto stream = raft::resource::get_cuda_stream(handle).get(); + log2_rows_ = 0; + while ((int64_t{1} << log2_rows_) < std::max(n_rows, 1)) { + ++log2_rows_; + } + max_rows_ = n_rows; + RAFT_CUDA_TRY(cudaMemsetAsync(col_max_.data(), 0, n_cols_ * sizeof(MathT), stream)); + if (n_rows == 0 || n_cols_ == 0) { return; } + constexpr int kThreads = 128; + const auto row_chunks = + static_cast(std::min(raft::ceildiv(n_rows, 1024), int64_t{1024})); + const dim3 grid(static_cast(raft::ceildiv(n_cols_, kThreads)), row_chunks); + column_max_abs_kernel<<>>(X, n_rows, n_cols_, col_max_.data()); + RAFT_CUDA_TRY(cudaPeekAtLastError()); + } + + /** + * sums[c, :] (+)= sum of the rows i of X with labels[i] == c. + * + * @param[in] X Row-major data [n_rows x n_cols], within the ranges given to set_scale + * @param[in] labels Label of each row, in [0, n_clusters) + * @param[in] n_rows Number of rows, at most the number given to set_scale + * @param[inout] sums Row-major output [n_clusters x n_cols] + * @param[in] reset_sums Overwrite `sums` if true, add to it otherwise + */ + template + void sums_by_label(raft::resources const& handle, + DataIterator X, + LabelsIterator labels, + int64_t n_rows, + MathT* sums, + bool reset_sums) + { + RAFT_EXPECTS(n_rows <= max_rows_, "fixed_point_sums: more rows than given to set_scale"); + auto stream = raft::resource::get_cuda_stream(handle).get(); + const int64_t cells = n_cols_ * n_clusters_; + if (cells == 0) { return; } + constexpr int kThreads = 256; + constexpr size_t kMaxSmem = 40 * 1024; + const size_t exp_bytes = sizeof(int) * n_cols_; + const bool use_shared = cells * sizeof(unsigned long long) + exp_bytes <= kMaxSmem; + const size_t smem = (use_shared ? cells * sizeof(unsigned long long) : 0) + exp_bytes; + if (n_rows > 0) { + const int64_t sms = raft::resource::get_device_properties(handle).multiProcessorCount; + const auto grid = static_cast(std::min( + raft::ceildiv(n_rows * n_cols_, kThreads), use_shared ? 4 * sms : 64 * sms)); + accumulate_kernel<<>>(X, + labels, + n_rows, + n_cols_, + n_clusters_, + col_max_.data(), + log2_rows_, + use_shared, + acc_.data()); + RAFT_CUDA_TRY(cudaPeekAtLastError()); + } + const auto grid = + static_cast(std::min(raft::ceildiv(cells, kThreads), 4096)); + convert_kernel<<>>( + acc_.data(), n_clusters_, n_cols_, col_max_.data(), log2_rows_, reset_sums, sums); + RAFT_CUDA_TRY(cudaPeekAtLastError()); + } + + private: + int64_t n_cols_; + int64_t n_clusters_; + int64_t max_rows_ = 0; + int log2_rows_ = 0; + rmm::device_uvector col_max_; + rmm::device_uvector acc_; +}; + +} // namespace cuvs::cluster::kmeans::detail::fixed_point_sums diff --git a/cpp/src/cluster/detail/kmeans_mg.cuh b/cpp/src/cluster/detail/kmeans_mg.cuh index a86c30a087..4b25221f63 100644 --- a/cpp/src/cluster/detail/kmeans_mg.cuh +++ b/cpp/src/cluster/detail/kmeans_mg.cuh @@ -605,7 +605,8 @@ void mnmg_fit( dev_res, rank_centroids_const, raft::make_const_mdspan(new_centroids.view()), - sqrd_norm_error_dev.view()); + sqrd_norm_error_dev.view(), + workspace); raft::copy( rank_centroids.data_handle(), new_centroids.data_handle(), n_clusters * n_features, stream); diff --git a/cpp/src/neighbors/ivf_pq/ivf_pq_build.cuh b/cpp/src/neighbors/ivf_pq/ivf_pq_build.cuh index 85024d2de2..8f6cdccd61 100644 --- a/cpp/src/neighbors/ivf_pq/ivf_pq_build.cuh +++ b/cpp/src/neighbors/ivf_pq/ivf_pq_build.cuh @@ -58,6 +58,8 @@ #include #include #include +#include +#include #include #include @@ -205,17 +207,6 @@ void flat_compute_residuals( stream.get()); } -template -__launch_bounds__(BlockDim) static __global__ void fill_indices_kernel(IdxT n_rows, - IdxT* data_indices, - IdxT* data_offsets, - const uint32_t* labels) -{ - const auto i = IdxT(BlockDim) * IdxT(blockIdx.x) + IdxT(threadIdx.x); - if (i >= n_rows) { return; } - data_indices[atomicAdd(data_offsets + labels[i], 1)] = i; -} - /** * @brief Calculate cluster offsets and arrange data indices into clusters. * @@ -249,13 +240,12 @@ auto calculate_offsets_and_indices(IdxT n_rows, stream.sync(); RAFT_EXPECTS(cumsum == n_rows, "cluster sizes do not add up."); RAFT_LOG_DEBUG("Max cluster size %d", max_cluster_size); - rmm::device_uvector data_offsets_buf(n_lists, stream); - auto data_offsets = data_offsets_buf.data(); - raft::copy(data_offsets, cluster_offsets, n_lists, stream); - constexpr uint32_t n_threads = 128; // NOLINT - const IdxT n_blocks = raft::div_rounding_up_unsafe(n_rows, n_threads); - fill_indices_kernel - <<>>(n_rows, data_indices, data_offsets, labels); + // Group the row indices by cluster with a stable sort, so that the order of the rows within a + // cluster (which decides the per-cluster training subsets) is the same in every build. + rmm::device_uvector sorted_labels(n_rows, stream); + raft::copy(sorted_labels.data(), labels, n_rows, stream); + thrust::sequence(exec_policy, data_indices, data_indices + n_rows); + thrust::stable_sort_by_key(exec_policy, sorted_labels.begin(), sorted_labels.end(), data_indices); return max_cluster_size; } diff --git a/cpp/tests/cluster/kmeans.cu b/cpp/tests/cluster/kmeans.cu index 23fc248917..a6e5b4c0f6 100644 --- a/cpp/tests/cluster/kmeans.cu +++ b/cpp/tests/cluster/kmeans.cu @@ -23,6 +23,7 @@ #include +#include #include #include @@ -495,6 +496,82 @@ class KmeansFitBatchedTest : public ::testing::TestWithParam(handle, n_clusters, n_features); + raft::random::RngState rng(p.rng_state.seed); + raft::random::uniform(handle, rng, d_init.data_handle(), n_clusters * n_features, T(-1), T(1)); + + auto d_sw = d_sw_view(); + std::optional> h_sw = std::nullopt; + auto h_sample_weight = raft::make_host_vector(testparams.weighted ? n_samples : 0); + if (testparams.weighted) { + std::fill_n(h_sample_weight.data_handle(), n_samples, T(1)); + h_sw = std::make_optional(raft::make_const_mdspan(h_sample_weight.view())); + } + + constexpr int n_runs = 3; + for (int host : {0, 1}) { + std::vector> centroids(n_runs, std::vector(n_clusters * n_features)); + std::vector inertias(n_runs); + std::vector n_iters(n_runs); + for (int r = 0; r < n_runs; ++r) { + auto d_c = raft::make_device_matrix(handle, n_clusters, n_features); + raft::copy(d_c.data_handle(), d_init.data_handle(), n_clusters * n_features, stream); + T inertia = 0; + if (host) { + auto pb = p; + pb.device_buffer_samples = testparams.device_buffer_samples; + int64_t n_iter = 0; + cuvs::cluster::kmeans::fit( + handle, + pb, + raft::make_const_mdspan(h_X->view()), + h_sw, + raft::make_device_matrix_view(d_c.data_handle(), n_clusters, n_features), + raft::make_host_scalar_view(&inertia), + raft::make_host_scalar_view(&n_iter)); + n_iters[r] = n_iter; + } else { + int n_iter = 0; + cuvs::cluster::kmeans::fit(handle, + p, + raft::make_const_mdspan(d_X->view()), + d_sw, + d_c.view(), + raft::make_host_scalar_view(&inertia), + raft::make_host_scalar_view(&n_iter)); + n_iters[r] = n_iter; + } + raft::copy(centroids[r].data(), d_c.data_handle(), n_clusters * n_features, stream); + raft::resource::sync_stream(handle, stream); + inertias[r] = inertia; + } + for (int r = 1; r < n_runs; ++r) { + EXPECT_EQ(n_iters[0], n_iters[r]) << (host ? "host" : "device") << " run " << r; + EXPECT_EQ(inertias[0], inertias[r]) << (host ? "host" : "device") << " run " << r; + EXPECT_EQ( + 0, std::memcmp(centroids[0].data(), centroids[r].data(), centroids[0].size() * sizeof(T))) + << (host ? "host" : "device") << " run " << r << ": centroids differ"; + } + } + } + T fitHostWithInitSize(int64_t init_size_value) { int n_features = testparams.n_col; @@ -709,6 +786,18 @@ TEST_P(KmeansFitBatchedTestD, Result) runZeroCost(); } +TEST_P(KmeansFitBatchedTestF, Reproducible) +{ + prepareBlobInputs(); + runReproducibleCheck(); +} + +TEST_P(KmeansFitBatchedTestD, Reproducible) +{ + prepareBlobInputs(); + runReproducibleCheck(); +} + INSTANTIATE_TEST_CASE_P(KmeansFitBatchedTests, KmeansFitBatchedTestF, ::testing::ValuesIn(batched_inputsf2)); diff --git a/cpp/tests/cluster/kmeans_balanced.cu b/cpp/tests/cluster/kmeans_balanced.cu index 3b7988e76b..6b9f0c6f79 100644 --- a/cpp/tests/cluster/kmeans_balanced.cu +++ b/cpp/tests/cluster/kmeans_balanced.cu @@ -10,8 +10,10 @@ #include #include #include +#include #include #include +#include #include #include #include @@ -22,6 +24,8 @@ #include +#include +#include #include #include @@ -301,4 +305,88 @@ KB_TEST((KmeansBalancedTest, KmeansBalancedTestFI8I32I32_SEP, inputsf_i32); +/* + * Identical fits must give bitwise-identical centers. With `skewed`, 95% of the rows are packed in + * a tiny region and n_lists divides n_rows, so most clusters start (nearly) empty and the balancing + * step must find a donor for many receivers. + */ + +template +void check_balanced_fit_reproducible( + cuvs::cluster::kmeans::balanced_donor_selection donor_selection, bool skewed = false) +{ + raft::resources handle; + auto stream = raft::resource::get_cuda_stream(handle); + const int64_t n_rows = 20000; + const int64_t n_cols = 24; + const int64_t n_lists = skewed ? 1000 : 200; + auto data_f = raft::make_device_matrix(handle, n_rows, n_cols); + auto data = raft::make_device_matrix(handle, n_rows, n_cols); + raft::random::RngState rng(1234ULL); + raft::random::uniform(handle, rng, data_f.data_handle(), n_rows * n_cols, -100.0f, 100.0f); + if (skewed) { + auto packed = raft::make_device_matrix_view( + data_f.data_handle(), n_rows * 95 / 100, n_cols); + raft::linalg::map( + handle, packed, raft::mul_const_op{1e-3f}, raft::make_const_mdspan(packed)); + } + raft::linalg::map( + handle, data.view(), raft::cast_op{}, raft::make_const_mdspan(data_f.view())); + + cuvs::cluster::kmeans::balanced_params params; + params.n_iters = 20; + params.donor_selection = donor_selection; + if (donor_selection == cuvs::cluster::kmeans::balanced_donor_selection::Random) { + // The IVF-PQ codebook settings + params.balance_lower_tolerance = 0.25f; + params.balance_upper_tolerance = 4.0f; + params.centroid_offset = 0.05f; + } + + std::vector> results; + for (int run = 0; run < 2; run++) { + auto centers = raft::make_device_matrix(handle, n_lists, n_cols); + ASSERT_NO_THROW(cuvs::cluster::kmeans::fit( + handle, params, raft::make_const_mdspan(data.view()), centers.view())); + std::vector h(centers.size()); + raft::update_host(h.data(), centers.data_handle(), h.size(), stream); + raft::resource::sync_stream(handle); + for (auto x : h) { + ASSERT_TRUE(std::isfinite(x)); + } + results.push_back(std::move(h)); + } + ASSERT_EQ(0, std::memcmp(results[0].data(), results[1].data(), results[0].size() * sizeof(float))) + << "balanced k-means centers differ between identical fits"; +} + +TEST(KmeansBalancedReproducible, FloatSizeSorted) +{ + check_balanced_fit_reproducible( + cuvs::cluster::kmeans::balanced_donor_selection::SizeSorted); +} + +TEST(KmeansBalancedReproducible, FloatRandomDonor) +{ + check_balanced_fit_reproducible(cuvs::cluster::kmeans::balanced_donor_selection::Random); +} + +TEST(KmeansBalancedReproducible, Int8SizeSorted) +{ + check_balanced_fit_reproducible( + cuvs::cluster::kmeans::balanced_donor_selection::SizeSorted); +} + +TEST(KmeansBalancedReproducible, SkewedSizeSorted) +{ + check_balanced_fit_reproducible( + cuvs::cluster::kmeans::balanced_donor_selection::SizeSorted, true); +} + +TEST(KmeansBalancedReproducible, SkewedRandomDonor) +{ + check_balanced_fit_reproducible(cuvs::cluster::kmeans::balanced_donor_selection::Random, + true); +} + } // namespace cuvs diff --git a/cpp/tests/neighbors/ann_ivf_flat/test_float_int64_t.cu b/cpp/tests/neighbors/ann_ivf_flat/test_float_int64_t.cu index fee22db3c5..0e7e47f20f 100644 --- a/cpp/tests/neighbors/ann_ivf_flat/test_float_int64_t.cu +++ b/cpp/tests/neighbors/ann_ivf_flat/test_float_int64_t.cu @@ -10,7 +10,9 @@ #include #include #include +#include +#include #include #include #include @@ -27,6 +29,30 @@ TEST_P(AnnIVFFlatTestF_float, AnnIVFFlat) INSTANTIATE_TEST_CASE_P(AnnIVFFlatTest, AnnIVFFlatTestF_float, ::testing::ValuesIn(inputs)); +TEST(IVFFlatBuild, Reproducible) +{ + raft::resources handle; + auto stream = raft::resource::get_cuda_stream(handle); + const int64_t n_rows = 100000; + const int64_t dim = 32; + auto dataset = raft::make_device_matrix(handle, n_rows, dim); + raft::random::RngState rng(42ULL); + raft::random::uniform(handle, rng, dataset.data_handle(), n_rows * dim, -1.0f, 1.0f); + + index_params params; + params.n_lists = 500; + std::vector> centers; + for (int run = 0; run < 2; run++) { + const auto index = build(handle, params, raft::make_const_mdspan(dataset.view())); + std::vector h(index.centers().size()); + raft::update_host(h.data(), index.centers().data_handle(), h.size(), stream); + raft::resource::sync_stream(handle); + centers.push_back(std::move(h)); + } + ASSERT_EQ(0, std::memcmp(centers[0].data(), centers[1].data(), centers[0].size() * sizeof(float))) + << "IVF-Flat centers differ between identical builds"; +} + TEST(IVFFlatSerialization, StreamRoundTrip) { raft::resources handle; diff --git a/cpp/tests/neighbors/ann_ivf_pq/test_float_int64_t.cu b/cpp/tests/neighbors/ann_ivf_pq/test_float_int64_t.cu index 47e5f8d0dd..8a89d45682 100644 --- a/cpp/tests/neighbors/ann_ivf_pq/test_float_int64_t.cu +++ b/cpp/tests/neighbors/ann_ivf_pq/test_float_int64_t.cu @@ -1,10 +1,16 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ #include "../ann_ivf_pq.cuh" +#include + +#include +#include +#include + namespace cuvs::neighbors::ivf_pq { using f32_f32_i64 = ivf_pq_test; @@ -28,4 +34,51 @@ INSTANTIATE(f32_f32_i64_filter, defaults() + small_dims() + big_dims_moderate_lut() + enum_variety_l2() + enum_variety_l2sqrt() + enum_variety_ip() + enum_variety_cosine()); +template +auto copy_to_host(raft::resources const& handle, MdspanT view) +{ + std::vector> h(view.size()); + raft::update_host( + h.data(), view.data_handle(), h.size(), raft::resource::get_cuda_stream(handle)); + raft::resource::sync_stream(handle); + return h; +} + +void check_ivf_pq_build_reproducible(codebook_gen codebook_kind) +{ + raft::resources handle; + const int64_t n_rows = 100000; + const int64_t dim = 32; + auto dataset = raft::make_device_matrix(handle, n_rows, dim); + raft::random::RngState rng(42ULL); + raft::random::uniform(handle, rng, dataset.data_handle(), n_rows * dim, -1.0f, 1.0f); + + index_params params; + params.n_lists = 500; + params.pq_dim = 16; + params.codebook_kind = codebook_kind; + std::vector> centers, pq_centers; + for (int run = 0; run < 2; run++) { + const auto index = build(handle, params, raft::make_const_mdspan(dataset.view())); + centers.push_back(copy_to_host(handle, index.centers())); + pq_centers.push_back(copy_to_host(handle, index.pq_centers())); + } + ASSERT_EQ(0, std::memcmp(centers[0].data(), centers[1].data(), centers[0].size() * sizeof(float))) + << "IVF-PQ cluster centers differ between identical builds"; + ASSERT_EQ( + 0, + std::memcmp(pq_centers[0].data(), pq_centers[1].data(), pq_centers[0].size() * sizeof(float))) + << "IVF-PQ codebooks differ between identical builds"; +} + +TEST(IvfPqBuild, ReproduciblePerSubspace) +{ + check_ivf_pq_build_reproducible(codebook_gen::PER_SUBSPACE); +} + +TEST(IvfPqBuild, ReproduciblePerCluster) +{ + check_ivf_pq_build_reproducible(codebook_gen::PER_CLUSTER); +} + } // namespace cuvs::neighbors::ivf_pq