diff --git a/.github/workflows/pr.yaml b/.github/workflows/pr.yaml index a6438c1a8d..b364a38cbf 100644 --- a/.github/workflows/pr.yaml +++ b/.github/workflows/pr.yaml @@ -11,450 +11,20 @@ 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 + - bench-build-base + - bench-build-pr2801 + - bench-build-pr2804 + - bench 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] + bench-build-base: + needs: [build-details] permissions: actions: read contents: read @@ -467,228 +37,10 @@ jobs: build_type: pull-request build-datetime: ${{ needs.build-details.outputs.build-datetime }} node_type: cpu16 - script: ci/build_cpp.sh - conda-cpp-tests: - needs: [conda-cpp-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-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 }}" - 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 - # save newly compiled entries; restore a compatible earlier archive first. - cache-key-files: | - cpp/CMakeLists.txt - cpp/cmake/** - cpp/include/cuvs/detail/jit_lto/** - cpp/src/detail/jit_lto/** - cpp/src/distance/detail/pairwise_matrix/jit_lto_kernels/** - cpp/src/neighbors/detail/cagra/jit_lto_kernels/** - cpp/src/neighbors/detail/jit_lto_kernels/** - cpp/src/neighbors/ivf_flat/detail/jit_lto_kernels/** - cpp/src/neighbors/ivf_pq/detail/jit_lto_kernels/** - cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/** - cpp/src/neighbors/ivf_sq/detail/jit_lto_kernels/** - cpp/tests/** - dependencies.yaml - # The CUDA driver's JIT cache is independent of the Python interpreter used to run C++ tests. - cache-key-matrix-fields: | - CUDA_VER - ARCH - LINUX_VER - GPU - DRIVER - DEPENDENCIES - # PRs restore a compatible main cache but do not create merge-ref-scoped entries. - 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] + script: ci/bench/build_base.sh + matrix_filter: map(select(.CUDA_VER == "13.3.0")) + bench-build-pr2801: + needs: [build-details] permissions: actions: read contents: read @@ -696,38 +48,15 @@ jobs: 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 + uses: rapidsai/shared-workflows/.github/workflows/conda-cpp-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] + script: ci/bench/build_pr2801.sh + matrix_filter: map(select(.CUDA_VER == "13.3.0")) + bench-build-pr2804: + needs: [build-details] permissions: actions: read contents: read @@ -735,18 +64,15 @@ jobs: packages: read pull-requests: read secrets: inherit # zizmor: ignore[secrets-inherit] - uses: rapidsai/shared-workflows/.github/workflows/wheels-build.yaml@main + uses: rapidsai/shared-workflows/.github/workflows/conda-cpp-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] + node_type: cpu16 + script: ci/bench/build_pr2804.sh + matrix_filter: map(select(.CUDA_VER == "13.3.0")) + bench: + needs: [bench-build-base, bench-build-pr2801, bench-build-pr2804] permissions: actions: read contents: read @@ -754,43 +80,11 @@ jobs: 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 + uses: rapidsai/shared-workflows/.github/workflows/conda-cpp-tests.yaml@main 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 }} + script: "ci/bench/run.sh" + matrix_type: nightly + # DEPENDENCIES selects the workload group of ci/bench/run.sh (kmeans, ivf, lite, tests, + # profile, profile2). + matrix_filter: '[{ARCH: "amd64", PY_VER: "3.13", CUDA_VER: "13.3.0", LINUX_VER: "ubuntu26.04", GPU: "h100", DRIVER: "latest", DEPENDENCIES: "kmeans"}, {ARCH: "amd64", PY_VER: "3.13", CUDA_VER: "13.3.0", LINUX_VER: "ubuntu26.04", GPU: "h100", DRIVER: "latest", DEPENDENCIES: "ivf"}, {ARCH: "arm64", PY_VER: "3.13", CUDA_VER: "13.3.0", LINUX_VER: "rockylinux8", GPU: "a100", DRIVER: "latest", DEPENDENCIES: "kmeans"}, {ARCH: "arm64", PY_VER: "3.13", CUDA_VER: "13.3.0", LINUX_VER: "rockylinux8", GPU: "a100", DRIVER: "latest", DEPENDENCIES: "ivf"}, {ARCH: "amd64", PY_VER: "3.14", CUDA_VER: "13.3.0", LINUX_VER: "ubuntu26.04", GPU: "rtxpro6000", DRIVER: "latest", DEPENDENCIES: "lite"}, {ARCH: "amd64", PY_VER: "3.13", CUDA_VER: "13.3.0", LINUX_VER: "ubuntu26.04", GPU: "l4", DRIVER: "latest", DEPENDENCIES: "lite"}, {ARCH: "arm64", PY_VER: "3.14", CUDA_VER: "13.3.0", LINUX_VER: "ubuntu26.04", GPU: "gh200", DRIVER: "latest", DEPENDENCIES: "lite"}, {ARCH: "arm64", PY_VER: "3.14", CUDA_VER: "13.3.0", LINUX_VER: "ubuntu26.04", GPU: "gb300", DRIVER: "latest", DEPENDENCIES: "lite"}, {ARCH: "amd64", PY_VER: "3.13", CUDA_VER: "13.3.0", LINUX_VER: "ubuntu26.04", GPU: "h100", DRIVER: "latest", DEPENDENCIES: "tests"}, {ARCH: "amd64", PY_VER: "3.13", CUDA_VER: "13.3.0", LINUX_VER: "ubuntu26.04", GPU: "l4", DRIVER: "latest", DEPENDENCIES: "tests"}]' diff --git a/ci/bench/bench.cu b/ci/bench/bench.cu new file mode 100644 index 0000000000..712c7c0c75 --- /dev/null +++ b/ci/bench/bench.cu @@ -0,0 +1,727 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +// End-to-end timing of k-means fit, balanced k-means and IVF / CAGRA builds through the public API. +// The same binary is run against several libcuvs builds (selected with LD_LIBRARY_PATH); one +// invocation runs one workload for one libcuvs build, `warmup` untimed + `reps` timed repetitions. +// +// Output: one line per timed repetition, +// RESULT | variant= ms= +#ifndef RAFT_SYSTEM_LITTLE_ENDIAN +#define RAFT_SYSTEM_LITTLE_ENDIAN 1 // as set by the cuvs CMake targets +#endif + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +using clk = std::chrono::steady_clock; + +static std::string g_variant = "?"; + +struct Args { + std::map kv; + std::string get(const std::string& k, const std::string& d) const + { + auto it = kv.find(k); + return it == kv.end() ? d : it->second; + } + int64_t geti(const std::string& k, int64_t d) const + { + auto it = kv.find(k); + return it == kv.end() ? d : std::stoll(it->second); + } + double getf(const std::string& k, double d) const + { + auto it = kv.find(k); + return it == kv.end() ? d : std::stod(it->second); + } +}; + +std::vector parse_list(const std::string& s) +{ + std::vector out; + std::stringstream ss(s); + std::string tok; + while (std::getline(ss, tok, ',')) { + if (!tok.empty()) { out.push_back(std::stoll(tok)); } + } + return out; +} + +std::vector split(const std::string& s, char c) +{ + std::vector out; + std::stringstream ss(s); + std::string tok; + while (std::getline(ss, tok, c)) { + out.push_back(tok); + } + return out; +} + +template +std::vector to_host(raft::resources const& res, const T* p, size_t n) +{ + std::vector h(n); + raft::copy(h.data(), p, n, raft::resource::get_cuda_stream(res)); + raft::resource::sync_stream(res); + return h; +} + +template +uint64_t fnv(const std::vector& v, uint64_t h = 1469598103934665603ULL) +{ + auto p = reinterpret_cast(v.data()); + for (size_t i = 0; i < v.size() * sizeof(T); i++) { + h = (h ^ p[i]) * 1099511628211ULL; + } + return h; +} + +bool fits(size_t bytes, const std::string& what) +{ + size_t free_b = 0, total_b = 0; + cudaMemGetInfo(&free_b, &total_b); + if (bytes > free_b * 0.85) { + printf("SKIP %s needs ~%.1f GB, free %.1f GB\n", what.c_str(), bytes / 1e9, free_b / 1e9); + fflush(stdout); + return false; + } + return true; +} + +template +double time_ms(raft::resources const& res, F&& f) +{ + raft::resource::sync_stream(res); + cudaDeviceSynchronize(); + auto t0 = clk::now(); + f(); + raft::resource::sync_stream(res); + cudaDeviceSynchronize(); + return std::chrono::duration(clk::now() - t0).count(); +} + +// Prints the libcuvs.so this process loaded and exits if it is not $BENCH_EXPECT_LIB. +void check_loaded_libcuvs() +{ + std::ifstream maps("/proc/self/maps"); + std::string line, path; + while (std::getline(maps, line)) { + if (line.find("libcuvs.so") != std::string::npos) { + path = line.substr(line.find('/')); + break; + } + } + printf("LIB variant=%s %s\n", g_variant.c_str(), path.c_str()); + const char* expect = std::getenv("BENCH_EXPECT_LIB"); + if (expect != nullptr && path != expect) { + printf("ERROR variant=%s loaded %s, expected %s\n", g_variant.c_str(), path.c_str(), expect); + fflush(stdout); + std::exit(3); + } +} + +// Dataset: Gaussian blobs, deterministic for a given (n, d, seed). +template +void make_data(raft::resources const& res, + T* out, + int64_t n, + int64_t d, + int64_t n_blobs, + T stddev, + uint64_t seed) +{ + auto labels = raft::make_device_vector(res, n); + raft::random::make_blobs(out, + labels.data_handle(), + n, + d, + n_blobs, + raft::resource::get_cuda_stream(res).get(), + true, + nullptr, + nullptr, + stddev, + true, + T(-10), + T(10), + seed); +} + +// ------------------------------------------------------------------------------------------------ +// k-means fit (regular). Array init with the first k rows, so every build starts from the same +// centroids; tol is tiny so that the iteration count is (nearly always) max_iter. +// ------------------------------------------------------------------------------------------------ +template +void suite_kmeans(const Args& a) +{ + const int64_t n = a.geti("n", 1000000); + const int64_t d = a.geti("d", 128); + const auto ks = parse_list(a.get("k", "100")); + const bool weights = a.geti("w", 0) != 0; + const int64_t batch = a.geti("batch", 0); // > 0: host data, batched fit + const std::string init = a.get("init", "array"); + const int reps = a.geti("reps", 2); + const int warmup = a.geti("warmup", 1); + const bool ref64 = a.geti("ref64", 0) != 0; // also fit in double from the same init + const char* dtype = sizeof(T) == 4 ? "f32" : "f64"; + + raft::resources res; + auto stream = raft::resource::get_cuda_stream(res); + size_t need = (batch > 0 ? 2 * batch : n) * d * sizeof(T) + n * 24 + (size_t(3) << 30); + if (ref64 && sizeof(T) == 4) { need += n * d * 8; } + if (!fits(need, "kmeans n=" + std::to_string(n) + " d=" + std::to_string(d))) { return; } + + auto X = raft::make_device_matrix(res, n, d); + // "uniform" (default): no cluster structure, so the cost keeps decreasing and every fit runs + // max_iter iterations. "blobs": 1000 Gaussian blobs. + const std::string data = a.get("data", "uniform"); + if (data == "blobs") { + make_data(res, X.data_handle(), n, d, 1000, T(2), 1234ULL); + } else { + raft::random::RngState rng(1234); + raft::random::uniform(res, rng, X.data_handle(), n * d, T(-1), T(1)); + } + const uint64_t xhash = + fnv(to_host(res, X.data_handle(), std::min(n * d, int64_t(1) << 22))); + auto W = raft::make_device_vector(res, n); + if (weights) { + raft::random::RngState rng(99); + raft::random::uniform(res, rng, W.data_handle(), n, T(0.5), T(1.5)); + } + std::vector hX, hW; + if (batch > 0) { + hX = to_host(res, X.data_handle(), n * d); + if (weights) { hW = to_host(res, W.data_handle(), n); } + } + std::optional> X64; + std::optional> W64; + if (ref64 && sizeof(T) == 4) { + X64.emplace(raft::make_device_matrix(res, n, d)); + thrust::transform(thrust::cuda::par.on(stream.get()), + X.data_handle(), + X.data_handle() + n * d, + X64->data_handle(), + [] __device__(T x) { return double(x); }); + if (weights) { + W64.emplace(raft::make_device_vector(res, n)); + thrust::transform(thrust::cuda::par.on(stream.get()), + W.data_handle(), + W.data_handle() + n, + W64->data_handle(), + [] __device__(T x) { return double(x); }); + } + } + + for (auto k : ks) { + const double flops_iter = 2.0 * n * k * d; + if (flops_iter > a.getf("max_flops_iter", 3e13)) { + printf("SKIP kmeans n=%ld d=%ld k=%ld (too slow)\n", n, d, k); + continue; + } + int max_iter = a.geti("iters", flops_iter > 3e12 ? 3 : 10); + cuvs::cluster::kmeans::params p; + p.n_clusters = k; + p.max_iter = max_iter; + p.tol = 1e-30; + p.init = init == "kpp" ? cuvs::cluster::kmeans::params::KMeansPlusPlus + : cuvs::cluster::kmeans::params::Array; + p.rng_state.seed = 42; + p.device_buffer_samples = batch; + auto C = raft::make_device_matrix(res, k, d); + for (int r = 0; r < warmup + reps; r++) { + raft::copy(C.data_handle(), X.data_handle(), k * d, stream); + T inertia = 0; + int64_t n_iter = 0; + double ms = time_ms(res, [&] { + if (batch > 0) { + std::optional> sw; + if (weights) { sw = raft::make_host_vector_view(hW.data(), n); } + cuvs::cluster::kmeans::fit(res, + p, + raft::make_host_matrix_view(hX.data(), n, d), + sw, + C.view(), + raft::make_host_scalar_view(&inertia), + raft::make_host_scalar_view(&n_iter)); + } else { + std::optional> sw; + if (weights) { sw = raft::make_const_mdspan(W.view()); } + cuvs::cluster::kmeans::fit(res, + p, + raft::make_const_mdspan(X.view()), + sw, + C.view(), + raft::make_host_scalar_view(&inertia), + raft::make_host_scalar_view(&n_iter)); + } + }); + if (r < warmup) { continue; } + auto hC = to_host(res, C.data_handle(), C.size()); + std::string ref_str; + if (X64.has_value() && r == warmup) { + // Reference: the same fit (same init, same iteration count) in double precision. + auto C64 = raft::make_device_matrix(res, k, d); + thrust::transform(thrust::cuda::par.on(stream.get()), + X.data_handle(), + X.data_handle() + k * d, + C64.data_handle(), + [] __device__(T x) { return double(x); }); + double inertia64 = 0; + int64_t n_iter64 = 0; + auto p64 = p; + p64.max_iter = n_iter; + std::optional> sw64; + if (weights) { sw64 = raft::make_const_mdspan(W64->view()); } + cuvs::cluster::kmeans::fit(res, + p64, + raft::make_const_mdspan(X64->view()), + sw64, + C64.view(), + raft::make_host_scalar_view(&inertia64), + raft::make_host_scalar_view(&n_iter64)); + auto hC64 = to_host(res, C64.data_handle(), C64.size()); + double max_abs = 0, scale = 0; + for (size_t i = 0; i < hC.size(); i++) { + max_abs = std::max(max_abs, std::abs(double(hC[i]) - hC64[i])); + scale = std::max(scale, std::abs(hC64[i])); + } + char buf[256]; + snprintf(buf, + sizeof(buf), + " inertia64=%.9g iters64=%ld center_maxabs_vs64=%.3g center_maxrel_vs64=%.3g", + inertia64, + n_iter64, + max_abs, + max_abs / scale); + ref_str = buf; + } + printf( + "RESULT kmeans_fit dtype=%s data=%s n=%ld d=%ld k=%ld w=%d batch=%ld init=%s max_iter=%d | " + "variant=%s ms=%.3f iters=%ld per_iter_ms=%.4f inertia=%.9g hash=%016lx xhash=%016lx%s\n", + dtype, + data.c_str(), + n, + d, + k, + int(weights), + batch, + init.c_str(), + max_iter, + g_variant.c_str(), + ms, + n_iter, + ms / std::max(n_iter, 1), + double(inertia), + fnv(hC), + xhash, + ref_str.c_str()); + fflush(stdout); + } + } +} + +// ------------------------------------------------------------------------------------------------ +// Balanced k-means fit and predict. +// ------------------------------------------------------------------------------------------------ +template +double host_inertia(const std::vector& X, + const std::vector& C, + const std::vector& labels, + int64_t n, + int64_t d) +{ + double s = 0; + for (int64_t i = 0; i < n; i++) { + const float* c = C.data() + size_t(labels[i]) * d; + const T* x = X.data() + size_t(i) * d; + double r = 0; + for (int64_t j = 0; j < d; j++) { + double diff = double(x[j]) - double(c[j]); + r += diff * diff; + } + s += r; + } + return s; +} + +template +void suite_balanced(const Args& a) +{ + const int64_t n = a.geti("n", 1000000); + const int64_t d = a.geti("d", 128); + const auto ks = parse_list(a.get("k", "1024")); + const int calls = a.geti("calls", 1); + const bool random = a.get("donor", "size") == "random"; + const int n_iters = a.geti("iters", 20); + const int reps = a.geti("reps", 2); + const int warmup = a.geti("warmup", 1); + const bool do_predict = a.geti("predict", 0) != 0; + const char* dtype = std::is_same_v ? "f32" : "i8"; + + raft::resources res; + auto stream = raft::resource::get_cuda_stream(res); + if (!fits(n * d * 5 + n * 16 + (size_t(2) << 30), "balanced n=" + std::to_string(n))) { return; } + auto Xf = raft::make_device_matrix(res, n, d); + make_data(res, Xf.data_handle(), n, d, 1000, 2.0f, 4321ULL); + auto X = raft::make_device_matrix(res, n, d); + if constexpr (std::is_same_v) { + raft::copy(X.data_handle(), Xf.data_handle(), n * d, stream); + } else { + thrust::transform(thrust::cuda::par.on(stream.get()), + Xf.data_handle(), + Xf.data_handle() + n * d, + X.data_handle(), + [] __device__(float x) { + return T(fmaxf(-127.f, fminf(127.f, rintf(x * 6.f)))); + }); + } + // distinct=1: every call gets its own dataset (like the per-subspace fits of IVF-PQ), instead + // of refitting the same data `calls` times. + const bool distinct = a.geti("distinct", 0) != 0 && calls > 1; + std::optional> XS; + if (distinct) { + XS.emplace(raft::make_device_matrix(res, calls * n, d)); + for (int i = 0; i < calls; i++) { + make_data(res, Xf.data_handle(), n, d, 1000, 2.0f, 4321ULL + i); + thrust::transform(thrust::cuda::par.on(stream.get()), + Xf.data_handle(), + Xf.data_handle() + n * d, + XS->data_handle() + i * n * d, + [] __device__(float x) { + return std::is_same_v + ? T(x) + : T(fmaxf(-127.f, fminf(127.f, rintf(x * 6.f)))); + }); + } + } + std::vector hX; + if (!distinct && n * d <= int64_t(1) << 28) { hX = to_host(res, X.data_handle(), n * d); } + + cuvs::cluster::kmeans::balanced_params params; + params.n_iters = n_iters; + if (random) { + params.donor_selection = cuvs::cluster::kmeans::balanced_donor_selection::Random; + params.balance_lower_tolerance = 0.25f; + params.balance_upper_tolerance = 4.0f; + params.centroid_offset = 0.05f; + } + for (auto k : ks) { + auto C = raft::make_device_matrix(res, k, d); + auto labels = raft::make_device_vector(res, n); + for (int r = 0; r < warmup + reps; r++) { + double ms = time_ms(res, [&] { + for (int i = 0; i < calls; i++) { + auto Xi = distinct ? raft::make_device_matrix_view( + XS->data_handle() + i * n * d, n, d) + : raft::make_const_mdspan(X.view()); + cuvs::cluster::kmeans::fit(res, params, Xi, C.view()); + } + }); + if (r < warmup) { continue; } + auto hC = to_host(res, C.data_handle(), C.size()); + cuvs::cluster::kmeans::predict(res, + params, + raft::make_const_mdspan(X.view()), + raft::make_const_mdspan(C.view()), + labels.view()); + auto hL = to_host(res, labels.data_handle(), n); + double inert = hX.empty() ? -1 : host_inertia(hX, hC, hL, n, d); + int64_t min_s = n, max_s = 0; + { + std::vector sizes(k, 0); + for (auto l : hL) { + sizes[l]++; + } + for (auto s : sizes) { + min_s = std::min(min_s, s); + max_s = std::max(max_s, s); + } + } + printf( + "RESULT balanced_fit dtype=%s n=%ld d=%ld k=%ld donor=%s n_iters=%d calls=%d%s | " + "variant=%s ms=%.3f inertia64=%.9g min_size=%ld max_size=%ld hash=%016lx\n", + dtype, + n, + d, + k, + random ? "random" : "size_sorted", + n_iters, + calls, + distinct ? " distinct=1" : "", + g_variant.c_str(), + ms, + inert, + min_s, + max_s, + fnv(hC)); + fflush(stdout); + } + if (do_predict) { + for (int r = 0; r < warmup + reps; r++) { + double ms = time_ms(res, [&] { + cuvs::cluster::kmeans::predict(res, + params, + raft::make_const_mdspan(X.view()), + raft::make_const_mdspan(C.view()), + labels.view()); + }); + if (r < warmup) { continue; } + printf("RESULT balanced_predict dtype=%s n=%ld d=%ld k=%ld | variant=%s ms=%.3f\n", + dtype, + n, + d, + k, + g_variant.c_str(), + ms); + fflush(stdout); + } + } + } +} + +// ------------------------------------------------------------------------------------------------ +// Index builds (+ recall@10 of each build, + one search timing). +// ------------------------------------------------------------------------------------------------ +double recall(const std::vector& found, const std::vector& gt, int64_t nq, int k) +{ + int64_t hits = 0; + for (int64_t q = 0; q < nq; q++) { + for (int i = 0; i < k; i++) { + for (int j = 0; j < k; j++) { + if (found[q * k + i] == gt[q * k + j]) { + hits++; + break; + } + } + } + } + return double(hits) / double(nq * k); +} + +void suite_ivf(const Args& a) +{ + const int64_t n = a.geti("n", 1000000); + const int64_t d = a.geti("d", 128); + const int64_t nq = a.geti("nq", 10000); + const int topk = 10; + const int reps = a.geti("reps", 2); + const int warmup = a.geti("warmup", 1); + const auto specs = split(a.get("index", "flat:1024"), ','); + const int n_probes = a.geti("probes", 50); + const int search_reps = a.geti("search_reps", 5); + + raft::resources res; + auto stream = raft::resource::get_cuda_stream(res); + if (!fits((n + nq) * d * 4 * 3 + (size_t(3) << 30), "ivf n=" + std::to_string(n))) { return; } + auto all = raft::make_device_matrix(res, n + nq, d); + make_data(res, all.data_handle(), n + nq, d, 10000, 2.0f, 777ULL); + auto X = raft::make_device_matrix_view(all.data_handle(), n, d); + auto Q = raft::make_device_matrix_view(all.data_handle() + n * d, nq, d); + + // Ground truth. + auto nbrs = raft::make_device_matrix(res, nq, topk); + auto dists = raft::make_device_matrix(res, nq, topk); + std::vector gt; + { + auto bf = cuvs::neighbors::brute_force::build(res, {}, X); + cuvs::neighbors::brute_force::search(res, {}, bf, Q, nbrs.view(), dists.view()); + gt = to_host(res, nbrs.data_handle(), nq * topk); + } + + for (const auto& spec : specs) { + auto f = split(spec, ':'); + const std::string kind = f[0]; + const uint32_t p1 = f.size() > 1 ? std::stoul(f[1]) : 1024; + const uint32_t p2 = f.size() > 2 ? std::stoul(f[2]) : 0; + std::string cfg = "index=" + kind + " n=" + std::to_string(n) + " d=" + std::to_string(d) + + (kind == "cagra" ? " graph_degree=" : " n_lists=") + std::to_string(p1) + + (p2 ? " pq_dim=" + std::to_string(p2) : ""); + // Build `warmup + reps` times; after each timed build, measure recall (untimed). After the last + // build, time `search_reps` searches. + std::function run; // (timed_build, time_search) + auto report = [&](double ms, double rec, uint64_t h) { + printf("RESULT build %s | variant=%s ms=%.3f recall=%.5f hash=%016lx\n", + cfg.c_str(), + g_variant.c_str(), + ms, + rec, + h); + fflush(stdout); + }; + auto search_timing = [&](auto&& search_fn) { + search_fn(); // warm (JIT, allocations) + std::vector t; + for (int i = 0; i < search_reps; i++) { + t.push_back(time_ms(res, search_fn)); + } + for (double ms : t) { + printf("RESULT search %s probes=%d nq=%ld | variant=%s ms=%.3f qps=%.0f\n", + cfg.c_str(), + n_probes, + nq, + g_variant.c_str(), + ms, + nq / (ms / 1000.0)); + } + fflush(stdout); + }; + for (int r = 0; r < warmup + reps; r++) { + const bool timed = r >= warmup; + const bool last = r == warmup + reps - 1; + if (kind == "flat") { + cuvs::neighbors::ivf_flat::index_params ip; + ip.n_lists = p1; + std::optional> idx; + double ms = time_ms(res, [&] { idx.emplace(cuvs::neighbors::ivf_flat::build(res, ip, X)); }); + if (!timed) { continue; } + cuvs::neighbors::ivf_flat::search_params sp; + sp.n_probes = n_probes; + auto s = [&] { + cuvs::neighbors::ivf_flat::search(res, sp, *idx, Q, nbrs.view(), dists.view()); + }; + s(); + report(ms, + recall(to_host(res, nbrs.data_handle(), nq * topk), gt, nq, topk), + fnv(to_host(res, idx->centers().data_handle(), idx->centers().size()))); + if (last) { search_timing(s); } + } else if (kind == "pq" || kind == "pqc") { + cuvs::neighbors::ivf_pq::index_params ip; + ip.n_lists = p1; + ip.pq_dim = p2; + if (kind == "pqc") { ip.codebook_kind = cuvs::neighbors::ivf_pq::codebook_gen::PER_CLUSTER; } + std::optional> idx; + double ms = time_ms(res, [&] { idx.emplace(cuvs::neighbors::ivf_pq::build(res, ip, X)); }); + if (!timed) { continue; } + cuvs::neighbors::ivf_pq::search_params sp; + sp.n_probes = n_probes; + auto s = [&] { + cuvs::neighbors::ivf_pq::search(res, sp, *idx, Q, nbrs.view(), dists.view()); + }; + s(); + auto hc = to_host(res, idx->centers().data_handle(), idx->centers().size()); + auto hp = to_host(res, idx->pq_centers().data_handle(), idx->pq_centers().size()); + report(ms, recall(to_host(res, nbrs.data_handle(), nq * topk), gt, nq, topk), fnv(hp, fnv(hc))); + if (last) { search_timing(s); } + } else if (kind == "sq") { + cuvs::neighbors::ivf_sq::index_params ip; + ip.n_lists = p1; + std::optional> idx; + double ms = time_ms(res, [&] { idx.emplace(cuvs::neighbors::ivf_sq::build(res, ip, X)); }); + if (!timed) { continue; } + cuvs::neighbors::ivf_sq::search_params sp; + sp.n_probes = n_probes; + cuvs::neighbors::ivf_sq::search(res, sp, *idx, Q, nbrs.view(), dists.view()); + auto c = idx->centers(); + report(ms, + recall(to_host(res, nbrs.data_handle(), nq * topk), gt, nq, topk), + fnv(to_host(res, c.data_handle(), c.size()))); + } else if (kind == "rbq") { + cuvs::neighbors::ivf_rabitq::index_params ip; + ip.n_lists = p1; + std::optional> idx; + double ms = + time_ms(res, [&] { idx.emplace(cuvs::neighbors::ivf_rabitq::build(res, ip, X)); }); + if (!timed) { continue; } + cuvs::neighbors::ivf_rabitq::search_params sp; + sp.n_probes = n_probes; + cuvs::neighbors::ivf_rabitq::search(res, sp, *idx, Q, nbrs.view(), dists.view()); + auto hn = to_host(res, nbrs.data_handle(), nq * topk); + report(ms, recall(hn, gt, nq, topk), fnv(hn)); + } else if (kind == "cagra") { + cuvs::neighbors::cagra::index_params ip; + ip.graph_degree = p1; + ip.intermediate_graph_degree = 2 * p1; + ip.graph_build_params = + cuvs::neighbors::cagra::graph_build_params::ivf_pq_params(X.extents()); + auto view = cuvs::neighbors::make_device_padded_dataset_view(res, X); + std::optional> idx; + double ms = time_ms(res, [&] { idx.emplace(cuvs::neighbors::cagra::build(res, ip, view)); }); + if (!timed) { continue; } + cuvs::neighbors::cagra::search_params sp; + sp.itopk_size = 64; + auto n32 = raft::make_device_matrix(res, nq, topk); + cuvs::neighbors::cagra::search(res, sp, *idx, Q, n32.view(), dists.view()); + auto h32 = to_host(res, n32.data_handle(), nq * topk); + std::vector hn(h32.begin(), h32.end()); + auto g = idx->graph(); + report(ms, recall(hn, gt, nq, topk), fnv(to_host(res, g.data_handle(), g.size()))); + } else { + fprintf(stderr, "unknown index %s\n", kind.c_str()); + exit(2); + } + } + } +} + +int main(int argc, char** argv) +{ + if (argc < 3) { + fprintf(stderr, "usage: %s [key=value...]\n", argv[0]); + return 2; + } + g_variant = argv[1]; + std::string suite = argv[2]; + Args a; + for (int i = 3; i < argc; i++) { + std::string s(argv[i]); + auto eq = s.find('='); + if (eq != std::string::npos) { a.kv[s.substr(0, eq)] = s.substr(eq + 1); } + } + check_loaded_libcuvs(); + try { + if (suite == "kmeans") { + if (a.get("dtype", "f32") == "f64") { + suite_kmeans(a); + } else { + suite_kmeans(a); + } + } else if (suite == "balanced") { + if (a.get("dtype", "f32") == "i8") { + suite_balanced(a); + } else { + suite_balanced(a); + } + } else if (suite == "ivf") { + suite_ivf(a); + } else { + fprintf(stderr, "unknown suite %s\n", suite.c_str()); + return 2; + } + } catch (std::exception& e) { + printf("ERROR variant=%s suite=%s: %s\n", g_variant.c_str(), suite.c_str(), e.what()); + fflush(stdout); + return 1; + } + return 0; +} diff --git a/ci/bench/build_base.sh b/ci/bench/build_base.sh new file mode 100755 index 0000000000..a722e77447 --- /dev/null +++ b/ci/bench/build_base.sh @@ -0,0 +1,5 @@ +#!/bin/bash +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +BENCH_VARIANT=base +source ci/bench/build_variant.sh diff --git a/ci/bench/build_pr2801.sh b/ci/bench/build_pr2801.sh new file mode 100755 index 0000000000..9e5b657820 --- /dev/null +++ b/ci/bench/build_pr2801.sh @@ -0,0 +1,5 @@ +#!/bin/bash +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +BENCH_VARIANT=pr2801 +source ci/bench/build_variant.sh diff --git a/ci/bench/build_pr2804.sh b/ci/bench/build_pr2804.sh new file mode 100755 index 0000000000..d03994fe44 --- /dev/null +++ b/ci/bench/build_pr2804.sh @@ -0,0 +1,5 @@ +#!/bin/bash +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +BENCH_VARIANT=pr2804 +source ci/bench/build_variant.sh diff --git a/ci/bench/build_variant.sh b/ci/bench/build_variant.sh new file mode 100755 index 0000000000..80ee3442e8 --- /dev/null +++ b/ci/bench/build_variant.sh @@ -0,0 +1,17 @@ +#!/bin/bash +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +# Sourced by ci/bench/build_.sh with BENCH_VARIANT set. Applies the variant's patch to +# this tree (upstream main), then builds libcuvs exactly like ci/build_cpp.sh and uploads it under +# a variant-specific artifact name. +set -euo pipefail +if [[ "${BENCH_VARIANT}" != "base" ]]; then + git apply "ci/bench/patches/${BENCH_VARIANT}.patch" +fi +rapids-logger "Building libcuvs variant ${BENCH_VARIANT}" +git status --short +source ci/build_cpp.sh +RAPIDS_PACKAGE_NAME="${RAPIDS_PACKAGE_NAME}_${BENCH_VARIANT}" +export RAPIDS_PACKAGE_NAME +echo "Artifact name: ${RAPIDS_PACKAGE_NAME}" diff --git a/ci/bench/patches/pr2801.patch b/ci/bench/patches/pr2801.patch new file mode 100644 index 0000000000..4447b6cba3 --- /dev/null +++ b/ci/bench/patches/pr2801.patch @@ -0,0 +1,826 @@ +diff --git a/cpp/src/cluster/detail/kmeans.cuh b/cpp/src/cluster/detail/kmeans.cuh +index 332fed16..4478cec2 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_centroid_sums.cuh b/cpp/src/cluster/detail/kmeans_centroid_sums.cuh +new file mode 100644 +index 00000000..6248e00a +--- /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 78edf875..efedd05b 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_mg.cuh b/cpp/src/cluster/detail/kmeans_mg.cuh +index a86c30a0..4b25221f 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/tests/cluster/kmeans.cu b/cpp/tests/cluster/kmeans.cu +index 23fc2489..a6e5b4c0 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/ci/bench/patches/pr2804.patch b/ci/bench/patches/pr2804.patch new file mode 100644 index 0000000000..0466486c0c --- /dev/null +++ b/ci/bench/patches/pr2804.patch @@ -0,0 +1,289 @@ +diff --git a/cpp/src/cluster/detail/kmeans_balanced.cuh b/cpp/src/cluster/detail/kmeans_balanced.cuh +index bb96c911..881c0bd9 100644 +--- a/cpp/src/cluster/detail/kmeans_balanced.cuh ++++ b/cpp/src/cluster/detail/kmeans_balanced.cuh +@@ -43,6 +43,7 @@ + + #include + #include ++#include + #include + #include + #include +@@ -640,7 +641,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 +653,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); +@@ -768,6 +774,8 @@ __launch_bounds__((raft::WarpSize * BlockDimY)) RAFT_KERNEL + * max_cluster_size > average_size * balance_upper_tolerance + * balance_upper_tolerance > 1 + * @param[in] centroid_offset offset from the donor cluster centroid towards a donor point ++ * @param[in] donor_selection how the donor points are chosen (see above) ++ * @param[in] balancing_step index of the current EM iteration; selects the pseudo-random offset + * @param[in] mapping_op Mapping operation from T to MathT + * @param[inout] device_memory memory resource to use for temporary allocations + * +@@ -791,6 +799,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 +811,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 +845,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 the donor points do not depend on how many fits ran before in the process. + 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 +861,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 +879,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 +@@ -979,6 +992,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) { +diff --git a/cpp/src/neighbors/ivf_pq/ivf_pq_build.cuh b/cpp/src/neighbors/ivf_pq/ivf_pq_build.cuh +index 85024d2d..8f6cdccd 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_balanced.cu b/cpp/tests/cluster/kmeans_balanced.cu +index 3b7988e7..185f400d 100644 +--- a/cpp/tests/cluster/kmeans_balanced.cu ++++ b/cpp/tests/cluster/kmeans_balanced.cu +@@ -3,6 +3,7 @@ + * SPDX-License-Identifier: Apache-2.0 + */ + ++#include "../../src/cluster/detail/kmeans_balanced.cuh" + #include "../test_utils.h" + + #include +@@ -12,16 +13,20 @@ + #include + #include + #include ++#include + #include + #include + #include + + #include ++#include + + #include + + #include + ++#include ++#include + #include + #include + +@@ -301,4 +306,90 @@ KB_TEST((KmeansBalancedTest, + KmeansBalancedTestFI8I32I32_SEP, + inputsf_i32); + ++/* ++ * The balancing step must pick the same donor points for identical inputs: they may depend neither ++ * on the order in which the GPU runs the warps nor on how many fits ran before in the process. ++ * Most rows are packed in a few clusters and many clusters are empty, so that many receivers need a ++ * donor. ++ */ ++void check_adjust_centers_reproducible(cuvs::cluster::kmeans::balanced_donor_selection donor) ++{ ++ raft::resources handle; ++ auto stream = raft::resource::get_cuda_stream(handle); ++ const int64_t n_rows = 10007; ++ const int64_t dim = 8; ++ const int64_t n_lists = 200; ++ auto dataset = raft::make_device_matrix(handle, n_rows, dim); ++ raft::random::RngState rng(1234ULL); ++ raft::random::uniform(handle, rng, dataset.data_handle(), n_rows * dim, -1.0f, 1.0f); ++ std::vector h_centers(n_lists * dim); ++ for (size_t i = 0; i < h_centers.size(); i++) { ++ h_centers[i] = static_cast(i % 97) / 97.0f; ++ } ++ ++ // 95% of the rows in clusters [0, 10), the rest spread over [10, 60); clusters [60, 200) empty. ++ std::vector h_labels(n_rows); ++ std::vector h_sizes(n_lists, 0); ++ for (int64_t i = 0; i < n_rows; i++) { ++ h_labels[i] = i < n_rows * 95 / 100 ? i % 10 : 10 + i % 50; ++ h_sizes[h_labels[i]]++; ++ } ++ auto labels = raft::make_device_vector(handle, n_rows); ++ auto sizes = raft::make_device_vector(handle, n_lists); ++ raft::update_device(labels.data_handle(), h_labels.data(), n_rows, stream); ++ raft::update_device(sizes.data_handle(), h_sizes.data(), n_lists, stream); ++ ++ const bool random = donor == cuvs::cluster::kmeans::balanced_donor_selection::Random; ++ const float lower_tol = random ? 0.25f : 0.7f; ++ const float upper_tol = random ? 4.0f : 1.5f; ++ const float offset = random ? 0.05f : 0.2f; ++ ++ std::vector> results; ++ for (int run = 0; run < 3; run++) { ++ auto centers = raft::make_device_matrix(handle, n_lists, dim); ++ raft::update_device(centers.data_handle(), h_centers.data(), h_centers.size(), stream); ++ bool adjusted = ++ cuvs::cluster::kmeans::detail::adjust_centers(handle, ++ centers.data_handle(), ++ n_lists, ++ dim, ++ dataset.data_handle(), ++ n_rows, ++ labels.data_handle(), ++ sizes.data_handle(), ++ lower_tol, ++ upper_tol, ++ offset, ++ donor, ++ /*balancing_step=*/3u, ++ raft::identity_op{}, ++ rmm::mr::get_current_device_resource_ref()); ++ ASSERT_TRUE(adjusted); ++ std::vector h(centers.size()); ++ raft::update_host(h.data(), centers.data_handle(), h.size(), stream); ++ raft::resource::sync_stream(handle); ++ results.push_back(std::move(h)); ++ } ++ // Both strategies move (at least) the first 60 empty clusters. ++ for (int64_t c = 60; c < 120; c++) { ++ ASSERT_NE(0, std::memcmp(&results[0][c * dim], &h_centers[c * dim], dim * sizeof(float))) ++ << "cluster " << c << " was not moved"; ++ } ++ for (int run = 1; run < 3; run++) { ++ ASSERT_EQ( ++ 0, std::memcmp(results[0].data(), results[run].data(), results[0].size() * sizeof(float))) ++ << "adjust_centers moved the centers differently for identical inputs (call " << run << ")"; ++ } ++} ++ ++TEST(KmeansBalancedAdjustCenters, ReproducibleSizeSorted) ++{ ++ check_adjust_centers_reproducible(cuvs::cluster::kmeans::balanced_donor_selection::SizeSorted); ++} ++ ++TEST(KmeansBalancedAdjustCenters, ReproducibleRandomDonor) ++{ ++ check_adjust_centers_reproducible(cuvs::cluster::kmeans::balanced_donor_selection::Random); ++} ++ + } // namespace cuvs diff --git a/ci/bench/run.sh b/ci/bench/run.sh new file mode 100755 index 0000000000..3594e70f87 --- /dev/null +++ b/ci/bench/run.sh @@ -0,0 +1,234 @@ +#!/bin/bash +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +# Benchmark three libcuvs builds made identically in this workflow run: +# base = upstream main (this tree) +# pr2801 = main + #2801 (deterministic k-means fit) +# pr2804 = main + #2804 (balanced k-means donor selection / IVF-PQ grouping) +# The benchmark is compiled once against the public headers (identical in all three); each variant +# gets a copy whose DT_RPATH points at that variant's libcuvs.so (checked at run time). Runs are interleaved: every round runs each variant +# once, rotating the order (ABC, BCA, CAB, ...); each run does 1 warmup + `reps` timed repetitions. +# +# The workload group comes from the matrix's DEPENDENCIES field (RAPIDS_DEPENDENCIES). +set -euo pipefail + +. /opt/conda/etc/profile.d/conda.sh +GROUP=${RAPIDS_DEPENDENCIES:-lite} +VARIANTS=(base pr2801 pr2804) +REPO=$(pwd) +W=/tmp/bench +mkdir -p "${W}" +RESULTS="${W}/results.txt" +: > "${RESULTS}" + +declare -A CHANNEL +for v in "${VARIANTS[@]}"; do + CHANNEL[${v}]=$(rapids-download-from-github \ + "$(rapids-artifact-name conda_cpp libcuvs cuvs --cuda "${RAPIDS_CUDA_VERSION}")_${v}") +done +CUDA_MM=${RAPIDS_CUDA_VERSION%.*} + +nvidia-smi +nvidia-smi -q -d CLOCK,POWER,PERFORMANCE | head -80 || true +lscpu | head -20 || true + +if [[ "${GROUP}" == "tests" ]]; then + rapids-logger "Run the C++ tests of the #2804 build" + rapids-mamba-retry create --yes -n tests --override-channels --strict-channel-priority \ + -c "${CHANNEL[pr2804]}" -c rapidsai-nightly -c conda-forge \ + libcuvs libcuvs-tests "cuda-version=${CUDA_MM}" + set +u + conda activate tests + set -u + conda list | grep -E "^(libcuvs|libraft|cuda-version) " + cd "${CONDA_PREFIX}/bin/gtests/libcuvs" + EXITCODE=0 + ./CLUSTER_TEST --gtest_filter='KmeansBalancedAdjustCenters.*' --gtest_repeat=50 \ + --gtest_brief=1 || EXITCODE=1 + ./CLUSTER_TEST --gtest_filter='*KmeansBalanced*' --gtest_brief=1 || EXITCODE=1 + ./NEIGHBORS_ANN_IVF_PQ_TEST --gtest_brief=1 || EXITCODE=1 + ./NEIGHBORS_ANN_IVF_FLAT_TEST --gtest_brief=1 || EXITCODE=1 + exit ${EXITCODE} +fi + +PKGS=(libcuvs "cuda-version=${CUDA_MM}" cuda-nvcc cuda-cudart-dev libcurand-dev cxx-compiler patchelf) +rapids-logger "Create the benchmark environment" +rapids-mamba-retry create --yes -n bench --override-channels --strict-channel-priority \ + -c "${CHANNEL[base]}" -c rapidsai-nightly -c conda-forge "${PKGS[@]}" +set +u +conda activate bench +set -u +conda list | grep -E "^(libcuvs|libraft|librmm|cuda-version|cuda-nvcc|cuda-cudart) " + +rapids-logger "Extract libcuvs.so of each variant" +for v in "${VARIANTS[@]}"; do + pkg=$(find "${CHANNEL[${v}]}" -name 'libcuvs-[0-9]*.conda' | head -n 1) + /opt/conda/bin/python -c "import sys; from conda_package_handling import api; api.extract(sys.argv[1], sys.argv[2])" \ + "${pkg}" "${W}/${v}" + echo "${v}: ${pkg}" +done +md5sum "${W}"/*/lib/libcuvs.so + +rapids-logger "Compile the benchmarks" +compile() { + local out=$1 + shift + # shellcheck disable=SC2046 + nvcc -std=c++20 -O3 -arch=native --extended-lambda --expt-relaxed-constexpr -ccbin "${CXX}" -w \ + -I"${CONDA_PREFIX}/include/rapids" -I"${CONDA_PREFIX}/include" \ + $(for d in "${CONDA_PREFIX}"/targets/*/include; do echo "-I$d"; done) \ + "$@" -o "${out}" \ + -L"${CONDA_PREFIX}/lib" $(for d in "${CONDA_PREFIX}"/targets/*/lib; do echo "-L$d"; done) \ + -lcuvs -lrmm -lrapids_logger -lcudart -lcurand +} +compile "${W}/bench" ci/bench/bench.cu +for v in "${VARIANTS[@]}"; do + cp "${W}/bench" "${W}/bench_${v}" + patchelf --force-rpath --set-rpath "${W}/${v}/lib:${CONDA_PREFIX}/lib" "${W}/bench_${v}" + echo "bench_${v} rpath: $(patchelf --print-rpath "${W}/bench_${v}")" +done +# The centroid-sum microbenchmark includes #2801's header directly. +mkdir -p "${W}/tree2801" +cp -r cpp "${W}/tree2801/" +(cd "${W}/tree2801" && git apply --include='cpp/src/*' "${REPO}/ci/bench/patches/pr2801.patch") +compile "${W}/sums_2801" -I"${W}/tree2801/cpp/src" ci/bench/sums_2801.cu + +run() { + local v=$1 + shift + echo "+ ${v} $*" + # libcuvs.so comes from the binary's DT_RPATH (searched before LD_LIBRARY_PATH); its own + # dependencies come from the environment. + BENCH_EXPECT_LIB="${W}/${v}/lib/libcuvs.so" LD_LIBRARY_PATH="${CONDA_PREFIX}/lib" \ + "${W}/bench_${v}" "${v}" "$@" 2>&1 | tee -a "${RESULTS}" | grep -E "^(RESULT|SKIP|ERROR|LIB)" || true +} +# rounds R args...: R interleaved rounds over all variants, rotating the order. +rounds() { + local R=$1 + shift + for ((r = 0; r < R; r++)); do + for i in 0 1 2; do + run "${VARIANTS[$(((r + i) % 3))]}" "$@" + done + done +} + +micro() { + rapids-logger "Centroid-sum microbenchmark (#2801 kernels vs. atomics)" + LD_LIBRARY_PATH="${CONDA_PREFIX}/lib" "${W}/sums_2801" | tee -a "${RESULTS}" +} + +R=5 +START=$(date +%s) +case "${GROUP}" in + kmeans) + micro + rapids-logger "k-means fit sweep (float, device data, Array init)" + for nd in "100000 32" "100000 128" "100000 1024" "1000000 32" "1000000 128" "1000000 1024" \ + "10000000 32" "10000000 128" "10000000 1024"; do + read -r n d <<< "${nd}" + rounds ${R} kmeans n="${n}" d="${d}" k=10,100,1000,10000 + done + rapids-logger "k-means fit with sample weights" + rounds ${R} kmeans n=1000000 d=128 k=100,1000 w=1 + rounds ${R} kmeans n=1000000 d=1024 k=100 w=1 + rounds ${R} kmeans n=10000000 d=32 k=1000 w=1 + rapids-logger "k-means fit, double" + rounds ${R} kmeans dtype=f64 n=1000000 d=128 k=100,1000 + rounds ${R} kmeans dtype=f64 n=100000 d=1024 k=1000 + rounds ${R} kmeans dtype=f64 n=10000000 d=32 k=100 + rounds ${R} kmeans dtype=f64 n=1000000 d=128 k=1000 w=1 + rapids-logger "k-means fit, batched host data" + rounds ${R} kmeans n=10000000 d=128 k=1000 batch=1000000 + rounds ${R} kmeans n=1000000 d=128 k=100 batch=250000 w=1 + rapids-logger "k-means fit, k-means++ init" + rounds ${R} kmeans n=1000000 d=128 k=1000 init=kpp + rapids-logger "k-means quality vs. a float64 fit" + rounds 2 kmeans data=blobs n=1000000 d=128 k=100,1000 ref64=1 reps=2 iters=20 + rounds 2 kmeans data=blobs n=100000 d=32 k=10,100 ref64=1 reps=2 iters=20 + rounds 2 kmeans data=blobs n=1000000 d=32 k=1000 w=1 ref64=1 reps=2 iters=20 + rounds 2 kmeans data=uniform n=1000000 d=128 k=1000 ref64=1 reps=2 iters=20 + ;; + ivf) + micro + rapids-logger "Balanced k-means fit / predict" + rounds ${R} balanced n=1000000 d=128 k=1024,4096 predict=1 + rounds ${R} balanced n=2500000 d=128 k=4096 + rounds ${R} balanced n=10000000 d=128 k=16384 predict=1 reps=1 + rounds ${R} balanced n=1000000 d=768 k=1024 + rounds ${R} balanced n=100000 d=32 k=4096 + rounds ${R} balanced dtype=i8 n=1000000 d=128 k=1024 + rounds ${R} balanced n=65536 d=2 k=256 calls=64 donor=random iters=25 + rounds ${R} balanced n=65536 d=2 k=256 calls=64 donor=size iters=25 + rapids-logger "Index builds (+ recall@10, search QPS)" + rounds ${R} ivf n=1000000 d=128 index=flat:1024,pq:1024:64,pqc:1024:32,sq:1024,rbq:1024,cagra:64 + rounds ${R} ivf n=5000000 d=128 index=flat:4096,pq:4096:64,sq:4096,rbq:4096 + rounds ${R} ivf n=10000000 d=96 index=flat:16384,pq:16384:48 reps=1 + rounds ${R} ivf n=1000000 d=768 index=flat:1024,pq:1024:96 + ;; + profile) + micro + rapids-logger "Balanced k-means, 64 small fits (IVF-PQ codebook-like)" + rounds ${R} balanced n=65536 d=2 k=256 calls=64 donor=size iters=25 distinct=1 + rounds ${R} balanced n=65536 d=2 k=256 calls=64 donor=random iters=25 distinct=1 + rounds ${R} balanced n=65536 d=2 k=256 calls=64 donor=size iters=25 + rounds ${R} balanced n=65536 d=2 k=256 calls=64 donor=random iters=25 + rapids-logger "Install Nsight Systems" + rapids-mamba-retry install --yes -n bench -c conda-forge nsight-systems + i=0 + for cfg in "n=100000 d=128 k=100" "n=100000 d=32 k=1000" "n=100000 d=1024 k=10" \ + "n=1000000 d=128 k=10"; do + i=$((i + 1)) + for v in base pr2801; do + rapids-logger "nsys: ${v} kmeans ${cfg}" + # shellcheck disable=SC2086 + BENCH_EXPECT_LIB="${W}/${v}/lib/libcuvs.so" LD_LIBRARY_PATH="${CONDA_PREFIX}/lib" \ + nsys profile -t cuda --stats=false -f true -o "${W}/prof_${v}_${i}" \ + "${W}/bench_${v}" "${v}" kmeans ${cfg} reps=3 warmup=1 | grep -E "^(RESULT|ERROR|LIB)" || true + nsys stats -r cuda_gpu_kern_sum,cuda_api_sum --format table "${W}/prof_${v}_${i}.nsys-rep" \ + 2>/dev/null | cut -c1-220 | head -n 70 || true + done + done + ;; + profile2) + rapids-logger "Install Nsight Systems" + rapids-mamba-retry install --yes -n bench -c conda-forge nsight-systems + i=0 + for cfg in "donor=size" "donor=random distinct=1"; do + i=$((i + 1)) + for v in base pr2804 base pr2804; do + rapids-logger "nsys: ${v} balanced 64 calls ${cfg}" + # shellcheck disable=SC2086 + BENCH_EXPECT_LIB="${W}/${v}/lib/libcuvs.so" LD_LIBRARY_PATH="${CONDA_PREFIX}/lib" \ + nsys profile -t cuda --stats=false -f true -o "${W}/prof_${v}_${i}" \ + "${W}/bench_${v}" "${v}" balanced n=65536 d=2 k=256 calls=64 iters=25 ${cfg} \ + reps=1 warmup=0 | grep -E "^(RESULT|ERROR|LIB)" || true + nsys stats -r cuda_gpu_kern_sum,cuda_api_sum --format table "${W}/prof_${v}_${i}.nsys-rep" \ + 2>/dev/null | cut -c1-250 | head -n 70 || true + done + done + ;; + lite) + micro + rounds ${R} kmeans n=1000000 d=128 k=100,1000,10000 reps=1 + rounds ${R} kmeans n=10000000 d=32 k=1000 reps=1 + rounds ${R} kmeans n=1000000 d=1024 k=1000 reps=1 + rounds ${R} kmeans n=1000000 d=128 k=1000 w=1 reps=1 + rounds ${R} kmeans dtype=f64 n=1000000 d=128 k=1000 reps=1 + rounds ${R} kmeans n=10000000 d=128 k=1000 batch=1000000 reps=1 + rounds ${R} balanced n=1000000 d=128 k=1024 reps=1 + rounds ${R} balanced n=65536 d=2 k=256 calls=64 donor=random iters=25 reps=1 + rounds ${R} ivf n=1000000 d=128 index=flat:1024,pq:1024:64,pqc:1024:32,cagra:64 reps=1 + rounds ${R} ivf n=5000000 d=128 index=flat:4096,pq:4096:64 reps=1 + ;; + *) + echo "unknown group ${GROUP}" + exit 1 + ;; +esac +echo "Benchmark time: $(($(date +%s) - START)) s" +nvidia-smi -q -d CLOCK,PERFORMANCE | head -60 || true + +rapids-logger "Summary (${GROUP})" +/opt/conda/bin/python ci/bench/summarize.py "${RESULTS}" diff --git a/ci/bench/summarize.py b/ci/bench/summarize.py new file mode 100644 index 0000000000..f0e362e0b3 --- /dev/null +++ b/ci/bench/summarize.py @@ -0,0 +1,96 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Summarize RESULT lines of ci/bench/bench.cu: per configuration and variant, the median and +min-max of the timings, the change of the median vs. base, and the quality metrics.""" + +import statistics +import sys +from collections import OrderedDict, defaultdict + +VARIANTS = ["base", "pr2801", "pr2804"] + + +def parse(path): + data = OrderedDict() + for line in open(path, errors="replace"): + if not line.startswith("RESULT "): + continue + left, _, right = line[len("RESULT ") :].strip().partition(" | ") + m = dict(kv.split("=", 1) for kv in right.split() if "=" in kv) + v = m.pop("variant") + data.setdefault(left, defaultdict(list))[v].append(m) + return data + + +def med(xs): + return statistics.median(xs) if xs else float("nan") + + +def fmt_t(xs): + if not xs: + return "-" + return f"{med(xs):.4g} ({min(xs):.4g}-{max(xs):.4g})" + + +def main(path): + data = parse(path) + print("| config | metric | " + " | ".join(VARIANTS) + " | pr2801 vs base | pr2804 vs base |") + print("|---|---|" + "---|" * len(VARIANTS) + "---|---|") + for cfg, per_v in data.items(): + key = "per_iter_ms" if cfg.startswith("kmeans_fit") else "ms" + if cfg.startswith("search"): + key = "qps" + t = {v: [float(m[key]) for m in per_v.get(v, []) if key in m] for v in VARIANTS} + base = med(t["base"]) + deltas = [] + for v in ("pr2801", "pr2804"): + if t[v] and t["base"]: + deltas.append(f"{100 * (med(t[v]) / base - 1):+.1f}%") + else: + deltas.append("-") + cells = [f"{fmt_t(t[v])} n={len(t[v])}" for v in VARIANTS] + print(f"| {cfg} | {key} | " + " | ".join(cells) + " | " + " | ".join(deltas) + " |") + extra = [] + for metric in ( + "ms", + "iters", + "inertia", + "recall", + "inertia64", + "max_size", + "center_maxrel_vs64", + "center_maxabs_vs64", + "iters64", + ): + if metric == key: + continue + vals = {v: [m[metric] for m in per_v.get(v, []) if metric in m] for v in VARIANTS} + if not any(vals.values()): + continue + row = [] + for v in VARIANTS: + xs = vals[v] + if not xs: + row.append("-") + continue + try: + f = [float(x) for x in xs] + row.append( + f"{med(f):.6g}" if min(f) == max(f) else f"{med(f):.6g} ({min(f):.6g}-{max(f):.6g})" + ) + except ValueError: + row.append(",".join(sorted(set(xs)))) + extra.append(f"| {cfg} | {metric} | " + " | ".join(row) + " | | |") + hashes = {v: len({m["hash"] for m in per_v.get(v, []) if "hash" in m}) for v in VARIANTS} + if any(hashes.values()): + extra.append( + f"| {cfg} | distinct results | " + + " | ".join(f"{hashes[v]} of {len(per_v.get(v, []))}" for v in VARIANTS) + + " | | |" + ) + for e in extra: + print(e) + + +if __name__ == "__main__": + main(sys.argv[1]) diff --git a/ci/bench/sums_2801.cu b/ci/bench/sums_2801.cu new file mode 100644 index 0000000000..80874beca4 --- /dev/null +++ b/ci/bench/sums_2801.cu @@ -0,0 +1,252 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +// Microbenchmarks for the k-means centroid update step: +// - raft::linalg::reduce_rows_by_key + reduce_cols_by_key (atomics) vs. +// sums_by_label::weighted_sums_by_label (fixed order), per call; +// - cub::DeviceReduce sum with run_to_run vs. gpu_to_gpu determinism. +#include "cluster/detail/kmeans_centroid_sums.cuh" + +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include + +#include +#include + +#include +#include +#include +#include +#include +#include +#include + +template +float time_ms(cudaStream_t stream, int reps, F f) +{ + cudaEvent_t e0, e1; + cudaEventCreate(&e0); + cudaEventCreate(&e1); + f(); + cudaStreamSynchronize(stream); + cudaEventRecord(e0, stream); + for (int r = 0; r < reps; ++r) { + f(); + } + cudaEventRecord(e1, stream); + cudaEventSynchronize(e1); + float ms = 0; + cudaEventElapsedTime(&ms, e0, e1); + cudaEventDestroy(e0); + cudaEventDestroy(e1); + return ms / reps; +} + +// Per-cluster weighted sums and weight sums: base (raft reduce_rows_by_key + reduce_cols_by_key, +// floating-point atomics) vs. #2801 (sums_by_label::weighted_sums_by_label, fixed order). Times +// both, checks run-to-run reproducibility (two calls each) and, when `check` is set, compares the +// cluster means (sum / weight) with a long double host reference. +// pattern 0: X ~ U(-1, 1); pattern 1: X ~ 100 + U(-1, 1) (large common offset) +// labels: uniform random, or skewed (half of the rows in cluster 0) +void bench_sums(int64_t n, int64_t d, int64_t k, int reps, int pattern, bool skew, bool check) +{ + using T = float; + using IndexT = int64_t; + raft::resources handle; + auto stream = raft::resource::get_cuda_stream(handle).get(); + rmm::device_uvector X(n * d, stream), W(n, stream), sums(k * d, stream), wsum(k, stream); + rmm::device_uvector labels(n, stream); + raft::random::RngState rng(42); + raft::random::uniform(handle, rng, X.data(), n * d, T(-1), T(1)); + if (pattern == 1) { + thrust::transform(thrust::cuda::par.on(stream), + X.begin(), + X.end(), + X.begin(), + [] __device__(T x) { return x + T(100); }); + } + raft::random::uniform(handle, rng, W.data(), n, T(0.5), T(1.5)); + raft::random::uniformInt(handle, rng, labels.data(), n, 0, int(k)); + if (skew) { + thrust::transform(thrust::cuda::par.on(stream), + labels.begin(), + labels.begin() + n / 2, + labels.begin(), + [] __device__(int) { return 0; }); + } + rmm::device_uvector ws(0, stream), keys_char(n, stream); + + auto run_atomic = [&] { + raft::linalg::reduce_rows_by_key( + X.data(), d, labels.data(), W.data(), keys_char.data(), n, d, k, sums.data(), stream, true); + raft::linalg::reduce_cols_by_key( + W.data(), labels.data(), wsum.data(), IndexT(1), n, k, stream, true); + }; + auto run_fixed = [&] { + cuvs::cluster::kmeans::detail::sums_by_label::weighted_sums_by_label( + handle, X.data(), W.data(), labels.data(), n, d, k, sums.data(), wsum.data(), ws, true); + }; + auto means = [&] { + std::vector hs(k * d), hw(k); + raft::copy(hs.data(), sums.data(), k * d, stream); + raft::copy(hw.data(), wsum.data(), k, stream); + cudaStreamSynchronize(stream); + std::vector m(k * d); + for (int64_t c = 0; c < k; c++) { + for (int64_t j = 0; j < d; j++) { + m[c * d + j] = hw[c] > 0 ? double(hs[c * d + j]) / double(hw[c]) : 0.0; + } + } + return std::make_pair(hs, m); + }; + + float t_atomic = time_ms(stream, reps, run_atomic); + auto [a1, ma] = means(); + run_atomic(); + auto [a2, ma2] = means(); + float t_fixed = time_ms(stream, reps, run_fixed); + auto [f1, mf] = means(); + run_fixed(); + auto [f2, mf2] = means(); + + std::string acc; + if (check) { + auto hX = std::vector(n * d); + auto hW = std::vector(n); + auto hL = std::vector(n); + raft::copy(hX.data(), X.data(), n * d, stream); + raft::copy(hW.data(), W.data(), n, stream); + raft::copy(hL.data(), labels.data(), n, stream); + cudaStreamSynchronize(stream); + std::vector rs(k * d, 0), rw(k, 0); + for (int64_t i = 0; i < n; i++) { + rw[hL[i]] += hW[i]; + for (int64_t j = 0; j < d; j++) { + rs[hL[i] * d + j] += (long double)hW[i] * hX[i * d + j]; + } + } + auto err = [&](const std::vector& m) { + double max_abs = 0, max_rel = 0; + for (int64_t c = 0; c < k; c++) { + if (rw[c] == 0) { continue; } + for (int64_t j = 0; j < d; j++) { + double ref = double(rs[c * d + j] / rw[c]); + double e = std::abs(m[c * d + j] - ref); + max_abs = std::max(max_abs, e); + max_rel = std::max(max_rel, e / std::max(std::abs(ref), 1e-30)); + } + } + return std::make_pair(max_abs, max_rel); + }; + auto [ea, ra] = err(ma); + auto [ef, rf] = err(mf); + char buf[256]; + snprintf(buf, + sizeof(buf), + " mean_maxabs_vs_f64: atomic=%.3g fixed=%.3g mean_maxrel_vs_f64: atomic=%.3g " + "fixed=%.3g", + ea, + ef, + ra, + rf); + acc = buf; + } + printf( + "MICRO centroid_sums n=%ld d=%ld k=%ld pattern=%s labels=%s base_atomic_ms=%.3f " + "pr2801_fixed_order_ms=%.3f ratio=%.2f reproducible: atomic=%d fixed=%d%s\n", + n, + d, + k, + pattern == 1 ? "offset100" : "uniform", + skew ? "skewed" : "uniform", + t_atomic, + t_fixed, + t_fixed / t_atomic, + int(std::memcmp(a1.data(), a2.data(), a1.size() * sizeof(T)) == 0), + int(std::memcmp(f1.data(), f2.data(), f1.size() * sizeof(T)) == 0), + acc.c_str()); + fflush(stdout); +} + +template +void reduce(raft::resources const& handle, const float* in, int64_t n, float* out, Determinism det) +{ + auto env = cuda::std::execution::env{ + cuda::stream_ref{raft::resource::get_cuda_stream(handle)}, + cuda::execution::require(det), + cuda::std::execution::prop{cuda::mr::get_memory_resource, + raft::resource::get_workspace_resource_ref(handle)}}; + RAFT_CUDA_TRY(cub::DeviceReduce::Reduce(in, out, n, cuda::std::plus<>{}, 0.0f, env)); +} + +void bench_reduce(int64_t n, int reps) +{ + raft::resources handle; + auto stream = raft::resource::get_cuda_stream(handle).get(); + rmm::device_uvector x(n, stream), out(1, stream); + raft::random::RngState rng(42); + raft::random::uniform(handle, rng, x.data(), n, -1.0f, 1.0f); + float t_rr = time_ms(stream, reps, [&] { + reduce(handle, x.data(), n, out.data(), cuda::execution::determinism::run_to_run); + }); + float t_g2g = time_ms(stream, reps, [&] { + reduce(handle, x.data(), n, out.data(), cuda::execution::determinism::gpu_to_gpu); + }); + printf("BENCH reduce_sum n=%ld run_to_run_ms=%.4f gpu_to_gpu_ms=%.4f ratio=%.2f\n", + n, + t_rr, + t_g2g, + t_g2g / t_rr); + fflush(stdout); +} + +int main() +{ + // Timing (and reproducibility) on k-means fit shapes. + for (auto [n, d, k] : std::vector>{{100000, 32, 10}, + {100000, 32, 1000}, + {100000, 128, 10}, + {100000, 128, 100}, + {100000, 128, 1000}, + {100000, 1024, 10}, + {100000, 1024, 100}, + {1000000, 32, 100}, + {1000000, 32, 1000}, + {1000000, 128, 10}, + {1000000, 128, 100}, + {1000000, 128, 1000}, + {1000000, 128, 10000}, + {1000000, 1024, 1000}, + {10000000, 32, 1000}, + {10000000, 128, 1000}, + {10000000, 2, 8}}) { + bench_sums(n, d, k, 20, 0, false, false); + } + bench_sums(1000000, 128, 100, 20, 0, true, false); + bench_sums(10000000, 32, 1000, 20, 0, true, false); + // Accuracy of the cluster means vs. a long double reference. + for (int pattern : {0, 1}) { + for (bool skew : {false, true}) { + bench_sums(1000000, 32, 10, 5, pattern, skew, true); + bench_sums(1000000, 32, 1000, 5, pattern, skew, true); + bench_sums(4000000, 16, 100, 5, pattern, skew, true); + } + } + for (int64_t n : {1000, 100000, 1000000, 10000000, 100000000}) { + bench_reduce(n, 50); + } + return 0; +}