From 41b0713736eb5a8062cc9ab7979ecf9aa933f83e Mon Sep 17 00:00:00 2001 From: divyegala Date: Wed, 17 Jun 2026 20:58:30 +0000 Subject: [PATCH 01/82] add example --- conda/recipes/libcuvs/recipe.yaml | 5 + cpp/tests/CMakeLists.txt | 2 + cpp/tests/cutile/CMakeLists.txt | 23 ++++ cpp/tests/cutile/cutile_vector_add.cu | 128 ++++++++++++++++++ cpp/tests/cutile/export_vector_add_cubin.py | 101 ++++++++++++++ cpp/tests/cutile/generate_cutile_cubins.cmake | 90 ++++++++++++ cpp/tests/cutile/vector_add_kernel.py | 17 +++ dependencies.yaml | 3 + 8 files changed, 369 insertions(+) create mode 100644 cpp/tests/cutile/CMakeLists.txt create mode 100644 cpp/tests/cutile/cutile_vector_add.cu create mode 100644 cpp/tests/cutile/export_vector_add_cubin.py create mode 100644 cpp/tests/cutile/generate_cutile_cubins.cmake create mode 100644 cpp/tests/cutile/vector_add_kernel.py diff --git a/conda/recipes/libcuvs/recipe.yaml b/conda/recipes/libcuvs/recipe.yaml index aa7a37db44..93f31f8cf2 100644 --- a/conda/recipes/libcuvs/recipe.yaml +++ b/conda/recipes/libcuvs/recipe.yaml @@ -80,6 +80,7 @@ cache: - cuda-cudart-dev - cuda-nvrtc-dev - cuda-profiler-api + - cutile-python - libcublas-dev - libcurand-dev - libcusolver-dev @@ -117,6 +118,7 @@ outputs: - cuda-cudart-dev - cuda-nvrtc-dev - cuda-profiler-api + - cutile-python - libcublas-dev - libcurand-dev - libcusolver-dev @@ -179,6 +181,7 @@ outputs: - cuda-cudart-dev - cuda-nvrtc-dev - cuda-profiler-api + - cutile-python - libcublas-dev - libcurand-dev - libcusolver-dev @@ -240,6 +243,7 @@ outputs: - cuda-cudart-dev - cuda-nvrtc-dev - cuda-profiler-api + - cutile-python - libcublas-dev - libcurand-dev - libcusolver-dev @@ -299,6 +303,7 @@ outputs: - openblas # required by some CPU algos in benchmarks - cuda-cudart-dev - cuda-profiler-api + - cutile-python - libcublas-dev - libcurand-dev - libcusolver-dev diff --git a/cpp/tests/CMakeLists.txt b/cpp/tests/CMakeLists.txt index 9b96f94bf0..ba6ed6e0e7 100644 --- a/cpp/tests/CMakeLists.txt +++ b/cpp/tests/CMakeLists.txt @@ -386,6 +386,8 @@ ConfigureTest( PERCENT 100 ) +add_subdirectory(cutile) + # ################################################################################################## # Install tests #################################################################################### # ################################################################################################## diff --git a/cpp/tests/cutile/CMakeLists.txt b/cpp/tests/cutile/CMakeLists.txt new file mode 100644 index 0000000000..989c8137d0 --- /dev/null +++ b/cpp/tests/cutile/CMakeLists.txt @@ -0,0 +1,23 @@ +# ============================================================================= +# cmake-format: off +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. +# SPDX-License-Identifier: Apache-2.0 +# cmake-format: on +# ============================================================================= + +include("${CMAKE_CURRENT_LIST_DIR}/generate_cutile_cubins.cmake") + +generate_cutile_vector_add_cubins(CUTILE_GENERATED_INCLUDE_DIR) + +ConfigureTest( + NAME CUTILE_VECTOR_ADD_TEST + PATH "${CMAKE_CURRENT_LIST_DIR}/cutile_vector_add.cu" + GPUS 1 + PERCENT 100 +) + +add_dependencies(CUTILE_VECTOR_ADD_TEST cutile_vector_add_cubins) + +target_include_directories( + CUTILE_VECTOR_ADD_TEST PRIVATE "${CUTILE_GENERATED_INCLUDE_DIR}" +) diff --git a/cpp/tests/cutile/cutile_vector_add.cu b/cpp/tests/cutile/cutile_vector_add.cu new file mode 100644 index 0000000000..77a5e51311 --- /dev/null +++ b/cpp/tests/cutile/cutile_vector_add.cu @@ -0,0 +1,128 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include "../test_utils.cuh" + +#include "vector_add_kernel_symbol.h" +#include "vector_add_sm_100_cubin.h" +#include "vector_add_sm_120_cubin.h" +#include "vector_add_sm_80_cubin.h" +#include "vector_add_sm_86_cubin.h" +#include "vector_add_sm_90_cubin.h" + +#include + +#include + +namespace cuvs { +namespace { + +struct EmbeddedCubin { + int cc_major; + int cc_minor; + const unsigned char* data; + size_t size; +}; + +// Lookup table for cubins built at configure time (see export_vector_add_cubin.py). +constexpr EmbeddedCubin kEmbeddedCubins[] = { + {8, 0, vector_add_sm_80_cubin, sizeof(vector_add_sm_80_cubin)}, + {8, 6, vector_add_sm_86_cubin, sizeof(vector_add_sm_86_cubin)}, + {9, 0, vector_add_sm_90_cubin, sizeof(vector_add_sm_90_cubin)}, + {10, 0, vector_add_sm_100_cubin, sizeof(vector_add_sm_100_cubin)}, + {12, 0, vector_add_sm_120_cubin, sizeof(vector_add_sm_120_cubin)}, +}; + +const EmbeddedCubin* find_embedded_cubin(int cc_major, int cc_minor) +{ + for (const auto& entry : kEmbeddedCubins) { + if (entry.cc_major == cc_major && entry.cc_minor == cc_minor) { return &entry; } + } + // Fall back to a cubin for the same major version (e.g. minor SKUs within a generation). + for (const auto& entry : kEmbeddedCubins) { + if (entry.cc_major == cc_major) { return &entry; } + } + return nullptr; +} + +class CutileVectorAddTest : public ::testing::Test { + protected: + void SetUp() override + { + int device = 0; + RAFT_CUDA_TRY(cudaGetDevice(&device)); + RAFT_CUDA_TRY( + cudaDeviceGetAttribute(&cc_major_, cudaDevAttrComputeCapabilityMajor, device)); + RAFT_CUDA_TRY( + cudaDeviceGetAttribute(&cc_minor_, cudaDevAttrComputeCapabilityMinor, device)); + } + + int cc_major_{}; + int cc_minor_{}; +}; + +} // namespace + +TEST_F(CutileVectorAddTest, EmbeddedCubinVectorAdd) +{ + const EmbeddedCubin* cubin = find_embedded_cubin(cc_major_, cc_minor_); + ASSERT_NE(cubin, nullptr) + << "No embedded cuTile cubin for compute capability " << cc_major_ << "." << cc_minor_; + + cudaLibrary_t library{}; + ASSERT_EQ(cudaSuccess, + cudaLibraryLoadData( + &library, cubin->data, nullptr, nullptr, 0, nullptr, nullptr, 0)) + << "cudaLibraryLoadData failed: " << cudaGetErrorString(cudaGetLastError()); + + cudaKernel_t kernel{}; + ASSERT_EQ(cudaSuccess, + cudaLibraryGetKernel(&kernel, library, CUTILE_VECTOR_ADD_KERNEL_SYMBOL)) + << "cudaLibraryGetKernel failed: " << cudaGetErrorString(cudaGetLastError()); + + constexpr int kN = 1024; + constexpr int kTile = 256; + constexpr int kGridDim = (kN + kTile - 1) / kTile; + + float *d_a = nullptr, *d_b = nullptr, *d_c = nullptr; + RAFT_CUDA_TRY(cudaMalloc(&d_a, kN * sizeof(float))); + RAFT_CUDA_TRY(cudaMalloc(&d_b, kN * sizeof(float))); + RAFT_CUDA_TRY(cudaMalloc(&d_c, kN * sizeof(float))); + + std::vector h_a(kN), h_b(kN); + for (int i = 0; i < kN; ++i) { + h_a[i] = static_cast(i); + h_b[i] = static_cast(i * 2); + } + RAFT_CUDA_TRY(cudaMemcpy(d_a, h_a.data(), kN * sizeof(float), cudaMemcpyHostToDevice)); + RAFT_CUDA_TRY(cudaMemcpy(d_b, h_b.data(), kN * sizeof(float), cudaMemcpyHostToDevice)); + RAFT_CUDA_TRY(cudaMemset(d_c, 0, kN * sizeof(float))); + + int64_t shape = kN; + int64_t stride = 1; + void* kernel_args[] = { + &d_a, &shape, &stride, &d_b, &shape, &stride, &d_c, &shape, &stride, + }; + + dim3 grid(kGridDim); + dim3 block(1); + ASSERT_EQ(cudaSuccess, cudaLaunchKernel(kernel, grid, block, kernel_args, 0, 0)) + << "cudaLaunchKernel failed: " << cudaGetErrorString(cudaGetLastError()); + RAFT_CUDA_TRY(cudaDeviceSynchronize()); + + std::vector h_c(kN); + RAFT_CUDA_TRY(cudaMemcpy(h_c.data(), d_c, kN * sizeof(float), cudaMemcpyDeviceToHost)); + + for (int i = 0; i < kN; ++i) { + ASSERT_FLOAT_EQ(h_a[i] + h_b[i], h_c[i]) << "@" << i; + } + + RAFT_CUDA_TRY(cudaFree(d_a)); + RAFT_CUDA_TRY(cudaFree(d_b)); + RAFT_CUDA_TRY(cudaFree(d_c)); + RAFT_CUDA_TRY(cudaLibraryUnload(library)); +} + +} // namespace cuvs diff --git a/cpp/tests/cutile/export_vector_add_cubin.py b/cpp/tests/cutile/export_vector_add_cubin.py new file mode 100644 index 0000000000..bf40a4ad80 --- /dev/null +++ b/cpp/tests/cutile/export_vector_add_cubin.py @@ -0,0 +1,101 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. +# SPDX-License-Identifier: Apache-2.0 +"""Export the cuTile vector-add kernel to a cubin for a single GPU target.""" + +from __future__ import annotations + +import argparse +import sys +from pathlib import Path + +import cuda.tile as ct +from cuda.tile.compilation import ( + ArrayConstraint, + CallingConvention, + ConstantConstraint, + KernelSignature, + export_kernel, +) + +from vector_add_kernel import TILE_SIZE, vector_add + +# cuTile / tileiras gpu_code values used at build time. These correspond to the +# cuvs library CUDA 13 real targets as follows (tileiras has no sm_*a/sm_*f names): +# sm_80 -> 80-real +# sm_86 -> 86-real +# sm_90 -> 90a-real +# sm_100 -> 100f-real +# sm_120 -> 120a-real +SUPPORTED_GPU_CODES = ("sm_80", "sm_86", "sm_90", "sm_100", "sm_120") + + +def _kernel_signature() -> KernelSignature: + array = ArrayConstraint( + ct.float32, + 1, + index_dtype=ct.int64, + stride_lower_bound_incl=0, + alias_groups=(), + may_alias_internally=False, + stride_constant=(1,), + ) + return KernelSignature( + parameters=[array, array, array, ConstantConstraint(TILE_SIZE)], + calling_convention=CallingConvention.cutile_python_v1(), + ).with_mangled_symbol("vector_add") + + +def export_cubin(output_file: Path, gpu_code: str, symbol_header: Path | None) -> str: + if gpu_code not in SUPPORTED_GPU_CODES: + raise ValueError( + f"Unsupported gpu_code {gpu_code!r}; expected one of {SUPPORTED_GPU_CODES}" + ) + + signature = _kernel_signature() + export_kernel( + vector_add, + signatures=[signature], + output_file=str(output_file), + gpu_code=gpu_code, + output_format="cubin", + ) + + if symbol_header is not None: + symbol_header.write_text( + "\n".join( + [ + "// Generated by export_vector_add_cubin.py; do not edit.", + "#pragma once", + f'#define CUTILE_VECTOR_ADD_KERNEL_SYMBOL "{signature.symbol}"', + "", + ] + ) + ) + + return signature.symbol + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("output_file", type=Path, help="Output cubin path") + parser.add_argument( + "--gpu-code", + required=True, + choices=SUPPORTED_GPU_CODES, + help="tileiras / export_kernel target (e.g. sm_120)", + ) + parser.add_argument( + "--symbol-header", + type=Path, + default=None, + help="Optional header that defines CUTILE_VECTOR_ADD_KERNEL_SYMBOL", + ) + args = parser.parse_args() + + symbol = export_cubin(args.output_file, args.gpu_code, args.symbol_header) + print(symbol) + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/cpp/tests/cutile/generate_cutile_cubins.cmake b/cpp/tests/cutile/generate_cutile_cubins.cmake new file mode 100644 index 0000000000..3425b03028 --- /dev/null +++ b/cpp/tests/cutile/generate_cutile_cubins.cmake @@ -0,0 +1,90 @@ +# ============================================================================= +# cmake-format: off +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. +# SPDX-License-Identifier: Apache-2.0 +# cmake-format: on +# ============================================================================= + +include_guard(GLOBAL) + +# Build-time cuTile cubin targets. Maps to cuvs CUDA 13 -real library arches (75-real omitted). +set(CUTILE_VECTOR_ADD_GPU_CODES sm_80 sm_86 sm_90 sm_100 sm_120) + +function(generate_cutile_vector_add_cubins output_include_dir_var) + find_package(Python3 REQUIRED COMPONENTS Interpreter) + find_package(CUDAToolkit REQUIRED) + + find_program( + CUTILE_BIN2C + NAMES bin2c + PATHS ${CUDAToolkit_BIN_DIR} + REQUIRED + ) + + execute_process( + COMMAND "${Python3_EXECUTABLE}" -c "import cuda.tile" + RESULT_VARIABLE _cutile_import_result + OUTPUT_QUIET + ERROR_QUIET + ) + if(NOT _cutile_import_result EQUAL 0) + message( + FATAL_ERROR + "cuda.tile (cuTile Python) is required to build CUTILE_VECTOR_ADD_TEST. " + "Install it in the active Python environment, e.g. pip install cuda-tile[tileiras]." + ) + endif() + + set(_cutile_source_dir "${CMAKE_CURRENT_FUNCTION_LIST_DIR}") + set(_cutile_binary_dir "${CMAKE_CURRENT_BINARY_DIR}/cutile_generated") + file(MAKE_DIRECTORY "${_cutile_binary_dir}") + + set(_symbol_header "${_cutile_binary_dir}/vector_add_kernel_symbol.h") + set(_first_gpu_code TRUE) + + foreach(_gpu_code IN LISTS CUTILE_VECTOR_ADD_GPU_CODES) + set(_cubin_file "${_cutile_binary_dir}/vector_add_${_gpu_code}.cubin") + set(_cubin_header "${_cutile_binary_dir}/vector_add_${_gpu_code}_cubin.h") + + if(_first_gpu_code) + set(_symbol_arg --symbol-header "${_symbol_header}") + set(_cubin_outputs "${_cubin_file}" "${_symbol_header}") + set(_first_gpu_code FALSE) + else() + set(_symbol_arg) + set(_cubin_outputs "${_cubin_file}") + endif() + + add_custom_command( + OUTPUT ${_cubin_outputs} + COMMAND + "${Python3_EXECUTABLE}" "${_cutile_source_dir}/export_vector_add_cubin.py" + "${_cubin_file}" --gpu-code "${_gpu_code}" ${_symbol_arg} + DEPENDS "${_cutile_source_dir}/export_vector_add_cubin.py" + "${_cutile_source_dir}/vector_add_kernel.py" + COMMENT "Exporting cuTile vector_add cubin for ${_gpu_code}" + VERBATIM + ) + + add_custom_command( + OUTPUT "${_cubin_header}" + COMMAND "${CUTILE_BIN2C}" --const --name "vector_add_${_gpu_code}_cubin" --static + "${_cubin_file}" > "${_cubin_header}" + DEPENDS "${_cubin_file}" + COMMENT "Embedding vector_add ${_gpu_code} cubin via bin2c" + VERBATIM + ) + + list(APPEND _generated_headers "${_cubin_header}") + endforeach() + + add_custom_target( + cutile_vector_add_cubins + DEPENDS "${_symbol_header}" ${_generated_headers} + ) + + set(${output_include_dir_var} + "${_cutile_binary_dir}" + PARENT_SCOPE + ) +endfunction() diff --git a/cpp/tests/cutile/vector_add_kernel.py b/cpp/tests/cutile/vector_add_kernel.py new file mode 100644 index 0000000000..46b7a607c6 --- /dev/null +++ b/cpp/tests/cutile/vector_add_kernel.py @@ -0,0 +1,17 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. +# SPDX-License-Identifier: Apache-2.0 +"""cuTile Python vector-add kernel used by the embedded-cubin example test.""" + +from __future__ import annotations + +import cuda.tile as ct + +TILE_SIZE = 256 + + +@ct.kernel +def vector_add(a, b, c, TILE_SIZE: ct.Constant): + bid = ct.bid(0) + ta = ct.load(a, bid, TILE_SIZE) + tb = ct.load(b, bid, TILE_SIZE) + ct.store(c, bid, ta + tb) diff --git a/dependencies.yaml b/dependencies.yaml index 744e4d9227..756041e60c 100644 --- a/dependencies.yaml +++ b/dependencies.yaml @@ -395,6 +395,7 @@ dependencies: - cuda-nvrtc-dev - cuda-nvtx-dev - cuda-profiler-api + - cutile-python - libcublas-dev - libcurand-dev - libcusolver-dev @@ -430,12 +431,14 @@ dependencies: packages: - &ctk_cu13 cuda-toolkit[cublas,curand,cusolver,cusparse,nvrtc]==13.* - &nvjitlink_cu13 nvidia-nvjitlink>=13.0,<14 + - &cutile_cu13 cuda-tile[tileiras] # if no matching matrix selectors passed, list the CUDA 13 requirement # (just as a source of documentation, as this populates pyproject.toml in source control) - matrix: packages: - *ctk_cu13 - *nvjitlink_cu13 + - *cutile_cu13 depends_on_cudart: common: - output_types: conda From b10c02ca5a5ef094ca892c6d32f6f14d6d63447f Mon Sep 17 00:00:00 2001 From: divyegala Date: Wed, 24 Jun 2026 16:34:11 +0000 Subject: [PATCH 02/82] initial integration --- cpp/CMakeLists.txt | 50 ++- .../modules/generate_cutile_kernels.cmake | 315 ++++++++++++++++++ cpp/cmake/modules/register_cubin.cpp.in | 22 ++ cpp/cmake/modules/register_tileir.cpp.in | 22 ++ .../cuvs/detail/jit_lto/AlgorithmPlanner.hpp | 63 +++- .../cuvs/detail/jit_lto/FragmentEntry.hpp | 63 ++++ .../cuvs/detail/jit_lto/cutile_arch_tags.hpp | 52 +++ .../cuvs/detail/jit_lto/cutile_module.hpp | 75 +++++ .../fused_distance_nn/fused_1nn_fragments.hpp | 21 ++ .../cuvs/detail/jit_lto/tileir_compat.hpp | 99 ++++++ cpp/src/detail/jit_lto/AlgorithmPlanner.cpp | 103 ++---- .../detail/jit_lto/LTOAlgorithmPlanner.cpp | 76 +++++ .../detail/jit_lto/TileAlgorithmPlanner.cpp | 38 +++ cpp/src/distance/detail/fused_distance_nn.cuh | 15 + .../cutile/export_fused_1nn.py | 136 ++++++++ .../cutile/fused_1nn_cutile_cubin_matrix.json | 40 +++ .../fused_1nn_cutile_tileir_matrix.json | 20 ++ .../cutile/fused_1nn_kernel.py | 68 ++++ .../cutile/fused_1nn_planner.hpp | 60 ++++ .../cutile/fused_1nn_tile.cu | 173 ++++++++++ .../cutile/fused_1nn_tile.hpp | 55 +++ .../pairwise_matrix_planner.hpp | 4 +- .../jit_lto_kernels/cagra_planner_base.hpp | 4 +- .../interleaved_scan_planner.hpp | 4 +- .../compute_similarity_planner.hpp | 4 +- .../detail/jit_lto_kernels/scan_planner.hpp | 4 +- cpp/tests/CMakeLists.txt | 3 +- cpp/tests/cutile/cutile_vector_add.cu | 176 ++++++++-- cpp/tests/cutile/export_vector_add_cubin.py | 58 +++- cpp/tests/cutile/generate_cutile_cubins.cmake | 27 ++ cpp/tests/neighbors/distance_nn.cu | 1 + cpp/tests/neighbors/distance_nn_helper.cuh | 45 ++- 32 files changed, 1743 insertions(+), 153 deletions(-) create mode 100644 cpp/cmake/modules/generate_cutile_kernels.cmake create mode 100644 cpp/cmake/modules/register_cubin.cpp.in create mode 100644 cpp/cmake/modules/register_tileir.cpp.in create mode 100644 cpp/include/cuvs/detail/jit_lto/cutile_arch_tags.hpp create mode 100644 cpp/include/cuvs/detail/jit_lto/cutile_module.hpp create mode 100644 cpp/include/cuvs/detail/jit_lto/fused_distance_nn/fused_1nn_fragments.hpp create mode 100644 cpp/include/cuvs/detail/jit_lto/tileir_compat.hpp create mode 100644 cpp/src/detail/jit_lto/LTOAlgorithmPlanner.cpp create mode 100644 cpp/src/detail/jit_lto/TileAlgorithmPlanner.cpp create mode 100644 cpp/src/distance/detail/fused_distance_nn/cutile/export_fused_1nn.py create mode 100644 cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_cutile_cubin_matrix.json create mode 100644 cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_cutile_tileir_matrix.json create mode 100644 cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_kernel.py create mode 100644 cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_planner.hpp create mode 100644 cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_tile.cu create mode 100644 cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_tile.hpp diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt index 227c2906cc..cc6e1975b3 100644 --- a/cpp/CMakeLists.txt +++ b/cpp/CMakeLists.txt @@ -957,6 +957,47 @@ if(NOT BUILD_CPU_ONLY) OUTPUT_FILE_FORMAT "${CMAKE_CURRENT_BINARY_DIR}/src/distance/detail/pairwise_matrix/dispatch_rbf_inst_data_@data_abbrev@_acc_@acc_abbrev@_out_@out_abbrev@_index_@index_abbrev@_op_@op_abbrev@.cu" ) + + include(cmake/modules/generate_cutile_kernels.cmake) + set(fused_1nn_cutile_dir + "${CMAKE_CURRENT_SOURCE_DIR}/src/distance/detail/fused_distance_nn/cutile") + set(cutile_fused_1nn_generated_dir + "${CMAKE_CURRENT_BINARY_DIR}/generated_kernels/distance/fused_1nn/cutile") + generate_cutile_cubin_kernels( + cutile_fused_1nn_files + KERNEL_DIR "${fused_1nn_cutile_dir}" + KERNEL_BASENAME "fused_1nn" + KERNEL_PYTHON "fused_1nn_kernel.py" + EXPORT_SCRIPT "export_fused_1nn.py" + OUTPUT_DIRECTORY "${cutile_fused_1nn_generated_dir}" + MATRIX_JSON_FILE "${fused_1nn_cutile_dir}/fused_1nn_cutile_cubin_matrix.json" + FRAGMENT_TAG_FORMAT + "cuvs::distance::detail::fragment_tag_fused_1nn_cubin" + FRAGMENT_TAG_HEADER_FILES + "" + "" + "" + ) + generate_cutile_tileir_kernels( + cutile_fused_1nn_files + KERNEL_DIR "${fused_1nn_cutile_dir}" + KERNEL_BASENAME "fused_1nn" + KERNEL_PYTHON "fused_1nn_kernel.py" + EXPORT_SCRIPT "export_fused_1nn.py" + OUTPUT_DIRECTORY "${cutile_fused_1nn_generated_dir}" + MATRIX_JSON_FILE "${fused_1nn_cutile_dir}/fused_1nn_cutile_tileir_matrix.json" + FRAGMENT_TAG_FORMAT + "cuvs::distance::detail::fragment_tag_fused_1nn_tileir" + FRAGMENT_TAG_HEADER_FILES + "" + "" + ) + if(NOT DEFINED CUVS_CUTILE_ENABLED) + set(CUVS_CUTILE_ENABLED 0) + endif() + target_compile_definitions( + cuvs_cpp_headers INTERFACE CUVS_CUTILE_ENABLED=${CUVS_CUTILE_ENABLED} + ) generate_inst_matrix( cagra_build_inst_files MATRIX_JSON_FILE "${CMAKE_CURRENT_SOURCE_DIR}/src/neighbors/cagra_build_matrix.json" @@ -1147,6 +1188,8 @@ if(NOT BUILD_CPU_ONLY) src/util/host_memory.cpp src/detail/jit_lto/AlgorithmLauncher.cpp src/detail/jit_lto/AlgorithmPlanner.cpp + src/detail/jit_lto/LTOAlgorithmPlanner.cpp + src/detail/jit_lto/TileAlgorithmPlanner.cpp src/detail/jit_lto/FragmentEntry.cpp src/detail/jit_lto/nvjitlink_checker.cpp src/detail/jit_lto/NVRTCLTOFragmentCompiler.cpp @@ -1234,6 +1277,8 @@ if(NOT BUILD_CPU_ONLY) src/stats/trustworthiness_score.cu ${CUVS_MG_ALGOS} ${jit_lto_files} + ${cutile_fused_1nn_files} + $<$:src/distance/detail/fused_distance_nn/cutile/fused_1nn_tile.cu> ) set_target_properties( @@ -1257,6 +1302,7 @@ if(NOT BUILD_CPU_ONLY) target_compile_definitions( cuvs_objs PRIVATE $<$:CUVS_BUILD_CAGRA_HNSWLIB> $<$:NVTX_ENABLED> + CUVS_CUTILE_ENABLED=${CUVS_CUTILE_ENABLED} ) target_link_libraries( @@ -1274,7 +1320,9 @@ if(NOT BUILD_CPU_ONLY) PUBLIC "$" "$" INTERFACE "$" - PRIVATE "${CMAKE_CURRENT_SOURCE_DIR}/src" "${CMAKE_CURRENT_BINARY_DIR}/src" + PRIVATE "${CMAKE_CURRENT_SOURCE_DIR}/src" + "${CMAKE_CURRENT_BINARY_DIR}/src" + "${cutile_fused_1nn_generated_dir}" ) # Endian detection diff --git a/cpp/cmake/modules/generate_cutile_kernels.cmake b/cpp/cmake/modules/generate_cutile_kernels.cmake new file mode 100644 index 0000000000..7b9c2521c4 --- /dev/null +++ b/cpp/cmake/modules/generate_cutile_kernels.cmake @@ -0,0 +1,315 @@ +# ============================================================================= +# cmake-format: off +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. +# SPDX-License-Identifier: Apache-2.0 +# cmake-format: on +# ============================================================================= + +include_guard(GLOBAL) + +include(${CMAKE_CURRENT_LIST_DIR}/compute_matrix_product.cmake) + +function(generate_cutile_kernels_stub) + set(CUVS_CUTILE_ENABLED 0 PARENT_SCOPE) +endfunction() + +function(_cutile_fragment_tag_header_files output_var) + set(${output_var} "") + foreach(_header IN LISTS ARGN) + if(NOT _header MATCHES "^(\".*\"|<.*>)$") + set(_header "\"${_header}\"") + endif() + string(APPEND ${output_var} "#include ${_header}\n") + endforeach() + set(${output_var} + "${${output_var}}" + PARENT_SCOPE + ) +endfunction() + +function(_cutile_kernels_setup) + set(options) + set(one_value MATRIX_JSON_FILE OUTPUT_DIRECTORY) + set(multi_value) + cmake_parse_arguments(_CUTILE "${options}" "${one_value}" "${multi_value}" ${ARGN}) + + find_package(Python3 REQUIRED COMPONENTS Interpreter) + find_package(CUDAToolkit REQUIRED) + + if(CUDAToolkit_VERSION VERSION_LESS 13.0) + message( + STATUS + "cuTile embedded kernels require CUDA 13.0+; skipping cuTile generation (found ${CUDAToolkit_VERSION})." + ) + set(_CUTILE_SETUP_OK + FALSE + PARENT_SCOPE + ) + return() + endif() + + find_program( + CUTILE_BIN2C + NAMES bin2c + PATHS ${CUDAToolkit_BIN_DIR} + REQUIRED + ) + + execute_process( + COMMAND "${Python3_EXECUTABLE}" -c "import cuda.tile" + RESULT_VARIABLE _cutile_import_result + OUTPUT_QUIET + ERROR_QUIET + ) + if(NOT _cutile_import_result EQUAL 0) + message( + FATAL_ERROR + "cuda.tile (cuTile Python) is required to build cuTile embedded kernels. " + "Install it in the active Python environment, e.g. pip install cuda-tile[tileiras]." + ) + endif() + + set_property( + DIRECTORY + PROPERTY CMAKE_CONFIGURE_DEPENDS "${_CUTILE_MATRIX_JSON_FILE}" + APPEND + ) + + file(MAKE_DIRECTORY "${_CUTILE_OUTPUT_DIRECTORY}") + + set(_CUTILE_SETUP_OK + TRUE + PARENT_SCOPE + ) +endfunction() + +function(process_cutile_cubin_matrix_entry source_list_var) + set(options) + set(one_value + KERNEL_DIR + KERNEL_BASENAME + KERNEL_PYTHON + EXPORT_SCRIPT + OUTPUT_DIRECTORY + FRAGMENT_TAG_FORMAT + MATRIX_JSON_ENTRY + ) + set(multi_value FRAGMENT_TAG_HEADER_FILES) + cmake_parse_arguments(_CUTILE "${options}" "${one_value}" "${multi_value}" ${ARGN}) + + populate_matrix_variables("${_CUTILE_MATRIX_JSON_ENTRY}") + _cutile_fragment_tag_header_files( + fragment_tag_header_files ${_CUTILE_FRAGMENT_TAG_HEADER_FILES} + ) + + string(CONFIGURE "${_CUTILE_FRAGMENT_TAG_FORMAT}" fragment_tag @ONLY) + + set(_artifact_basename "${_CUTILE_KERNEL_BASENAME}_${data_type}_${gpu_code}") + set(_cubin_file "${_CUTILE_OUTPUT_DIRECTORY}/${_artifact_basename}.cubin") + set(_cubin_header "${_CUTILE_OUTPUT_DIRECTORY}/${_artifact_basename}_cubin.h") + set(_cubin_cpp "${_CUTILE_OUTPUT_DIRECTORY}/${_artifact_basename}_cubin.cpp") + set(cubin_header_file "${_artifact_basename}_cubin.h") + + add_custom_command( + OUTPUT "${_cubin_file}" + COMMAND + "${Python3_EXECUTABLE}" "${_CUTILE_KERNEL_DIR}/${_CUTILE_EXPORT_SCRIPT}" "${_cubin_file}" + --format cubin --data-type "${data_type}" --gpu-code "${gpu_code}" + DEPENDS "${_CUTILE_KERNEL_DIR}/${_CUTILE_EXPORT_SCRIPT}" + "${_CUTILE_KERNEL_DIR}/${_CUTILE_KERNEL_PYTHON}" + COMMENT "Exporting cuTile ${_CUTILE_KERNEL_BASENAME} cubin ${data_type} ${gpu_code}" + VERBATIM + ) + + add_custom_command( + OUTPUT "${_cubin_header}" + COMMAND "${CUTILE_BIN2C}" --const --name embedded_cubin --static "${_cubin_file}" + > "${_cubin_header}" + DEPENDS "${_cubin_file}" + VERBATIM + ) + + configure_file( + "${CMAKE_CURRENT_FUNCTION_LIST_DIR}/register_cubin.cpp.in" "${_cubin_cpp}" @ONLY + ) + list(APPEND ${source_list_var} "${_cubin_header}" "${_cubin_cpp}") + set(${source_list_var} + "${${source_list_var}}" + PARENT_SCOPE + ) +endfunction() + +function(process_cutile_tileir_matrix_entry source_list_var) + set(options) + set(one_value + KERNEL_DIR + KERNEL_BASENAME + KERNEL_PYTHON + EXPORT_SCRIPT + OUTPUT_DIRECTORY + FRAGMENT_TAG_FORMAT + MATRIX_JSON_ENTRY + ) + set(multi_value FRAGMENT_TAG_HEADER_FILES) + cmake_parse_arguments(_CUTILE "${options}" "${one_value}" "${multi_value}" ${ARGN}) + + populate_matrix_variables("${_CUTILE_MATRIX_JSON_ENTRY}") + _cutile_fragment_tag_header_files( + fragment_tag_header_files ${_CUTILE_FRAGMENT_TAG_HEADER_FILES} + ) + + string(CONFIGURE "${_CUTILE_FRAGMENT_TAG_FORMAT}" fragment_tag @ONLY) + set(_tileir_file "${_CUTILE_OUTPUT_DIRECTORY}/${_CUTILE_KERNEL_BASENAME}_${data_type}.tilebc") + set(_tileir_header "${_CUTILE_OUTPUT_DIRECTORY}/${_CUTILE_KERNEL_BASENAME}_${data_type}_tileir.h") + set(_tileir_cpp "${_CUTILE_OUTPUT_DIRECTORY}/${_CUTILE_KERNEL_BASENAME}_${data_type}_tileir.cpp") + set(tileir_header_file "${_CUTILE_KERNEL_BASENAME}_${data_type}_tileir.h") + + add_custom_command( + OUTPUT "${_tileir_file}" + COMMAND + "${Python3_EXECUTABLE}" "${_CUTILE_KERNEL_DIR}/${_CUTILE_EXPORT_SCRIPT}" "${_tileir_file}" + --format tileir_bytecode --data-type "${data_type}" --gpu-code "${export_gpu_code}" + --bytecode-version "${bytecode_version}" + DEPENDS "${_CUTILE_KERNEL_DIR}/${_CUTILE_EXPORT_SCRIPT}" + "${_CUTILE_KERNEL_DIR}/${_CUTILE_KERNEL_PYTHON}" + COMMENT "Exporting cuTile ${_CUTILE_KERNEL_BASENAME} TileIR bytecode ${data_type}" + VERBATIM + ) + + add_custom_command( + OUTPUT "${_tileir_header}" + COMMAND "${CUTILE_BIN2C}" --const --name embedded_tileir --static "${_tileir_file}" + > "${_tileir_header}" + DEPENDS "${_tileir_file}" + VERBATIM + ) + + configure_file( + "${CMAKE_CURRENT_FUNCTION_LIST_DIR}/register_tileir.cpp.in" "${_tileir_cpp}" @ONLY + ) + list(APPEND ${source_list_var} "${_tileir_header}" "${_tileir_cpp}") + set(${source_list_var} + "${${source_list_var}}" + PARENT_SCOPE + ) +endfunction() + +function(generate_cutile_cubin_kernels source_list_var) + set(options) + set(one_value + KERNEL_DIR + KERNEL_BASENAME + KERNEL_PYTHON + EXPORT_SCRIPT + OUTPUT_DIRECTORY + MATRIX_JSON_FILE + FRAGMENT_TAG_FORMAT + ) + set(multi_value FRAGMENT_TAG_HEADER_FILES) + cmake_parse_arguments(_CUTILE "${options}" "${one_value}" "${multi_value}" ${ARGN}) + + if(NOT _CUTILE_KERNEL_BASENAME) + message(FATAL_ERROR "generate_cutile_cubin_kernels: KERNEL_BASENAME is required") + endif() + if(NOT _CUTILE_KERNEL_PYTHON) + set(_CUTILE_KERNEL_PYTHON "fused_1nn_kernel.py") + endif() + + _cutile_kernels_setup( + MATRIX_JSON_FILE "${_CUTILE_MATRIX_JSON_FILE}" + OUTPUT_DIRECTORY "${_CUTILE_OUTPUT_DIRECTORY}" + ) + if(NOT _CUTILE_SETUP_OK) + generate_cutile_kernels_stub() + set(${source_list_var} + "" + PARENT_SCOPE + ) + return() + endif() + + compute_matrix_product(matrix_product MATRIX_JSON_FILE "${_CUTILE_MATRIX_JSON_FILE}") + + string(JSON len LENGTH "${matrix_product}") + math(EXPR last "${len} - 1") + + # cmake-lint: disable=C0103,E1120 + foreach(i RANGE "${last}") + string(JSON matrix_json_entry GET "${matrix_product}" "${i}") + process_cutile_cubin_matrix_entry( + "${source_list_var}" + KERNEL_DIR "${_CUTILE_KERNEL_DIR}" + KERNEL_BASENAME "${_CUTILE_KERNEL_BASENAME}" + KERNEL_PYTHON "${_CUTILE_KERNEL_PYTHON}" + EXPORT_SCRIPT "${_CUTILE_EXPORT_SCRIPT}" + OUTPUT_DIRECTORY "${_CUTILE_OUTPUT_DIRECTORY}" + FRAGMENT_TAG_FORMAT "${_CUTILE_FRAGMENT_TAG_FORMAT}" + FRAGMENT_TAG_HEADER_FILES ${_CUTILE_FRAGMENT_TAG_HEADER_FILES} + MATRIX_JSON_ENTRY "${matrix_json_entry}" + ) + endforeach() + + set(CUVS_CUTILE_ENABLED 1 PARENT_SCOPE) + set(${source_list_var} + "${${source_list_var}}" + PARENT_SCOPE + ) +endfunction() + +function(generate_cutile_tileir_kernels source_list_var) + set(options) + set(one_value + KERNEL_DIR + KERNEL_BASENAME + KERNEL_PYTHON + EXPORT_SCRIPT + OUTPUT_DIRECTORY + MATRIX_JSON_FILE + FRAGMENT_TAG_FORMAT + ) + set(multi_value FRAGMENT_TAG_HEADER_FILES) + cmake_parse_arguments(_CUTILE "${options}" "${one_value}" "${multi_value}" ${ARGN}) + + if(NOT _CUTILE_KERNEL_BASENAME) + message(FATAL_ERROR "generate_cutile_tileir_kernels: KERNEL_BASENAME is required") + endif() + if(NOT _CUTILE_KERNEL_PYTHON) + set(_CUTILE_KERNEL_PYTHON "fused_1nn_kernel.py") + endif() + + _cutile_kernels_setup( + MATRIX_JSON_FILE "${_CUTILE_MATRIX_JSON_FILE}" + OUTPUT_DIRECTORY "${_CUTILE_OUTPUT_DIRECTORY}" + ) + if(NOT _CUTILE_SETUP_OK) + generate_cutile_kernels_stub() + return() + endif() + + compute_matrix_product(matrix_product MATRIX_JSON_FILE "${_CUTILE_MATRIX_JSON_FILE}") + + string(JSON len LENGTH "${matrix_product}") + math(EXPR last "${len} - 1") + + # cmake-lint: disable=C0103,E1120 + foreach(i RANGE "${last}") + string(JSON matrix_json_entry GET "${matrix_product}" "${i}") + process_cutile_tileir_matrix_entry( + "${source_list_var}" + KERNEL_DIR "${_CUTILE_KERNEL_DIR}" + KERNEL_BASENAME "${_CUTILE_KERNEL_BASENAME}" + KERNEL_PYTHON "${_CUTILE_KERNEL_PYTHON}" + EXPORT_SCRIPT "${_CUTILE_EXPORT_SCRIPT}" + OUTPUT_DIRECTORY "${_CUTILE_OUTPUT_DIRECTORY}" + FRAGMENT_TAG_FORMAT "${_CUTILE_FRAGMENT_TAG_FORMAT}" + FRAGMENT_TAG_HEADER_FILES ${_CUTILE_FRAGMENT_TAG_HEADER_FILES} + MATRIX_JSON_ENTRY "${matrix_json_entry}" + ) + endforeach() + + set(CUVS_CUTILE_ENABLED 1 PARENT_SCOPE) + set(${source_list_var} + "${${source_list_var}}" + PARENT_SCOPE + ) +endfunction() diff --git a/cpp/cmake/modules/register_cubin.cpp.in b/cpp/cmake/modules/register_cubin.cpp.in new file mode 100644 index 0000000000..c27d6829ee --- /dev/null +++ b/cpp/cmake/modules/register_cubin.cpp.in @@ -0,0 +1,22 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include "@cubin_header_file@" +#include + +@fragment_tag_header_files@ + +namespace { + +using fragment_tag = @fragment_tag@; +using fragment_entry = StaticCubinFragmentEntry; + +} // namespace + +template <> +const uint8_t* const fragment_entry::data = embedded_cubin; + +template <> +const size_t fragment_entry::length = sizeof(embedded_cubin); diff --git a/cpp/cmake/modules/register_tileir.cpp.in b/cpp/cmake/modules/register_tileir.cpp.in new file mode 100644 index 0000000000..fb81acedbc --- /dev/null +++ b/cpp/cmake/modules/register_tileir.cpp.in @@ -0,0 +1,22 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include "@tileir_header_file@" +#include + +@fragment_tag_header_files@ + +namespace { + +using fragment_tag = @fragment_tag@; +using fragment_entry = StaticTileIrBytecodeFragmentEntry; + +} // namespace + +template <> +const uint8_t* const fragment_entry::data = embedded_tileir; + +template <> +const size_t fragment_entry::length = sizeof(embedded_tileir); diff --git a/cpp/include/cuvs/detail/jit_lto/AlgorithmPlanner.hpp b/cpp/include/cuvs/detail/jit_lto/AlgorithmPlanner.hpp index 7f275b1285..d727c73b9d 100644 --- a/cpp/include/cuvs/detail/jit_lto/AlgorithmPlanner.hpp +++ b/cpp/include/cuvs/detail/jit_lto/AlgorithmPlanner.hpp @@ -10,6 +10,7 @@ #include #include #include +#include #include #include @@ -19,6 +20,7 @@ struct LauncherJitCache { std::shared_mutex mutex; std::unordered_map> launchers; + std::unordered_set build_failed; }; struct AlgorithmPlanner { @@ -27,9 +29,32 @@ struct AlgorithmPlanner { { } + virtual ~AlgorithmPlanner() = default; + std::shared_ptr get_launcher(); + /** Returns nullptr when no module can be loaded for the current device (does not RAFT_FAIL). */ + std::shared_ptr try_get_launcher(); + std::string entrypoint; + + protected: + virtual std::shared_ptr build() = 0; + + virtual std::string get_planner_key() const = 0; + + std::shared_ptr read_cache(std::string const& launch_key) const; + + LauncherJitCache& jit_cache_; +}; + +/** Links embedded LTO fatbin fragments at runtime via nvJitLink. */ +struct LTOAlgorithmPlanner : AlgorithmPlanner { + LTOAlgorithmPlanner(std::string entrypoint, LauncherJitCache& jit_cache) + : AlgorithmPlanner(std::move(entrypoint), jit_cache) + { + } + std::vector> fragments; template >> @@ -45,16 +70,38 @@ struct AlgorithmPlanner { } protected: - /** Extra link-time option strings passed to nvJitLink. Base build() - * always passes "-lto" and "-arch=sm_XX" first; derived planners may append here in their - * constructor body. */ + /** Extra link-time option strings passed to nvJitLink. */ std::vector linktime_extra_options; - private: - std::string get_fragments_key() const; - std::shared_ptr build(); + std::string get_planner_key() const override; - std::shared_ptr read_cache(std::string const& launch_key) const; + std::shared_ptr build() override; +}; - LauncherJitCache& jit_cache_; +/** Loads prebuilt cubins or TileIR bytecode via cudaLibraryLoadData. */ +struct TileAlgorithmPlanner : AlgorithmPlanner { + TileAlgorithmPlanner(std::string entrypoint, LauncherJitCache& jit_cache) + : AlgorithmPlanner(std::move(entrypoint), jit_cache) + { + } + + template + void add_static_fragment() + { + cubin_fragments_.push_back(std::make_unique>()); + } + + template + void add_static_tileir_fragment() + { + tileir_fragment_ = std::make_unique>(); + } + + protected: + std::vector> cubin_fragments_; + std::unique_ptr tileir_fragment_; + + std::string get_planner_key() const override; + + std::shared_ptr build() override; }; diff --git a/cpp/include/cuvs/detail/jit_lto/FragmentEntry.hpp b/cpp/include/cuvs/detail/jit_lto/FragmentEntry.hpp index 35aa46633c..df69ec1d7b 100644 --- a/cpp/include/cuvs/detail/jit_lto/FragmentEntry.hpp +++ b/cpp/include/cuvs/detail/jit_lto/FragmentEntry.hpp @@ -62,3 +62,66 @@ struct UDFFatbinFragment final : FatbinFragmentEntry { std::string key_; std::vector bytes_; }; + +/** Embedded CUDA binary module (cubin), loaded directly via cudaLibraryLoadData. */ +struct CubinFragmentEntry { + virtual ~CubinFragmentEntry() = default; + + virtual const uint8_t* get_data() const = 0; + + virtual size_t get_length() const = 0; + + virtual const char* get_key() const = 0; + + virtual int get_cc_major() const = 0; + + virtual int get_cc_minor() const = 0; +}; + +template +struct StaticCubinFragmentEntry final : CubinFragmentEntry { + const uint8_t* get_data() const override { return StaticCubinFragmentEntry::data; } + + size_t get_length() const override { return StaticCubinFragmentEntry::length; } + + const char* get_key() const override + { + return typeid(StaticCubinFragmentEntry).name(); + } + + int get_cc_major() const override { return FragmentTag::cc_major; } + + int get_cc_minor() const override { return FragmentTag::cc_minor; } + + static const uint8_t* const data; + static const size_t length; +}; + +/** Embedded TileIR bytecode, JIT-compiled by the driver when no matching cubin exists. */ +struct TileIrBytecodeFragmentEntry { + virtual ~TileIrBytecodeFragmentEntry() = default; + + virtual const uint8_t* get_data() const = 0; + + virtual size_t get_length() const = 0; + + virtual const char* get_key() const = 0; +}; + +template +struct StaticTileIrBytecodeFragmentEntry final : TileIrBytecodeFragmentEntry { + const uint8_t* get_data() const override + { + return StaticTileIrBytecodeFragmentEntry::data; + } + + size_t get_length() const override { return StaticTileIrBytecodeFragmentEntry::length; } + + const char* get_key() const override + { + return typeid(StaticTileIrBytecodeFragmentEntry).name(); + } + + static const uint8_t* const data; + static const size_t length; +}; diff --git a/cpp/include/cuvs/detail/jit_lto/cutile_arch_tags.hpp b/cpp/include/cuvs/detail/jit_lto/cutile_arch_tags.hpp new file mode 100644 index 0000000000..2c915a278b --- /dev/null +++ b/cpp/include/cuvs/detail/jit_lto/cutile_arch_tags.hpp @@ -0,0 +1,52 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +#ifndef CUVS_CUTILE_ENABLED +#define CUVS_CUTILE_ENABLED 0 +#endif + +namespace cuvs::detail::jit_lto { + +#if CUVS_CUTILE_ENABLED + +/** Must stay in sync with cuTile matrix _arch entries and planner add_static_fragment calls. */ +struct cutile_arch_8_0 { + static constexpr int cc_major = 8; + static constexpr int cc_minor = 0; +}; + +struct cutile_arch_8_6 { + static constexpr int cc_major = 8; + static constexpr int cc_minor = 6; +}; + +struct cutile_arch_9_0 { + static constexpr int cc_major = 9; + static constexpr int cc_minor = 0; +}; + +struct cutile_arch_12_0 { + static constexpr int cc_major = 12; + static constexpr int cc_minor = 0; +}; + +inline bool is_embedded_cubin_arch(int cc_major, int cc_minor) +{ + if (cc_major == 8 && cc_minor == 0) { return true; } + if (cc_major == 8 && cc_minor == 6) { return true; } + if (cc_major == 9 && cc_minor == 0) { return true; } + if (cc_major == 12 && cc_minor == 0) { return true; } + return false; +} + +#else + +inline bool is_embedded_cubin_arch(int, int) { return false; } + +#endif + +} // namespace cuvs::detail::jit_lto diff --git a/cpp/include/cuvs/detail/jit_lto/cutile_module.hpp b/cpp/include/cuvs/detail/jit_lto/cutile_module.hpp new file mode 100644 index 0000000000..dff0f472a7 --- /dev/null +++ b/cpp/include/cuvs/detail/jit_lto/cutile_module.hpp @@ -0,0 +1,75 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +#include +#include +#include +#include +#include +#include + +#include + +#include +#include +#include + +#include + +namespace cuvs::detail::jit_lto { + +struct CutileModuleImage { + const uint8_t* data; + size_t size; +}; + +inline bool get_device_compute_capability(int& cc_major, int& cc_minor) +{ + int device = 0; + if (cudaGetDevice(&device) != cudaSuccess) { return false; } + if (cudaDeviceGetAttribute(&cc_major, cudaDevAttrComputeCapabilityMajor, device) != cudaSuccess) { + return false; + } + if (cudaDeviceGetAttribute(&cc_minor, cudaDevAttrComputeCapabilityMinor, device) != cudaSuccess) { + return false; + } + return true; +} + +/** Selects a prebuilt cubin for the device CC, or embedded TileIR when the driver can JIT it. */ +inline std::optional resolve_cutile_module_image( + int cc_major, + int cc_minor, + int driver_version, + const std::vector>& cubin_fragments, + const TileIrBytecodeFragmentEntry* tileir_fragment) +{ + for (const auto& fragment : cubin_fragments) { + if (fragment->get_cc_major() == cc_major && fragment->get_cc_minor() == cc_minor) { + return CutileModuleImage{fragment->get_data(), fragment->get_length()}; + } + } + if (tileir_fragment != nullptr && tileir_fallback_available(driver_version)) { + return CutileModuleImage{tileir_fragment->get_data(), tileir_fragment->get_length()}; + } + return std::nullopt; +} + +inline std::shared_ptr load_cutile_launcher(const CutileModuleImage& image, + const std::string& kernel_symbol) +{ + cudaLibrary_t library{}; + RAFT_CUDA_TRY( + cudaLibraryLoadData(&library, image.data, nullptr, nullptr, 0, nullptr, nullptr, 0)); + + cudaKernel_t kernel{}; + RAFT_CUDA_TRY(cudaLibraryGetKernel(&kernel, library, kernel_symbol.c_str())); + + return std::make_shared(kernel, library); +} + +} // namespace cuvs::detail::jit_lto diff --git a/cpp/include/cuvs/detail/jit_lto/fused_distance_nn/fused_1nn_fragments.hpp b/cpp/include/cuvs/detail/jit_lto/fused_distance_nn/fused_1nn_fragments.hpp new file mode 100644 index 0000000000..517118bbe2 --- /dev/null +++ b/cpp/include/cuvs/detail/jit_lto/fused_distance_nn/fused_1nn_fragments.hpp @@ -0,0 +1,21 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +#include + +namespace cuvs::distance::detail { + +template +struct fragment_tag_fused_1nn_cubin { + static constexpr int cc_major = ArchTag::cc_major; + static constexpr int cc_minor = ArchTag::cc_minor; +}; + +template +struct fragment_tag_fused_1nn_tileir {}; + +} // namespace cuvs::distance::detail diff --git a/cpp/include/cuvs/detail/jit_lto/tileir_compat.hpp b/cpp/include/cuvs/detail/jit_lto/tileir_compat.hpp new file mode 100644 index 0000000000..d63759fb36 --- /dev/null +++ b/cpp/include/cuvs/detail/jit_lto/tileir_compat.hpp @@ -0,0 +1,99 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +#ifndef CUVS_CUTILE_ENABLED +#define CUVS_CUTILE_ENABLED 0 +#endif + +#include +#include + +#include + +namespace cuvs::detail::jit_lto { + +/** Minimum CUDA driver version (from cudaDriverGetVersion) for TileIR JIT of embedded bytecode. */ +inline constexpr int kMinTileIrJitDriverVersion = 13010; // CUDA 13.1 / driver >= 590.44 + +/** Minimum CUDA runtime version (from cudaRuntimeGetVersion) for cuTile integration. */ +inline constexpr int kMinCutileRuntimeVersion = 13000; + +inline constexpr bool library_built_with_cutile() +{ +#if CUVS_CUTILE_ENABLED + return true; +#else + return false; +#endif +} + +inline bool runtime_cuda13_or_newer() +{ + int runtime_version = 0; + if (cudaRuntimeGetVersion(&runtime_version) != cudaSuccess) { return false; } + return runtime_version >= kMinCutileRuntimeVersion; +} + +/** True when this build embeds cuTile artifacts and the runtime is CUDA 13+. */ +inline bool cutile_integration_enabled() +{ + return library_built_with_cutile() && runtime_cuda13_or_newer(); +} + +/** True when this build embeds a prebuilt cubin for the given compute capability. */ +inline bool has_embedded_cubin_for_arch(int cc_major, int cc_minor) +{ + return is_embedded_cubin_arch(cc_major, cc_minor); +} + +/** True when the driver can JIT-compile embedded TileIR bytecode at load time. */ +inline bool tileir_fallback_available(int driver_version) +{ + return driver_version >= kMinTileIrJitDriverVersion; +} + +/** + * True when a cuTile launch may be attempted for the given device: cuTile is enabled, the runtime + * is CUDA 13+, and either a matching embedded cubin exists (no driver JIT required) or the driver + * can JIT the embedded TileIR bytecode fallback. + */ +inline bool cutile_launch_available_for_arch(int cc_major, int cc_minor, int driver_version) +{ + if (!cutile_integration_enabled()) { return false; } + if (has_embedded_cubin_for_arch(cc_major, cc_minor)) { return true; } + return tileir_fallback_available(driver_version); +} + +inline bool query_driver_version(int& driver_version) +{ + return cudaDriverGetVersion(&driver_version) == cudaSuccess; +} + +inline bool query_current_device_arch(int& cc_major, int& cc_minor) +{ + int device = 0; + if (cudaGetDevice(&device) != cudaSuccess) { return false; } + if (cudaDeviceGetAttribute(&cc_major, cudaDevAttrComputeCapabilityMajor, device) != cudaSuccess) { + return false; + } + if (cudaDeviceGetAttribute(&cc_minor, cudaDevAttrComputeCapabilityMinor, device) != cudaSuccess) { + return false; + } + return true; +} + +inline bool cutile_launch_available_on_current_device() +{ + int cc_major = 0; + int cc_minor = 0; + int driver_version = 0; + if (!query_current_device_arch(cc_major, cc_minor)) { return false; } + if (!query_driver_version(driver_version)) { return false; } + return cutile_launch_available_for_arch(cc_major, cc_minor, driver_version); +} + +} // namespace cuvs::detail::jit_lto diff --git a/cpp/src/detail/jit_lto/AlgorithmPlanner.cpp b/cpp/src/detail/jit_lto/AlgorithmPlanner.cpp index 7416ea396d..486d6f1aa5 100644 --- a/cpp/src/detail/jit_lto/AlgorithmPlanner.cpp +++ b/cpp/src/detail/jit_lto/AlgorithmPlanner.cpp @@ -3,33 +3,16 @@ * SPDX-License-Identifier: Apache-2.0 */ -#include -#include #include #include -#include #include #include -#include #include -#include - -#include "cuda_runtime.h" -#include "nvJitLink.h" #include #include -std::string AlgorithmPlanner::get_fragments_key() const -{ - std::string key = ""; - for (const auto& fragment : this->fragments) { - key += fragment->get_key(); - } - return key; -} - std::shared_ptr AlgorithmPlanner::read_cache(std::string const& launch_key) const { auto& launchers = jit_cache_.launchers; @@ -38,79 +21,37 @@ std::shared_ptr AlgorithmPlanner::read_cache(std::string cons return nullptr; } -std::shared_ptr AlgorithmPlanner::get_launcher() +std::shared_ptr AlgorithmPlanner::try_get_launcher() { - auto& launchers = jit_cache_.launchers; - auto launch_key = this->get_fragments_key(); + auto launch_key = this->get_planner_key(); - if (auto hit = read_cache(launch_key)) { return hit; } + { + std::shared_lock read_lock(jit_cache_.mutex); + if (jit_cache_.build_failed.count(launch_key)) { return nullptr; } + if (auto hit = read_cache(launch_key)) { return hit; } + } std::unique_lock write_lock(jit_cache_.mutex); - if (auto it = launchers.find(launch_key); it != launchers.end()) { return it->second; } + if (jit_cache_.build_failed.count(launch_key)) { return nullptr; } + if (auto it = jit_cache_.launchers.find(launch_key); it != jit_cache_.launchers.end()) { + return it->second; + } - std::string log_message = - "JIT compiling launcher for kernel: " + this->entrypoint + " and device functions: "; - for (const auto& fragment : this->fragments) { - log_message += std::string{fragment->get_key()} + ","; + RAFT_LOG_DEBUG("Building launcher for kernel entrypoint: %s", this->entrypoint.c_str()); + auto launcher = this->build(); + if (!launcher) { + jit_cache_.build_failed.insert(launch_key); + return nullptr; } - log_message.pop_back(); - RAFT_LOG_DEBUG("%s", log_message.c_str()); - auto launcher = this->build(); - launchers[launch_key] = launcher; + jit_cache_.launchers[launch_key] = launcher; return launcher; } -std::shared_ptr AlgorithmPlanner::build() +std::shared_ptr AlgorithmPlanner::get_launcher() { - int device = 0; - int major = 0; - int minor = 0; - RAFT_CUDA_TRY(cudaGetDevice(&device)); - RAFT_CUDA_TRY(cudaDeviceGetAttribute(&major, cudaDevAttrComputeCapabilityMajor, device)); - RAFT_CUDA_TRY(cudaDeviceGetAttribute(&minor, cudaDevAttrComputeCapabilityMinor, device)); - - std::string archs = "-arch=sm_" + std::to_string((major * 10 + minor)); - - // Load the generated LTO IR and link them together - nvJitLinkHandle handle; - std::vector lopts; - lopts.reserve(2 + linktime_extra_options.size()); - lopts.push_back("-lto"); - lopts.push_back(archs.c_str()); - for (auto const& opt : linktime_extra_options) { - lopts.push_back(opt.c_str()); - } - auto result = nvJitLinkCreate(&handle, static_cast(lopts.size()), lopts.data()); - check_nvjitlink_result(handle, result); - - for (const auto& frag : this->fragments) { - frag->add_to(handle); + auto launcher = try_get_launcher(); + if (!launcher) { + RAFT_FAIL("Failed to build launcher for kernel entrypoint: %s", this->entrypoint.c_str()); } - - // Call to nvJitLinkComplete causes linker to link together all the LTO-IR - // modules perform any optimizations and generate cubin from it. - result = nvJitLinkComplete(handle); - check_nvjitlink_result(handle, result); - - // get cubin from nvJitLink - size_t cubin_size; - result = nvJitLinkGetLinkedCubinSize(handle, &cubin_size); - check_nvjitlink_result(handle, result); - - std::unique_ptr cubin{new char[cubin_size]}; - result = nvJitLinkGetLinkedCubin(handle, cubin.get()); - check_nvjitlink_result(handle, result); - - result = nvJitLinkDestroy(&handle); - RAFT_EXPECTS(result == NVJITLINK_SUCCESS, "nvJitLinkDestroy failed"); - - // cubin is linked, so now load it - cudaLibrary_t library; - RAFT_CUDA_TRY( - cudaLibraryLoadData(&library, cubin.get(), nullptr, nullptr, 0, nullptr, nullptr, 0)); - - cudaKernel_t kernel; - RAFT_CUDA_TRY(cudaLibraryGetKernel(&kernel, library, this->entrypoint.c_str())); - - return std::make_shared(kernel, library); + return launcher; } diff --git a/cpp/src/detail/jit_lto/LTOAlgorithmPlanner.cpp b/cpp/src/detail/jit_lto/LTOAlgorithmPlanner.cpp new file mode 100644 index 0000000000..da7c0408b4 --- /dev/null +++ b/cpp/src/detail/jit_lto/LTOAlgorithmPlanner.cpp @@ -0,0 +1,76 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include +#include +#include +#include + +#include +#include + +#include "cuda_runtime.h" +#include "nvJitLink.h" + +#include + +std::string LTOAlgorithmPlanner::get_planner_key() const +{ + std::string key; + for (const auto& fragment : this->fragments) { + key += fragment->get_key(); + } + return key; +} + +std::shared_ptr LTOAlgorithmPlanner::build() +{ + int device = 0; + int major = 0; + int minor = 0; + RAFT_CUDA_TRY(cudaGetDevice(&device)); + RAFT_CUDA_TRY(cudaDeviceGetAttribute(&major, cudaDevAttrComputeCapabilityMajor, device)); + RAFT_CUDA_TRY(cudaDeviceGetAttribute(&minor, cudaDevAttrComputeCapabilityMinor, device)); + + std::string archs = "-arch=sm_" + std::to_string((major * 10 + minor)); + + nvJitLinkHandle handle; + std::vector lopts; + lopts.reserve(2 + linktime_extra_options.size()); + lopts.push_back("-lto"); + lopts.push_back(archs.c_str()); + for (auto const& opt : linktime_extra_options) { + lopts.push_back(opt.c_str()); + } + auto result = nvJitLinkCreate(&handle, static_cast(lopts.size()), lopts.data()); + check_nvjitlink_result(handle, result); + + for (const auto& frag : this->fragments) { + frag->add_to(handle); + } + + result = nvJitLinkComplete(handle); + check_nvjitlink_result(handle, result); + + size_t cubin_size; + result = nvJitLinkGetLinkedCubinSize(handle, &cubin_size); + check_nvjitlink_result(handle, result); + + std::unique_ptr cubin{new char[cubin_size]}; + result = nvJitLinkGetLinkedCubin(handle, cubin.get()); + check_nvjitlink_result(handle, result); + + result = nvJitLinkDestroy(&handle); + RAFT_EXPECTS(result == NVJITLINK_SUCCESS, "nvJitLinkDestroy failed"); + + cudaLibrary_t library; + RAFT_CUDA_TRY( + cudaLibraryLoadData(&library, cubin.get(), nullptr, nullptr, 0, nullptr, nullptr, 0)); + + cudaKernel_t kernel; + RAFT_CUDA_TRY(cudaLibraryGetKernel(&kernel, library, this->entrypoint.c_str())); + + return std::make_shared(kernel, library); +} diff --git a/cpp/src/detail/jit_lto/TileAlgorithmPlanner.cpp b/cpp/src/detail/jit_lto/TileAlgorithmPlanner.cpp new file mode 100644 index 0000000000..edb6269213 --- /dev/null +++ b/cpp/src/detail/jit_lto/TileAlgorithmPlanner.cpp @@ -0,0 +1,38 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include +#include + +#include +#include + +std::string TileAlgorithmPlanner::get_planner_key() const +{ + std::string key = this->entrypoint; + for (const auto& fragment : cubin_fragments_) { + key += fragment->get_key(); + } + if (tileir_fragment_) { key += tileir_fragment_->get_key(); } + return key; +} + +std::shared_ptr TileAlgorithmPlanner::build() +{ + int cc_major = 0; + int cc_minor = 0; + if (!cuvs::detail::jit_lto::get_device_compute_capability(cc_major, cc_minor)) { + return nullptr; + } + + int driver_version = 0; + if (cudaDriverGetVersion(&driver_version) != cudaSuccess) { return nullptr; } + + auto image = cuvs::detail::jit_lto::resolve_cutile_module_image( + cc_major, cc_minor, driver_version, cubin_fragments_, tileir_fragment_.get()); + if (!image) { return nullptr; } + + return cuvs::detail::jit_lto::load_cutile_launcher(*image, this->entrypoint); +} diff --git a/cpp/src/distance/detail/fused_distance_nn.cuh b/cpp/src/distance/detail/fused_distance_nn.cuh index f9dbd968ec..8b47092b58 100644 --- a/cpp/src/distance/detail/fused_distance_nn.cuh +++ b/cpp/src/distance/detail/fused_distance_nn.cuh @@ -5,14 +5,22 @@ #pragma once +#ifndef CUVS_CUTILE_ENABLED +#define CUVS_CUTILE_ENABLED 0 +#endif + #include "distance_ops/l2_exp.cuh" // ops::l2_exp_distance_op #include "fused_distance_nn/cutlass_base.cuh" +#if CUVS_CUTILE_ENABLED +#include "fused_distance_nn/cutile/fused_1nn_tile.hpp" +#endif #include "fused_distance_nn/fused_cosine_nn.cuh" #include "fused_distance_nn/fused_l2_nn.cuh" #include "fused_distance_nn/helper_structs.cuh" #include "fused_distance_nn/simt_kernel.cuh" #include "pairwise_distance_base.cuh" // PairwiseDistances #include +#include #include // raft::KeyValuePair #include // raft::identity_op #include // Policy @@ -54,6 +62,13 @@ void fusedDistanceNNImpl(OutT* min, // The kernel policy is determined by fusedDistanceNN. typedef Policy P; +#if CUVS_CUTILE_ENABLED + if (cuvs::detail::jit_lto::cutile_launch_available_on_current_device() && + try_fused_1nn_tile(min, x, y, m, n, k, metric, stream)) { + return; + } +#endif + dim3 blk(P::Nthreads); auto nblks = raft::ceildiv(m, P::Nthreads); constexpr auto maxVal = std::numeric_limits::max(); diff --git a/cpp/src/distance/detail/fused_distance_nn/cutile/export_fused_1nn.py b/cpp/src/distance/detail/fused_distance_nn/cutile/export_fused_1nn.py new file mode 100644 index 0000000000..6a20be24ef --- /dev/null +++ b/cpp/src/distance/detail/fused_distance_nn/cutile/export_fused_1nn.py @@ -0,0 +1,136 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. +# SPDX-License-Identifier: Apache-2.0 +"""Export fused 1-NN cuTile kernels to cubin or TileIR bytecode.""" + +from __future__ import annotations + +import argparse +import sys +from pathlib import Path +from typing import Literal + +import cuda.tile as ct +from cuda.tile.compilation import ( + ArrayConstraint, + CallingConvention, + ConstantConstraint, + KernelSignature, + ScalarConstraint, + export_kernel, +) + +from fused_1nn_kernel import KERNELS, KERNEL_SYMBOLS, TILE_CONSTANTS + +DEFAULT_TILEIR_BYTECODE_VERSION = "13.1" +# cuTile requires a gpu_code even for TileIR bytecode export: it selects the compilation +# target / feature set for lowering, not the runtime architecture (the driver JITs at load). +DEFAULT_TILEIR_EXPORT_GPU_CODE = "sm_80" + + +def _dtype_for(data_type: str): + if data_type == "half": + return ct.float16 + if data_type == "float": + return ct.float32 + raise ValueError(f"Unsupported data_type {data_type!r}") + + +def _kernel_signature(data_type: str) -> KernelSignature: + elem = _dtype_for(data_type) + array = ArrayConstraint( + elem, + 2, + index_dtype=ct.int64, + stride_lower_bound_incl=0, + alias_groups=(), + may_alias_internally=False, + ) + idx_array = ArrayConstraint( + ct.int64, + 1, + index_dtype=ct.int64, + stride_lower_bound_incl=0, + alias_groups=(), + may_alias_internally=False, + stride_constant=(1,), + ) + dist_array = ArrayConstraint( + ct.float32, + 1, + index_dtype=ct.int64, + stride_lower_bound_incl=0, + alias_groups=(), + may_alias_internally=False, + stride_constant=(1,), + ) + tm, tn, tk = TILE_CONSTANTS + return KernelSignature( + parameters=[ + array, + array, + idx_array, + dist_array, + ScalarConstraint(ct.int64), + ScalarConstraint(ct.int64), + ScalarConstraint(ct.int64), + ConstantConstraint(tm), + ConstantConstraint(tn), + ConstantConstraint(tk), + ], + calling_convention=CallingConvention.cutile_python_v1(), + ).with_symbol(KERNEL_SYMBOLS[data_type]) + + +def export_binary( + output_file: Path, + *, + output_format: Literal["cubin", "tileir_bytecode"], + data_type: str, + gpu_code: str, + bytecode_version: str | None = None, +) -> str: + kernel = KERNELS[data_type] + signature = _kernel_signature(data_type) + + export_kwargs = { + "kernel": kernel, + "signatures": [signature], + "output_file": str(output_file), + "gpu_code": gpu_code, + "output_format": output_format, + } + if output_format == "tileir_bytecode": + export_kwargs["bytecode_version"] = bytecode_version or DEFAULT_TILEIR_BYTECODE_VERSION + + export_kernel(**export_kwargs) + + return signature.symbol + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("output_file", type=Path) + parser.add_argument("--format", choices=("cubin", "tileir_bytecode"), default="cubin") + parser.add_argument("--data-type", choices=tuple(KERNELS.keys()), required=True) + parser.add_argument( + "--gpu-code", + default=DEFAULT_TILEIR_EXPORT_GPU_CODE, + help="Target SM for cubin export, or compile hint for TileIR bytecode export", + ) + parser.add_argument("--bytecode-version", default=DEFAULT_TILEIR_BYTECODE_VERSION) + args = parser.parse_args() + + print( + export_binary( + args.output_file, + output_format=args.format, + data_type=args.data_type, + gpu_code=args.gpu_code, + bytecode_version=args.bytecode_version, + ) + ) + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_cutile_cubin_matrix.json b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_cutile_cubin_matrix.json new file mode 100644 index 0000000000..fbd4bfdd64 --- /dev/null +++ b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_cutile_cubin_matrix.json @@ -0,0 +1,40 @@ +[ + { + "_data": [ + { + "data_type": "half", + "data_abbrev": "h" + }, + { + "data_type": "float", + "data_abbrev": "f" + } + ], + "_arch": [ + { + "gpu_code": "sm_80", + "cc_major": 8, + "cc_minor": 0, + "arch_tag": "cutile_arch_8_0" + }, + { + "gpu_code": "sm_86", + "cc_major": 8, + "cc_minor": 6, + "arch_tag": "cutile_arch_8_6" + }, + { + "gpu_code": "sm_90", + "cc_major": 9, + "cc_minor": 0, + "arch_tag": "cutile_arch_9_0" + }, + { + "gpu_code": "sm_120", + "cc_major": 12, + "cc_minor": 0, + "arch_tag": "cutile_arch_12_0" + } + ] + } +] diff --git a/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_cutile_tileir_matrix.json b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_cutile_tileir_matrix.json new file mode 100644 index 0000000000..364c94594c --- /dev/null +++ b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_cutile_tileir_matrix.json @@ -0,0 +1,20 @@ +[ + { + "_data": [ + { + "data_type": "half", + "data_abbrev": "h" + }, + { + "data_type": "float", + "data_abbrev": "f" + } + ], + "_tileir": [ + { + "export_gpu_code": "sm_80", + "bytecode_version": "13.1" + } + ] + } +] diff --git a/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_kernel.py b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_kernel.py new file mode 100644 index 0000000000..232b9506af --- /dev/null +++ b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_kernel.py @@ -0,0 +1,68 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. +# SPDX-License-Identifier: Apache-2.0 +"""cuTile fused GEMM + inner-product 1-NN (argmax dot product) for cuVS.""" + +from __future__ import annotations + +import cuda.tile as ct + +ConstInt = ct.Constant[int] + +TILE_M = 128 +TILE_N = 256 +TILE_K = 64 + + +def _make_kernel(data_type: str): + if data_type == "half": + dtype = ct.float16 + acc_dtype = ct.float32 + elif data_type == "float": + dtype = ct.float32 + acc_dtype = ct.float32 + else: + raise ValueError(f"Unsupported data_type {data_type!r}") + + @ct.kernel + def fused_1nn_kernel(A, B, OutIdx, OutDist, M, N, K, tm: ConstInt, tn: ConstInt, tk: ConstInt): + bidm = ct.bid(0) + + best_dist = ct.full((tm,), -3.4e38, acc_dtype) + best_idx = ct.zeros((tm,), ct.int64) + + num_tiles_k = ct.num_tiles(A, axis=1, shape=(tm, tk)) + num_tiles_n = ct.num_tiles(B, axis=0, shape=(tn, tk)) + zero_pad = ct.PaddingMode.ZERO + + for n in range(num_tiles_n): + accumulator = ct.full((tm, tn), 0, dtype=acc_dtype) + + for k in range(num_tiles_k): + a = ct.load(A, index=(bidm, k), shape=(tm, tk), padding_mode=zero_pad) + b_T = ct.load(B, index=(n, k), shape=(tn, tk), padding_mode=zero_pad) + accumulator = ct.mma(a, ct.transpose(b_T), accumulator) + + curr_max = ct.max(accumulator, axis=1) + curr_idx = ct.argmax(accumulator, axis=1) + + update = curr_max > best_dist + best_dist = ct.where(update, curr_max, best_dist) + best_idx = ct.where(update, n * tn + curr_idx, best_idx) + + ct.store(OutIdx, index=(bidm,), tile=best_idx) + ct.store(OutDist, index=(bidm,), tile=best_dist) + + return fused_1nn_kernel + + +KERNELS = { + "half": _make_kernel("half"), + "float": _make_kernel("float"), +} + +KERNEL_SYMBOLS = { + "half": "fused_1nn_half", + "float": "fused_1nn_float", +} + +TILE_CONSTANTS = (TILE_M, TILE_N, TILE_K) diff --git a/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_planner.hpp b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_planner.hpp new file mode 100644 index 0000000000..dd2a539528 --- /dev/null +++ b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_planner.hpp @@ -0,0 +1,60 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +#include + +#include +#include +#include +#include + +namespace cuvs::distance::detail { + +/** Must match KERNEL_SYMBOLS in fused_1nn_kernel.py (export uses with_symbol). */ +template +inline const char* fused_1nn_kernel_entrypoint() +{ + if constexpr (std::is_same_v) { + return "fused_1nn_half"; + } else if constexpr (std::is_same_v) { + return "fused_1nn_float"; + } else { + static_assert(sizeof(DataTag) == 0, "unsupported fused 1-NN cuTile data type"); + return ""; + } +} + +template +struct Fused1nnTilePlanner : TileAlgorithmPlanner { + inline static LauncherJitCache launcher_jit_cache{}; + + Fused1nnTilePlanner() + : TileAlgorithmPlanner(fused_1nn_kernel_entrypoint(), launcher_jit_cache) + { + } + + /** Registers embedded cubin modules (one per SM); see register_cubin.cpp object files. */ + void add_entrypoint() + { + using cuvs::detail::jit_lto::cutile_arch_12_0; + using cuvs::detail::jit_lto::cutile_arch_8_0; + using cuvs::detail::jit_lto::cutile_arch_8_6; + using cuvs::detail::jit_lto::cutile_arch_9_0; + + this->add_static_fragment>(); + this->add_static_fragment>(); + this->add_static_fragment>(); + this->add_static_fragment>(); + } + + void add_tileir_fallback() + { + this->add_static_tileir_fragment>(); + } +}; + +} // namespace cuvs::distance::detail diff --git a/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_tile.cu b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_tile.cu new file mode 100644 index 0000000000..af8b0b181f --- /dev/null +++ b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_tile.cu @@ -0,0 +1,173 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include "fused_1nn_tile.hpp" + +#include "fused_1nn_planner.hpp" + +#include +#include + +namespace cuvs { +namespace distance { +namespace detail { + +namespace { + +template +__global__ void pack_fused_1nn_kvp(OutT* out, const int64_t* idx, const float* dist, IdxT len) +{ + IdxT i = blockIdx.x * blockDim.x + threadIdx.x; + if (i < len) { + out[i].key = static_cast(idx[i]); + out[i].value = static_cast(dist[i]); + } +} + +template +bool launch_fused_1nn_tile(const DataT* x, + const DataT* y, + OutT* out, + IdxT m, + IdxT n, + IdxT k, + cudaStream_t stream) +{ + Fused1nnTilePlanner planner; + planner.add_entrypoint(); + planner.add_tileir_fallback(); + auto launcher = planner.try_get_launcher(); + if (!launcher) { return false; } + + int64_t* d_idx = nullptr; + float* d_dist = nullptr; + RAFT_CUDA_TRY(cudaMallocAsync(&d_idx, m * sizeof(int64_t), stream)); + RAFT_CUDA_TRY(cudaMallocAsync(&d_dist, m * sizeof(float), stream)); + + int64_t shape_x[2] = {m, k}; + int64_t stride_x[2] = {k, 1}; + int64_t shape_y[2] = {n, k}; + int64_t stride_y[2] = {k, 1}; + int64_t shape_idx[1] = {m}; + int64_t stride_idx[1] = {1}; + int64_t shape_dist[1] = {m}; + int64_t stride_dist[1] = {1}; + + int64_t M = m, N = n, K = k; + constexpr int64_t tm = 128, tn = 256, tk = 64; + + void* x_ptr = const_cast(x); + void* y_ptr = const_cast(y); + void* idx_ptr = d_idx; + void* dist_ptr = d_dist; + + dim3 grid((m + tm - 1) / tm, 1, 1); + dim3 block(1, 1, 1); + + using fused_1nn_cutile_kernel_t = void(void*, + int64_t*, + int64_t*, + void*, + int64_t*, + int64_t*, + void*, + int64_t*, + int64_t*, + void*, + int64_t*, + int64_t*, + int64_t, + int64_t, + int64_t, + int64_t, + int64_t, + int64_t); + launcher->template dispatch( + stream, + grid, + block, + 0, + x_ptr, + shape_x, + stride_x, + y_ptr, + shape_y, + stride_y, + idx_ptr, + shape_idx, + stride_idx, + dist_ptr, + shape_dist, + stride_dist, + M, + N, + K, + tm, + tn, + tk); + + pack_fused_1nn_kvp<<<(m + 255) / 256, 256, 0, stream>>>(out, d_idx, d_dist, m); + RAFT_CUDA_TRY(cudaGetLastError()); + RAFT_CUDA_TRY(cudaFreeAsync(d_idx, stream)); + RAFT_CUDA_TRY(cudaFreeAsync(d_dist, stream)); + return true; +} + +} // namespace + +template , int>> +bool try_fused_1nn_tile(OutT* min, + const DataT* x, + const DataT* y, + IdxT m, + IdxT n, + IdxT k, + cuvs::distance::DistanceType metric, + cudaStream_t stream) +{ + if (metric != cuvs::distance::DistanceType::InnerProduct) { return false; } + + if constexpr (std::is_same_v) { + return launch_fused_1nn_tile( + x, y, min, m, n, k, stream); + } else if constexpr (std::is_same_v) { + return launch_fused_1nn_tile( + x, y, min, m, n, k, stream); + } else { + return false; + } +} + +using kvp_i_f = raft::KeyValuePair; +using kvp_i64_f = raft::KeyValuePair; +using kvp_i_h = raft::KeyValuePair; +using kvp_i64_h = raft::KeyValuePair; + +#define CUVS_INST_TRY_FUSED_1NN_TILE(DataT, OutT, IdxT) \ + template CUVS_EXPORT bool try_fused_1nn_tile(OutT*, \ + const DataT*, \ + const DataT*, \ + IdxT, \ + IdxT, \ + IdxT, \ + cuvs::distance::DistanceType, \ + cudaStream_t) + +// int and int32_t are the same on LP64; one instantiation covers both. +CUVS_INST_TRY_FUSED_1NN_TILE(float, kvp_i_f, int); +CUVS_INST_TRY_FUSED_1NN_TILE(float, kvp_i64_f, int64_t); +CUVS_INST_TRY_FUSED_1NN_TILE(half, kvp_i_f, int); +CUVS_INST_TRY_FUSED_1NN_TILE(half, kvp_i64_f, int64_t); +CUVS_INST_TRY_FUSED_1NN_TILE(half, kvp_i_h, int); +CUVS_INST_TRY_FUSED_1NN_TILE(half, kvp_i64_h, int64_t); + +#undef CUVS_INST_TRY_FUSED_1NN_TILE + +} // namespace detail +} // namespace distance +} // namespace cuvs diff --git a/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_tile.hpp b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_tile.hpp new file mode 100644 index 0000000000..30f804d399 --- /dev/null +++ b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_tile.hpp @@ -0,0 +1,55 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +#include + +#include +#include + +#include + +namespace cuvs { +namespace distance { +namespace detail { + +template +inline constexpr bool is_fused_1nn_kvp_output_v = + std::is_same_v> || + std::is_same_v>; + +template , int> = 0> +bool try_fused_1nn_tile(OutT* min, + const DataT* x, + const DataT* y, + IdxT m, + IdxT n, + IdxT k, + cuvs::distance::DistanceType metric, + cudaStream_t stream); + +template , int> = 0> +bool try_fused_1nn_tile(OutT*, + const DataT*, + const DataT*, + IdxT, + IdxT, + IdxT, + cuvs::distance::DistanceType, + cudaStream_t) +{ + return false; +} + +} // namespace detail +} // namespace distance +} // namespace cuvs diff --git a/cpp/src/distance/detail/pairwise_matrix/jit_lto_kernels/pairwise_matrix_planner.hpp b/cpp/src/distance/detail/pairwise_matrix/jit_lto_kernels/pairwise_matrix_planner.hpp index 0d00b3eca6..f89a383596 100644 --- a/cpp/src/distance/detail/pairwise_matrix/jit_lto_kernels/pairwise_matrix_planner.hpp +++ b/cpp/src/distance/detail/pairwise_matrix/jit_lto_kernels/pairwise_matrix_planner.hpp @@ -20,7 +20,7 @@ template -struct PairwiseMatrixPlanner : AlgorithmPlanner { +struct PairwiseMatrixPlanner : LTOAlgorithmPlanner { using DistanceTag = DistanceTag_; using DataTag = DataTag_; using AccTag = AccTag_; @@ -33,7 +33,7 @@ struct PairwiseMatrixPlanner : AlgorithmPlanner { inline static LauncherJitCache launcher_jit_cache{}; - PairwiseMatrixPlanner() : AlgorithmPlanner(kPairwiseMatrixJitEntrypoint, launcher_jit_cache) {} + PairwiseMatrixPlanner() : LTOAlgorithmPlanner(kPairwiseMatrixJitEntrypoint, launcher_jit_cache) {} void add_entrypoint() { diff --git a/cpp/src/neighbors/detail/cagra/jit_lto_kernels/cagra_planner_base.hpp b/cpp/src/neighbors/detail/cagra/jit_lto_kernels/cagra_planner_base.hpp index 0c3ed64d13..b44a7f044e 100644 --- a/cpp/src/neighbors/detail/cagra/jit_lto_kernels/cagra_planner_base.hpp +++ b/cpp/src/neighbors/detail/cagra/jit_lto_kernels/cagra_planner_base.hpp @@ -25,7 +25,7 @@ template -struct CagraPlannerBase : AlgorithmPlanner { +struct CagraPlannerBase : LTOAlgorithmPlanner { using DataTag = DataTag_; using IndexTag = IndexTag_; using DistanceTag = DistanceTag_; @@ -34,7 +34,7 @@ struct CagraPlannerBase : AlgorithmPlanner { using SampleFilterJitTag = SampleFilterJitTag_; explicit CagraPlannerBase(std::string entrypoint, LauncherJitCache& jit_cache) - : AlgorithmPlanner(std::move(entrypoint), jit_cache) + : LTOAlgorithmPlanner(std::move(entrypoint), jit_cache) { } diff --git a/cpp/src/neighbors/ivf_flat/detail/jit_lto_kernels/interleaved_scan_planner.hpp b/cpp/src/neighbors/ivf_flat/detail/jit_lto_kernels/interleaved_scan_planner.hpp index ed8191016b..7899d970ab 100644 --- a/cpp/src/neighbors/ivf_flat/detail/jit_lto_kernels/interleaved_scan_planner.hpp +++ b/cpp/src/neighbors/ivf_flat/detail/jit_lto_kernels/interleaved_scan_planner.hpp @@ -14,10 +14,10 @@ namespace cuvs::neighbors::ivf_flat::detail { -struct InterleavedScanPlanner : AlgorithmPlanner { +struct InterleavedScanPlanner : LTOAlgorithmPlanner { inline static LauncherJitCache launcher_jit_cache{}; - InterleavedScanPlanner() : AlgorithmPlanner("interleaved_scan", launcher_jit_cache) {} + InterleavedScanPlanner() : LTOAlgorithmPlanner("interleaved_scan", launcher_jit_cache) {} template void add_entrypoint() diff --git a/cpp/src/neighbors/ivf_pq/detail/jit_lto_kernels/compute_similarity_planner.hpp b/cpp/src/neighbors/ivf_pq/detail/jit_lto_kernels/compute_similarity_planner.hpp index 0621966cad..7152aaeebd 100644 --- a/cpp/src/neighbors/ivf_pq/detail/jit_lto_kernels/compute_similarity_planner.hpp +++ b/cpp/src/neighbors/ivf_pq/detail/jit_lto_kernels/compute_similarity_planner.hpp @@ -12,10 +12,10 @@ namespace cuvs::neighbors::ivf_pq::detail { -struct ComputeSimilarityPlanner : AlgorithmPlanner { +struct ComputeSimilarityPlanner : LTOAlgorithmPlanner { inline static LauncherJitCache launcher_jit_cache{}; - ComputeSimilarityPlanner() : AlgorithmPlanner("compute_similarity", launcher_jit_cache) {} + ComputeSimilarityPlanner() : LTOAlgorithmPlanner("compute_similarity", launcher_jit_cache) {} template void add_entrypoint() diff --git a/cpp/src/neighbors/ivf_sq/detail/jit_lto_kernels/scan_planner.hpp b/cpp/src/neighbors/ivf_sq/detail/jit_lto_kernels/scan_planner.hpp index 05ea34532e..5dc47dc612 100644 --- a/cpp/src/neighbors/ivf_sq/detail/jit_lto_kernels/scan_planner.hpp +++ b/cpp/src/neighbors/ivf_sq/detail/jit_lto_kernels/scan_planner.hpp @@ -13,10 +13,10 @@ namespace cuvs::neighbors::ivf_sq::detail { -struct IvfSqScanPlanner : AlgorithmPlanner { +struct IvfSqScanPlanner : LTOAlgorithmPlanner { inline static LauncherJitCache launcher_jit_cache{}; - IvfSqScanPlanner() : AlgorithmPlanner("ivf_sq_scan", launcher_jit_cache) {} + IvfSqScanPlanner() : LTOAlgorithmPlanner("ivf_sq_scan", launcher_jit_cache) {} template void add_entrypoint() diff --git a/cpp/tests/CMakeLists.txt b/cpp/tests/CMakeLists.txt index ba6ed6e0e7..006b35b5c4 100644 --- a/cpp/tests/CMakeLists.txt +++ b/cpp/tests/CMakeLists.txt @@ -386,7 +386,8 @@ ConfigureTest( PERCENT 100 ) -add_subdirectory(cutile) +# cuTile vector-add example test disabled; fused 1-NN cuTile is covered via libcuvs integration. +# add_subdirectory(cutile) # ################################################################################################## # Install tests #################################################################################### diff --git a/cpp/tests/cutile/cutile_vector_add.cu b/cpp/tests/cutile/cutile_vector_add.cu index 77a5e51311..07d694bef1 100644 --- a/cpp/tests/cutile/cutile_vector_add.cu +++ b/cpp/tests/cutile/cutile_vector_add.cu @@ -11,10 +11,15 @@ #include "vector_add_sm_80_cubin.h" #include "vector_add_sm_86_cubin.h" #include "vector_add_sm_90_cubin.h" +#include "vector_add_tileir_bytecode.h" + +#include #include #include +#include +#include namespace cuvs { namespace { @@ -26,7 +31,7 @@ struct EmbeddedCubin { size_t size; }; -// Lookup table for cubins built at configure time (see export_vector_add_cubin.py). +// Prebuilt cubins for known library targets (see export_vector_add_cubin.py). constexpr EmbeddedCubin kEmbeddedCubins[] = { {8, 0, vector_add_sm_80_cubin, sizeof(vector_add_sm_80_cubin)}, {8, 6, vector_add_sm_86_cubin, sizeof(vector_add_sm_86_cubin)}, @@ -35,53 +40,128 @@ constexpr EmbeddedCubin kEmbeddedCubins[] = { {12, 0, vector_add_sm_120_cubin, sizeof(vector_add_sm_120_cubin)}, }; -const EmbeddedCubin* find_embedded_cubin(int cc_major, int cc_minor) +constexpr EmbeddedCubin kTileIrBytecode = { + -1, + -1, + vector_add_tileir_bytecode, + sizeof(vector_add_tileir_bytecode), +}; + +struct CutileModuleImage { + const uint8_t* data; + size_t size; +}; + +std::optional resolve_vector_add_module(int cc_major, int cc_minor) { for (const auto& entry : kEmbeddedCubins) { - if (entry.cc_major == cc_major && entry.cc_minor == cc_minor) { return &entry; } + if (entry.cc_major == cc_major && entry.cc_minor == cc_minor) { + return CutileModuleImage{reinterpret_cast(entry.data), entry.size}; + } } - // Fall back to a cubin for the same major version (e.g. minor SKUs within a generation). - for (const auto& entry : kEmbeddedCubins) { - if (entry.cc_major == cc_major) { return &entry; } + + int driver_version = 0; + if (cudaDriverGetVersion(&driver_version) != cudaSuccess) { return std::nullopt; } + if (!cuvs::detail::jit_lto::tileir_fallback_available(driver_version)) { + return std::nullopt; } - return nullptr; + return CutileModuleImage{ + reinterpret_cast(kTileIrBytecode.data), kTileIrBytecode.size}; } -class CutileVectorAddTest : public ::testing::Test { - protected: - void SetUp() override +struct LoadedKernel { + cudaLibrary_t library = nullptr; + cudaKernel_t kernel = nullptr; + bool used_tileir_jit{false}; + const char* skip_reason{nullptr}; + + LoadedKernel() = default; + + LoadedKernel(LoadedKernel&& other) noexcept { *this = std::move(other); } + + LoadedKernel& operator=(LoadedKernel&& other) noexcept { - int device = 0; - RAFT_CUDA_TRY(cudaGetDevice(&device)); - RAFT_CUDA_TRY( - cudaDeviceGetAttribute(&cc_major_, cudaDevAttrComputeCapabilityMajor, device)); - RAFT_CUDA_TRY( - cudaDeviceGetAttribute(&cc_minor_, cudaDevAttrComputeCapabilityMinor, device)); + if (this != &other) { + unload(); + library = other.library; + kernel = other.kernel; + used_tileir_jit = other.used_tileir_jit; + skip_reason = other.skip_reason; + other.library = nullptr; + other.kernel = nullptr; + } + return *this; } - int cc_major_{}; - int cc_minor_{}; -}; + LoadedKernel(const LoadedKernel&) = delete; + LoadedKernel& operator=(const LoadedKernel&) = delete; -} // namespace + ~LoadedKernel() { unload(); } -TEST_F(CutileVectorAddTest, EmbeddedCubinVectorAdd) + explicit operator bool() const { return kernel != nullptr; } + + private: + void unload() + { + if (library != nullptr) { + RAFT_CUDA_TRY(cudaLibraryUnload(library)); + library = nullptr; + kernel = nullptr; + } + } +}; + +LoadedKernel load_vector_add_kernel(int cc_major, int cc_minor) { - const EmbeddedCubin* cubin = find_embedded_cubin(cc_major_, cc_minor_); - ASSERT_NE(cubin, nullptr) - << "No embedded cuTile cubin for compute capability " << cc_major_ << "." << cc_minor_; + LoadedKernel result{}; + result.used_tileir_jit = !cuvs::detail::jit_lto::is_embedded_cubin_arch(cc_major, cc_minor); + + auto image = resolve_vector_add_module(cc_major, cc_minor); + if (!image) { + if (result.used_tileir_jit) { + result.skip_reason = + "TileIR driver JIT unavailable for this GPU. Requires CUDA 13.1+ driver (>= 590.44)."; + } else { + ADD_FAILURE() << "No embedded cuTile module for compute capability " << cc_major << "." + << cc_minor; + } + return result; + } - cudaLibrary_t library{}; - ASSERT_EQ(cudaSuccess, - cudaLibraryLoadData( - &library, cubin->data, nullptr, nullptr, 0, nullptr, nullptr, 0)) - << "cudaLibraryLoadData failed: " << cudaGetErrorString(cudaGetLastError()); + const cudaError_t load_status = + cudaLibraryLoadData(&result.library, image->data, nullptr, nullptr, 0, nullptr, nullptr, 0); + if (load_status != cudaSuccess) { + if (result.used_tileir_jit) { + result.skip_reason = + "TileIR driver JIT unavailable for this GPU (requires CUDA 13.1+ driver >= 590.44)."; + SCOPED_TRACE(cudaGetErrorString(load_status)); + } else { + ADD_FAILURE() << "cudaLibraryLoadData failed: " << cudaGetErrorString(load_status); + } + return result; + } - cudaKernel_t kernel{}; - ASSERT_EQ(cudaSuccess, - cudaLibraryGetKernel(&kernel, library, CUTILE_VECTOR_ADD_KERNEL_SYMBOL)) - << "cudaLibraryGetKernel failed: " << cudaGetErrorString(cudaGetLastError()); + const cudaError_t kernel_status = + cudaLibraryGetKernel(&result.kernel, result.library, CUTILE_VECTOR_ADD_KERNEL_SYMBOL); + if (kernel_status != cudaSuccess) { + if (result.library != nullptr) { + RAFT_CUDA_TRY(cudaLibraryUnload(result.library)); + result.library = nullptr; + } + result.kernel = nullptr; + if (result.used_tileir_jit) { + result.skip_reason = + "TileIR driver JIT unavailable for this GPU (requires CUDA 13.1+ driver >= 590.44)."; + SCOPED_TRACE(cudaGetErrorString(kernel_status)); + } else { + ADD_FAILURE() << "cudaLibraryGetKernel failed: " << cudaGetErrorString(kernel_status); + } + } + return result; +} +void run_vector_add(cudaKernel_t kernel) +{ constexpr int kN = 1024; constexpr int kTile = 256; constexpr int kGridDim = (kN + kTile - 1) / kTile; @@ -122,7 +202,35 @@ TEST_F(CutileVectorAddTest, EmbeddedCubinVectorAdd) RAFT_CUDA_TRY(cudaFree(d_a)); RAFT_CUDA_TRY(cudaFree(d_b)); RAFT_CUDA_TRY(cudaFree(d_c)); - RAFT_CUDA_TRY(cudaLibraryUnload(library)); +} + +class CutileVectorAddTest : public ::testing::Test { + protected: + void SetUp() override + { + int device = 0; + RAFT_CUDA_TRY(cudaGetDevice(&device)); + RAFT_CUDA_TRY( + cudaDeviceGetAttribute(&cc_major_, cudaDevAttrComputeCapabilityMajor, device)); + RAFT_CUDA_TRY( + cudaDeviceGetAttribute(&cc_minor_, cudaDevAttrComputeCapabilityMinor, device)); + } + + int cc_major_{}; + int cc_minor_{}; +}; + +} // namespace + +TEST_F(CutileVectorAddTest, EmbeddedCubinVectorAdd) +{ + LoadedKernel loaded = load_vector_add_kernel(cc_major_, cc_minor_); + if (loaded.skip_reason) { GTEST_SKIP() << loaded.skip_reason; } + if (!loaded) { return; } + + SCOPED_TRACE(loaded.used_tileir_jit ? "loaded via TileIR driver JIT" + : "loaded via prebuilt cubin"); + run_vector_add(loaded.kernel); } } // namespace cuvs diff --git a/cpp/tests/cutile/export_vector_add_cubin.py b/cpp/tests/cutile/export_vector_add_cubin.py index bf40a4ad80..fa099189cd 100644 --- a/cpp/tests/cutile/export_vector_add_cubin.py +++ b/cpp/tests/cutile/export_vector_add_cubin.py @@ -1,12 +1,13 @@ # SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. # SPDX-License-Identifier: Apache-2.0 -"""Export the cuTile vector-add kernel to a cubin for a single GPU target.""" +"""Export the cuTile vector-add kernel to cubin or TileIR bytecode.""" from __future__ import annotations import argparse import sys from pathlib import Path +from typing import Literal import cuda.tile as ct from cuda.tile.compilation import ( @@ -28,6 +29,9 @@ # sm_120 -> 120a-real SUPPORTED_GPU_CODES = ("sm_80", "sm_86", "sm_90", "sm_100", "sm_120") +# Minimum TileIR bytecode version supported by cuTile; also the most portable choice. +DEFAULT_TILEIR_BYTECODE_VERSION = "13.1" + def _kernel_signature() -> KernelSignature: array = ArrayConstraint( @@ -45,20 +49,31 @@ def _kernel_signature() -> KernelSignature: ).with_mangled_symbol("vector_add") -def export_cubin(output_file: Path, gpu_code: str, symbol_header: Path | None) -> str: - if gpu_code not in SUPPORTED_GPU_CODES: +def export_kernel_binary( + output_file: Path, + *, + output_format: Literal["cubin", "tileir_bytecode"], + gpu_code: str, + bytecode_version: str | None = None, + symbol_header: Path | None = None, +) -> str: + if output_format == "cubin" and gpu_code not in SUPPORTED_GPU_CODES: raise ValueError( f"Unsupported gpu_code {gpu_code!r}; expected one of {SUPPORTED_GPU_CODES}" ) signature = _kernel_signature() - export_kernel( - vector_add, - signatures=[signature], - output_file=str(output_file), - gpu_code=gpu_code, - output_format="cubin", - ) + export_kwargs: dict = { + "kernel": vector_add, + "signatures": [signature], + "output_file": str(output_file), + "gpu_code": gpu_code, + "output_format": output_format, + } + if output_format == "tileir_bytecode": + export_kwargs["bytecode_version"] = bytecode_version or DEFAULT_TILEIR_BYTECODE_VERSION + + export_kernel(**export_kwargs) if symbol_header is not None: symbol_header.write_text( @@ -77,12 +92,23 @@ def export_cubin(output_file: Path, gpu_code: str, symbol_header: Path | None) - def main() -> int: parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument("output_file", type=Path, help="Output cubin path") + parser.add_argument("output_file", type=Path, help="Output cubin or .tilebc path") + parser.add_argument( + "--format", + choices=("cubin", "tileir_bytecode"), + default="cubin", + help="Export format (default: cubin)", + ) parser.add_argument( "--gpu-code", required=True, choices=SUPPORTED_GPU_CODES, - help="tileiras / export_kernel target (e.g. sm_120)", + help="tileiras / export_kernel compile target (e.g. sm_120)", + ) + parser.add_argument( + "--bytecode-version", + default=DEFAULT_TILEIR_BYTECODE_VERSION, + help="TileIR bytecode version when --format=tileir_bytecode (default: 13.1)", ) parser.add_argument( "--symbol-header", @@ -92,7 +118,13 @@ def main() -> int: ) args = parser.parse_args() - symbol = export_cubin(args.output_file, args.gpu_code, args.symbol_header) + symbol = export_kernel_binary( + args.output_file, + output_format=args.format, + gpu_code=args.gpu_code, + bytecode_version=args.bytecode_version, + symbol_header=args.symbol_header, + ) print(symbol) return 0 diff --git a/cpp/tests/cutile/generate_cutile_cubins.cmake b/cpp/tests/cutile/generate_cutile_cubins.cmake index 3425b03028..766d3167c6 100644 --- a/cpp/tests/cutile/generate_cutile_cubins.cmake +++ b/cpp/tests/cutile/generate_cutile_cubins.cmake @@ -78,6 +78,33 @@ function(generate_cutile_vector_add_cubins output_include_dir_var) list(APPEND _generated_headers "${_cubin_header}") endforeach() + # Portable TileIR bytecode for driver JIT on architectures without a prebuilt cubin. + # Requires a CUDA 13.1+ driver (>= 590.44); see Tile IR bytecode docs. + set(_tileir_file "${_cutile_binary_dir}/vector_add.tilebc") + set(_tileir_header "${_cutile_binary_dir}/vector_add_tileir_bytecode.h") + + add_custom_command( + OUTPUT "${_tileir_file}" + COMMAND + "${Python3_EXECUTABLE}" "${_cutile_source_dir}/export_vector_add_cubin.py" + "${_tileir_file}" --format tileir_bytecode --gpu-code sm_80 --bytecode-version 13.1 + DEPENDS "${_cutile_source_dir}/export_vector_add_cubin.py" + "${_cutile_source_dir}/vector_add_kernel.py" + COMMENT "Exporting cuTile vector_add TileIR bytecode (v13.1)" + VERBATIM + ) + + add_custom_command( + OUTPUT "${_tileir_header}" + COMMAND "${CUTILE_BIN2C}" --const --name vector_add_tileir_bytecode --static "${_tileir_file}" + > "${_tileir_header}" + DEPENDS "${_tileir_file}" + COMMENT "Embedding vector_add TileIR bytecode via bin2c" + VERBATIM + ) + + list(APPEND _generated_headers "${_tileir_header}") + add_custom_target( cutile_vector_add_cubins DEPENDS "${_symbol_header}" ${_generated_headers} diff --git a/cpp/tests/neighbors/distance_nn.cu b/cpp/tests/neighbors/distance_nn.cu index f31f3ebacf..f5efaa5bec 100644 --- a/cpp/tests/neighbors/distance_nn.cu +++ b/cpp/tests/neighbors/distance_nn.cu @@ -187,6 +187,7 @@ const std::vector> input_fp32 = { {4096, 16384, 128, DistanceType::L2Expanded, true, uint64_t(31415926), 0.1}, {4096, 4096, 64, DistanceType::L2SqrtExpanded, false, uint64_t(31415926), 0.1}, {4096, 16384, 128, DistanceType::L2SqrtExpanded, false, uint64_t(31415926), 0.1}, + {512, 1024, 64, DistanceType::InnerProduct, false, uint64_t(31415926), 0.1}, {4096, 4096, 64, DistanceType::CosineExpanded, false, uint64_t(31415926), 0.1}, {8192, 4096, 64, DistanceType::CosineExpanded, false, uint64_t(31415926), 0.1}, // Fused implementation for cosine distance ignores the sqrt parameter, therefore diff --git a/cpp/tests/neighbors/distance_nn_helper.cuh b/cpp/tests/neighbors/distance_nn_helper.cuh index fda7b76573..422879918f 100644 --- a/cpp/tests/neighbors/distance_nn_helper.cuh +++ b/cpp/tests/neighbors/distance_nn_helper.cuh @@ -66,6 +66,16 @@ __device__ AccT cosine_distance(const DataT* v1, const DataT* v2, IdxT K) } // This is a naive implementation of 1-NN computation +template +__device__ AccT inner_product_score(const DataT* v1, const DataT* v2, IdxT K) +{ + AccT score = AccT(0.0); + for (IdxT i = 0; i < K; i++) { + score += AccT(v1[i]) * AccT(v2[i]); + } + return score; +} + template RAFT_KERNEL ref_nn_kernel( OutT* out, const DataT* A, const DataT* B, IdxT M, IdxT N, IdxT K, bool sqrt, DistanceType metric) @@ -73,22 +83,47 @@ RAFT_KERNEL ref_nn_kernel( IdxT tid = threadIdx.x + blockIdx.x * IdxT(blockDim.x); for (IdxT m = tid; m < M; m += (blockDim.x * gridDim.x)) { - IdxT min_index = N + 1; - AccT min_dist = max_val(); + IdxT best_index = N + 1; + AccT best_score = min_val(); + AccT best_dist = max_val(); for (IdxT n = 0; n < N; n++) { + if (metric == DistanceType::InnerProduct) { + AccT score = inner_product_score(&A[m * K], &B[n * K], K); + if (score > best_score) { + best_score = score; + best_index = n; + } + continue; + } + AccT dist; if (metric == DistanceType::L2SqrtExpanded || metric == DistanceType::L2Expanded) { dist = l2_distance(&A[m * K], &B[n * K], K); } else if (metric == DistanceType::CosineExpanded) { dist = cosine_distance(&A[m * K], &B[n * K], K); + } else { + continue; + } + if (dist < best_dist) { + best_dist = dist; + best_index = n; } - if (dist < min_dist) { - min_dist = dist; - min_index = n; + } + + if (metric == DistanceType::InnerProduct) { + if constexpr (std::is_fundamental::value) { + out[m] = AccT(best_score); + } else { + out[m].key = IdxT(best_index); + out[m].value = AccT(best_score); } + continue; } + IdxT min_index = best_index; + AccT min_dist = best_dist; + if constexpr (std::is_fundamental::value) { static_assert(std::is_same::value, "OutT and AccT are not same type"); out[m] = AccT(min_dist); From 1b934ddaa684bb09d2df9e611e6c7482b7ee975b Mon Sep 17 00:00:00 2001 From: divyegala Date: Wed, 24 Jun 2026 19:33:18 +0000 Subject: [PATCH 03/82] attempt to fix tile linkage --- cpp/CMakeLists.txt | 22 +-- .../modules/generate_cutile_kernels.cmake | 183 +++++------------- ...cpp.in => register_cutile_fragment.cpp.in} | 8 +- cpp/cmake/modules/register_tileir.cpp.in | 22 --- .../cutile/fused_1nn_cutile_cubin_matrix.json | 40 ---- .../cutile/fused_1nn_cutile_matrix.json | 64 ++++++ .../fused_1nn_cutile_tileir_matrix.json | 20 -- 7 files changed, 120 insertions(+), 239 deletions(-) rename cpp/cmake/modules/{register_cubin.cpp.in => register_cutile_fragment.cpp.in} (57%) delete mode 100644 cpp/cmake/modules/register_tileir.cpp.in delete mode 100644 cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_cutile_cubin_matrix.json create mode 100644 cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_cutile_matrix.json delete mode 100644 cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_cutile_tileir_matrix.json diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt index cc6e1975b3..a1f3f3973c 100644 --- a/cpp/CMakeLists.txt +++ b/cpp/CMakeLists.txt @@ -963,33 +963,21 @@ if(NOT BUILD_CPU_ONLY) "${CMAKE_CURRENT_SOURCE_DIR}/src/distance/detail/fused_distance_nn/cutile") set(cutile_fused_1nn_generated_dir "${CMAKE_CURRENT_BINARY_DIR}/generated_kernels/distance/fused_1nn/cutile") - generate_cutile_cubin_kernels( + generate_cutile_kernels( cutile_fused_1nn_files KERNEL_DIR "${fused_1nn_cutile_dir}" KERNEL_BASENAME "fused_1nn" KERNEL_PYTHON "fused_1nn_kernel.py" EXPORT_SCRIPT "export_fused_1nn.py" OUTPUT_DIRECTORY "${cutile_fused_1nn_generated_dir}" - MATRIX_JSON_FILE "${fused_1nn_cutile_dir}/fused_1nn_cutile_cubin_matrix.json" - FRAGMENT_TAG_FORMAT + MATRIX_JSON_FILE "${fused_1nn_cutile_dir}/fused_1nn_cutile_matrix.json" + FRAGMENT_TAG_FORMAT_CUBIN "cuvs::distance::detail::fragment_tag_fused_1nn_cubin" - FRAGMENT_TAG_HEADER_FILES - "" - "" - "" - ) - generate_cutile_tileir_kernels( - cutile_fused_1nn_files - KERNEL_DIR "${fused_1nn_cutile_dir}" - KERNEL_BASENAME "fused_1nn" - KERNEL_PYTHON "fused_1nn_kernel.py" - EXPORT_SCRIPT "export_fused_1nn.py" - OUTPUT_DIRECTORY "${cutile_fused_1nn_generated_dir}" - MATRIX_JSON_FILE "${fused_1nn_cutile_dir}/fused_1nn_cutile_tileir_matrix.json" - FRAGMENT_TAG_FORMAT + FRAGMENT_TAG_FORMAT_TILEIR "cuvs::distance::detail::fragment_tag_fused_1nn_tileir" FRAGMENT_TAG_HEADER_FILES "" + "" "" ) if(NOT DEFINED CUVS_CUTILE_ENABLED) diff --git a/cpp/cmake/modules/generate_cutile_kernels.cmake b/cpp/cmake/modules/generate_cutile_kernels.cmake index 7b9c2521c4..f0219dc842 100644 --- a/cpp/cmake/modules/generate_cutile_kernels.cmake +++ b/cpp/cmake/modules/generate_cutile_kernels.cmake @@ -77,13 +77,15 @@ function(_cutile_kernels_setup) file(MAKE_DIRECTORY "${_CUTILE_OUTPUT_DIRECTORY}") + set(Python3_EXECUTABLE "${Python3_EXECUTABLE}" PARENT_SCOPE) + set(CUTILE_BIN2C "${CUTILE_BIN2C}" PARENT_SCOPE) set(_CUTILE_SETUP_OK TRUE PARENT_SCOPE ) endfunction() -function(process_cutile_cubin_matrix_entry source_list_var) +function(process_cutile_matrix_entry source_list_var) set(options) set(one_value KERNEL_DIR @@ -91,110 +93,75 @@ function(process_cutile_cubin_matrix_entry source_list_var) KERNEL_PYTHON EXPORT_SCRIPT OUTPUT_DIRECTORY - FRAGMENT_TAG_FORMAT + FRAGMENT_TAG_FORMAT_CUBIN + FRAGMENT_TAG_FORMAT_TILEIR MATRIX_JSON_ENTRY ) set(multi_value FRAGMENT_TAG_HEADER_FILES) cmake_parse_arguments(_CUTILE "${options}" "${one_value}" "${multi_value}" ${ARGN}) - populate_matrix_variables("${_CUTILE_MATRIX_JSON_ENTRY}") - _cutile_fragment_tag_header_files( - fragment_tag_header_files ${_CUTILE_FRAGMENT_TAG_HEADER_FILES} - ) - - string(CONFIGURE "${_CUTILE_FRAGMENT_TAG_FORMAT}" fragment_tag @ONLY) - - set(_artifact_basename "${_CUTILE_KERNEL_BASENAME}_${data_type}_${gpu_code}") - set(_cubin_file "${_CUTILE_OUTPUT_DIRECTORY}/${_artifact_basename}.cubin") - set(_cubin_header "${_CUTILE_OUTPUT_DIRECTORY}/${_artifact_basename}_cubin.h") - set(_cubin_cpp "${_CUTILE_OUTPUT_DIRECTORY}/${_artifact_basename}_cubin.cpp") - set(cubin_header_file "${_artifact_basename}_cubin.h") - - add_custom_command( - OUTPUT "${_cubin_file}" - COMMAND - "${Python3_EXECUTABLE}" "${_CUTILE_KERNEL_DIR}/${_CUTILE_EXPORT_SCRIPT}" "${_cubin_file}" - --format cubin --data-type "${data_type}" --gpu-code "${gpu_code}" - DEPENDS "${_CUTILE_KERNEL_DIR}/${_CUTILE_EXPORT_SCRIPT}" - "${_CUTILE_KERNEL_DIR}/${_CUTILE_KERNEL_PYTHON}" - COMMENT "Exporting cuTile ${_CUTILE_KERNEL_BASENAME} cubin ${data_type} ${gpu_code}" - VERBATIM - ) - - add_custom_command( - OUTPUT "${_cubin_header}" - COMMAND "${CUTILE_BIN2C}" --const --name embedded_cubin --static "${_cubin_file}" - > "${_cubin_header}" - DEPENDS "${_cubin_file}" - VERBATIM - ) + find_package(Python3 REQUIRED COMPONENTS Interpreter) - configure_file( - "${CMAKE_CURRENT_FUNCTION_LIST_DIR}/register_cubin.cpp.in" "${_cubin_cpp}" @ONLY - ) - list(APPEND ${source_list_var} "${_cubin_header}" "${_cubin_cpp}") - set(${source_list_var} - "${${source_list_var}}" - PARENT_SCOPE - ) -endfunction() + populate_matrix_variables("${_CUTILE_MATRIX_JSON_ENTRY}") -function(process_cutile_tileir_matrix_entry source_list_var) - set(options) - set(one_value - KERNEL_DIR - KERNEL_BASENAME - KERNEL_PYTHON - EXPORT_SCRIPT - OUTPUT_DIRECTORY - FRAGMENT_TAG_FORMAT - MATRIX_JSON_ENTRY - ) - set(multi_value FRAGMENT_TAG_HEADER_FILES) - cmake_parse_arguments(_CUTILE "${options}" "${one_value}" "${multi_value}" ${ARGN}) + if(register STREQUAL "cubin") + string(CONFIGURE "${_CUTILE_FRAGMENT_TAG_FORMAT_CUBIN}" fragment_tag @ONLY) + set(bin2c_symbol embedded_cubin) + set(fragment_entry_type "StaticCubinFragmentEntry") + elseif(register STREQUAL "tileir") + string(CONFIGURE "${_CUTILE_FRAGMENT_TAG_FORMAT_TILEIR}" fragment_tag @ONLY) + set(bin2c_symbol embedded_tileir) + set(fragment_entry_type "StaticTileIrBytecodeFragmentEntry") + else() + message(FATAL_ERROR "Unknown cuTile register kind '${register}'") + endif() - populate_matrix_variables("${_CUTILE_MATRIX_JSON_ENTRY}") _cutile_fragment_tag_header_files( fragment_tag_header_files ${_CUTILE_FRAGMENT_TAG_HEADER_FILES} ) - string(CONFIGURE "${_CUTILE_FRAGMENT_TAG_FORMAT}" fragment_tag @ONLY) - set(_tileir_file "${_CUTILE_OUTPUT_DIRECTORY}/${_CUTILE_KERNEL_BASENAME}_${data_type}.tilebc") - set(_tileir_header "${_CUTILE_OUTPUT_DIRECTORY}/${_CUTILE_KERNEL_BASENAME}_${data_type}_tileir.h") - set(_tileir_cpp "${_CUTILE_OUTPUT_DIRECTORY}/${_CUTILE_KERNEL_BASENAME}_${data_type}_tileir.cpp") - set(tileir_header_file "${_CUTILE_KERNEL_BASENAME}_${data_type}_tileir.h") + string(CONFIGURE "${artifact_basename}" _artifact_basename @ONLY) + set(_artifact_stem "${_CUTILE_KERNEL_BASENAME}_${_artifact_basename}") + set(_artifact_file "${_CUTILE_OUTPUT_DIRECTORY}/${_artifact_stem}.${artifact_ext}") + set(_embedded_header "${_CUTILE_OUTPUT_DIRECTORY}/${_artifact_stem}_${register}.h") + set(_fragment_cpp "${_CUTILE_OUTPUT_DIRECTORY}/${_artifact_stem}_${register}.cpp") + set(embedded_header_file "${_artifact_stem}_${register}.h") + + set(_python_args --format "${output_format}" --data-type "${data_type}" --gpu-code "${gpu_code}") + if(DEFINED bytecode_version AND NOT "${bytecode_version}" STREQUAL "") + list(APPEND _python_args --bytecode-version "${bytecode_version}") + endif() add_custom_command( - OUTPUT "${_tileir_file}" - COMMAND - "${Python3_EXECUTABLE}" "${_CUTILE_KERNEL_DIR}/${_CUTILE_EXPORT_SCRIPT}" "${_tileir_file}" - --format tileir_bytecode --data-type "${data_type}" --gpu-code "${export_gpu_code}" - --bytecode-version "${bytecode_version}" + OUTPUT "${_artifact_file}" + COMMAND "${Python3_EXECUTABLE}" "${_CUTILE_KERNEL_DIR}/${_CUTILE_EXPORT_SCRIPT}" + "${_artifact_file}" ${_python_args} + WORKING_DIRECTORY "${_CUTILE_KERNEL_DIR}" DEPENDS "${_CUTILE_KERNEL_DIR}/${_CUTILE_EXPORT_SCRIPT}" "${_CUTILE_KERNEL_DIR}/${_CUTILE_KERNEL_PYTHON}" - COMMENT "Exporting cuTile ${_CUTILE_KERNEL_BASENAME} TileIR bytecode ${data_type}" + COMMENT "Exporting cuTile ${_CUTILE_KERNEL_BASENAME} ${output_format} ${data_type}" VERBATIM ) add_custom_command( - OUTPUT "${_tileir_header}" - COMMAND "${CUTILE_BIN2C}" --const --name embedded_tileir --static "${_tileir_file}" - > "${_tileir_header}" - DEPENDS "${_tileir_file}" + OUTPUT "${_embedded_header}" + COMMAND "${CUTILE_BIN2C}" --const --name ${bin2c_symbol} --static "${_artifact_file}" + > "${_embedded_header}" + DEPENDS "${_artifact_file}" VERBATIM ) configure_file( - "${CMAKE_CURRENT_FUNCTION_LIST_DIR}/register_tileir.cpp.in" "${_tileir_cpp}" @ONLY + "${CMAKE_CURRENT_FUNCTION_LIST_DIR}/register_cutile_fragment.cpp.in" "${_fragment_cpp}" @ONLY ) - list(APPEND ${source_list_var} "${_tileir_header}" "${_tileir_cpp}") + list(APPEND ${source_list_var} "${_embedded_header}" "${_fragment_cpp}") set(${source_list_var} "${${source_list_var}}" PARENT_SCOPE ) endfunction() -function(generate_cutile_cubin_kernels source_list_var) +function(generate_cutile_kernels source_list_var) set(options) set(one_value KERNEL_DIR @@ -203,13 +170,14 @@ function(generate_cutile_cubin_kernels source_list_var) EXPORT_SCRIPT OUTPUT_DIRECTORY MATRIX_JSON_FILE - FRAGMENT_TAG_FORMAT + FRAGMENT_TAG_FORMAT_CUBIN + FRAGMENT_TAG_FORMAT_TILEIR ) set(multi_value FRAGMENT_TAG_HEADER_FILES) cmake_parse_arguments(_CUTILE "${options}" "${one_value}" "${multi_value}" ${ARGN}) if(NOT _CUTILE_KERNEL_BASENAME) - message(FATAL_ERROR "generate_cutile_cubin_kernels: KERNEL_BASENAME is required") + message(FATAL_ERROR "generate_cutile_kernels: KERNEL_BASENAME is required") endif() if(NOT _CUTILE_KERNEL_PYTHON) set(_CUTILE_KERNEL_PYTHON "fused_1nn_kernel.py") @@ -236,72 +204,15 @@ function(generate_cutile_cubin_kernels source_list_var) # cmake-lint: disable=C0103,E1120 foreach(i RANGE "${last}") string(JSON matrix_json_entry GET "${matrix_product}" "${i}") - process_cutile_cubin_matrix_entry( - "${source_list_var}" - KERNEL_DIR "${_CUTILE_KERNEL_DIR}" - KERNEL_BASENAME "${_CUTILE_KERNEL_BASENAME}" - KERNEL_PYTHON "${_CUTILE_KERNEL_PYTHON}" - EXPORT_SCRIPT "${_CUTILE_EXPORT_SCRIPT}" - OUTPUT_DIRECTORY "${_CUTILE_OUTPUT_DIRECTORY}" - FRAGMENT_TAG_FORMAT "${_CUTILE_FRAGMENT_TAG_FORMAT}" - FRAGMENT_TAG_HEADER_FILES ${_CUTILE_FRAGMENT_TAG_HEADER_FILES} - MATRIX_JSON_ENTRY "${matrix_json_entry}" - ) - endforeach() - - set(CUVS_CUTILE_ENABLED 1 PARENT_SCOPE) - set(${source_list_var} - "${${source_list_var}}" - PARENT_SCOPE - ) -endfunction() - -function(generate_cutile_tileir_kernels source_list_var) - set(options) - set(one_value - KERNEL_DIR - KERNEL_BASENAME - KERNEL_PYTHON - EXPORT_SCRIPT - OUTPUT_DIRECTORY - MATRIX_JSON_FILE - FRAGMENT_TAG_FORMAT - ) - set(multi_value FRAGMENT_TAG_HEADER_FILES) - cmake_parse_arguments(_CUTILE "${options}" "${one_value}" "${multi_value}" ${ARGN}) - - if(NOT _CUTILE_KERNEL_BASENAME) - message(FATAL_ERROR "generate_cutile_tileir_kernels: KERNEL_BASENAME is required") - endif() - if(NOT _CUTILE_KERNEL_PYTHON) - set(_CUTILE_KERNEL_PYTHON "fused_1nn_kernel.py") - endif() - - _cutile_kernels_setup( - MATRIX_JSON_FILE "${_CUTILE_MATRIX_JSON_FILE}" - OUTPUT_DIRECTORY "${_CUTILE_OUTPUT_DIRECTORY}" - ) - if(NOT _CUTILE_SETUP_OK) - generate_cutile_kernels_stub() - return() - endif() - - compute_matrix_product(matrix_product MATRIX_JSON_FILE "${_CUTILE_MATRIX_JSON_FILE}") - - string(JSON len LENGTH "${matrix_product}") - math(EXPR last "${len} - 1") - - # cmake-lint: disable=C0103,E1120 - foreach(i RANGE "${last}") - string(JSON matrix_json_entry GET "${matrix_product}" "${i}") - process_cutile_tileir_matrix_entry( + process_cutile_matrix_entry( "${source_list_var}" KERNEL_DIR "${_CUTILE_KERNEL_DIR}" KERNEL_BASENAME "${_CUTILE_KERNEL_BASENAME}" KERNEL_PYTHON "${_CUTILE_KERNEL_PYTHON}" EXPORT_SCRIPT "${_CUTILE_EXPORT_SCRIPT}" OUTPUT_DIRECTORY "${_CUTILE_OUTPUT_DIRECTORY}" - FRAGMENT_TAG_FORMAT "${_CUTILE_FRAGMENT_TAG_FORMAT}" + FRAGMENT_TAG_FORMAT_CUBIN "${_CUTILE_FRAGMENT_TAG_FORMAT_CUBIN}" + FRAGMENT_TAG_FORMAT_TILEIR "${_CUTILE_FRAGMENT_TAG_FORMAT_TILEIR}" FRAGMENT_TAG_HEADER_FILES ${_CUTILE_FRAGMENT_TAG_HEADER_FILES} MATRIX_JSON_ENTRY "${matrix_json_entry}" ) diff --git a/cpp/cmake/modules/register_cubin.cpp.in b/cpp/cmake/modules/register_cutile_fragment.cpp.in similarity index 57% rename from cpp/cmake/modules/register_cubin.cpp.in rename to cpp/cmake/modules/register_cutile_fragment.cpp.in index c27d6829ee..0fc074bdbb 100644 --- a/cpp/cmake/modules/register_cubin.cpp.in +++ b/cpp/cmake/modules/register_cutile_fragment.cpp.in @@ -3,7 +3,7 @@ * SPDX-License-Identifier: Apache-2.0 */ -#include "@cubin_header_file@" +#include "@embedded_header_file@" #include @fragment_tag_header_files@ @@ -11,12 +11,12 @@ namespace { using fragment_tag = @fragment_tag@; -using fragment_entry = StaticCubinFragmentEntry; +using fragment_entry = @fragment_entry_type@; } // namespace template <> -const uint8_t* const fragment_entry::data = embedded_cubin; +const uint8_t* const fragment_entry::data = @bin2c_symbol@; template <> -const size_t fragment_entry::length = sizeof(embedded_cubin); +const size_t fragment_entry::length = sizeof(@bin2c_symbol@); diff --git a/cpp/cmake/modules/register_tileir.cpp.in b/cpp/cmake/modules/register_tileir.cpp.in deleted file mode 100644 index fb81acedbc..0000000000 --- a/cpp/cmake/modules/register_tileir.cpp.in +++ /dev/null @@ -1,22 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. - * SPDX-License-Identifier: Apache-2.0 - */ - -#include "@tileir_header_file@" -#include - -@fragment_tag_header_files@ - -namespace { - -using fragment_tag = @fragment_tag@; -using fragment_entry = StaticTileIrBytecodeFragmentEntry; - -} // namespace - -template <> -const uint8_t* const fragment_entry::data = embedded_tileir; - -template <> -const size_t fragment_entry::length = sizeof(embedded_tileir); diff --git a/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_cutile_cubin_matrix.json b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_cutile_cubin_matrix.json deleted file mode 100644 index fbd4bfdd64..0000000000 --- a/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_cutile_cubin_matrix.json +++ /dev/null @@ -1,40 +0,0 @@ -[ - { - "_data": [ - { - "data_type": "half", - "data_abbrev": "h" - }, - { - "data_type": "float", - "data_abbrev": "f" - } - ], - "_arch": [ - { - "gpu_code": "sm_80", - "cc_major": 8, - "cc_minor": 0, - "arch_tag": "cutile_arch_8_0" - }, - { - "gpu_code": "sm_86", - "cc_major": 8, - "cc_minor": 6, - "arch_tag": "cutile_arch_8_6" - }, - { - "gpu_code": "sm_90", - "cc_major": 9, - "cc_minor": 0, - "arch_tag": "cutile_arch_9_0" - }, - { - "gpu_code": "sm_120", - "cc_major": 12, - "cc_minor": 0, - "arch_tag": "cutile_arch_12_0" - } - ] - } -] diff --git a/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_cutile_matrix.json b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_cutile_matrix.json new file mode 100644 index 0000000000..52955863c5 --- /dev/null +++ b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_cutile_matrix.json @@ -0,0 +1,64 @@ +[ + { + "_data": [ + { + "data_type": "half", + "data_abbrev": "h" + }, + { + "data_type": "float", + "data_abbrev": "f" + } + ], + "_export": [ + { + "output_format": "cubin", + "artifact_ext": "cubin", + "artifact_basename": "@data_type@_@gpu_code@", + "register": "cubin", + "gpu_code": "sm_80", + "cc_major": 8, + "cc_minor": 0, + "arch_tag": "cutile_arch_8_0" + }, + { + "output_format": "cubin", + "artifact_ext": "cubin", + "artifact_basename": "@data_type@_@gpu_code@", + "register": "cubin", + "gpu_code": "sm_86", + "cc_major": 8, + "cc_minor": 6, + "arch_tag": "cutile_arch_8_6" + }, + { + "output_format": "cubin", + "artifact_ext": "cubin", + "artifact_basename": "@data_type@_@gpu_code@", + "register": "cubin", + "gpu_code": "sm_90", + "cc_major": 9, + "cc_minor": 0, + "arch_tag": "cutile_arch_9_0" + }, + { + "output_format": "cubin", + "artifact_ext": "cubin", + "artifact_basename": "@data_type@_@gpu_code@", + "register": "cubin", + "gpu_code": "sm_120", + "cc_major": 12, + "cc_minor": 0, + "arch_tag": "cutile_arch_12_0" + }, + { + "output_format": "tileir_bytecode", + "artifact_ext": "tilebc", + "artifact_basename": "@data_type@", + "register": "tileir", + "gpu_code": "sm_80", + "bytecode_version": "13.1" + } + ] + } +] diff --git a/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_cutile_tileir_matrix.json b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_cutile_tileir_matrix.json deleted file mode 100644 index 364c94594c..0000000000 --- a/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_cutile_tileir_matrix.json +++ /dev/null @@ -1,20 +0,0 @@ -[ - { - "_data": [ - { - "data_type": "half", - "data_abbrev": "h" - }, - { - "data_type": "float", - "data_abbrev": "f" - } - ], - "_tileir": [ - { - "export_gpu_code": "sm_80", - "bytecode_version": "13.1" - } - ] - } -] From c7f7cbd30bcf408c4823606d7f93f9e0df1526cc Mon Sep 17 00:00:00 2001 From: divyegala Date: Wed, 24 Jun 2026 21:10:59 +0000 Subject: [PATCH 04/82] working test, remove example --- .../cutile/fused_1nn_tile.cu | 110 ++++---- .../cutile/fused_1nn_tile.hpp | 15 +- cpp/tests/CMakeLists.txt | 3 - cpp/tests/cutile/CMakeLists.txt | 23 -- cpp/tests/cutile/cutile_vector_add.cu | 236 ------------------ cpp/tests/cutile/export_vector_add_cubin.py | 133 ---------- cpp/tests/cutile/generate_cutile_cubins.cmake | 117 --------- cpp/tests/cutile/vector_add_kernel.py | 17 -- 8 files changed, 60 insertions(+), 594 deletions(-) delete mode 100644 cpp/tests/cutile/CMakeLists.txt delete mode 100644 cpp/tests/cutile/cutile_vector_add.cu delete mode 100644 cpp/tests/cutile/export_vector_add_cubin.py delete mode 100644 cpp/tests/cutile/generate_cutile_cubins.cmake delete mode 100644 cpp/tests/cutile/vector_add_kernel.py diff --git a/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_tile.cu b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_tile.cu index af8b0b181f..0ad4ee62a5 100644 --- a/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_tile.cu +++ b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_tile.cu @@ -16,6 +16,8 @@ namespace detail { namespace { +constexpr int64_t TILE_M = 128; + template __global__ void pack_fused_1nn_kvp(OutT* out, const int64_t* idx, const float* dist, IdxT len) { @@ -27,13 +29,8 @@ __global__ void pack_fused_1nn_kvp(OutT* out, const int64_t* idx, const float* d } template -bool launch_fused_1nn_tile(const DataT* x, - const DataT* y, - OutT* out, - IdxT m, - IdxT n, - IdxT k, - cudaStream_t stream) +bool launch_fused_1nn_tile( + const DataT* x, const DataT* y, OutT* out, IdxT m, IdxT n, IdxT k, cudaStream_t stream) { Fused1nnTilePlanner planner; planner.add_entrypoint(); @@ -46,67 +43,70 @@ bool launch_fused_1nn_tile(const DataT* x, RAFT_CUDA_TRY(cudaMallocAsync(&d_idx, m * sizeof(int64_t), stream)); RAFT_CUDA_TRY(cudaMallocAsync(&d_dist, m * sizeof(float), stream)); - int64_t shape_x[2] = {m, k}; - int64_t stride_x[2] = {k, 1}; - int64_t shape_y[2] = {n, k}; - int64_t stride_y[2] = {k, 1}; - int64_t shape_idx[1] = {m}; - int64_t stride_idx[1] = {1}; - int64_t shape_dist[1] = {m}; - int64_t stride_dist[1] = {1}; + int64_t shape_x[2] = {m, k}; + int64_t stride_x[2] = {k, 1}; + int64_t shape_y[2] = {n, k}; + int64_t stride_y[2] = {k, 1}; + int64_t shape_idx = m; + int64_t stride_idx = 1; + int64_t shape_dist = m; + int64_t stride_dist = 1; int64_t M = m, N = n, K = k; - constexpr int64_t tm = 128, tn = 256, tk = 64; void* x_ptr = const_cast(x); void* y_ptr = const_cast(y); void* idx_ptr = d_idx; void* dist_ptr = d_dist; - dim3 grid((m + tm - 1) / tm, 1, 1); + dim3 grid((m + TILE_M - 1) / TILE_M, 1, 1); dim3 block(1, 1, 1); + // cutile_python_v1 (see fused_1nn_float PTX): each 2D array is (ptr, shape0, shape1, + // stride0, stride1); each 1D array is (ptr, shape, stride); ConstantConstraint tile sizes + // are embedded in the module. using fused_1nn_cutile_kernel_t = void(void*, - int64_t*, - int64_t*, - void*, - int64_t*, - int64_t*, + int64_t, + int64_t, + int64_t, + int64_t, void*, - int64_t*, - int64_t*, + int64_t, + int64_t, + int64_t, + int64_t, void*, - int64_t*, - int64_t*, int64_t, int64_t, + void*, + int64_t, int64_t, int64_t, int64_t, int64_t); - launcher->template dispatch( - stream, - grid, - block, - 0, - x_ptr, - shape_x, - stride_x, - y_ptr, - shape_y, - stride_y, - idx_ptr, - shape_idx, - stride_idx, - dist_ptr, - shape_dist, - stride_dist, - M, - N, - K, - tm, - tn, - tk); + launcher->template dispatch(stream, + grid, + block, + 0, + x_ptr, + shape_x[0], + shape_x[1], + stride_x[0], + stride_x[1], + y_ptr, + shape_y[0], + shape_y[1], + stride_y[0], + stride_y[1], + idx_ptr, + shape_idx, + stride_idx, + dist_ptr, + shape_dist, + stride_dist, + M, + N, + K); pack_fused_1nn_kvp<<<(m + 255) / 256, 256, 0, stream>>>(out, d_idx, d_dist, m); RAFT_CUDA_TRY(cudaGetLastError()); @@ -148,13 +148,13 @@ using kvp_i64_f = raft::KeyValuePair; using kvp_i_h = raft::KeyValuePair; using kvp_i64_h = raft::KeyValuePair; -#define CUVS_INST_TRY_FUSED_1NN_TILE(DataT, OutT, IdxT) \ +#define CUVS_INST_TRY_FUSED_1NN_TILE(DataT, OutT, IdxT) \ template CUVS_EXPORT bool try_fused_1nn_tile(OutT*, \ - const DataT*, \ - const DataT*, \ - IdxT, \ - IdxT, \ - IdxT, \ + const DataT*, \ + const DataT*, \ + IdxT, \ + IdxT, \ + IdxT, \ cuvs::distance::DistanceType, \ cudaStream_t) diff --git a/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_tile.hpp b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_tile.hpp index 30f804d399..d72a020ba7 100644 --- a/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_tile.hpp +++ b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_tile.hpp @@ -18,8 +18,9 @@ namespace detail { template inline constexpr bool is_fused_1nn_kvp_output_v = - std::is_same_v> || - std::is_same_v>; + (std::is_same_v || std::is_same_v) && + (std::is_same_v> || + std::is_same_v>); template , int> = 0> -bool try_fused_1nn_tile(OutT*, - const DataT*, - const DataT*, - IdxT, - IdxT, - IdxT, - cuvs::distance::DistanceType, - cudaStream_t) +bool try_fused_1nn_tile( + OutT*, const DataT*, const DataT*, IdxT, IdxT, IdxT, cuvs::distance::DistanceType, cudaStream_t) { return false; } diff --git a/cpp/tests/CMakeLists.txt b/cpp/tests/CMakeLists.txt index 006b35b5c4..9b96f94bf0 100644 --- a/cpp/tests/CMakeLists.txt +++ b/cpp/tests/CMakeLists.txt @@ -386,9 +386,6 @@ ConfigureTest( PERCENT 100 ) -# cuTile vector-add example test disabled; fused 1-NN cuTile is covered via libcuvs integration. -# add_subdirectory(cutile) - # ################################################################################################## # Install tests #################################################################################### # ################################################################################################## diff --git a/cpp/tests/cutile/CMakeLists.txt b/cpp/tests/cutile/CMakeLists.txt deleted file mode 100644 index 989c8137d0..0000000000 --- a/cpp/tests/cutile/CMakeLists.txt +++ /dev/null @@ -1,23 +0,0 @@ -# ============================================================================= -# cmake-format: off -# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. -# SPDX-License-Identifier: Apache-2.0 -# cmake-format: on -# ============================================================================= - -include("${CMAKE_CURRENT_LIST_DIR}/generate_cutile_cubins.cmake") - -generate_cutile_vector_add_cubins(CUTILE_GENERATED_INCLUDE_DIR) - -ConfigureTest( - NAME CUTILE_VECTOR_ADD_TEST - PATH "${CMAKE_CURRENT_LIST_DIR}/cutile_vector_add.cu" - GPUS 1 - PERCENT 100 -) - -add_dependencies(CUTILE_VECTOR_ADD_TEST cutile_vector_add_cubins) - -target_include_directories( - CUTILE_VECTOR_ADD_TEST PRIVATE "${CUTILE_GENERATED_INCLUDE_DIR}" -) diff --git a/cpp/tests/cutile/cutile_vector_add.cu b/cpp/tests/cutile/cutile_vector_add.cu deleted file mode 100644 index 07d694bef1..0000000000 --- a/cpp/tests/cutile/cutile_vector_add.cu +++ /dev/null @@ -1,236 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. - * SPDX-License-Identifier: Apache-2.0 - */ - -#include "../test_utils.cuh" - -#include "vector_add_kernel_symbol.h" -#include "vector_add_sm_100_cubin.h" -#include "vector_add_sm_120_cubin.h" -#include "vector_add_sm_80_cubin.h" -#include "vector_add_sm_86_cubin.h" -#include "vector_add_sm_90_cubin.h" -#include "vector_add_tileir_bytecode.h" - -#include - -#include - -#include -#include -#include - -namespace cuvs { -namespace { - -struct EmbeddedCubin { - int cc_major; - int cc_minor; - const unsigned char* data; - size_t size; -}; - -// Prebuilt cubins for known library targets (see export_vector_add_cubin.py). -constexpr EmbeddedCubin kEmbeddedCubins[] = { - {8, 0, vector_add_sm_80_cubin, sizeof(vector_add_sm_80_cubin)}, - {8, 6, vector_add_sm_86_cubin, sizeof(vector_add_sm_86_cubin)}, - {9, 0, vector_add_sm_90_cubin, sizeof(vector_add_sm_90_cubin)}, - {10, 0, vector_add_sm_100_cubin, sizeof(vector_add_sm_100_cubin)}, - {12, 0, vector_add_sm_120_cubin, sizeof(vector_add_sm_120_cubin)}, -}; - -constexpr EmbeddedCubin kTileIrBytecode = { - -1, - -1, - vector_add_tileir_bytecode, - sizeof(vector_add_tileir_bytecode), -}; - -struct CutileModuleImage { - const uint8_t* data; - size_t size; -}; - -std::optional resolve_vector_add_module(int cc_major, int cc_minor) -{ - for (const auto& entry : kEmbeddedCubins) { - if (entry.cc_major == cc_major && entry.cc_minor == cc_minor) { - return CutileModuleImage{reinterpret_cast(entry.data), entry.size}; - } - } - - int driver_version = 0; - if (cudaDriverGetVersion(&driver_version) != cudaSuccess) { return std::nullopt; } - if (!cuvs::detail::jit_lto::tileir_fallback_available(driver_version)) { - return std::nullopt; - } - return CutileModuleImage{ - reinterpret_cast(kTileIrBytecode.data), kTileIrBytecode.size}; -} - -struct LoadedKernel { - cudaLibrary_t library = nullptr; - cudaKernel_t kernel = nullptr; - bool used_tileir_jit{false}; - const char* skip_reason{nullptr}; - - LoadedKernel() = default; - - LoadedKernel(LoadedKernel&& other) noexcept { *this = std::move(other); } - - LoadedKernel& operator=(LoadedKernel&& other) noexcept - { - if (this != &other) { - unload(); - library = other.library; - kernel = other.kernel; - used_tileir_jit = other.used_tileir_jit; - skip_reason = other.skip_reason; - other.library = nullptr; - other.kernel = nullptr; - } - return *this; - } - - LoadedKernel(const LoadedKernel&) = delete; - LoadedKernel& operator=(const LoadedKernel&) = delete; - - ~LoadedKernel() { unload(); } - - explicit operator bool() const { return kernel != nullptr; } - - private: - void unload() - { - if (library != nullptr) { - RAFT_CUDA_TRY(cudaLibraryUnload(library)); - library = nullptr; - kernel = nullptr; - } - } -}; - -LoadedKernel load_vector_add_kernel(int cc_major, int cc_minor) -{ - LoadedKernel result{}; - result.used_tileir_jit = !cuvs::detail::jit_lto::is_embedded_cubin_arch(cc_major, cc_minor); - - auto image = resolve_vector_add_module(cc_major, cc_minor); - if (!image) { - if (result.used_tileir_jit) { - result.skip_reason = - "TileIR driver JIT unavailable for this GPU. Requires CUDA 13.1+ driver (>= 590.44)."; - } else { - ADD_FAILURE() << "No embedded cuTile module for compute capability " << cc_major << "." - << cc_minor; - } - return result; - } - - const cudaError_t load_status = - cudaLibraryLoadData(&result.library, image->data, nullptr, nullptr, 0, nullptr, nullptr, 0); - if (load_status != cudaSuccess) { - if (result.used_tileir_jit) { - result.skip_reason = - "TileIR driver JIT unavailable for this GPU (requires CUDA 13.1+ driver >= 590.44)."; - SCOPED_TRACE(cudaGetErrorString(load_status)); - } else { - ADD_FAILURE() << "cudaLibraryLoadData failed: " << cudaGetErrorString(load_status); - } - return result; - } - - const cudaError_t kernel_status = - cudaLibraryGetKernel(&result.kernel, result.library, CUTILE_VECTOR_ADD_KERNEL_SYMBOL); - if (kernel_status != cudaSuccess) { - if (result.library != nullptr) { - RAFT_CUDA_TRY(cudaLibraryUnload(result.library)); - result.library = nullptr; - } - result.kernel = nullptr; - if (result.used_tileir_jit) { - result.skip_reason = - "TileIR driver JIT unavailable for this GPU (requires CUDA 13.1+ driver >= 590.44)."; - SCOPED_TRACE(cudaGetErrorString(kernel_status)); - } else { - ADD_FAILURE() << "cudaLibraryGetKernel failed: " << cudaGetErrorString(kernel_status); - } - } - return result; -} - -void run_vector_add(cudaKernel_t kernel) -{ - constexpr int kN = 1024; - constexpr int kTile = 256; - constexpr int kGridDim = (kN + kTile - 1) / kTile; - - float *d_a = nullptr, *d_b = nullptr, *d_c = nullptr; - RAFT_CUDA_TRY(cudaMalloc(&d_a, kN * sizeof(float))); - RAFT_CUDA_TRY(cudaMalloc(&d_b, kN * sizeof(float))); - RAFT_CUDA_TRY(cudaMalloc(&d_c, kN * sizeof(float))); - - std::vector h_a(kN), h_b(kN); - for (int i = 0; i < kN; ++i) { - h_a[i] = static_cast(i); - h_b[i] = static_cast(i * 2); - } - RAFT_CUDA_TRY(cudaMemcpy(d_a, h_a.data(), kN * sizeof(float), cudaMemcpyHostToDevice)); - RAFT_CUDA_TRY(cudaMemcpy(d_b, h_b.data(), kN * sizeof(float), cudaMemcpyHostToDevice)); - RAFT_CUDA_TRY(cudaMemset(d_c, 0, kN * sizeof(float))); - - int64_t shape = kN; - int64_t stride = 1; - void* kernel_args[] = { - &d_a, &shape, &stride, &d_b, &shape, &stride, &d_c, &shape, &stride, - }; - - dim3 grid(kGridDim); - dim3 block(1); - ASSERT_EQ(cudaSuccess, cudaLaunchKernel(kernel, grid, block, kernel_args, 0, 0)) - << "cudaLaunchKernel failed: " << cudaGetErrorString(cudaGetLastError()); - RAFT_CUDA_TRY(cudaDeviceSynchronize()); - - std::vector h_c(kN); - RAFT_CUDA_TRY(cudaMemcpy(h_c.data(), d_c, kN * sizeof(float), cudaMemcpyDeviceToHost)); - - for (int i = 0; i < kN; ++i) { - ASSERT_FLOAT_EQ(h_a[i] + h_b[i], h_c[i]) << "@" << i; - } - - RAFT_CUDA_TRY(cudaFree(d_a)); - RAFT_CUDA_TRY(cudaFree(d_b)); - RAFT_CUDA_TRY(cudaFree(d_c)); -} - -class CutileVectorAddTest : public ::testing::Test { - protected: - void SetUp() override - { - int device = 0; - RAFT_CUDA_TRY(cudaGetDevice(&device)); - RAFT_CUDA_TRY( - cudaDeviceGetAttribute(&cc_major_, cudaDevAttrComputeCapabilityMajor, device)); - RAFT_CUDA_TRY( - cudaDeviceGetAttribute(&cc_minor_, cudaDevAttrComputeCapabilityMinor, device)); - } - - int cc_major_{}; - int cc_minor_{}; -}; - -} // namespace - -TEST_F(CutileVectorAddTest, EmbeddedCubinVectorAdd) -{ - LoadedKernel loaded = load_vector_add_kernel(cc_major_, cc_minor_); - if (loaded.skip_reason) { GTEST_SKIP() << loaded.skip_reason; } - if (!loaded) { return; } - - SCOPED_TRACE(loaded.used_tileir_jit ? "loaded via TileIR driver JIT" - : "loaded via prebuilt cubin"); - run_vector_add(loaded.kernel); -} - -} // namespace cuvs diff --git a/cpp/tests/cutile/export_vector_add_cubin.py b/cpp/tests/cutile/export_vector_add_cubin.py deleted file mode 100644 index fa099189cd..0000000000 --- a/cpp/tests/cutile/export_vector_add_cubin.py +++ /dev/null @@ -1,133 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. -# SPDX-License-Identifier: Apache-2.0 -"""Export the cuTile vector-add kernel to cubin or TileIR bytecode.""" - -from __future__ import annotations - -import argparse -import sys -from pathlib import Path -from typing import Literal - -import cuda.tile as ct -from cuda.tile.compilation import ( - ArrayConstraint, - CallingConvention, - ConstantConstraint, - KernelSignature, - export_kernel, -) - -from vector_add_kernel import TILE_SIZE, vector_add - -# cuTile / tileiras gpu_code values used at build time. These correspond to the -# cuvs library CUDA 13 real targets as follows (tileiras has no sm_*a/sm_*f names): -# sm_80 -> 80-real -# sm_86 -> 86-real -# sm_90 -> 90a-real -# sm_100 -> 100f-real -# sm_120 -> 120a-real -SUPPORTED_GPU_CODES = ("sm_80", "sm_86", "sm_90", "sm_100", "sm_120") - -# Minimum TileIR bytecode version supported by cuTile; also the most portable choice. -DEFAULT_TILEIR_BYTECODE_VERSION = "13.1" - - -def _kernel_signature() -> KernelSignature: - array = ArrayConstraint( - ct.float32, - 1, - index_dtype=ct.int64, - stride_lower_bound_incl=0, - alias_groups=(), - may_alias_internally=False, - stride_constant=(1,), - ) - return KernelSignature( - parameters=[array, array, array, ConstantConstraint(TILE_SIZE)], - calling_convention=CallingConvention.cutile_python_v1(), - ).with_mangled_symbol("vector_add") - - -def export_kernel_binary( - output_file: Path, - *, - output_format: Literal["cubin", "tileir_bytecode"], - gpu_code: str, - bytecode_version: str | None = None, - symbol_header: Path | None = None, -) -> str: - if output_format == "cubin" and gpu_code not in SUPPORTED_GPU_CODES: - raise ValueError( - f"Unsupported gpu_code {gpu_code!r}; expected one of {SUPPORTED_GPU_CODES}" - ) - - signature = _kernel_signature() - export_kwargs: dict = { - "kernel": vector_add, - "signatures": [signature], - "output_file": str(output_file), - "gpu_code": gpu_code, - "output_format": output_format, - } - if output_format == "tileir_bytecode": - export_kwargs["bytecode_version"] = bytecode_version or DEFAULT_TILEIR_BYTECODE_VERSION - - export_kernel(**export_kwargs) - - if symbol_header is not None: - symbol_header.write_text( - "\n".join( - [ - "// Generated by export_vector_add_cubin.py; do not edit.", - "#pragma once", - f'#define CUTILE_VECTOR_ADD_KERNEL_SYMBOL "{signature.symbol}"', - "", - ] - ) - ) - - return signature.symbol - - -def main() -> int: - parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument("output_file", type=Path, help="Output cubin or .tilebc path") - parser.add_argument( - "--format", - choices=("cubin", "tileir_bytecode"), - default="cubin", - help="Export format (default: cubin)", - ) - parser.add_argument( - "--gpu-code", - required=True, - choices=SUPPORTED_GPU_CODES, - help="tileiras / export_kernel compile target (e.g. sm_120)", - ) - parser.add_argument( - "--bytecode-version", - default=DEFAULT_TILEIR_BYTECODE_VERSION, - help="TileIR bytecode version when --format=tileir_bytecode (default: 13.1)", - ) - parser.add_argument( - "--symbol-header", - type=Path, - default=None, - help="Optional header that defines CUTILE_VECTOR_ADD_KERNEL_SYMBOL", - ) - args = parser.parse_args() - - symbol = export_kernel_binary( - args.output_file, - output_format=args.format, - gpu_code=args.gpu_code, - bytecode_version=args.bytecode_version, - symbol_header=args.symbol_header, - ) - print(symbol) - return 0 - - -if __name__ == "__main__": - sys.exit(main()) diff --git a/cpp/tests/cutile/generate_cutile_cubins.cmake b/cpp/tests/cutile/generate_cutile_cubins.cmake deleted file mode 100644 index 766d3167c6..0000000000 --- a/cpp/tests/cutile/generate_cutile_cubins.cmake +++ /dev/null @@ -1,117 +0,0 @@ -# ============================================================================= -# cmake-format: off -# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. -# SPDX-License-Identifier: Apache-2.0 -# cmake-format: on -# ============================================================================= - -include_guard(GLOBAL) - -# Build-time cuTile cubin targets. Maps to cuvs CUDA 13 -real library arches (75-real omitted). -set(CUTILE_VECTOR_ADD_GPU_CODES sm_80 sm_86 sm_90 sm_100 sm_120) - -function(generate_cutile_vector_add_cubins output_include_dir_var) - find_package(Python3 REQUIRED COMPONENTS Interpreter) - find_package(CUDAToolkit REQUIRED) - - find_program( - CUTILE_BIN2C - NAMES bin2c - PATHS ${CUDAToolkit_BIN_DIR} - REQUIRED - ) - - execute_process( - COMMAND "${Python3_EXECUTABLE}" -c "import cuda.tile" - RESULT_VARIABLE _cutile_import_result - OUTPUT_QUIET - ERROR_QUIET - ) - if(NOT _cutile_import_result EQUAL 0) - message( - FATAL_ERROR - "cuda.tile (cuTile Python) is required to build CUTILE_VECTOR_ADD_TEST. " - "Install it in the active Python environment, e.g. pip install cuda-tile[tileiras]." - ) - endif() - - set(_cutile_source_dir "${CMAKE_CURRENT_FUNCTION_LIST_DIR}") - set(_cutile_binary_dir "${CMAKE_CURRENT_BINARY_DIR}/cutile_generated") - file(MAKE_DIRECTORY "${_cutile_binary_dir}") - - set(_symbol_header "${_cutile_binary_dir}/vector_add_kernel_symbol.h") - set(_first_gpu_code TRUE) - - foreach(_gpu_code IN LISTS CUTILE_VECTOR_ADD_GPU_CODES) - set(_cubin_file "${_cutile_binary_dir}/vector_add_${_gpu_code}.cubin") - set(_cubin_header "${_cutile_binary_dir}/vector_add_${_gpu_code}_cubin.h") - - if(_first_gpu_code) - set(_symbol_arg --symbol-header "${_symbol_header}") - set(_cubin_outputs "${_cubin_file}" "${_symbol_header}") - set(_first_gpu_code FALSE) - else() - set(_symbol_arg) - set(_cubin_outputs "${_cubin_file}") - endif() - - add_custom_command( - OUTPUT ${_cubin_outputs} - COMMAND - "${Python3_EXECUTABLE}" "${_cutile_source_dir}/export_vector_add_cubin.py" - "${_cubin_file}" --gpu-code "${_gpu_code}" ${_symbol_arg} - DEPENDS "${_cutile_source_dir}/export_vector_add_cubin.py" - "${_cutile_source_dir}/vector_add_kernel.py" - COMMENT "Exporting cuTile vector_add cubin for ${_gpu_code}" - VERBATIM - ) - - add_custom_command( - OUTPUT "${_cubin_header}" - COMMAND "${CUTILE_BIN2C}" --const --name "vector_add_${_gpu_code}_cubin" --static - "${_cubin_file}" > "${_cubin_header}" - DEPENDS "${_cubin_file}" - COMMENT "Embedding vector_add ${_gpu_code} cubin via bin2c" - VERBATIM - ) - - list(APPEND _generated_headers "${_cubin_header}") - endforeach() - - # Portable TileIR bytecode for driver JIT on architectures without a prebuilt cubin. - # Requires a CUDA 13.1+ driver (>= 590.44); see Tile IR bytecode docs. - set(_tileir_file "${_cutile_binary_dir}/vector_add.tilebc") - set(_tileir_header "${_cutile_binary_dir}/vector_add_tileir_bytecode.h") - - add_custom_command( - OUTPUT "${_tileir_file}" - COMMAND - "${Python3_EXECUTABLE}" "${_cutile_source_dir}/export_vector_add_cubin.py" - "${_tileir_file}" --format tileir_bytecode --gpu-code sm_80 --bytecode-version 13.1 - DEPENDS "${_cutile_source_dir}/export_vector_add_cubin.py" - "${_cutile_source_dir}/vector_add_kernel.py" - COMMENT "Exporting cuTile vector_add TileIR bytecode (v13.1)" - VERBATIM - ) - - add_custom_command( - OUTPUT "${_tileir_header}" - COMMAND "${CUTILE_BIN2C}" --const --name vector_add_tileir_bytecode --static "${_tileir_file}" - > "${_tileir_header}" - DEPENDS "${_tileir_file}" - COMMENT "Embedding vector_add TileIR bytecode via bin2c" - VERBATIM - ) - - list(APPEND _generated_headers "${_tileir_header}") - - add_custom_target( - cutile_vector_add_cubins - DEPENDS "${_symbol_header}" ${_generated_headers} - ) - - set(${output_include_dir_var} - "${_cutile_binary_dir}" - PARENT_SCOPE - ) -endfunction() diff --git a/cpp/tests/cutile/vector_add_kernel.py b/cpp/tests/cutile/vector_add_kernel.py deleted file mode 100644 index 46b7a607c6..0000000000 --- a/cpp/tests/cutile/vector_add_kernel.py +++ /dev/null @@ -1,17 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. -# SPDX-License-Identifier: Apache-2.0 -"""cuTile Python vector-add kernel used by the embedded-cubin example test.""" - -from __future__ import annotations - -import cuda.tile as ct - -TILE_SIZE = 256 - - -@ct.kernel -def vector_add(a, b, c, TILE_SIZE: ct.Constant): - bid = ct.bid(0) - ta = ct.load(a, bid, TILE_SIZE) - tb = ct.load(b, bid, TILE_SIZE) - ct.store(c, bid, ta + tb) From 86c9311b87fae027be808b1a4d04f174e0bfbbb2 Mon Sep 17 00:00:00 2001 From: divyegala Date: Wed, 24 Jun 2026 21:16:48 +0000 Subject: [PATCH 05/82] style check --- cpp/CMakeLists.txt | 47 +++++----- .../modules/generate_cutile_kernels.cmake | 86 ++++++++++--------- .../modules/register_cutile_fragment.cpp.in | 8 +- .../cuvs/detail/jit_lto/FragmentEntry.hpp | 5 +- .../cuvs/detail/jit_lto/tileir_compat.hpp | 4 +- .../detail/jit_lto/TileAlgorithmPlanner.cpp | 4 +- cpp/src/distance/detail/fused_distance_nn.cuh | 4 +- .../cutile/export_fused_1nn.py | 16 +++- .../cutile/fused_1nn_kernel.py | 30 ++++--- cpp/tests/neighbors/distance_nn_helper.cuh | 4 +- python/libcuvs/pyproject.toml | 1 + 11 files changed, 118 insertions(+), 91 deletions(-) diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt index a1f3f3973c..70e2509a88 100644 --- a/cpp/CMakeLists.txt +++ b/cpp/CMakeLists.txt @@ -960,32 +960,38 @@ if(NOT BUILD_CPU_ONLY) include(cmake/modules/generate_cutile_kernels.cmake) set(fused_1nn_cutile_dir - "${CMAKE_CURRENT_SOURCE_DIR}/src/distance/detail/fused_distance_nn/cutile") + "${CMAKE_CURRENT_SOURCE_DIR}/src/distance/detail/fused_distance_nn/cutile" + ) set(cutile_fused_1nn_generated_dir - "${CMAKE_CURRENT_BINARY_DIR}/generated_kernels/distance/fused_1nn/cutile") + "${CMAKE_CURRENT_BINARY_DIR}/generated_kernels/distance/fused_1nn/cutile" + ) generate_cutile_kernels( cutile_fused_1nn_files - KERNEL_DIR "${fused_1nn_cutile_dir}" - KERNEL_BASENAME "fused_1nn" - KERNEL_PYTHON "fused_1nn_kernel.py" - EXPORT_SCRIPT "export_fused_1nn.py" - OUTPUT_DIRECTORY "${cutile_fused_1nn_generated_dir}" - MATRIX_JSON_FILE "${fused_1nn_cutile_dir}/fused_1nn_cutile_matrix.json" + KERNEL_DIR + "${fused_1nn_cutile_dir}" + KERNEL_BASENAME + "fused_1nn" + KERNEL_PYTHON + "fused_1nn_kernel.py" + EXPORT_SCRIPT + "export_fused_1nn.py" + OUTPUT_DIRECTORY + "${cutile_fused_1nn_generated_dir}" + MATRIX_JSON_FILE + "${fused_1nn_cutile_dir}/fused_1nn_cutile_matrix.json" FRAGMENT_TAG_FORMAT_CUBIN - "cuvs::distance::detail::fragment_tag_fused_1nn_cubin" + "cuvs::distance::detail::fragment_tag_fused_1nn_cubin" FRAGMENT_TAG_FORMAT_TILEIR - "cuvs::distance::detail::fragment_tag_fused_1nn_tileir" + "cuvs::distance::detail::fragment_tag_fused_1nn_tileir" FRAGMENT_TAG_HEADER_FILES - "" - "" - "" + "" + "" + "" ) if(NOT DEFINED CUVS_CUTILE_ENABLED) set(CUVS_CUTILE_ENABLED 0) endif() - target_compile_definitions( - cuvs_cpp_headers INTERFACE CUVS_CUTILE_ENABLED=${CUVS_CUTILE_ENABLED} - ) + target_compile_definitions(cuvs_cpp_headers INTERFACE CUVS_CUTILE_ENABLED=${CUVS_CUTILE_ENABLED}) generate_inst_matrix( cagra_build_inst_files MATRIX_JSON_FILE "${CMAKE_CURRENT_SOURCE_DIR}/src/neighbors/cagra_build_matrix.json" @@ -1288,9 +1294,9 @@ if(NOT BUILD_CPU_ONLY) ) target_compile_definitions( - cuvs_objs PRIVATE $<$:CUVS_BUILD_CAGRA_HNSWLIB> - $<$:NVTX_ENABLED> - CUVS_CUTILE_ENABLED=${CUVS_CUTILE_ENABLED} + cuvs_objs + PRIVATE $<$:CUVS_BUILD_CAGRA_HNSWLIB> + $<$:NVTX_ENABLED> CUVS_CUTILE_ENABLED=${CUVS_CUTILE_ENABLED} ) target_link_libraries( @@ -1308,8 +1314,7 @@ if(NOT BUILD_CPU_ONLY) PUBLIC "$" "$" INTERFACE "$" - PRIVATE "${CMAKE_CURRENT_SOURCE_DIR}/src" - "${CMAKE_CURRENT_BINARY_DIR}/src" + PRIVATE "${CMAKE_CURRENT_SOURCE_DIR}/src" "${CMAKE_CURRENT_BINARY_DIR}/src" "${cutile_fused_1nn_generated_dir}" ) diff --git a/cpp/cmake/modules/generate_cutile_kernels.cmake b/cpp/cmake/modules/generate_cutile_kernels.cmake index f0219dc842..ac8d369cdc 100644 --- a/cpp/cmake/modules/generate_cutile_kernels.cmake +++ b/cpp/cmake/modules/generate_cutile_kernels.cmake @@ -10,7 +10,10 @@ include_guard(GLOBAL) include(${CMAKE_CURRENT_LIST_DIR}/compute_matrix_product.cmake) function(generate_cutile_kernels_stub) - set(CUVS_CUTILE_ENABLED 0 PARENT_SCOPE) + set(CUVS_CUTILE_ENABLED + 0 + PARENT_SCOPE + ) endfunction() function(_cutile_fragment_tag_header_files output_var) @@ -51,15 +54,13 @@ function(_cutile_kernels_setup) find_program( CUTILE_BIN2C NAMES bin2c - PATHS ${CUDAToolkit_BIN_DIR} - REQUIRED + PATHS ${CUDAToolkit_BIN_DIR} REQUIRED ) execute_process( COMMAND "${Python3_EXECUTABLE}" -c "import cuda.tile" RESULT_VARIABLE _cutile_import_result - OUTPUT_QUIET - ERROR_QUIET + OUTPUT_QUIET ERROR_QUIET ) if(NOT _cutile_import_result EQUAL 0) message( @@ -77,8 +78,14 @@ function(_cutile_kernels_setup) file(MAKE_DIRECTORY "${_CUTILE_OUTPUT_DIRECTORY}") - set(Python3_EXECUTABLE "${Python3_EXECUTABLE}" PARENT_SCOPE) - set(CUTILE_BIN2C "${CUTILE_BIN2C}" PARENT_SCOPE) + set(Python3_EXECUTABLE + "${Python3_EXECUTABLE}" + PARENT_SCOPE + ) + set(CUTILE_BIN2C + "${CUTILE_BIN2C}" + PARENT_SCOPE + ) set(_CUTILE_SETUP_OK TRUE PARENT_SCOPE @@ -87,15 +94,8 @@ endfunction() function(process_cutile_matrix_entry source_list_var) set(options) - set(one_value - KERNEL_DIR - KERNEL_BASENAME - KERNEL_PYTHON - EXPORT_SCRIPT - OUTPUT_DIRECTORY - FRAGMENT_TAG_FORMAT_CUBIN - FRAGMENT_TAG_FORMAT_TILEIR - MATRIX_JSON_ENTRY + set(one_value KERNEL_DIR KERNEL_BASENAME KERNEL_PYTHON EXPORT_SCRIPT OUTPUT_DIRECTORY + FRAGMENT_TAG_FORMAT_CUBIN FRAGMENT_TAG_FORMAT_TILEIR MATRIX_JSON_ENTRY ) set(multi_value FRAGMENT_TAG_HEADER_FILES) cmake_parse_arguments(_CUTILE "${options}" "${one_value}" "${multi_value}" ${ARGN}) @@ -116,9 +116,7 @@ function(process_cutile_matrix_entry source_list_var) message(FATAL_ERROR "Unknown cuTile register kind '${register}'") endif() - _cutile_fragment_tag_header_files( - fragment_tag_header_files ${_CUTILE_FRAGMENT_TAG_HEADER_FILES} - ) + _cutile_fragment_tag_header_files(fragment_tag_header_files ${_CUTILE_FRAGMENT_TAG_HEADER_FILES}) string(CONFIGURE "${artifact_basename}" _artifact_basename @ONLY) set(_artifact_stem "${_CUTILE_KERNEL_BASENAME}_${_artifact_basename}") @@ -145,8 +143,8 @@ function(process_cutile_matrix_entry source_list_var) add_custom_command( OUTPUT "${_embedded_header}" - COMMAND "${CUTILE_BIN2C}" --const --name ${bin2c_symbol} --static "${_artifact_file}" - > "${_embedded_header}" + COMMAND "${CUTILE_BIN2C}" --const --name ${bin2c_symbol} --static "${_artifact_file}" > + "${_embedded_header}" DEPENDS "${_artifact_file}" VERBATIM ) @@ -163,15 +161,8 @@ endfunction() function(generate_cutile_kernels source_list_var) set(options) - set(one_value - KERNEL_DIR - KERNEL_BASENAME - KERNEL_PYTHON - EXPORT_SCRIPT - OUTPUT_DIRECTORY - MATRIX_JSON_FILE - FRAGMENT_TAG_FORMAT_CUBIN - FRAGMENT_TAG_FORMAT_TILEIR + set(one_value KERNEL_DIR KERNEL_BASENAME KERNEL_PYTHON EXPORT_SCRIPT OUTPUT_DIRECTORY + MATRIX_JSON_FILE FRAGMENT_TAG_FORMAT_CUBIN FRAGMENT_TAG_FORMAT_TILEIR ) set(multi_value FRAGMENT_TAG_HEADER_FILES) cmake_parse_arguments(_CUTILE "${options}" "${one_value}" "${multi_value}" ${ARGN}) @@ -184,8 +175,7 @@ function(generate_cutile_kernels source_list_var) endif() _cutile_kernels_setup( - MATRIX_JSON_FILE "${_CUTILE_MATRIX_JSON_FILE}" - OUTPUT_DIRECTORY "${_CUTILE_OUTPUT_DIRECTORY}" + MATRIX_JSON_FILE "${_CUTILE_MATRIX_JSON_FILE}" OUTPUT_DIRECTORY "${_CUTILE_OUTPUT_DIRECTORY}" ) if(NOT _CUTILE_SETUP_OK) generate_cutile_kernels_stub() @@ -206,19 +196,31 @@ function(generate_cutile_kernels source_list_var) string(JSON matrix_json_entry GET "${matrix_product}" "${i}") process_cutile_matrix_entry( "${source_list_var}" - KERNEL_DIR "${_CUTILE_KERNEL_DIR}" - KERNEL_BASENAME "${_CUTILE_KERNEL_BASENAME}" - KERNEL_PYTHON "${_CUTILE_KERNEL_PYTHON}" - EXPORT_SCRIPT "${_CUTILE_EXPORT_SCRIPT}" - OUTPUT_DIRECTORY "${_CUTILE_OUTPUT_DIRECTORY}" - FRAGMENT_TAG_FORMAT_CUBIN "${_CUTILE_FRAGMENT_TAG_FORMAT_CUBIN}" - FRAGMENT_TAG_FORMAT_TILEIR "${_CUTILE_FRAGMENT_TAG_FORMAT_TILEIR}" - FRAGMENT_TAG_HEADER_FILES ${_CUTILE_FRAGMENT_TAG_HEADER_FILES} - MATRIX_JSON_ENTRY "${matrix_json_entry}" + KERNEL_DIR + "${_CUTILE_KERNEL_DIR}" + KERNEL_BASENAME + "${_CUTILE_KERNEL_BASENAME}" + KERNEL_PYTHON + "${_CUTILE_KERNEL_PYTHON}" + EXPORT_SCRIPT + "${_CUTILE_EXPORT_SCRIPT}" + OUTPUT_DIRECTORY + "${_CUTILE_OUTPUT_DIRECTORY}" + FRAGMENT_TAG_FORMAT_CUBIN + "${_CUTILE_FRAGMENT_TAG_FORMAT_CUBIN}" + FRAGMENT_TAG_FORMAT_TILEIR + "${_CUTILE_FRAGMENT_TAG_FORMAT_TILEIR}" + FRAGMENT_TAG_HEADER_FILES + ${_CUTILE_FRAGMENT_TAG_HEADER_FILES} + MATRIX_JSON_ENTRY + "${matrix_json_entry}" ) endforeach() - set(CUVS_CUTILE_ENABLED 1 PARENT_SCOPE) + set(CUVS_CUTILE_ENABLED + 1 + PARENT_SCOPE + ) set(${source_list_var} "${${source_list_var}}" PARENT_SCOPE diff --git a/cpp/cmake/modules/register_cutile_fragment.cpp.in b/cpp/cmake/modules/register_cutile_fragment.cpp.in index 0fc074bdbb..3ffd5c0d0c 100644 --- a/cpp/cmake/modules/register_cutile_fragment.cpp.in +++ b/cpp/cmake/modules/register_cutile_fragment.cpp.in @@ -8,10 +8,10 @@ @fragment_tag_header_files@ -namespace { - -using fragment_tag = @fragment_tag@; -using fragment_entry = @fragment_entry_type@; + namespace +{ + using fragment_tag = @fragment_tag@; + using fragment_entry = @fragment_entry_type@; } // namespace diff --git a/cpp/include/cuvs/detail/jit_lto/FragmentEntry.hpp b/cpp/include/cuvs/detail/jit_lto/FragmentEntry.hpp index df69ec1d7b..6c399d860a 100644 --- a/cpp/include/cuvs/detail/jit_lto/FragmentEntry.hpp +++ b/cpp/include/cuvs/detail/jit_lto/FragmentEntry.hpp @@ -115,7 +115,10 @@ struct StaticTileIrBytecodeFragmentEntry final : TileIrBytecodeFragmentEntry { return StaticTileIrBytecodeFragmentEntry::data; } - size_t get_length() const override { return StaticTileIrBytecodeFragmentEntry::length; } + size_t get_length() const override + { + return StaticTileIrBytecodeFragmentEntry::length; + } const char* get_key() const override { diff --git a/cpp/include/cuvs/detail/jit_lto/tileir_compat.hpp b/cpp/include/cuvs/detail/jit_lto/tileir_compat.hpp index d63759fb36..f15407fd4c 100644 --- a/cpp/include/cuvs/detail/jit_lto/tileir_compat.hpp +++ b/cpp/include/cuvs/detail/jit_lto/tileir_compat.hpp @@ -88,8 +88,8 @@ inline bool query_current_device_arch(int& cc_major, int& cc_minor) inline bool cutile_launch_available_on_current_device() { - int cc_major = 0; - int cc_minor = 0; + int cc_major = 0; + int cc_minor = 0; int driver_version = 0; if (!query_current_device_arch(cc_major, cc_minor)) { return false; } if (!query_driver_version(driver_version)) { return false; } diff --git a/cpp/src/detail/jit_lto/TileAlgorithmPlanner.cpp b/cpp/src/detail/jit_lto/TileAlgorithmPlanner.cpp index edb6269213..e0ce77e789 100644 --- a/cpp/src/detail/jit_lto/TileAlgorithmPlanner.cpp +++ b/cpp/src/detail/jit_lto/TileAlgorithmPlanner.cpp @@ -23,9 +23,7 @@ std::shared_ptr TileAlgorithmPlanner::build() { int cc_major = 0; int cc_minor = 0; - if (!cuvs::detail::jit_lto::get_device_compute_capability(cc_major, cc_minor)) { - return nullptr; - } + if (!cuvs::detail::jit_lto::get_device_compute_capability(cc_major, cc_minor)) { return nullptr; } int driver_version = 0; if (cudaDriverGetVersion(&driver_version) != cudaSuccess) { return nullptr; } diff --git a/cpp/src/distance/detail/fused_distance_nn.cuh b/cpp/src/distance/detail/fused_distance_nn.cuh index 8b47092b58..b1b18e58f6 100644 --- a/cpp/src/distance/detail/fused_distance_nn.cuh +++ b/cpp/src/distance/detail/fused_distance_nn.cuh @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2024, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ @@ -19,8 +19,8 @@ #include "fused_distance_nn/helper_structs.cuh" #include "fused_distance_nn/simt_kernel.cuh" #include "pairwise_distance_base.cuh" // PairwiseDistances -#include #include +#include #include // raft::KeyValuePair #include // raft::identity_op #include // Policy diff --git a/cpp/src/distance/detail/fused_distance_nn/cutile/export_fused_1nn.py b/cpp/src/distance/detail/fused_distance_nn/cutile/export_fused_1nn.py index 6a20be24ef..10a4fa9ec1 100644 --- a/cpp/src/distance/detail/fused_distance_nn/cutile/export_fused_1nn.py +++ b/cpp/src/distance/detail/fused_distance_nn/cutile/export_fused_1nn.py @@ -100,7 +100,9 @@ def export_binary( "output_format": output_format, } if output_format == "tileir_bytecode": - export_kwargs["bytecode_version"] = bytecode_version or DEFAULT_TILEIR_BYTECODE_VERSION + export_kwargs["bytecode_version"] = ( + bytecode_version or DEFAULT_TILEIR_BYTECODE_VERSION + ) export_kernel(**export_kwargs) @@ -110,14 +112,20 @@ def export_binary( def main() -> int: parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("output_file", type=Path) - parser.add_argument("--format", choices=("cubin", "tileir_bytecode"), default="cubin") - parser.add_argument("--data-type", choices=tuple(KERNELS.keys()), required=True) + parser.add_argument( + "--format", choices=("cubin", "tileir_bytecode"), default="cubin" + ) + parser.add_argument( + "--data-type", choices=tuple(KERNELS.keys()), required=True + ) parser.add_argument( "--gpu-code", default=DEFAULT_TILEIR_EXPORT_GPU_CODE, help="Target SM for cubin export, or compile hint for TileIR bytecode export", ) - parser.add_argument("--bytecode-version", default=DEFAULT_TILEIR_BYTECODE_VERSION) + parser.add_argument( + "--bytecode-version", default=DEFAULT_TILEIR_BYTECODE_VERSION + ) args = parser.parse_args() print( diff --git a/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_kernel.py b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_kernel.py index 232b9506af..65fe165b70 100644 --- a/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_kernel.py +++ b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_kernel.py @@ -14,17 +14,23 @@ def _make_kernel(data_type: str): - if data_type == "half": - dtype = ct.float16 - acc_dtype = ct.float32 - elif data_type == "float": - dtype = ct.float32 - acc_dtype = ct.float32 - else: + if data_type not in ("half", "float"): raise ValueError(f"Unsupported data_type {data_type!r}") + acc_dtype = ct.float32 @ct.kernel - def fused_1nn_kernel(A, B, OutIdx, OutDist, M, N, K, tm: ConstInt, tn: ConstInt, tk: ConstInt): + def fused_1nn_kernel( + A, + B, + OutIdx, + OutDist, + M, + N, + K, + tm: ConstInt, + tn: ConstInt, + tk: ConstInt, + ): bidm = ct.bid(0) best_dist = ct.full((tm,), -3.4e38, acc_dtype) @@ -38,8 +44,12 @@ def fused_1nn_kernel(A, B, OutIdx, OutDist, M, N, K, tm: ConstInt, tn: ConstInt, accumulator = ct.full((tm, tn), 0, dtype=acc_dtype) for k in range(num_tiles_k): - a = ct.load(A, index=(bidm, k), shape=(tm, tk), padding_mode=zero_pad) - b_T = ct.load(B, index=(n, k), shape=(tn, tk), padding_mode=zero_pad) + a = ct.load( + A, index=(bidm, k), shape=(tm, tk), padding_mode=zero_pad + ) + b_T = ct.load( + B, index=(n, k), shape=(tn, tk), padding_mode=zero_pad + ) accumulator = ct.mma(a, ct.transpose(b_T), accumulator) curr_max = ct.max(accumulator, axis=1) diff --git a/cpp/tests/neighbors/distance_nn_helper.cuh b/cpp/tests/neighbors/distance_nn_helper.cuh index 422879918f..ea440387b4 100644 --- a/cpp/tests/neighbors/distance_nn_helper.cuh +++ b/cpp/tests/neighbors/distance_nn_helper.cuh @@ -91,8 +91,8 @@ RAFT_KERNEL ref_nn_kernel( if (metric == DistanceType::InnerProduct) { AccT score = inner_product_score(&A[m * K], &B[n * K], K); if (score > best_score) { - best_score = score; - best_index = n; + best_score = score; + best_index = n; } continue; } diff --git a/python/libcuvs/pyproject.toml b/python/libcuvs/pyproject.toml index 5025daa66d..b4e848304f 100644 --- a/python/libcuvs/pyproject.toml +++ b/python/libcuvs/pyproject.toml @@ -19,6 +19,7 @@ authors = [ license = "Apache-2.0" requires-python = ">=3.11" dependencies = [ + "cuda-tile[tileiras]", "cuda-toolkit[cublas,curand,cusolver,cusparse,nvrtc]==13.*", "libraft==26.8.*,>=0.0.0a0", "librmm==26.8.*,>=0.0.0a0", From e352629e5398e33dc780888f3b96c8c610f4dafd Mon Sep 17 00:00:00 2001 From: divyegala Date: Thu, 25 Jun 2026 21:16:05 +0000 Subject: [PATCH 06/82] start integrating other metrics --- cpp/CMakeLists.txt | 4 +- .../modules/generate_cutile_kernels.cmake | 32 ++++- .../modules/register_cutile_fragment.cpp.in | 9 ++ .../cuvs/detail/jit_lto/AlgorithmPlanner.hpp | 3 + .../cuvs/detail/jit_lto/FragmentEntry.hpp | 39 ++++++ .../fused_distance_nn/fused_1nn_fragments.hpp | 70 +++++++++- .../detail/jit_lto/TileAlgorithmPlanner.cpp | 46 +++++++ cpp/src/distance/detail/fused_distance_nn.cuh | 2 +- .../cutile/export_fused_1nn.py | 97 +++++++++----- .../cutile/fused_1nn_cutile_matrix.json | 31 ++++- .../cutile/fused_1nn_kernel.py | 85 ++++++++++--- .../cutile/fused_1nn_planner.hpp | 56 +++++--- .../cutile/fused_1nn_tile.cu | 120 +++++++++++++----- .../cutile/fused_1nn_tile.hpp | 32 +++-- 14 files changed, 508 insertions(+), 118 deletions(-) diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt index 70e2509a88..87716cd296 100644 --- a/cpp/CMakeLists.txt +++ b/cpp/CMakeLists.txt @@ -980,9 +980,9 @@ if(NOT BUILD_CPU_ONLY) MATRIX_JSON_FILE "${fused_1nn_cutile_dir}/fused_1nn_cutile_matrix.json" FRAGMENT_TAG_FORMAT_CUBIN - "cuvs::distance::detail::fragment_tag_fused_1nn_cubin" + "cuvs::distance::detail::fragment_tag_fused_1nn_cubin, cuvs::detail::jit_lto::@arch_tag@>" FRAGMENT_TAG_FORMAT_TILEIR - "cuvs::distance::detail::fragment_tag_fused_1nn_tileir" + "cuvs::distance::detail::fragment_tag_fused_1nn_tileir>" FRAGMENT_TAG_HEADER_FILES "" "" diff --git a/cpp/cmake/modules/generate_cutile_kernels.cmake b/cpp/cmake/modules/generate_cutile_kernels.cmake index ac8d369cdc..abdca118ad 100644 --- a/cpp/cmake/modules/generate_cutile_kernels.cmake +++ b/cpp/cmake/modules/generate_cutile_kernels.cmake @@ -92,6 +92,30 @@ function(_cutile_kernels_setup) ) endfunction() +function(_cutile_generate_matrix_tiles_header header_path matrix_json_file) + file(READ "${matrix_json_file}" _matrix_json) + string(JSON _tile0 GET "${_matrix_json}" 0 "_tile" 0) + string(JSON _tile_m GET "${_tile0}" "tile_m") + string(JSON _tile_n GET "${_tile0}" "tile_n") + string(JSON _tile_k GET "${_tile0}" "tile_k") + file( + WRITE "${header_path}" + "/* + * Generated from ${matrix_json_file} by generate_cutile_kernels.cmake — do not edit. + */ +#pragma once + +#include + +namespace cuvs::distance::detail { + +using fused_1nn_matrix_tile = cutile_tile_config<${_tile_m}, ${_tile_n}, ${_tile_k}>; + +} // namespace cuvs::distance::detail +" + ) +endfunction() + function(process_cutile_matrix_entry source_list_var) set(options) set(one_value KERNEL_DIR KERNEL_BASENAME KERNEL_PYTHON EXPORT_SCRIPT OUTPUT_DIRECTORY @@ -125,7 +149,10 @@ function(process_cutile_matrix_entry source_list_var) set(_fragment_cpp "${_CUTILE_OUTPUT_DIRECTORY}/${_artifact_stem}_${register}.cpp") set(embedded_header_file "${_artifact_stem}_${register}.h") - set(_python_args --format "${output_format}" --data-type "${data_type}" --gpu-code "${gpu_code}") + set(_python_args + --format "${output_format}" --data-type "${data_type}" --metric "${metric}" --tile-m + "${tile_m}" --tile-n "${tile_n}" --tile-k "${tile_k}" --gpu-code "${gpu_code}" + ) if(DEFINED bytecode_version AND NOT "${bytecode_version}" STREQUAL "") list(APPEND _python_args --bytecode-version "${bytecode_version}") endif() @@ -188,6 +215,9 @@ function(generate_cutile_kernels source_list_var) compute_matrix_product(matrix_product MATRIX_JSON_FILE "${_CUTILE_MATRIX_JSON_FILE}") + set(_matrix_tiles_header "${_CUTILE_OUTPUT_DIRECTORY}/fused_1nn_cutile_tiles.hpp") + _cutile_generate_matrix_tiles_header("${_matrix_tiles_header}" "${_CUTILE_MATRIX_JSON_FILE}") + string(JSON len LENGTH "${matrix_product}") math(EXPR last "${len} - 1") diff --git a/cpp/cmake/modules/register_cutile_fragment.cpp.in b/cpp/cmake/modules/register_cutile_fragment.cpp.in index 3ffd5c0d0c..de0472a779 100644 --- a/cpp/cmake/modules/register_cutile_fragment.cpp.in +++ b/cpp/cmake/modules/register_cutile_fragment.cpp.in @@ -20,3 +20,12 @@ const uint8_t* const fragment_entry::data = @bin2c_symbol@; template <> const size_t fragment_entry::length = sizeof(@bin2c_symbol@); + +template <> +const int fragment_entry::tile_m = @tile_m@; + +template <> +const int fragment_entry::tile_n = @tile_n@; + +template <> +const int fragment_entry::tile_k = @tile_k@; diff --git a/cpp/include/cuvs/detail/jit_lto/AlgorithmPlanner.hpp b/cpp/include/cuvs/detail/jit_lto/AlgorithmPlanner.hpp index d727c73b9d..7ff8487d20 100644 --- a/cpp/include/cuvs/detail/jit_lto/AlgorithmPlanner.hpp +++ b/cpp/include/cuvs/detail/jit_lto/AlgorithmPlanner.hpp @@ -97,6 +97,9 @@ struct TileAlgorithmPlanner : AlgorithmPlanner { tileir_fragment_ = std::make_unique>(); } + /** Tile geometry from the cubin or TileIR fragment that would load on this device. */ + CutileTileConfig tile_config() const; + protected: std::vector> cubin_fragments_; std::unique_ptr tileir_fragment_; diff --git a/cpp/include/cuvs/detail/jit_lto/FragmentEntry.hpp b/cpp/include/cuvs/detail/jit_lto/FragmentEntry.hpp index 6c399d860a..0961595f8d 100644 --- a/cpp/include/cuvs/detail/jit_lto/FragmentEntry.hpp +++ b/cpp/include/cuvs/detail/jit_lto/FragmentEntry.hpp @@ -63,6 +63,13 @@ struct UDFFatbinFragment final : FatbinFragmentEntry { std::vector bytes_; }; +/** cuTile GEMM-style block geometry embedded in generated Static*FragmentEntry specializations. */ +struct CutileTileConfig { + int tile_m; + int tile_n; + int tile_k; +}; + /** Embedded CUDA binary module (cubin), loaded directly via cudaLibraryLoadData. */ struct CubinFragmentEntry { virtual ~CubinFragmentEntry() = default; @@ -76,6 +83,12 @@ struct CubinFragmentEntry { virtual int get_cc_major() const = 0; virtual int get_cc_minor() const = 0; + + virtual int get_tile_m() const { return 0; } + + virtual int get_tile_n() const { return 0; } + + virtual int get_tile_k() const { return 0; } }; template @@ -93,6 +106,16 @@ struct StaticCubinFragmentEntry final : CubinFragmentEntry { int get_cc_minor() const override { return FragmentTag::cc_minor; } + int get_tile_m() const override { return tile_m; } + + int get_tile_n() const override { return tile_n; } + + int get_tile_k() const override { return tile_k; } + + static const int tile_m; + static const int tile_n; + static const int tile_k; + static const uint8_t* const data; static const size_t length; }; @@ -106,6 +129,12 @@ struct TileIrBytecodeFragmentEntry { virtual size_t get_length() const = 0; virtual const char* get_key() const = 0; + + virtual int get_tile_m() const { return 0; } + + virtual int get_tile_n() const { return 0; } + + virtual int get_tile_k() const { return 0; } }; template @@ -125,6 +154,16 @@ struct StaticTileIrBytecodeFragmentEntry final : TileIrBytecodeFragmentEntry { return typeid(StaticTileIrBytecodeFragmentEntry).name(); } + int get_tile_m() const override { return tile_m; } + + int get_tile_n() const override { return tile_n; } + + int get_tile_k() const override { return tile_k; } + + static const int tile_m; + static const int tile_n; + static const int tile_k; + static const uint8_t* const data; static const size_t length; }; diff --git a/cpp/include/cuvs/detail/jit_lto/fused_distance_nn/fused_1nn_fragments.hpp b/cpp/include/cuvs/detail/jit_lto/fused_distance_nn/fused_1nn_fragments.hpp index 517118bbe2..658c6e882b 100644 --- a/cpp/include/cuvs/detail/jit_lto/fused_distance_nn/fused_1nn_fragments.hpp +++ b/cpp/include/cuvs/detail/jit_lto/fused_distance_nn/fused_1nn_fragments.hpp @@ -6,16 +6,82 @@ #pragma once #include +#include namespace cuvs::distance::detail { -template +struct metric_tag_ip {}; +struct metric_tag_l2 {}; +struct metric_tag_cos {}; + +template +struct cutile_tile_config { + static constexpr int tile_m = TileM; + static constexpr int tile_n = TileN; + static constexpr int tile_k = TileK; +}; + +template +struct fused_1nn_metric_tag; + +template <> +struct fused_1nn_metric_tag { + using type = metric_tag_ip; +}; + +template <> +struct fused_1nn_metric_tag { + using type = metric_tag_l2; +}; + +template <> +struct fused_1nn_metric_tag { + using type = metric_tag_l2; +}; + +template <> +struct fused_1nn_metric_tag { + using type = metric_tag_cos; +}; + +/** Whether sqrt is applied when packing distance into KVP output. */ +template +constexpr bool fused_1nn_apply_sqrt_at_pack(bool is_sqrt) +{ + if constexpr (Metric == cuvs::distance::DistanceType::L2Expanded || + Metric == cuvs::distance::DistanceType::L2SqrtExpanded) { + return is_sqrt; + } else { + return false; + } +} + +template +using fused_1nn_metric_tag_t = typename fused_1nn_metric_tag::type; + +template +struct fused_1nn_data_tag; + +template <> +struct fused_1nn_data_tag { + using type = cuvs::neighbors::detail::tag_f; +}; + +template <> +struct fused_1nn_data_tag { + using type = cuvs::neighbors::detail::tag_h; +}; + +template +using fused_1nn_data_tag_t = typename fused_1nn_data_tag::type; + +template struct fragment_tag_fused_1nn_cubin { static constexpr int cc_major = ArchTag::cc_major; static constexpr int cc_minor = ArchTag::cc_minor; }; -template +template struct fragment_tag_fused_1nn_tileir {}; } // namespace cuvs::distance::detail diff --git a/cpp/src/detail/jit_lto/TileAlgorithmPlanner.cpp b/cpp/src/detail/jit_lto/TileAlgorithmPlanner.cpp index e0ce77e789..1487abb239 100644 --- a/cpp/src/detail/jit_lto/TileAlgorithmPlanner.cpp +++ b/cpp/src/detail/jit_lto/TileAlgorithmPlanner.cpp @@ -9,6 +9,31 @@ #include #include +#include +#include + +namespace { + +template +CutileTileConfig tile_config_from_fragment(const FragmentT* fragment, const std::string& entrypoint) +{ + if (fragment == nullptr) { + RAFT_FAIL("cuTile planner '%s' has no registered fragments", entrypoint.c_str()); + } + const int tile_m = fragment->get_tile_m(); + const int tile_n = fragment->get_tile_n(); + const int tile_k = fragment->get_tile_k(); + if (tile_m <= 0 || tile_n <= 0 || tile_k <= 0) { + RAFT_FAIL( + "cuTile planner '%s' is missing tile geometry in its static fragment (check " + "register_cutile_fragment.cpp generation)", + entrypoint.c_str()); + } + return CutileTileConfig{tile_m, tile_n, tile_k}; +} + +} // namespace + std::string TileAlgorithmPlanner::get_planner_key() const { std::string key = this->entrypoint; @@ -19,6 +44,27 @@ std::string TileAlgorithmPlanner::get_planner_key() const return key; } +CutileTileConfig TileAlgorithmPlanner::tile_config() const +{ + int cc_major = 0; + int cc_minor = 0; + if (cuvs::detail::jit_lto::get_device_compute_capability(cc_major, cc_minor)) { + for (const auto& fragment : cubin_fragments_) { + if (fragment->get_cc_major() == cc_major && fragment->get_cc_minor() == cc_minor) { + return tile_config_from_fragment(fragment.get(), entrypoint); + } + } + } + + if (tileir_fragment_) { return tile_config_from_fragment(tileir_fragment_.get(), entrypoint); } + + if (!cubin_fragments_.empty()) { + return tile_config_from_fragment(cubin_fragments_.front().get(), entrypoint); + } + + RAFT_FAIL("cuTile planner '%s' has no registered fragments", entrypoint.c_str()); +} + std::shared_ptr TileAlgorithmPlanner::build() { int cc_major = 0; diff --git a/cpp/src/distance/detail/fused_distance_nn.cuh b/cpp/src/distance/detail/fused_distance_nn.cuh index b1b18e58f6..63ef5396ac 100644 --- a/cpp/src/distance/detail/fused_distance_nn.cuh +++ b/cpp/src/distance/detail/fused_distance_nn.cuh @@ -64,7 +64,7 @@ void fusedDistanceNNImpl(OutT* min, #if CUVS_CUTILE_ENABLED if (cuvs::detail::jit_lto::cutile_launch_available_on_current_device() && - try_fused_1nn_tile(min, x, y, m, n, k, metric, stream)) { + try_fused_1nn_tile(min, x, y, xn, yn, m, n, k, metric, sqrt, stream)) { return; } #endif diff --git a/cpp/src/distance/detail/fused_distance_nn/cutile/export_fused_1nn.py b/cpp/src/distance/detail/fused_distance_nn/cutile/export_fused_1nn.py index 10a4fa9ec1..fcefe7a027 100644 --- a/cpp/src/distance/detail/fused_distance_nn/cutile/export_fused_1nn.py +++ b/cpp/src/distance/detail/fused_distance_nn/cutile/export_fused_1nn.py @@ -19,7 +19,7 @@ export_kernel, ) -from fused_1nn_kernel import KERNELS, KERNEL_SYMBOLS, TILE_CONSTANTS +from fused_1nn_kernel import METRICS, kernel_symbol, make_kernel, metric_abbrev DEFAULT_TILEIR_BYTECODE_VERSION = "13.1" # cuTile requires a gpu_code even for TileIR bytecode export: it selects the compilation @@ -35,50 +35,75 @@ def _dtype_for(data_type: str): raise ValueError(f"Unsupported data_type {data_type!r}") -def _kernel_signature(data_type: str) -> KernelSignature: - elem = _dtype_for(data_type) - array = ArrayConstraint( - elem, - 2, - index_dtype=ct.int64, - stride_lower_bound_incl=0, - alias_groups=(), - may_alias_internally=False, - ) - idx_array = ArrayConstraint( - ct.int64, - 1, +def _data_abbrev(data_type: str) -> str: + return {"half": "h", "float": "f"}[data_type] + + +def _relaxed_matrix_constraint(elem_dtype): + """Array constraints matching the relaxed TMA-friendly layout from gemm_nn_cutile.""" + return ArrayConstraint( + elem_dtype, + ndim=2, index_dtype=ct.int64, - stride_lower_bound_incl=0, + stride_lower_bound_incl=(0, None), alias_groups=(), may_alias_internally=False, - stride_constant=(1,), + stride_constant=(None, 1), + stride_divisible_by=(8, 1), + shape_divisible_by=(1, 1), + base_addr_divisible_by=16, ) - dist_array = ArrayConstraint( - ct.float32, - 1, + + +def _relaxed_vector_constraint(elem_dtype, *, tma_friendly: bool = False): + base_div = 16 if tma_friendly else 1 + return ArrayConstraint( + elem_dtype, + ndim=1, index_dtype=ct.int64, - stride_lower_bound_incl=0, + stride_lower_bound_incl=(None,), alias_groups=(), may_alias_internally=False, stride_constant=(1,), + stride_divisible_by=(1,), + shape_divisible_by=(1,), + base_addr_divisible_by=base_div, ) - tm, tn, tk = TILE_CONSTANTS + + +def _kernel_signature( + data_type: str, + metric: str, + tile_m: int, + tile_n: int, + tile_k: int, +) -> KernelSignature: + elem = _dtype_for(data_type) + matrix = _relaxed_matrix_constraint(elem) + norm_array = _relaxed_vector_constraint(elem, tma_friendly=True) + idx_array = _relaxed_vector_constraint(ct.int64) + dist_array = _relaxed_vector_constraint(ct.float32) + + abbrev = _data_abbrev(data_type) + symbol = kernel_symbol(abbrev, metric_abbrev(metric)) + return KernelSignature( parameters=[ - array, - array, + matrix, + matrix, + norm_array, + norm_array, idx_array, dist_array, ScalarConstraint(ct.int64), ScalarConstraint(ct.int64), ScalarConstraint(ct.int64), - ConstantConstraint(tm), - ConstantConstraint(tn), - ConstantConstraint(tk), + ConstantConstraint(tile_m), + ConstantConstraint(tile_n), + ConstantConstraint(tile_k), ], calling_convention=CallingConvention.cutile_python_v1(), - ).with_symbol(KERNEL_SYMBOLS[data_type]) + ).with_symbol(symbol) def export_binary( @@ -86,11 +111,15 @@ def export_binary( *, output_format: Literal["cubin", "tileir_bytecode"], data_type: str, + metric: str, + tile_m: int, + tile_n: int, + tile_k: int, gpu_code: str, bytecode_version: str | None = None, ) -> str: - kernel = KERNELS[data_type] - signature = _kernel_signature(data_type) + kernel = make_kernel(data_type, metric, tile_m, tile_n, tile_k) + signature = _kernel_signature(data_type, metric, tile_m, tile_n, tile_k) export_kwargs = { "kernel": kernel, @@ -116,8 +145,12 @@ def main() -> int: "--format", choices=("cubin", "tileir_bytecode"), default="cubin" ) parser.add_argument( - "--data-type", choices=tuple(KERNELS.keys()), required=True + "--data-type", choices=("half", "float"), required=True ) + parser.add_argument("--metric", choices=METRICS, required=True) + parser.add_argument("--tile-m", type=int, required=True) + parser.add_argument("--tile-n", type=int, required=True) + parser.add_argument("--tile-k", type=int, required=True) parser.add_argument( "--gpu-code", default=DEFAULT_TILEIR_EXPORT_GPU_CODE, @@ -133,6 +166,10 @@ def main() -> int: args.output_file, output_format=args.format, data_type=args.data_type, + metric=args.metric, + tile_m=args.tile_m, + tile_n=args.tile_n, + tile_k=args.tile_k, gpu_code=args.gpu_code, bytecode_version=args.bytecode_version, ) diff --git a/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_cutile_matrix.json b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_cutile_matrix.json index 52955863c5..3aa9dffd8a 100644 --- a/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_cutile_matrix.json +++ b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_cutile_matrix.json @@ -10,11 +10,32 @@ "data_abbrev": "f" } ], + "_metric": [ + { + "metric": "inner_product", + "metric_abbrev": "ip" + }, + { + "metric": "l2_expanded", + "metric_abbrev": "l2" + }, + { + "metric": "cosine_expanded", + "metric_abbrev": "cos" + } + ], + "_tile": [ + { + "tile_m": 128, + "tile_n": 128, + "tile_k": 64 + } + ], "_export": [ { "output_format": "cubin", "artifact_ext": "cubin", - "artifact_basename": "@data_type@_@gpu_code@", + "artifact_basename": "@data_type@_@metric_abbrev@_@gpu_code@", "register": "cubin", "gpu_code": "sm_80", "cc_major": 8, @@ -24,7 +45,7 @@ { "output_format": "cubin", "artifact_ext": "cubin", - "artifact_basename": "@data_type@_@gpu_code@", + "artifact_basename": "@data_type@_@metric_abbrev@_@gpu_code@", "register": "cubin", "gpu_code": "sm_86", "cc_major": 8, @@ -34,7 +55,7 @@ { "output_format": "cubin", "artifact_ext": "cubin", - "artifact_basename": "@data_type@_@gpu_code@", + "artifact_basename": "@data_type@_@metric_abbrev@_@gpu_code@", "register": "cubin", "gpu_code": "sm_90", "cc_major": 9, @@ -44,7 +65,7 @@ { "output_format": "cubin", "artifact_ext": "cubin", - "artifact_basename": "@data_type@_@gpu_code@", + "artifact_basename": "@data_type@_@metric_abbrev@_@gpu_code@", "register": "cubin", "gpu_code": "sm_120", "cc_major": 12, @@ -54,7 +75,7 @@ { "output_format": "tileir_bytecode", "artifact_ext": "tilebc", - "artifact_basename": "@data_type@", + "artifact_basename": "@data_type@_@metric_abbrev@", "register": "tileir", "gpu_code": "sm_80", "bytecode_version": "13.1" diff --git a/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_kernel.py b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_kernel.py index 65fe165b70..162e6ceb6b 100644 --- a/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_kernel.py +++ b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_kernel.py @@ -1,6 +1,6 @@ # SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. # SPDX-License-Identifier: Apache-2.0 -"""cuTile fused GEMM + inner-product 1-NN (argmax dot product) for cuVS.""" +"""cuTile fused GEMM + 1-NN kernels (InnerProduct, L2Expanded, CosineExpanded).""" from __future__ import annotations @@ -8,20 +8,38 @@ ConstInt = ct.Constant[int] -TILE_M = 128 -TILE_N = 256 -TILE_K = 64 +# Default tile geometry; overridden per export via make_kernel(..., tile_m, tile_n, tile_k). +DEFAULT_TILE_M = 128 +DEFAULT_TILE_N = 128 +DEFAULT_TILE_K = 64 +METRICS = ("inner_product", "l2_expanded", "cosine_expanded") -def _make_kernel(data_type: str): + +def make_kernel( + data_type: str, + metric: str, + tile_m: int = DEFAULT_TILE_M, + tile_n: int = DEFAULT_TILE_N, + tile_k: int = DEFAULT_TILE_K, +): + """Build a cuTile kernel with metric and tile sizes baked in at compile time.""" if data_type not in ("half", "float"): raise ValueError(f"Unsupported data_type {data_type!r}") + if metric not in METRICS: + raise ValueError(f"Unsupported metric {metric!r}") + acc_dtype = ct.float32 + is_ip = metric == "inner_product" + is_l2 = metric == "l2_expanded" + is_cos = metric == "cosine_expanded" @ct.kernel def fused_1nn_kernel( A, B, + A_norm, + B_norm, OutIdx, OutDist, M, @@ -33,7 +51,10 @@ def fused_1nn_kernel( ): bidm = ct.bid(0) - best_dist = ct.full((tm,), -3.4e38, acc_dtype) + if is_ip: + best_dist = ct.full((tm,), -3.4e38, acc_dtype) + else: + best_dist = ct.full((tm,), 3.4e38, acc_dtype) best_idx = ct.zeros((tm,), ct.int64) num_tiles_k = ct.num_tiles(A, axis=1, shape=(tm, tk)) @@ -52,11 +73,37 @@ def fused_1nn_kernel( ) accumulator = ct.mma(a, ct.transpose(b_T), accumulator) - curr_max = ct.max(accumulator, axis=1) - curr_idx = ct.argmax(accumulator, axis=1) + if is_ip: + score = accumulator + elif is_l2 or is_cos: + a_norm = ct.load( + A_norm, index=(bidm,), shape=(tm,), padding_mode=zero_pad + ) + b_norm = ct.load( + B_norm, index=(n,), shape=(tn,), padding_mode=zero_pad + ) + if is_l2: + # L2 expanded: ||x||^2 + ||y||^2 - 2 * dot(x, y); norms are squared. + score = ( + a_norm[:, None] + b_norm[None, :] - (2.0 * accumulator) + ) + elif is_cos: + # Cosine expanded distance: 1 - dot / (||x|| * ||y||); norms are L2 (not squared). + # No sqrt during the reduction — only arithmetic on stored distance if needed. + denom = a_norm[:, None] * b_norm[None, :] + score = 1.0 - (accumulator / denom) + + if is_ip: + curr_best = ct.max(score, axis=1) + curr_idx = ct.argmax(score, axis=1) + update = curr_best > best_dist + best_dist = ct.where(update, curr_best, best_dist) + else: + curr_best = ct.min(score, axis=1) + curr_idx = ct.argmin(score, axis=1) + update = curr_best < best_dist + best_dist = ct.where(update, curr_best, best_dist) - update = curr_max > best_dist - best_dist = ct.where(update, curr_max, best_dist) best_idx = ct.where(update, n * tn + curr_idx, best_idx) ct.store(OutIdx, index=(bidm,), tile=best_idx) @@ -65,14 +112,14 @@ def fused_1nn_kernel( return fused_1nn_kernel -KERNELS = { - "half": _make_kernel("half"), - "float": _make_kernel("float"), -} +def kernel_symbol(data_abbrev: str, metric_abbrev: str) -> str: + """Must stay in sync with fused_1nn_kernel_entrypoint() in fused_1nn_planner.hpp.""" + return f"fused_1nn_{data_abbrev}_{metric_abbrev}" -KERNEL_SYMBOLS = { - "half": "fused_1nn_half", - "float": "fused_1nn_float", -} -TILE_CONSTANTS = (TILE_M, TILE_N, TILE_K) +def metric_abbrev(metric: str) -> str: + return { + "inner_product": "ip", + "l2_expanded": "l2", + "cosine_expanded": "cos", + }[metric] diff --git a/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_planner.hpp b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_planner.hpp index dd2a539528..ae0ae118bd 100644 --- a/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_planner.hpp +++ b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_planner.hpp @@ -12,32 +12,49 @@ #include #include +#include "fused_1nn_cutile_tiles.hpp" + namespace cuvs::distance::detail { -/** Must match KERNEL_SYMBOLS in fused_1nn_kernel.py (export uses with_symbol). */ -template +/** Must match kernel_symbol() in fused_1nn_kernel.py (export uses with_symbol). */ +template inline const char* fused_1nn_kernel_entrypoint() { - if constexpr (std::is_same_v) { - return "fused_1nn_half"; - } else if constexpr (std::is_same_v) { - return "fused_1nn_float"; - } else { - static_assert(sizeof(DataTag) == 0, "unsupported fused 1-NN cuTile data type"); - return ""; + if constexpr (std::is_same_v) { + if constexpr (std::is_same_v) { + return "fused_1nn_f_ip"; + } else if constexpr (std::is_same_v) { + return "fused_1nn_f_l2"; + } else if constexpr (std::is_same_v) { + return "fused_1nn_f_cos"; + } + } else if constexpr (std::is_same_v) { + if constexpr (std::is_same_v) { + return "fused_1nn_h_ip"; + } else if constexpr (std::is_same_v) { + return "fused_1nn_h_l2"; + } else if constexpr (std::is_same_v) { + return "fused_1nn_h_cos"; + } } + static_assert(sizeof(DataTag) == 0, "unsupported fused 1-NN cuTile data/metric combination"); + return ""; } -template +template struct Fused1nnTilePlanner : TileAlgorithmPlanner { + using DataTag = fused_1nn_data_tag_t; + using MetricTag = fused_1nn_metric_tag_t; + inline static LauncherJitCache launcher_jit_cache{}; Fused1nnTilePlanner() - : TileAlgorithmPlanner(fused_1nn_kernel_entrypoint(), launcher_jit_cache) + : TileAlgorithmPlanner(fused_1nn_kernel_entrypoint(), launcher_jit_cache) { } - /** Registers embedded cubin modules (one per SM); see register_cubin.cpp object files. */ + /** Registers embedded cubin modules (one per SM); see register_cutile_fragment.cpp object files. + */ void add_entrypoint() { using cuvs::detail::jit_lto::cutile_arch_12_0; @@ -45,15 +62,20 @@ struct Fused1nnTilePlanner : TileAlgorithmPlanner { using cuvs::detail::jit_lto::cutile_arch_8_6; using cuvs::detail::jit_lto::cutile_arch_9_0; - this->add_static_fragment>(); - this->add_static_fragment>(); - this->add_static_fragment>(); - this->add_static_fragment>(); + this->add_static_fragment< + fragment_tag_fused_1nn_cubin>(); + this->add_static_fragment< + fragment_tag_fused_1nn_cubin>(); + this->add_static_fragment< + fragment_tag_fused_1nn_cubin>(); + this->add_static_fragment< + fragment_tag_fused_1nn_cubin>(); } void add_tileir_fallback() { - this->add_static_tileir_fragment>(); + this->add_static_tileir_fragment< + fragment_tag_fused_1nn_tileir>(); } }; diff --git a/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_tile.cu b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_tile.cu index 0ad4ee62a5..e343afca30 100644 --- a/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_tile.cu +++ b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_tile.cu @@ -16,28 +16,42 @@ namespace detail { namespace { -constexpr int64_t TILE_M = 128; - template -__global__ void pack_fused_1nn_kvp(OutT* out, const int64_t* idx, const float* dist, IdxT len) +__global__ void pack_fused_1nn_kvp( + OutT* out, const int64_t* idx, const float* dist, IdxT len, bool apply_sqrt) { IdxT i = blockIdx.x * blockDim.x + threadIdx.x; if (i < len) { - out[i].key = static_cast(idx[i]); - out[i].value = static_cast(dist[i]); + out[i].key = static_cast(idx[i]); + float value = dist[i]; + if (apply_sqrt) { value = sqrtf(value); } + out[i].value = static_cast(value); } } -template -bool launch_fused_1nn_tile( - const DataT* x, const DataT* y, OutT* out, IdxT m, IdxT n, IdxT k, cudaStream_t stream) +template +bool launch_fused_1nn_tile(const DataT* x, + const DataT* y, + const DataT* xn, + const DataT* yn, + OutT* out, + IdxT m, + IdxT n, + IdxT k, + bool is_sqrt, + cudaStream_t stream) { - Fused1nnTilePlanner planner; + if constexpr (!std::is_same_v && !std::is_same_v) { return false; } + + Fused1nnTilePlanner planner; planner.add_entrypoint(); planner.add_tileir_fallback(); - auto launcher = planner.try_get_launcher(); + const CutileTileConfig tile_cfg = planner.tile_config(); + auto launcher = planner.try_get_launcher(); if (!launcher) { return false; } + const bool apply_sqrt = fused_1nn_apply_sqrt_at_pack(is_sqrt); + int64_t* d_idx = nullptr; float* d_dist = nullptr; RAFT_CUDA_TRY(cudaMallocAsync(&d_idx, m * sizeof(int64_t), stream)); @@ -47,6 +61,10 @@ bool launch_fused_1nn_tile( int64_t stride_x[2] = {k, 1}; int64_t shape_y[2] = {n, k}; int64_t stride_y[2] = {k, 1}; + int64_t shape_xn = m; + int64_t stride_xn = 1; + int64_t shape_yn = n; + int64_t stride_yn = 1; int64_t shape_idx = m; int64_t stride_idx = 1; int64_t shape_dist = m; @@ -56,15 +74,17 @@ bool launch_fused_1nn_tile( void* x_ptr = const_cast(x); void* y_ptr = const_cast(y); + void* xn_ptr = const_cast(xn); + void* yn_ptr = const_cast(yn); void* idx_ptr = d_idx; void* dist_ptr = d_dist; - dim3 grid((m + TILE_M - 1) / TILE_M, 1, 1); + const int64_t tile_m = tile_cfg.tile_m; + dim3 grid((m + tile_m - 1) / tile_m, 1, 1); dim3 block(1, 1, 1); - // cutile_python_v1 (see fused_1nn_float PTX): each 2D array is (ptr, shape0, shape1, - // stride0, stride1); each 1D array is (ptr, shape, stride); ConstantConstraint tile sizes - // are embedded in the module. + // cutile_python_v1: 2D array (ptr, shape0, shape1, stride0, stride1); + // 1D array (ptr, shape, stride); tile sizes are embedded constants. using fused_1nn_cutile_kernel_t = void(void*, int64_t, int64_t, @@ -81,6 +101,12 @@ bool launch_fused_1nn_tile( void*, int64_t, int64_t, + void*, + int64_t, + int64_t, + void*, + int64_t, + int64_t, int64_t, int64_t, int64_t); @@ -98,6 +124,12 @@ bool launch_fused_1nn_tile( shape_y[1], stride_y[0], stride_y[1], + xn_ptr, + shape_xn, + stride_xn, + yn_ptr, + shape_yn, + stride_yn, idx_ptr, shape_idx, stride_idx, @@ -108,39 +140,62 @@ bool launch_fused_1nn_tile( N, K); - pack_fused_1nn_kvp<<<(m + 255) / 256, 256, 0, stream>>>(out, d_idx, d_dist, m); + pack_fused_1nn_kvp + <<<(m + 255) / 256, 256, 0, stream>>>(out, d_idx, d_dist, m, apply_sqrt); RAFT_CUDA_TRY(cudaGetLastError()); RAFT_CUDA_TRY(cudaFreeAsync(d_idx, stream)); RAFT_CUDA_TRY(cudaFreeAsync(d_dist, stream)); return true; } +template +bool try_fused_1nn_tile_dispatch(OutT* min, + const DataT* x, + const DataT* y, + const DataT* xn, + const DataT* yn, + IdxT m, + IdxT n, + IdxT k, + cuvs::distance::DistanceType metric, + bool is_sqrt, + cudaStream_t stream) +{ + switch (metric) { + case cuvs::distance::DistanceType::InnerProduct: + return launch_fused_1nn_tile( + x, y, xn, yn, min, m, n, k, is_sqrt, stream); + case cuvs::distance::DistanceType::L2Expanded: + return launch_fused_1nn_tile( + x, y, xn, yn, min, m, n, k, is_sqrt, stream); + case cuvs::distance::DistanceType::L2SqrtExpanded: + return launch_fused_1nn_tile( + x, y, xn, yn, min, m, n, k, is_sqrt, stream); + case cuvs::distance::DistanceType::CosineExpanded: + return launch_fused_1nn_tile( + x, y, xn, yn, min, m, n, k, is_sqrt, stream); + default: return false; + } +} + } // namespace -template , int>> +template + requires Fused1nnKvpOutput bool try_fused_1nn_tile(OutT* min, const DataT* x, const DataT* y, + const DataT* xn, + const DataT* yn, IdxT m, IdxT n, IdxT k, cuvs::distance::DistanceType metric, + bool is_sqrt, cudaStream_t stream) { - if (metric != cuvs::distance::DistanceType::InnerProduct) { return false; } - - if constexpr (std::is_same_v) { - return launch_fused_1nn_tile( - x, y, min, m, n, k, stream); - } else if constexpr (std::is_same_v) { - return launch_fused_1nn_tile( - x, y, min, m, n, k, stream); - } else { - return false; - } + return try_fused_1nn_tile_dispatch( + min, x, y, xn, yn, m, n, k, metric, is_sqrt, stream); } using kvp_i_f = raft::KeyValuePair; @@ -150,15 +205,18 @@ using kvp_i64_h = raft::KeyValuePair; #define CUVS_INST_TRY_FUSED_1NN_TILE(DataT, OutT, IdxT) \ template CUVS_EXPORT bool try_fused_1nn_tile(OutT*, \ + const DataT*, \ + const DataT*, \ const DataT*, \ const DataT*, \ IdxT, \ IdxT, \ IdxT, \ cuvs::distance::DistanceType, \ + bool, \ cudaStream_t) -// int and int32_t are the same on LP64; one instantiation covers both. +// int and int64_t are the same on LP64; one instantiation covers both. CUVS_INST_TRY_FUSED_1NN_TILE(float, kvp_i_f, int); CUVS_INST_TRY_FUSED_1NN_TILE(float, kvp_i64_f, int64_t); CUVS_INST_TRY_FUSED_1NN_TILE(half, kvp_i_f, int); diff --git a/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_tile.hpp b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_tile.hpp index d72a020ba7..807ecdb233 100644 --- a/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_tile.hpp +++ b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_tile.hpp @@ -5,6 +5,7 @@ #pragma once +#include #include #include @@ -22,25 +23,36 @@ inline constexpr bool is_fused_1nn_kvp_output_v = (std::is_same_v> || std::is_same_v>); -template , int> = 0> +template +concept Fused1nnKvpOutput = is_fused_1nn_kvp_output_v; + +template + requires Fused1nnKvpOutput bool try_fused_1nn_tile(OutT* min, const DataT* x, const DataT* y, + const DataT* xn, + const DataT* yn, IdxT m, IdxT n, IdxT k, cuvs::distance::DistanceType metric, + bool is_sqrt, cudaStream_t stream); -template , int> = 0> -bool try_fused_1nn_tile( - OutT*, const DataT*, const DataT*, IdxT, IdxT, IdxT, cuvs::distance::DistanceType, cudaStream_t) +template + requires(!Fused1nnKvpOutput) +bool try_fused_1nn_tile(OutT*, + const DataT*, + const DataT*, + const DataT*, + const DataT*, + IdxT, + IdxT, + IdxT, + cuvs::distance::DistanceType, + bool, + cudaStream_t) { return false; } From d6560fca9e678c9eabcbcf519506c000f15c993d Mon Sep 17 00:00:00 2001 From: divyegala Date: Thu, 25 Jun 2026 21:37:40 +0000 Subject: [PATCH 07/82] if constexpr exit --- .../cutile/fused_1nn_planner.hpp | 39 ++++++++++--------- 1 file changed, 21 insertions(+), 18 deletions(-) diff --git a/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_planner.hpp b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_planner.hpp index ae0ae118bd..c70ab3f87b 100644 --- a/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_planner.hpp +++ b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_planner.hpp @@ -20,25 +20,28 @@ namespace cuvs::distance::detail { template inline const char* fused_1nn_kernel_entrypoint() { - if constexpr (std::is_same_v) { - if constexpr (std::is_same_v) { - return "fused_1nn_f_ip"; - } else if constexpr (std::is_same_v) { - return "fused_1nn_f_l2"; - } else if constexpr (std::is_same_v) { - return "fused_1nn_f_cos"; - } - } else if constexpr (std::is_same_v) { - if constexpr (std::is_same_v) { - return "fused_1nn_h_ip"; - } else if constexpr (std::is_same_v) { - return "fused_1nn_h_l2"; - } else if constexpr (std::is_same_v) { - return "fused_1nn_h_cos"; - } + if constexpr (std::is_same_v && + std::is_same_v) { + return "fused_1nn_f_ip"; + } else if constexpr (std::is_same_v && + std::is_same_v) { + return "fused_1nn_f_l2"; + } else if constexpr (std::is_same_v && + std::is_same_v) { + return "fused_1nn_f_cos"; + } else if constexpr (std::is_same_v && + std::is_same_v) { + return "fused_1nn_h_ip"; + } else if constexpr (std::is_same_v && + std::is_same_v) { + return "fused_1nn_h_l2"; + } else if constexpr (std::is_same_v && + std::is_same_v) { + return "fused_1nn_h_cos"; + } else { + static_assert(sizeof(DataTag) == 0, "unsupported fused 1-NN cuTile data/metric combination"); + return ""; } - static_assert(sizeof(DataTag) == 0, "unsupported fused 1-NN cuTile data/metric combination"); - return ""; } template From 674285321fd5c37d9d79aea86da4b9a9a670c1ff Mon Sep 17 00:00:00 2001 From: divyegala Date: Thu, 25 Jun 2026 22:50:48 +0000 Subject: [PATCH 08/82] passing KMeans tests --- .../cuvs/detail/jit_lto/tileir_compat.hpp | 11 ++- cpp/src/cluster/detail/kmeans_balanced.cuh | 85 ++++++++++++++----- cpp/src/cluster/detail/kmeans_common.cuh | 47 +++++++--- .../detail/minClusterDistanceCompute.cu | 44 ++++++---- cpp/src/distance/detail/fused_distance_nn.cuh | 27 +++--- .../cutile/fused_1nn_kernel.py | 10 +++ .../cutile/fused_1nn_tile.cu | 2 + .../cutile/fused_1nn_tile.hpp | 33 ++++++- 8 files changed, 187 insertions(+), 72 deletions(-) diff --git a/cpp/include/cuvs/detail/jit_lto/tileir_compat.hpp b/cpp/include/cuvs/detail/jit_lto/tileir_compat.hpp index f15407fd4c..f114233179 100644 --- a/cpp/include/cuvs/detail/jit_lto/tileir_compat.hpp +++ b/cpp/include/cuvs/detail/jit_lto/tileir_compat.hpp @@ -61,12 +61,16 @@ inline bool tileir_fallback_available(int driver_version) * is CUDA 13+, and either a matching embedded cubin exists (no driver JIT required) or the driver * can JIT the embedded TileIR bytecode fallback. */ +#if CUVS_CUTILE_ENABLED inline bool cutile_launch_available_for_arch(int cc_major, int cc_minor, int driver_version) { - if (!cutile_integration_enabled()) { return false; } + if (!runtime_cuda13_or_newer()) { return false; } if (has_embedded_cubin_for_arch(cc_major, cc_minor)) { return true; } return tileir_fallback_available(driver_version); } +#else +inline constexpr bool cutile_launch_available_for_arch(int, int, int) { return false; } +#endif inline bool query_driver_version(int& driver_version) { @@ -86,6 +90,7 @@ inline bool query_current_device_arch(int& cc_major, int& cc_minor) return true; } +#if CUVS_CUTILE_ENABLED inline bool cutile_launch_available_on_current_device() { int cc_major = 0; @@ -95,5 +100,9 @@ inline bool cutile_launch_available_on_current_device() if (!query_driver_version(driver_version)) { return false; } return cutile_launch_available_for_arch(cc_major, cc_minor, driver_version); } +#else +/** Compile-time false when cuTile is not built; use in if constexpr to skip cuTile-only paths. */ +inline constexpr bool cutile_launch_available_on_current_device() { return false; } +#endif } // namespace cuvs::detail::jit_lto diff --git a/cpp/src/cluster/detail/kmeans_balanced.cuh b/cpp/src/cluster/detail/kmeans_balanced.cuh index 7fac255810..86e254f473 100644 --- a/cpp/src/cluster/detail/kmeans_balanced.cuh +++ b/cpp/src/cluster/detail/kmeans_balanced.cuh @@ -121,32 +121,63 @@ inline std::enable_if_t> predict_core( break; } case cuvs::distance::DistanceType::InnerProduct: { - // TODO: pass buffer - rmm::device_uvector distances(n_rows * n_clusters, stream, mr); + if (use_cutile_fused_nn(handle, n_rows, n_clusters, dim)) { + rmm::device_uvector L2NormBuf_OR_DistBuf(0, stream, mr); + rmm::device_uvector workspace(0, stream, mr); + + auto X_view = raft::make_device_matrix_view(dataset, n_rows, dim); + auto centroids_view = + raft::make_device_matrix_view(centers, n_clusters, dim); + auto X_norm_view = raft::make_device_vector_view(dataset_norm, n_rows); + + auto minClusterAndDistance = + raft::make_device_mdarray, IdxT>( + handle, mr, raft::make_extents(n_rows)); + + cuvs::cluster::kmeans::detail::minClusterAndDistanceCompute( + handle, + X_view, + centroids_view, + minClusterAndDistance.view(), + X_norm_view, + L2NormBuf_OR_DistBuf, + params.metric, + 0, + 0, + workspace); + + raft::linalg::map(handle, + raft::make_const_mdspan(minClusterAndDistance.view()), + raft::make_device_vector_view(labels, n_rows), + raft::compose_op, raft::key_op>()); + } else { + rmm::device_uvector distances(n_rows * n_clusters, stream, mr); - MathT alpha = -1.0; - MathT beta = 0.0; + MathT alpha = -1.0; + MathT beta = 0.0; - raft::linalg::gemm(handle, - true, - false, - n_clusters, - n_rows, - dim, - &alpha, - centers, - dim, - dataset, - dim, - &beta, - distances.data(), - n_clusters, - stream); + raft::linalg::gemm(handle, + true, + false, + n_clusters, + n_rows, + dim, + &alpha, + centers, + dim, + dataset, + dim, + &beta, + distances.data(), + n_clusters, + stream); - auto distances_const_view = raft::make_device_matrix_view( - distances.data(), n_rows, n_clusters); - auto labels_view = raft::make_device_vector_view(labels, n_rows); - raft::matrix::argmin(handle, distances_const_view, labels_view); + auto distances_const_view = + raft::make_device_matrix_view( + distances.data(), n_rows, n_clusters); + auto labels_view = raft::make_device_vector_view(labels, n_rows); + raft::matrix::argmin(handle, distances_const_view, labels_view); + } break; } default: { @@ -195,6 +226,14 @@ auto calc_minibatch_size(const raft::resources& handle, mem_per_row += sizeof(MathT) * n_clusters; } } break; + case distance::DistanceType::InnerProduct: { + if (use_cutile_fused_nn(handle, n_rows, n_clusters, dim)) { + mem_per_row += sizeof(int); + mem_per_row += sizeof(raft::KeyValuePair); + } else { + mem_per_row += sizeof(MathT) * n_clusters; + } + } break; // Other metrics require storing a distance matrix. default: { mem_per_row += sizeof(MathT) * n_clusters; diff --git a/cpp/src/cluster/detail/kmeans_common.cuh b/cpp/src/cluster/detail/kmeans_common.cuh index ba98dadca6..0606d77dec 100644 --- a/cpp/src/cluster/detail/kmeans_common.cuh +++ b/cpp/src/cluster/detail/kmeans_common.cuh @@ -7,6 +7,7 @@ #include "../../distance/distance.cuh" #include #include +#include #include #include @@ -57,31 +58,51 @@ namespace cuvs::cluster::kmeans::detail { +template +inline constexpr bool is_cutile_fused_data_type_v = + std::is_same_v || std::is_same_v; + /** - * @brief Returns true if the fused distance NN implementation should be used. + * @brief Returns true if the fused distance NN implementation should be used (CUTLASS and/or + * cuTile). + * + * Float/half: use fused whenever cuTile can launch (any architecture and problem size). If cuTile + * is unavailable, fall back to legacy CUTLASS fused on Ampere and Hopper only. Double and other + * types never use cuTile; they keep the historical CUTLASS/unfused heuristics on pre-Blackwell + * GPUs. * - * On Ampere (SM <= 8.x) always use fused. - * On Hopper (SM 9.x) use fused when m or n >= 4096. - * On Blackwell (SM >= 10.x) use unfused. + * Callers route through fusedDistanceNNMinReduce when this returns true; cuTile dispatch inside + * that API is gated separately by dtype (see fusedDistanceNNImpl). */ template bool use_fused(const raft::resources& handle, IdxT m, IdxT n, IdxT k) { + (void)k; cudaDeviceProp prop; prop = raft::resource::get_device_properties(handle); - if (prop.major <= 8) { - // Use fused for Ampere or before - return true; - } else if (prop.major == 9 && (m >= 4096 || n >= 4096)) { - // On Hopper if m, n are bigger than 4096, use fused - return true; - } else if (prop.major >= 10) { - // On Blackwell onwards, use unfused - return false; + + if constexpr (is_cutile_fused_data_type_v) { + if constexpr (cuvs::detail::jit_lto::library_built_with_cutile()) { + if (cuvs::detail::jit_lto::cutile_launch_available_on_current_device()) { return true; } + } + return prop.major <= 9; } + + if (prop.major >= 10) { return false; } + if (prop.major <= 8) { return true; } + if (prop.major == 9 && (m >= 4096 || n >= 4096)) { return true; } return false; } +/** True when assignment should use the cuTile fused 1-NN kernel (float/half only). */ +template +bool use_cutile_fused_nn(const raft::resources& /*handle*/, IdxT /*m*/, IdxT /*n*/, IdxT /*k*/) +{ + if constexpr (!is_cutile_fused_data_type_v) { return false; } + if constexpr (!cuvs::detail::jit_lto::library_built_with_cutile()) { return false; } + return cuvs::detail::jit_lto::cutile_launch_available_on_current_device(); +} + template struct SamplingOp { DataT* rnd; diff --git a/cpp/src/cluster/detail/minClusterDistanceCompute.cu b/cpp/src/cluster/detail/minClusterDistanceCompute.cu index b15119599e..65678faa08 100644 --- a/cpp/src/cluster/detail/minClusterDistanceCompute.cu +++ b/cpp/src/cluster/detail/minClusterDistanceCompute.cu @@ -27,33 +27,41 @@ void minClusterAndDistanceCompute( int batch_centroids, rmm::device_uvector& workspace) { - cudaStream_t stream = raft::resource::get_cuda_stream(handle); - auto n_samples = X.extent(0); - auto n_features = X.extent(1); - auto n_clusters = centroids.extent(0); - bool is_fused = metric == cuvs::distance::DistanceType::L2Expanded || - metric == cuvs::distance::DistanceType::L2SqrtExpanded || - metric == cuvs::distance::DistanceType::CosineExpanded; - - if (is_fused) { + cudaStream_t stream = raft::resource::get_cuda_stream(handle); + auto n_samples = X.extent(0); + auto n_features = X.extent(1); + auto n_clusters = centroids.extent(0); + bool is_l2_cos_fused = metric == cuvs::distance::DistanceType::L2Expanded || + metric == cuvs::distance::DistanceType::L2SqrtExpanded || + metric == cuvs::distance::DistanceType::CosineExpanded; + const bool is_ip_cutile = + metric == cuvs::distance::DistanceType::InnerProduct && + use_cutile_fused_nn(handle, n_samples, n_clusters, n_features); + + if (is_l2_cos_fused || is_ip_cutile) { L2NormBuf_OR_DistBuf.resize(n_clusters, stream); auto centroidsNorm = raft::make_device_vector_view(L2NormBuf_OR_DistBuf.data(), n_clusters); - if (metric == cuvs::distance::DistanceType::CosineExpanded) { - raft::linalg::norm( - handle, centroids, centroidsNorm, raft::sqrt_op{}); - } else { - raft::linalg::norm( - handle, centroids, centroidsNorm); + if (is_l2_cos_fused) { + if (metric == cuvs::distance::DistanceType::CosineExpanded) { + raft::linalg::norm( + handle, centroids, centroidsNorm, raft::sqrt_op{}); + } else { + raft::linalg::norm( + handle, centroids, centroidsNorm); + } } raft::KeyValuePair initial_value(0, std::numeric_limits::max()); raft::matrix::fill(handle, minClusterAndDistance, initial_value); - bool should_use_fused = + const bool should_use_fused = use_fused(handle, n_samples, n_clusters, n_features); + auto centroidsNormConst = + raft::make_device_vector_view(L2NormBuf_OR_DistBuf.data(), n_clusters); + if (should_use_fused) { workspace.resize((sizeof(int)) * n_samples, stream); @@ -62,7 +70,7 @@ void minClusterAndDistanceCompute( X.data_handle(), centroids.data_handle(), L2NormX.data_handle(), - centroidsNorm.data_handle(), + centroidsNormConst.data_handle(), n_samples, n_clusters, n_features, @@ -83,7 +91,7 @@ void minClusterAndDistanceCompute( X.data_handle(), centroids.data_handle(), L2NormX.data_handle(), - centroidsNorm.data_handle(), + centroidsNormConst.data_handle(), n_samples, n_clusters, n_features, diff --git a/cpp/src/distance/detail/fused_distance_nn.cuh b/cpp/src/distance/detail/fused_distance_nn.cuh index 63ef5396ac..476ab9c2be 100644 --- a/cpp/src/distance/detail/fused_distance_nn.cuh +++ b/cpp/src/distance/detail/fused_distance_nn.cuh @@ -5,21 +5,14 @@ #pragma once -#ifndef CUVS_CUTILE_ENABLED -#define CUVS_CUTILE_ENABLED 0 -#endif - #include "distance_ops/l2_exp.cuh" // ops::l2_exp_distance_op -#include "fused_distance_nn/cutlass_base.cuh" -#if CUVS_CUTILE_ENABLED #include "fused_distance_nn/cutile/fused_1nn_tile.hpp" -#endif +#include "fused_distance_nn/cutlass_base.cuh" #include "fused_distance_nn/fused_cosine_nn.cuh" #include "fused_distance_nn/fused_l2_nn.cuh" #include "fused_distance_nn/helper_structs.cuh" #include "fused_distance_nn/simt_kernel.cuh" #include "pairwise_distance_base.cuh" // PairwiseDistances -#include #include #include // raft::KeyValuePair #include // raft::identity_op @@ -62,12 +55,16 @@ void fusedDistanceNNImpl(OutT* min, // The kernel policy is determined by fusedDistanceNN. typedef Policy P; -#if CUVS_CUTILE_ENABLED - if (cuvs::detail::jit_lto::cutile_launch_available_on_current_device() && - try_fused_1nn_tile(min, x, y, xn, yn, m, n, k, metric, sqrt, stream)) { - return; + // Callers (e.g. use_fused) enable this API for CUTLASS fused as well as cuTile; only try cuTile + // for float/half KVP output so double and other types never instantiate cuTile symbols here. + if constexpr (is_fused_1nn_cutile_data_v) { + if constexpr (cuvs::detail::jit_lto::library_built_with_cutile() && + is_fused_1nn_kvp_output_v) { + if (try_fused_1nn_tile(min, x, y, xn, yn, m, n, k, metric, sqrt, stream)) { + return; + } + } } -#endif dim3 blk(P::Nthreads); auto nblks = raft::ceildiv(m, P::Nthreads); @@ -88,10 +85,12 @@ void fusedDistanceNNImpl(OutT* min, break; case cuvs::distance::DistanceType::L2SqrtExpanded: case cuvs::distance::DistanceType::L2Expanded: - // initOutBuffer is take care by fusedDistanceNNImpl() so we set it false to fusedL2NNImpl. fusedL2NNImpl( min, x, y, xn, yn, m, n, k, workspace, redOp, pairRedOp, sqrt, false, stream); break; + case cuvs::distance::DistanceType::InnerProduct: + // cuTile is the only fused InnerProduct implementation; callers must gate on availability. + break; default: assert("only cosine/l2 metric is supported with fusedDistanceNN\n"); break; } } diff --git a/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_kernel.py b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_kernel.py index 162e6ceb6b..7d78525869 100644 --- a/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_kernel.py +++ b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_kernel.py @@ -93,6 +93,16 @@ def fused_1nn_kernel( denom = a_norm[:, None] * b_norm[None, :] score = 1.0 - (accumulator / denom) + # Only the final N-tile can include zero-padded centroid columns. + if n == num_tiles_n - 1: + col = ct.arange(tn, dtype=ct.int64) + global_col = n * tn + col + valid = global_col < N + if is_ip: + score = ct.where(valid[None, :], score, -3.4e38) + else: + score = ct.where(valid[None, :], score, 3.4e38) + if is_ip: curr_best = ct.max(score, axis=1) curr_idx = ct.argmax(score, axis=1) diff --git a/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_tile.cu b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_tile.cu index e343afca30..d292f0522b 100644 --- a/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_tile.cu +++ b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_tile.cu @@ -110,6 +110,7 @@ bool launch_fused_1nn_tile(const DataT* x, int64_t, int64_t, int64_t); + std::cout << "Launching cuTile kernel" << std::endl; launcher->template dispatch(stream, grid, block, @@ -194,6 +195,7 @@ bool try_fused_1nn_tile(OutT* min, bool is_sqrt, cudaStream_t stream) { + if (!cuvs::detail::jit_lto::cutile_launch_available_on_current_device()) { return false; } return try_fused_1nn_tile_dispatch( min, x, y, xn, yn, m, n, k, metric, is_sqrt, stream); } diff --git a/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_tile.hpp b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_tile.hpp index 807ecdb233..563d2583d8 100644 --- a/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_tile.hpp +++ b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_tile.hpp @@ -11,21 +11,30 @@ #include #include +#include #include +#ifndef CUVS_CUTILE_ENABLED +#define CUVS_CUTILE_ENABLED 0 +#endif + namespace cuvs { namespace distance { namespace detail { +template +inline constexpr bool is_fused_1nn_cutile_data_v = + std::is_same_v || std::is_same_v; + template inline constexpr bool is_fused_1nn_kvp_output_v = - (std::is_same_v || std::is_same_v) && - (std::is_same_v> || - std::is_same_v>); + is_fused_1nn_cutile_data_v && (std::is_same_v> || + std::is_same_v>); template concept Fused1nnKvpOutput = is_fused_1nn_kvp_output_v; +#if CUVS_CUTILE_ENABLED template requires Fused1nnKvpOutput bool try_fused_1nn_tile(OutT* min, @@ -39,6 +48,24 @@ bool try_fused_1nn_tile(OutT* min, cuvs::distance::DistanceType metric, bool is_sqrt, cudaStream_t stream); +#else +template + requires Fused1nnKvpOutput +bool try_fused_1nn_tile(OutT*, + const DataT*, + const DataT*, + const DataT*, + const DataT*, + IdxT, + IdxT, + IdxT, + cuvs::distance::DistanceType, + bool, + cudaStream_t) +{ + return false; +} +#endif template requires(!Fused1nnKvpOutput) From db6b38524cc260dffc2b8d73ec61ebd887f68e43 Mon Sep 17 00:00:00 2001 From: divyegala Date: Tue, 30 Jun 2026 04:12:33 +0000 Subject: [PATCH 09/82] undo kvp, add constrains and alignments to tile export --- cpp/CMakeLists.txt | 4 +- .../modules/generate_cutile_kernels.cmake | 18 +- .../fused_distance_nn/fused_1nn_fragments.hpp | 26 +- cpp/src/cluster/detail/kmeans.cuh | 67 ++--- cpp/src/cluster/detail/kmeans_balanced.cuh | 139 +++++++---- cpp/src/cluster/detail/kmeans_common.cuh | 207 ++++++++-------- cpp/src/cluster/detail/kmeans_mg.cuh | 48 ++-- cpp/src/cluster/detail/kmeans_mg_batched.cuh | 16 +- .../detail/minClusterDistanceCompute.cu | 234 ++++++++++-------- cpp/src/cluster/kmeans.cuh | 25 +- cpp/src/distance/detail/fused_distance_nn.cuh | 81 +++--- .../cutile/export_fused_1nn.py | 77 ++++-- .../cutile/fused_1nn_cutile_matrix.json | 26 +- .../cutile/fused_1nn_kernel.py | 61 +++-- .../cutile/fused_1nn_planner.hpp | 54 ++-- .../cutile/fused_1nn_tile.cu | 135 +++++----- .../cutile/fused_1nn_tile.hpp | 40 +-- .../fused_distance_nn/fused_cosine_nn.cuh | 37 ++- .../detail/fused_distance_nn/fused_l2_nn.cuh | 37 +-- .../fused_distance_nn/helper_structs.cuh | 105 ++++++-- .../predicated_tile_iterator_reduced_vec.h | 4 +- cpp/src/distance/fused_distance_nn-inl.cuh | 157 ++++-------- cpp/tests/CMakeLists.txt | 2 +- cpp/tests/neighbors/distance_nn.cu | 73 +++--- cpp/tests/neighbors/distance_nn_helper.cuh | 28 +++ 25 files changed, 950 insertions(+), 751 deletions(-) diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt index 87716cd296..84979d05c1 100644 --- a/cpp/CMakeLists.txt +++ b/cpp/CMakeLists.txt @@ -980,9 +980,9 @@ if(NOT BUILD_CPU_ONLY) MATRIX_JSON_FILE "${fused_1nn_cutile_dir}/fused_1nn_cutile_matrix.json" FRAGMENT_TAG_FORMAT_CUBIN - "cuvs::distance::detail::fragment_tag_fused_1nn_cubin, cuvs::detail::jit_lto::@arch_tag@>" + "cuvs::distance::detail::fragment_tag_fused_1nn_cubin, cuvs::detail::jit_lto::@arch_tag@>" FRAGMENT_TAG_FORMAT_TILEIR - "cuvs::distance::detail::fragment_tag_fused_1nn_tileir>" + "cuvs::distance::detail::fragment_tag_fused_1nn_tileir>" FRAGMENT_TAG_HEADER_FILES "" "" diff --git a/cpp/cmake/modules/generate_cutile_kernels.cmake b/cpp/cmake/modules/generate_cutile_kernels.cmake index abdca118ad..9cd8a207c8 100644 --- a/cpp/cmake/modules/generate_cutile_kernels.cmake +++ b/cpp/cmake/modules/generate_cutile_kernels.cmake @@ -150,8 +150,22 @@ function(process_cutile_matrix_entry source_list_var) set(embedded_header_file "${_artifact_stem}_${register}.h") set(_python_args - --format "${output_format}" --data-type "${data_type}" --metric "${metric}" --tile-m - "${tile_m}" --tile-n "${tile_n}" --tile-k "${tile_k}" --gpu-code "${gpu_code}" + --format + "${output_format}" + --data-type + "${data_type}" + --metric + "${metric}" + --index-type + "${index_type}" + --tile-m + "${tile_m}" + --tile-n + "${tile_n}" + --tile-k + "${tile_k}" + --gpu-code + "${gpu_code}" ) if(DEFINED bytecode_version AND NOT "${bytecode_version}" STREQUAL "") list(APPEND _python_args --bytecode-version "${bytecode_version}") diff --git a/cpp/include/cuvs/detail/jit_lto/fused_distance_nn/fused_1nn_fragments.hpp b/cpp/include/cuvs/detail/jit_lto/fused_distance_nn/fused_1nn_fragments.hpp index 658c6e882b..c6afe16b5c 100644 --- a/cpp/include/cuvs/detail/jit_lto/fused_distance_nn/fused_1nn_fragments.hpp +++ b/cpp/include/cuvs/detail/jit_lto/fused_distance_nn/fused_1nn_fragments.hpp @@ -5,6 +5,8 @@ #pragma once +#include + #include #include @@ -75,13 +77,33 @@ struct fused_1nn_data_tag { template using fused_1nn_data_tag_t = typename fused_1nn_data_tag::type; -template +template +struct fused_1nn_index_tag; + +template <> +struct fused_1nn_index_tag { + using type = cuvs::neighbors::detail::tag_index_i32; +}; + +template <> +struct fused_1nn_index_tag { + using type = cuvs::neighbors::detail::tag_index_i64; +}; + +template +using fused_1nn_index_tag_t = typename fused_1nn_index_tag::type; + +template struct fragment_tag_fused_1nn_cubin { static constexpr int cc_major = ArchTag::cc_major; static constexpr int cc_minor = ArchTag::cc_minor; }; -template +template struct fragment_tag_fused_1nn_tileir {}; } // namespace cuvs::distance::detail diff --git a/cpp/src/cluster/detail/kmeans.cuh b/cpp/src/cluster/detail/kmeans.cuh index 635e8813bd..757108312f 100644 --- a/cpp/src/cluster/detail/kmeans.cuh +++ b/cpp/src/cluster/detail/kmeans.cuh @@ -682,8 +682,8 @@ void kmeans_fit( DataT* cur_centroids_ptr = cur_centroids_buf.data(); DataT* new_centroids_ptr = new_centroids_buf.data(); - auto minClusterAndDistance = raft::make_device_vector, IndexT>( - handle, streaming_batch_size); + auto nearest_idx = raft::make_device_vector(handle, streaming_batch_size); + auto nearest_dist = raft::make_device_vector(handle, streaming_batch_size); auto L2NormBatch = raft::make_device_vector(handle, streaming_batch_size); auto batch_weights_buf = raft::make_device_vector(handle, streaming_batch_size); rmm::device_uvector L2NormBuf_OR_DistBuf(0, stream); @@ -853,8 +853,10 @@ void kmeans_fit( auto batch_weights_view = cur_batch_weights(static_cast(data_batch.offset()), wt_data, cur_batch_size); - auto minCAD_view = raft::make_device_vector_view, IndexT>( - minClusterAndDistance.data_handle(), cur_batch_size); + auto nearest_idx_view = + raft::make_device_vector_view(nearest_idx.data_handle(), cur_batch_size); + auto nearest_dist_view = + raft::make_device_vector_view(nearest_dist.data_handle(), cur_batch_size); if constexpr (!data_on_device) { if (need_compute_norms) { @@ -883,7 +885,8 @@ void kmeans_fit( metric, iter_params.batch_samples, iter_params.batch_centroids, - minCAD_view, + nearest_idx_view, + nearest_dist_view, l2_const_view, L2NormBuf_OR_DistBuf, ws, @@ -1071,8 +1074,7 @@ void kmeans_predict(raft::resources const& handle, raft::make_const_mdspan(weight.view())); } - auto minClusterAndDistance = - raft::make_device_vector, IndexT>(handle, n_samples); + auto nearest_dist = raft::make_device_vector(handle, n_samples); rmm::device_uvector L2NormBuf_OR_DistBuf(0, stream); // L2 norm of X: ||x||^2 @@ -1082,50 +1084,35 @@ void kmeans_predict(raft::resources const& handle, raft::linalg::norm(handle, X, L2NormX.view()); } - // computes minClusterAndDistance[0:n_samples) where minClusterAndDistance[i] - // is a pair where - // 'key' is index to a sample in 'centroids' (index of the nearest - // centroid) and 'value' is the distance between the sample 'X[i]' and the - // 'centroid[key]' auto l2normx_view = raft::make_device_vector_view(L2NormX.data_handle(), n_samples); - cuvs::cluster::kmeans::detail::minClusterAndDistanceCompute( - handle, - X, - centroids, - minClusterAndDistance.view(), - l2normx_view, - L2NormBuf_OR_DistBuf, - pams.metric, - pams.batch_samples, - pams.batch_centroids, - workspace); + cuvs::cluster::kmeans::detail::minClusterAndDistanceCompute(handle, + X, + centroids, + labels, + nearest_dist.view(), + l2normx_view, + L2NormBuf_OR_DistBuf, + pams.metric, + pams.batch_samples, + pams.batch_centroids, + workspace); - // calculate cluster cost phi_x(C) rmm::device_scalar clusterCostD(stream); - raft::linalg::map( - handle, - minClusterAndDistance.view(), - [=] __device__(const raft::KeyValuePair kvp, DataT wt) { - raft::KeyValuePair res; - res.value = kvp.value * wt; - res.key = kvp.key; - return res; - }, - raft::make_const_mdspan(minClusterAndDistance.view()), - raft::make_const_mdspan(weight.view())); + raft::linalg::map(handle, + nearest_dist.view(), + raft::mul_op{}, + raft::make_const_mdspan(nearest_dist.view()), + raft::make_const_mdspan(weight.view())); cuvs::cluster::kmeans::detail::computeClusterCost( handle, - minClusterAndDistance.view(), + nearest_dist.view(), workspace, raft::make_device_scalar_view(clusterCostD.data()), - raft::value_op{}, + raft::identity_op{}, raft::add_op{}); - raft::linalg::map( - handle, labels, raft::key_op{}, raft::make_const_mdspan(minClusterAndDistance.view())); - inertia[0] = clusterCostD.value(stream); } diff --git a/cpp/src/cluster/detail/kmeans_balanced.cuh b/cpp/src/cluster/detail/kmeans_balanced.cuh index 86e254f473..007c462247 100644 --- a/cpp/src/cluster/detail/kmeans_balanced.cuh +++ b/cpp/src/cluster/detail/kmeans_balanced.cuh @@ -98,58 +98,90 @@ inline std::enable_if_t> predict_core( raft::make_device_matrix_view(centers, n_clusters, dim); auto X_norm_view = raft::make_device_vector_view(dataset_norm, n_rows); - auto minClusterAndDistance = raft::make_device_mdarray, IdxT>( - handle, mr, raft::make_extents(n_rows)); - - cuvs::cluster::kmeans::detail::minClusterAndDistanceCompute( - handle, - X_view, - centroids_view, - minClusterAndDistance.view(), - X_norm_view, - L2NormBuf_OR_DistBuf, - params.metric, - 0, // batch_samples (unused for fused reduction) - 0, // batch_centroids (unused for fused reduction) - workspace); - - // Copy keys to output labels - raft::linalg::map(handle, - raft::make_const_mdspan(minClusterAndDistance.view()), - raft::make_device_vector_view(labels, n_rows), - raft::compose_op, raft::key_op>()); - break; - } - case cuvs::distance::DistanceType::InnerProduct: { - if (use_cutile_fused_nn(handle, n_rows, n_clusters, dim)) { - rmm::device_uvector L2NormBuf_OR_DistBuf(0, stream, mr); - rmm::device_uvector workspace(0, stream, mr); - - auto X_view = raft::make_device_matrix_view(dataset, n_rows, dim); - auto centroids_view = - raft::make_device_matrix_view(centers, n_clusters, dim); - auto X_norm_view = raft::make_device_vector_view(dataset_norm, n_rows); - - auto minClusterAndDistance = - raft::make_device_mdarray, IdxT>( - handle, mr, raft::make_extents(n_rows)); + auto nearest_dist = + raft::make_device_mdarray(handle, mr, raft::make_extents(n_rows)); + if constexpr (std::is_same_v) { + auto labels_view = raft::make_device_vector_view(labels, n_rows); cuvs::cluster::kmeans::detail::minClusterAndDistanceCompute( handle, X_view, centroids_view, - minClusterAndDistance.view(), + labels_view, + nearest_dist.view(), + X_norm_view, + L2NormBuf_OR_DistBuf, + params.metric, + 0, // batch_samples (unused for fused reduction) + 0, // batch_centroids (unused for fused reduction) + workspace); + } else { + auto nearest_idx = + raft::make_device_mdarray(handle, mr, raft::make_extents(n_rows)); + cuvs::cluster::kmeans::detail::minClusterAndDistanceCompute( + handle, + X_view, + centroids_view, + nearest_idx.view(), + nearest_dist.view(), X_norm_view, L2NormBuf_OR_DistBuf, params.metric, 0, 0, workspace); + raft::copy( + handle, raft::make_device_vector_view(labels, n_rows), nearest_idx.view()); + } + break; + } + case cuvs::distance::DistanceType::InnerProduct: { + if (uses_fused_distance_nn( + use_fused(handle, n_rows, n_clusters, dim, params.metric))) { + rmm::device_uvector L2NormBuf_OR_DistBuf(0, stream, mr); + rmm::device_uvector workspace(0, stream, mr); - raft::linalg::map(handle, - raft::make_const_mdspan(minClusterAndDistance.view()), - raft::make_device_vector_view(labels, n_rows), - raft::compose_op, raft::key_op>()); + auto X_view = raft::make_device_matrix_view(dataset, n_rows, dim); + auto centroids_view = + raft::make_device_matrix_view(centers, n_clusters, dim); + auto X_norm_view = raft::make_device_vector_view(dataset_norm, n_rows); + + auto nearest_dist = + raft::make_device_mdarray(handle, mr, raft::make_extents(n_rows)); + + if constexpr (std::is_same_v) { + auto labels_view = raft::make_device_vector_view(labels, n_rows); + cuvs::cluster::kmeans::detail::minClusterAndDistanceCompute( + handle, + X_view, + centroids_view, + labels_view, + nearest_dist.view(), + X_norm_view, + L2NormBuf_OR_DistBuf, + params.metric, + 0, + 0, + workspace); + } else { + auto nearest_idx = + raft::make_device_mdarray(handle, mr, raft::make_extents(n_rows)); + cuvs::cluster::kmeans::detail::minClusterAndDistanceCompute( + handle, + X_view, + centroids_view, + nearest_idx.view(), + nearest_dist.view(), + X_norm_view, + L2NormBuf_OR_DistBuf, + params.metric, + 0, + 0, + workspace); + raft::copy(handle, + raft::make_device_vector_view(labels, n_rows), + nearest_idx.view()); + } } else { rmm::device_uvector distances(n_rows * n_clusters, stream, mr); @@ -216,22 +248,19 @@ auto calc_minibatch_size(const raft::resources& handle, size_t mem_per_row = 0; switch (metric) { case distance::DistanceType::L2Expanded: - case distance::DistanceType::L2SqrtExpanded: { - if (use_fused(handle, n_rows, n_clusters, dim)) { - // fusedL2NN needs a mutex and a key-value pair for each row. - mem_per_row += sizeof(int); - mem_per_row += sizeof(raft::KeyValuePair); - } else { - // unfused path needs a full GEMM output (distance matrix row). - mem_per_row += sizeof(MathT) * n_clusters; - } - } break; + case distance::DistanceType::L2SqrtExpanded: case distance::DistanceType::InnerProduct: { - if (use_cutile_fused_nn(handle, n_rows, n_clusters, dim)) { - mem_per_row += sizeof(int); - mem_per_row += sizeof(raft::KeyValuePair); - } else { - mem_per_row += sizeof(MathT) * n_clusters; + switch (use_fused(handle, n_rows, n_clusters, dim, metric)) { + case FusedDistancePath::FusedCutile: break; + case FusedDistancePath::FusedCutlass: + // fusedDistanceNNMinReduce CUTLASS fallback: mutex workspace + scratch KVP per row. + mem_per_row += sizeof(int); + mem_per_row += sizeof(raft::KeyValuePair); + break; + case FusedDistancePath::Unfused: + // unfused / GEMM+argmin path needs a full distance matrix row. + mem_per_row += sizeof(MathT) * n_clusters; + break; } } break; // Other metrics require storing a distance matrix. diff --git a/cpp/src/cluster/detail/kmeans_common.cuh b/cpp/src/cluster/detail/kmeans_common.cuh index 0606d77dec..f0d9bde801 100644 --- a/cpp/src/cluster/detail/kmeans_common.cuh +++ b/cpp/src/cluster/detail/kmeans_common.cuh @@ -62,20 +62,42 @@ template inline constexpr bool is_cutile_fused_data_type_v = std::is_same_v || std::is_same_v; +/** Which fused-distance implementation minCluster* will use (or Unfused). */ +enum class FusedDistancePath : std::uint8_t { + /** unfusedDistanceNNMinReduce or batched pairwise distance. */ + Unfused = 0, + /** fusedDistanceNNMinReduce via cuTile; no CUTLASS mutex / KVP scratch. */ + FusedCutile, + /** fusedDistanceNNMinReduce via legacy CUTLASS; needs mutex workspace + KVP scratch. */ + FusedCutlass, +}; + +inline constexpr bool uses_fused_distance_nn(FusedDistancePath path) +{ + return path != FusedDistancePath::Unfused; +} + +inline constexpr bool needs_cutlass_kvp_scratch(FusedDistancePath path) +{ + return path == FusedDistancePath::FusedCutlass; +} + +inline constexpr bool needs_fused_mutex_workspace(FusedDistancePath path) +{ + return path == FusedDistancePath::FusedCutlass; +} + /** - * @brief Returns true if the fused distance NN implementation should be used (CUTLASS and/or - * cuTile). + * @brief Selects the fused-distance assignment path for KMeans. * - * Float/half: use fused whenever cuTile can launch (any architecture and problem size). If cuTile - * is unavailable, fall back to legacy CUTLASS fused on Ampere and Hopper only. Double and other - * types never use cuTile; they keep the historical CUTLASS/unfused heuristics on pre-Blackwell + * Float/half: cuTile when the build and device support it. Otherwise L2/L2Sqrt/Cosine may use + * legacy CUTLASS fused on Ampere/Hopper (large enough problems). InnerProduct without cuTile uses + * Unfused. Double never uses cuTile; keeps historical CUTLASS/unfused heuristics on pre-Blackwell * GPUs. - * - * Callers route through fusedDistanceNNMinReduce when this returns true; cuTile dispatch inside - * that API is gated separately by dtype (see fusedDistanceNNImpl). */ template -bool use_fused(const raft::resources& handle, IdxT m, IdxT n, IdxT k) +FusedDistancePath use_fused( + const raft::resources& handle, IdxT m, IdxT n, IdxT k, cuvs::distance::DistanceType metric) { (void)k; cudaDeviceProp prop; @@ -83,24 +105,20 @@ bool use_fused(const raft::resources& handle, IdxT m, IdxT n, IdxT k) if constexpr (is_cutile_fused_data_type_v) { if constexpr (cuvs::detail::jit_lto::library_built_with_cutile()) { - if (cuvs::detail::jit_lto::cutile_launch_available_on_current_device()) { return true; } + if (cuvs::detail::jit_lto::cutile_launch_available_on_current_device()) { + return FusedDistancePath::FusedCutile; + } } - return prop.major <= 9; + if (metric == cuvs::distance::DistanceType::InnerProduct) { return FusedDistancePath::Unfused; } + if (prop.major <= 8) { return FusedDistancePath::FusedCutlass; } + if (prop.major == 9 && (m >= 4096 || n >= 4096)) { return FusedDistancePath::FusedCutlass; } + return FusedDistancePath::Unfused; } - if (prop.major >= 10) { return false; } - if (prop.major <= 8) { return true; } - if (prop.major == 9 && (m >= 4096 || n >= 4096)) { return true; } - return false; -} - -/** True when assignment should use the cuTile fused 1-NN kernel (float/half only). */ -template -bool use_cutile_fused_nn(const raft::resources& /*handle*/, IdxT /*m*/, IdxT /*n*/, IdxT /*k*/) -{ - if constexpr (!is_cutile_fused_data_type_v) { return false; } - if constexpr (!cuvs::detail::jit_lto::library_built_with_cutile()) { return false; } - return cuvs::detail::jit_lto::cutile_launch_available_on_current_device(); + if (prop.major >= 10) { return FusedDistancePath::Unfused; } + if (prop.major <= 8) { return FusedDistancePath::FusedCutlass; } + if (prop.major == 9 && (m >= 4096 || n >= 4096)) { return FusedDistancePath::FusedCutlass; } + return FusedDistancePath::Unfused; } template @@ -391,33 +409,32 @@ void shuffleAndGather(raft::resources const& handle, stream); } -// Calculates a pair for every sample in input 'X' where key is an -// index to an sample in 'centroids' (index of the nearest centroid) and 'value' -// is the distance between the sample and the 'centroid[key]' +// Calculates nearest centroid index and distance for every sample in input 'X'. template -void minClusterAndDistanceCompute( - raft::resources const& handle, - raft::device_matrix_view X, - raft::device_matrix_view centroids, - raft::device_vector_view, IndexT> minClusterAndDistance, - raft::device_vector_view L2NormX, - rmm::device_uvector& L2NormBuf_OR_DistBuf, - cuvs::distance::DistanceType metric, - int batch_samples, - int batch_centroids, - rmm::device_uvector& workspace); - -#define EXTERN_TEMPLATE_MIN_CLUSTER_AND_DISTANCE(DataT, IndexT) \ - extern template void minClusterAndDistanceCompute( \ - raft::resources const& handle, \ - raft::device_matrix_view X, \ - raft::device_matrix_view centroids, \ - raft::device_vector_view, IndexT> minClusterAndDistance, \ - raft::device_vector_view L2NormX, \ - rmm::device_uvector& L2NormBuf_OR_DistBuf, \ - cuvs::distance::DistanceType metric, \ - int batch_samples, \ - int batch_centroids, \ +void minClusterAndDistanceCompute(raft::resources const& handle, + raft::device_matrix_view X, + raft::device_matrix_view centroids, + raft::device_vector_view nearest_idx, + raft::device_vector_view nearest_dist, + raft::device_vector_view L2NormX, + rmm::device_uvector& L2NormBuf_OR_DistBuf, + cuvs::distance::DistanceType metric, + int batch_samples, + int batch_centroids, + rmm::device_uvector& workspace); + +#define EXTERN_TEMPLATE_MIN_CLUSTER_AND_DISTANCE(DataT, IndexT) \ + extern template void minClusterAndDistanceCompute( \ + raft::resources const& handle, \ + raft::device_matrix_view X, \ + raft::device_matrix_view centroids, \ + raft::device_vector_view nearest_idx, \ + raft::device_vector_view nearest_dist, \ + raft::device_vector_view L2NormX, \ + rmm::device_uvector& L2NormBuf_OR_DistBuf, \ + cuvs::distance::DistanceType metric, \ + int batch_samples, \ + int batch_centroids, \ rmm::device_uvector& workspace); EXTERN_TEMPLATE_MIN_CLUSTER_AND_DISTANCE(float, int64_t) @@ -476,22 +493,16 @@ void countSamplesInCluster(raft::resources const& handle, // stores (key, value) pair corresponding to each sample where // - key is the index of nearest cluster // - value is the distance to the nearest cluster - auto minClusterAndDistance = - raft::make_device_vector, IndexT>(handle, n_samples); - - // temporary buffer to store distance matrix, destructor releases the resource + auto nearest_idx = raft::make_device_vector(handle, n_samples); + auto nearest_dist = raft::make_device_vector(handle, n_samples); rmm::device_uvector L2NormBuf_OR_DistBuf(0, stream); - // computes minClusterAndDistance[0:n_samples) where minClusterAndDistance[i] - // is a pair where - // 'key' is index to an sample in 'centroids' (index of the nearest - // centroid) and 'value' is the distance between the sample 'X[i]' and the - // 'centroid[key]' cuvs::cluster::kmeans::detail::minClusterAndDistanceCompute( handle, X, (raft::device_matrix_view)centroids, - minClusterAndDistance.view(), + nearest_idx.view(), + nearest_dist.view(), L2NormX, L2NormBuf_OR_DistBuf, params.metric, @@ -499,12 +510,8 @@ void countSamplesInCluster(raft::resources const& handle, params.batch_centroids, workspace); - cuda::transform_iterator itr(minClusterAndDistance.data_handle(), - cuvs::cluster::kmeans::detail::KeyValueIndexOp{}); - - // count # of samples in each cluster countLabels(handle, - itr, + nearest_idx.data_handle(), sampleCountInCluster.data_handle(), (IndexT)n_samples, (IndexT)n_clusters, @@ -689,7 +696,8 @@ __device__ void check_convergence(raft::device_scalar_view clusteri * @param[in] batch_samples_param Batch-samples param forwarded to minClusterAndDistanceCompute * @param[in] batch_centroids_param Batch-centroids param forwarded to * minClusterAndDistanceCompute - * @param[inout] minClusterAndDistance Work buffer [batch_size] + * @param[inout] nearest_idx Nearest cluster index per sample [batch_size] + * @param[inout] nearest_dist Nearest distance per sample [batch_size] * @param[in] L2NormBatch Precomputed data norms [batch_size] * @param[inout] L2NormBuf_OR_DistBuf Resizable scratch * @param[inout] workspace Resizable scratch @@ -698,29 +706,30 @@ __device__ void check_convergence(raft::device_scalar_view clusteri * @param[inout] clustering_cost Running cost scalar (device) (added into) */ template -void process_batch( - raft::resources const& handle, - raft::device_matrix_view batch_data, - raft::device_vector_view batch_weights, - raft::device_matrix_view centroids, - cuvs::distance::DistanceType metric, - int batch_samples_param, - int batch_centroids_param, - raft::device_vector_view, IndexT> minClusterAndDistance, - raft::device_vector_view L2NormBatch, - rmm::device_uvector& L2NormBuf_OR_DistBuf, - rmm::device_uvector& workspace, - raft::device_matrix_view centroid_sums, - raft::device_vector_view weight_per_cluster, - raft::device_scalar_view clustering_cost, - rmm::device_uvector& batch_workspace) +void process_batch(raft::resources const& handle, + raft::device_matrix_view batch_data, + raft::device_vector_view batch_weights, + raft::device_matrix_view centroids, + cuvs::distance::DistanceType metric, + int batch_samples_param, + int batch_centroids_param, + raft::device_vector_view nearest_idx, + raft::device_vector_view nearest_dist, + raft::device_vector_view L2NormBatch, + rmm::device_uvector& L2NormBuf_OR_DistBuf, + rmm::device_uvector& workspace, + raft::device_matrix_view centroid_sums, + raft::device_vector_view weight_per_cluster, + raft::device_scalar_view clustering_cost, + rmm::device_uvector& batch_workspace) { cudaStream_t stream = raft::resource::get_cuda_stream(handle); minClusterAndDistanceCompute(handle, batch_data, centroids, - minClusterAndDistance, + nearest_idx, + nearest_dist, L2NormBatch, L2NormBuf_OR_DistBuf, metric, @@ -728,36 +737,30 @@ void process_batch( batch_centroids_param, workspace); - KeyValueIndexOp conversion_op; - thrust::transform_iterator, - const raft::KeyValuePair*> - labels_itr(minClusterAndDistance.data_handle(), conversion_op); - compute_centroid_adjustments(handle, batch_data, batch_weights, - labels_itr, + nearest_idx.data_handle(), static_cast(centroid_sums.extent(0)), centroid_sums, weight_per_cluster, batch_workspace, /*reset_sums=*/false); - raft::linalg::map( - handle, - minClusterAndDistance, - [=] __device__(const raft::KeyValuePair kvp, DataT wt) { - raft::KeyValuePair res; - res.value = kvp.value * wt; - res.key = kvp.key; - return res; - }, - raft::make_const_mdspan(minClusterAndDistance), - batch_weights); + auto weighted_dist = raft::make_device_vector(handle, nearest_dist.extent(0)); + raft::linalg::map(handle, + weighted_dist.view(), + raft::mul_op{}, + raft::make_const_mdspan(nearest_dist), + raft::make_const_mdspan(batch_weights)); auto batch_cost = raft::make_device_scalar(handle, DataT{0}); - computeClusterCost( - handle, minClusterAndDistance, workspace, batch_cost.view(), raft::value_op{}, raft::add_op{}); + computeClusterCost(handle, + weighted_dist.view(), + workspace, + batch_cost.view(), + raft::identity_op{}, + raft::add_op{}); raft::linalg::add(clustering_cost.data_handle(), clustering_cost.data_handle(), batch_cost.data_handle(), diff --git a/cpp/src/cluster/detail/kmeans_mg.cuh b/cpp/src/cluster/detail/kmeans_mg.cuh index ce3ca5a1fe..955d51c2a9 100644 --- a/cpp/src/cluster/detail/kmeans_mg.cuh +++ b/cpp/src/cluster/detail/kmeans_mg.cuh @@ -539,11 +539,8 @@ void fit(const raft::resources& handle, THROW("unknown initialization method to select initial centers"); } - // stores (key, value) pair corresponding to each sample where - // - key is the index of nearest cluster - // - value is the distance to the nearest cluster - auto minClusterAndDistance = - raft::make_device_vector, IndexT>(handle, n_samples); + auto nearest_idx = raft::make_device_vector(handle, n_samples); + auto nearest_dist = raft::make_device_vector(handle, n_samples); // temporary buffer to store L2 norm of centroids or distance matrix, // destructor releases the resource @@ -577,15 +574,11 @@ void fit(const raft::resources& handle, auto const_centroids = raft::make_device_matrix_view( centroids.data_handle(), centroids.extent(0), centroids.extent(1)); - // computes minClusterAndDistance[0:n_samples) where - // minClusterAndDistance[i] is a pair where - // 'key' is index to an sample in 'centroids' (index of the nearest - // centroid) and 'value' is the distance between the sample 'X[i]' and the - // 'centroid[key]' cuvs::cluster::kmeans::min_cluster_and_distance(handle, X, const_centroids, - minClusterAndDistance.view(), + nearest_idx.view(), + nearest_dist.view(), L2NormX.view(), L2NormBuf_OR_DistBuf, params.metric, @@ -595,9 +588,7 @@ void fit(const raft::resources& handle, workspace.resize(n_samples, stream); - cuda::transform_iterator keys_itr( - minClusterAndDistance.data_handle(), - cuvs::cluster::kmeans::detail::KeyValueIndexOp{}); + const IndexT* keys_itr = nearest_idx.data_handle(); raft::linalg::reduce_rows_by_key((DataT*)X.data_handle(), X.extent(1), keys_itr, @@ -696,35 +687,24 @@ void fit(const raft::resources& handle, raft::make_device_vector_view(centroids.data_handle(), newCentroids.size()), raft::make_device_vector_view(newCentroids.data_handle(), newCentroids.size())); - bool done = false; - rmm::device_scalar> clusterCostD(stream); + bool done = false; + auto clusterCostD = raft::make_device_scalar(handle, DataT{0}); // calculate cluster cost phi_x(C) cuvs::cluster::kmeans::cluster_cost( handle, - minClusterAndDistance.view(), + nearest_dist.view(), workspace, - raft::make_device_scalar_view(clusterCostD.data()), - cuda::proclaim_return_type>( - [] __device__(const raft::KeyValuePair& a, - const raft::KeyValuePair& b) { - raft::KeyValuePair res; - res.key = 0; - res.value = a.value + b.value; - return res; - })); + clusterCostD.view(), + cuda::proclaim_return_type( + [] __device__(const DataT& a, const DataT& b) { return a + b; })); // Cluster cost phi_x(C) from all ranks - comm.allreduce(&(clusterCostD.data()->value), - &(clusterCostD.data()->value), - 1, - raft::comms::op_t::SUM, - stream); + comm.allreduce( + clusterCostD.data_handle(), clusterCostD.data_handle(), 1, raft::comms::op_t::SUM, stream); DataT curClusteringCost = 0; - raft::copy(handle, - raft::make_host_scalar_view(&curClusteringCost), - raft::make_device_scalar_view(&(clusterCostD.data()->value))); + raft::copy(handle, raft::make_host_scalar_view(&curClusteringCost), clusterCostD.view()); ASSERT(comm.sync_stream(stream) == raft::comms::status_t::SUCCESS, "An error occurred in the distributed operation. This can result " diff --git a/cpp/src/cluster/detail/kmeans_mg_batched.cuh b/cpp/src/cluster/detail/kmeans_mg_batched.cuh index 98fed41636..ccc89991a6 100644 --- a/cpp/src/cluster/detail/kmeans_mg_batched.cuh +++ b/cpp/src/cluster/detail/kmeans_mg_batched.cuh @@ -157,9 +157,9 @@ void mnmg_fit(const raft::resources& handle, auto sqrd_norm_error_dev = raft::make_device_scalar(dev_res, DataT{0}); IndexT alloc_batch_size = has_data ? streaming_batch_size : IndexT{1}; auto batch_weights = raft::make_device_vector(dev_res, alloc_batch_size); - auto minClusterAndDistance = - raft::make_device_vector, IndexT>(dev_res, alloc_batch_size); - auto L2NormBatch = raft::make_device_vector(dev_res, alloc_batch_size); + auto nearest_idx = raft::make_device_vector(dev_res, alloc_batch_size); + auto nearest_dist = raft::make_device_vector(dev_res, alloc_batch_size); + auto L2NormBatch = raft::make_device_vector(dev_res, alloc_batch_size); rmm::device_uvector L2NormBuf_OR_DistBuf(0, stream); rmm::device_uvector workspace(0, stream); rmm::device_uvector batch_workspace(0, stream); @@ -353,9 +353,10 @@ void mnmg_fit(const raft::resources& handle, auto L2NormBatch_const = raft::make_const_mdspan(L2NormBatch_view); - auto minClusterAndDistance_view = - raft::make_device_vector_view, IndexT>( - minClusterAndDistance.data_handle(), current_batch_size); + auto nearest_idx_view = raft::make_device_vector_view( + nearest_idx.data_handle(), current_batch_size); + auto nearest_dist_view = raft::make_device_vector_view( + nearest_dist.data_handle(), current_batch_size); cuvs::cluster::kmeans::detail::process_batch( dev_res, @@ -365,7 +366,8 @@ void mnmg_fit(const raft::resources& handle, metric, params.batch_samples, params.batch_centroids, - minClusterAndDistance_view, + nearest_idx_view, + nearest_dist_view, L2NormBatch_const, L2NormBuf_OR_DistBuf, workspace, diff --git a/cpp/src/cluster/detail/minClusterDistanceCompute.cu b/cpp/src/cluster/detail/minClusterDistanceCompute.cu index 65678faa08..d01f48fc0a 100644 --- a/cpp/src/cluster/detail/minClusterDistanceCompute.cu +++ b/cpp/src/cluster/detail/minClusterDistanceCompute.cu @@ -11,39 +11,66 @@ namespace cuvs::cluster::kmeans::detail { -// Calculates a pair for every sample in input 'X' where key is an -// index to an sample in 'centroids' (index of the nearest centroid) and 'value' -// is the distance between the sample and the 'centroids[key]'. +namespace { + +template +__global__ void unpack_kvp_to_soa(IndexT* nearest_idx, + DataT* nearest_dist, + const raft::KeyValuePair* kvp, + IndexT n) +{ + IndexT i = blockIdx.x * blockDim.x + threadIdx.x; + if (i < n) { + if (nearest_idx != nullptr) { nearest_idx[i] = kvp[i].key; } + if (nearest_dist != nullptr) { nearest_dist[i] = kvp[i].value; } + } +} + +template +void unpack_kvp(raft::resources const& handle, + raft::device_vector_view nearest_idx, + raft::device_vector_view nearest_dist, + raft::device_vector_view, IndexT> kvp) +{ + auto stream = raft::resource::get_cuda_stream(handle); + auto n = static_cast(kvp.extent(0)); + int blks = static_cast((n + 255) / 256); + unpack_kvp_to_soa<<>>( + nearest_idx.data_handle(), nearest_dist.data_handle(), kvp.data_handle(), n); + RAFT_CUDA_TRY(cudaGetLastError()); +} + +} // namespace + template -void minClusterAndDistanceCompute( - raft::resources const& handle, - raft::device_matrix_view X, - raft::device_matrix_view centroids, - raft::device_vector_view, IndexT> minClusterAndDistance, - raft::device_vector_view L2NormX, - rmm::device_uvector& L2NormBuf_OR_DistBuf, - cuvs::distance::DistanceType metric, - int batch_samples, - int batch_centroids, - rmm::device_uvector& workspace) +void minClusterAndDistanceCompute(raft::resources const& handle, + raft::device_matrix_view X, + raft::device_matrix_view centroids, + raft::device_vector_view nearest_idx, + raft::device_vector_view nearest_dist, + raft::device_vector_view L2NormX, + rmm::device_uvector& L2NormBuf_OR_DistBuf, + cuvs::distance::DistanceType metric, + int batch_samples, + int batch_centroids, + rmm::device_uvector& workspace) { cudaStream_t stream = raft::resource::get_cuda_stream(handle); auto n_samples = X.extent(0); auto n_features = X.extent(1); auto n_clusters = centroids.extent(0); - bool is_l2_cos_fused = metric == cuvs::distance::DistanceType::L2Expanded || + const bool is_l2_cos = metric == cuvs::distance::DistanceType::L2Expanded || metric == cuvs::distance::DistanceType::L2SqrtExpanded || metric == cuvs::distance::DistanceType::CosineExpanded; - const bool is_ip_cutile = - metric == cuvs::distance::DistanceType::InnerProduct && - use_cutile_fused_nn(handle, n_samples, n_clusters, n_features); + const FusedDistancePath fused_path = + use_fused(handle, n_samples, n_clusters, n_features, metric); - if (is_l2_cos_fused || is_ip_cutile) { + if (uses_fused_distance_nn(fused_path)) { L2NormBuf_OR_DistBuf.resize(n_clusters, stream); auto centroidsNorm = raft::make_device_vector_view(L2NormBuf_OR_DistBuf.data(), n_clusters); - if (is_l2_cos_fused) { + if (is_l2_cos) { if (metric == cuvs::distance::DistanceType::CosineExpanded) { raft::linalg::norm( handle, centroids, centroidsNorm, raft::sqrt_op{}); @@ -53,20 +80,61 @@ void minClusterAndDistanceCompute( } } - raft::KeyValuePair initial_value(0, std::numeric_limits::max()); - raft::matrix::fill(handle, minClusterAndDistance, initial_value); + auto centroidsNormConst = + raft::make_device_vector_view(L2NormBuf_OR_DistBuf.data(), n_clusters); + + raft::KeyValuePair* cutlass_kvp_scratch = nullptr; + rmm::device_uvector> temp_kvp(0, stream); + if (needs_cutlass_kvp_scratch(fused_path)) { + temp_kvp.resize(n_samples, stream); + cutlass_kvp_scratch = temp_kvp.data(); + workspace.resize(sizeof(int) * n_samples, stream); + } + + cuvs::distance::fusedDistanceNNMinReduce( + nearest_idx.data_handle(), + nearest_dist.data_handle(), + X.data_handle(), + centroids.data_handle(), + L2NormX.data_handle(), + centroidsNormConst.data_handle(), + n_samples, + n_clusters, + n_features, + needs_fused_mutex_workspace(fused_path) ? (void*)workspace.data() : nullptr, + metric != cuvs::distance::DistanceType::L2Expanded, + true, + true, + metric, + 0.0f, + cutlass_kvp_scratch, + stream); + } else if (is_l2_cos) { + L2NormBuf_OR_DistBuf.resize(n_clusters, stream); + auto centroidsNorm = + raft::make_device_vector_view(L2NormBuf_OR_DistBuf.data(), n_clusters); - const bool should_use_fused = - use_fused(handle, n_samples, n_clusters, n_features); + if (metric == cuvs::distance::DistanceType::CosineExpanded) { + raft::linalg::norm( + handle, centroids, centroidsNorm, raft::sqrt_op{}); + } else { + raft::linalg::norm( + handle, centroids, centroidsNorm); + } auto centroidsNormConst = raft::make_device_vector_view(L2NormBuf_OR_DistBuf.data(), n_clusters); - if (should_use_fused) { - workspace.resize((sizeof(int)) * n_samples, stream); + workspace.resize(sizeof(DataT) * n_samples * n_clusters, stream); + auto temp_kvp = + raft::make_device_vector, IndexT>(handle, n_samples); + raft::KeyValuePair initial_value(0, std::numeric_limits::max()); + raft::matrix::fill(handle, temp_kvp.view(), initial_value); - cuvs::distance::fusedDistanceNNMinReduce, IndexT>( - minClusterAndDistance.data_handle(), + cuvs::distance:: + unfusedDistanceNNMinReduce, IndexT>( + handle, + temp_kvp.data_handle(), X.data_handle(), centroids.data_handle(), L2NormX.data_handle(), @@ -81,84 +149,44 @@ void minClusterAndDistanceCompute( metric, 0.0f, stream); - } else { - workspace.resize(sizeof(DataT) * n_samples * n_clusters, stream); - - cuvs::distance:: - unfusedDistanceNNMinReduce, IndexT>( - handle, - minClusterAndDistance.data_handle(), - X.data_handle(), - centroids.data_handle(), - L2NormX.data_handle(), - centroidsNormConst.data_handle(), - n_samples, - n_clusters, - n_features, - (void*)workspace.data(), - metric != cuvs::distance::DistanceType::L2Expanded, - false, - true, - metric, - 0.0f, - stream); - } + unpack_kvp(handle, nearest_idx, nearest_dist, raft::make_const_mdspan(temp_kvp.view())); } else { auto dataBatchSize = getDataBatchSize(batch_samples, n_samples); auto centroidsBatchSize = getCentroidsBatchSize(batch_centroids, n_clusters); - // TODO: Unless pool allocator is used, passing in a workspace for this - // isn't really increasing performance because this needs to do a re-allocation - // anyways. ref https://github.com/rapidsai/raft/issues/930 L2NormBuf_OR_DistBuf.resize(dataBatchSize * centroidsBatchSize, stream); - // pairwiseDistance[ns x nc] - tensor wrapper around the distance buffer auto pairwiseDistance = raft::make_device_matrix_view( L2NormBuf_OR_DistBuf.data(), dataBatchSize, centroidsBatchSize); + auto temp_kvp = + raft::make_device_vector, IndexT>(handle, n_samples); raft::KeyValuePair initial_value(0, std::numeric_limits::max()); - raft::matrix::fill(handle, minClusterAndDistance, initial_value); + raft::matrix::fill(handle, temp_kvp.view(), initial_value); - // tile over the input dataset for (IndexT dIdx = 0; dIdx < n_samples; dIdx += dataBatchSize) { - // # of samples for the current batch auto ns = std::min((IndexT)dataBatchSize, n_samples - dIdx); - // datasetView [ns x n_features] - view representing the current batch of - // input dataset auto datasetView = raft::make_device_matrix_view( X.data_handle() + (dIdx * n_features), ns, n_features); - // minClusterAndDistanceView [ns x n_clusters] - auto minClusterAndDistanceView = - raft::make_device_vector_view, IndexT>( - minClusterAndDistance.data_handle() + dIdx, ns); + auto temp_kvp_view = raft::make_device_vector_view, IndexT>( + temp_kvp.data_handle() + dIdx, ns); - // tile over the centroids for (IndexT cIdx = 0; cIdx < n_clusters; cIdx += centroidsBatchSize) { - // # of centroids for the current batch auto nc = std::min((IndexT)centroidsBatchSize, n_clusters - cIdx); - // centroidsView [nc x n_features] - view representing the current batch - // of centroids auto centroidsView = raft::make_device_matrix_view( centroids.data_handle() + (cIdx * n_features), nc, n_features); - // pairwiseDistanceView [ns x nc] - view representing the pairwise - // distance for current batch auto pairwiseDistanceView = raft::make_device_matrix_view(pairwiseDistance.data_handle(), ns, nc); - // calculate pairwise distance between current tile of cluster centroids - // and input dataset pairwise_distance_kmeans( handle, datasetView, centroidsView, pairwiseDistanceView, metric); - // argmin reduction returning pair - // calculates the closest centroid and the distance to the closest - // centroid raft::linalg::coalescedReduction( - minClusterAndDistanceView.data_handle(), + temp_kvp_view.data_handle(), pairwiseDistanceView.data_handle(), pairwiseDistanceView.extent(1), pairwiseDistanceView.extent(0), @@ -175,20 +203,23 @@ void minClusterAndDistanceCompute( raft::identity_op{}); } } + + unpack_kvp(handle, nearest_idx, nearest_dist, raft::make_const_mdspan(temp_kvp.view())); } } -#define INSTANTIATE_MIN_CLUSTER_AND_DISTANCE(DataT, IndexT) \ - template void minClusterAndDistanceCompute( \ - raft::resources const& handle, \ - raft::device_matrix_view X, \ - raft::device_matrix_view centroids, \ - raft::device_vector_view, IndexT> minClusterAndDistance, \ - raft::device_vector_view L2NormX, \ - rmm::device_uvector& L2NormBuf_OR_DistBuf, \ - cuvs::distance::DistanceType metric, \ - int batch_samples, \ - int batch_centroids, \ +#define INSTANTIATE_MIN_CLUSTER_AND_DISTANCE(DataT, IndexT) \ + template void minClusterAndDistanceCompute( \ + raft::resources const& handle, \ + raft::device_matrix_view X, \ + raft::device_matrix_view centroids, \ + raft::device_vector_view nearest_idx, \ + raft::device_vector_view nearest_dist, \ + raft::device_vector_view L2NormX, \ + rmm::device_uvector& L2NormBuf_OR_DistBuf, \ + cuvs::distance::DistanceType metric, \ + int batch_samples, \ + int batch_centroids, \ rmm::device_uvector& workspace); INSTANTIATE_MIN_CLUSTER_AND_DISTANCE(float, int64_t) @@ -215,13 +246,17 @@ void minClusterDistanceCompute(raft::resources const& handle, auto n_features = X.extent(1); auto n_clusters = centroids.extent(0); - bool is_fused = metric == cuvs::distance::DistanceType::L2Expanded || - metric == cuvs::distance::DistanceType::L2SqrtExpanded || - metric == cuvs::distance::DistanceType::CosineExpanded; + const bool is_l2_cos = metric == cuvs::distance::DistanceType::L2Expanded || + metric == cuvs::distance::DistanceType::L2SqrtExpanded || + metric == cuvs::distance::DistanceType::CosineExpanded; raft::matrix::fill(handle, minClusterDistance, std::numeric_limits::max()); - if (is_fused) { + const FusedDistancePath fused_path = + is_l2_cos ? use_fused(handle, n_samples, n_clusters, n_features, metric) + : FusedDistancePath::Unfused; + + if (uses_fused_distance_nn(fused_path)) { L2NormBuf_OR_DistBuf.resize(n_clusters, stream); auto centroidsNorm = raft::make_device_vector_view(L2NormBuf_OR_DistBuf.data(), n_clusters); @@ -241,9 +276,16 @@ void minClusterDistanceCompute(raft::resources const& handle, centroidsNorm); } - workspace.resize(sizeof(int) * n_samples, stream); + raft::KeyValuePair* cutlass_kvp_scratch = nullptr; + rmm::device_uvector> temp_kvp(0, stream); + if (needs_cutlass_kvp_scratch(fused_path)) { + temp_kvp.resize(n_samples, stream); + cutlass_kvp_scratch = temp_kvp.data(); + workspace.resize(sizeof(int) * n_samples, stream); + } - cuvs::distance::fusedDistanceNNMinReduce( + cuvs::distance::fusedDistanceNNMinReduce( + nullptr, minClusterDistance.data_handle(), X.data_handle(), centroids.data_handle(), @@ -252,12 +294,13 @@ void minClusterDistanceCompute(raft::resources const& handle, n_samples, n_clusters, n_features, - (void*)workspace.data(), + needs_fused_mutex_workspace(fused_path) ? (void*)workspace.data() : nullptr, metric != cuvs::distance::DistanceType::L2Expanded, - false, + true, true, metric, 0.0f, + cutlass_kvp_scratch, stream); } else { auto dataBatchSize = getDataBatchSize(batch_samples, n_samples); @@ -268,8 +311,6 @@ void minClusterDistanceCompute(raft::resources const& handle, auto pairwiseDistance = raft::make_device_matrix_view( L2NormBuf_OR_DistBuf.data(), dataBatchSize, centroidsBatchSize); - // tile over the input data and calculate distance matrix [n_samples x - // n_clusters] for (IndexT dIdx = 0; dIdx < n_samples; dIdx += dataBatchSize) { auto ns = std::min((IndexT)dataBatchSize, n_samples - dIdx); @@ -279,7 +320,6 @@ void minClusterDistanceCompute(raft::resources const& handle, auto minClusterDistanceView = raft::make_device_vector_view(minClusterDistance.data_handle() + dIdx, ns); - // tile over the centroids for (IndexT cIdx = 0; cIdx < n_clusters; cIdx += centroidsBatchSize) { auto nc = std::min((IndexT)centroidsBatchSize, n_clusters - cIdx); diff --git a/cpp/src/cluster/kmeans.cuh b/cpp/src/cluster/kmeans.cuh index 06da1fc1de..003604769b 100644 --- a/cpp/src/cluster/kmeans.cuh +++ b/cpp/src/cluster/kmeans.cuh @@ -486,22 +486,23 @@ void cluster_cost( * */ template -void min_cluster_and_distance( - raft::resources const& handle, - raft::device_matrix_view X, - raft::device_matrix_view centroids, - raft::device_vector_view, IndexT> minClusterAndDistance, - raft::device_vector_view L2NormX, - rmm::device_uvector& L2NormBuf_OR_DistBuf, - cuvs::distance::DistanceType metric, - int batch_samples, - int batch_centroids, - rmm::device_uvector& workspace) +void min_cluster_and_distance(raft::resources const& handle, + raft::device_matrix_view X, + raft::device_matrix_view centroids, + raft::device_vector_view nearest_idx, + raft::device_vector_view nearest_dist, + raft::device_vector_view L2NormX, + rmm::device_uvector& L2NormBuf_OR_DistBuf, + cuvs::distance::DistanceType metric, + int batch_samples, + int batch_centroids, + rmm::device_uvector& workspace) { cuvs::cluster::kmeans::detail::minClusterAndDistanceCompute(handle, X, centroids, - minClusterAndDistance, + nearest_idx, + nearest_dist, L2NormX, L2NormBuf_OR_DistBuf, metric, diff --git a/cpp/src/distance/detail/fused_distance_nn.cuh b/cpp/src/distance/detail/fused_distance_nn.cuh index 476ab9c2be..a2ec5422dd 100644 --- a/cpp/src/distance/detail/fused_distance_nn.cuh +++ b/cpp/src/distance/detail/fused_distance_nn.cuh @@ -14,6 +14,7 @@ #include "fused_distance_nn/simt_kernel.cuh" #include "pairwise_distance_base.cuh" // PairwiseDistances #include +#include #include // raft::KeyValuePair #include // raft::identity_op #include // Policy @@ -28,13 +29,9 @@ namespace distance { namespace detail { -template -void fusedDistanceNNImpl(OutT* min, +template +void fusedDistanceNNImpl(IdxT* nearest_idx, + DataT* nearest_dist, const DataT* x, const DataT* y, const DataT* xn, @@ -50,49 +47,77 @@ void fusedDistanceNNImpl(OutT* min, bool isRowMajor, cuvs::distance::DistanceType metric, float metric_arg, + raft::KeyValuePair* cutlass_kvp_scratch, cudaStream_t stream) { - // The kernel policy is determined by fusedDistanceNN. typedef Policy P; + typedef raft::KeyValuePair KVP; + constexpr auto maxVal = std::numeric_limits::max(); - // Callers (e.g. use_fused) enable this API for CUTLASS fused as well as cuTile; only try cuTile - // for float/half KVP output so double and other types never instantiate cuTile symbols here. if constexpr (is_fused_1nn_cutile_data_v) { - if constexpr (cuvs::detail::jit_lto::library_built_with_cutile() && - is_fused_1nn_kvp_output_v) { - if (try_fused_1nn_tile(min, x, y, xn, yn, m, n, k, metric, sqrt, stream)) { + if constexpr (cuvs::detail::jit_lto::library_built_with_cutile()) { + if (try_fused_1nn_tile( + nearest_idx, nearest_dist, x, y, xn, yn, m, n, k, metric, sqrt, stream)) { return; } } } - dim3 blk(P::Nthreads); - auto nblks = raft::ceildiv(m, P::Nthreads); - constexpr auto maxVal = std::numeric_limits::max(); - typedef raft::KeyValuePair KVPair; + RAFT_EXPECTS(cutlass_kvp_scratch != nullptr, "CUTLASS fused 1-NN requires a scratch KVP buffer"); - RAFT_CUDA_TRY(cudaMemsetAsync(workspace, 0, sizeof(int) * m, stream)); if (initOutBuffer) { - initKernel - <<>>(min, m, maxVal, redOp); - RAFT_CUDA_TRY(cudaGetLastError()); + initFused1nnOutput(nearest_idx, nearest_dist, m, std::numeric_limits::max(), stream); } + MinAndDistanceReduceOpImpl cutlass_redOp; + cutlass_redOp.out_kvp = cutlass_kvp_scratch; + initialize( + cutlass_kvp_scratch, m, maxVal, cutlass_redOp, stream); + + RAFT_CUDA_TRY(cudaMemsetAsync(workspace, 0, sizeof(int) * m, stream)); + switch (metric) { case cuvs::distance::DistanceType::CosineExpanded: - fusedCosineNN( - min, x, y, xn, yn, m, n, k, workspace, redOp, pairRedOp, sqrt, stream); + fusedCosineNN(nearest_idx, + nearest_dist, + x, + y, + xn, + yn, + m, + n, + k, + workspace, + cutlass_redOp, + pairRedOp, + sqrt, + cutlass_kvp_scratch, + stream); break; case cuvs::distance::DistanceType::L2SqrtExpanded: case cuvs::distance::DistanceType::L2Expanded: - fusedL2NNImpl( - min, x, y, xn, yn, m, n, k, workspace, redOp, pairRedOp, sqrt, false, stream); - break; - case cuvs::distance::DistanceType::InnerProduct: - // cuTile is the only fused InnerProduct implementation; callers must gate on availability. + fusedL2NNImpl(nearest_idx, + nearest_dist, + x, + y, + xn, + yn, + m, + n, + k, + workspace, + cutlass_redOp, + pairRedOp, + sqrt, + false, + cutlass_kvp_scratch, + stream); break; + case cuvs::distance::DistanceType::InnerProduct: break; default: assert("only cosine/l2 metric is supported with fusedDistanceNN\n"); break; } + + unpackFused1nnKvpToSoa(nearest_idx, nearest_dist, cutlass_kvp_scratch, m, stream); } } // namespace detail diff --git a/cpp/src/distance/detail/fused_distance_nn/cutile/export_fused_1nn.py b/cpp/src/distance/detail/fused_distance_nn/cutile/export_fused_1nn.py index fcefe7a027..1211456c9e 100644 --- a/cpp/src/distance/detail/fused_distance_nn/cutile/export_fused_1nn.py +++ b/cpp/src/distance/detail/fused_distance_nn/cutile/export_fused_1nn.py @@ -19,7 +19,14 @@ export_kernel, ) -from fused_1nn_kernel import METRICS, kernel_symbol, make_kernel, metric_abbrev +from fused_1nn_kernel import ( + INDEX_TYPES, + METRICS, + index_abbrev, + kernel_symbol, + make_kernel, + metric_abbrev, +) DEFAULT_TILEIR_BYTECODE_VERSION = "13.1" # cuTile requires a gpu_code even for TileIR bytecode export: it selects the compilation @@ -39,53 +46,82 @@ def _data_abbrev(data_type: str) -> str: return {"half": "h", "float": "f"}[data_type] -def _relaxed_matrix_constraint(elem_dtype): - """Array constraints matching the relaxed TMA-friendly layout from gemm_nn_cutile.""" +def _elem_stride_divisible_for_tma(elem_dtype) -> tuple[int, int]: + """Row stride (dim 0) divisible enough for 16-byte TMA access; last dim stride 1.""" + bytes_per_elem = 2 if elem_dtype == ct.float16 else 4 + return (16 // bytes_per_elem, 1) + + +def _cuvs_matrix_constraint(elem_dtype): + """Row-major device matrices for cuVS KMeans benchmarks. + + Assumes raft/cupy-style contiguous layout: stride[-1]==1, stride[0]==D, + 16-byte base alignment, and row pitch 16-byte aligned (float32 D%4==0, + float16 D%8==0). Applies to both points and centroids matrices. + + shape_divisible_by is (1, 1); tail tiles are masked in the kernel. + Odd D or general layouts need a separate relaxed export profile. + """ return ArrayConstraint( elem_dtype, ndim=2, - index_dtype=ct.int64, + index_dtype=ct.int32, stride_lower_bound_incl=(0, None), alias_groups=(), may_alias_internally=False, stride_constant=(None, 1), - stride_divisible_by=(8, 1), + stride_divisible_by=_elem_stride_divisible_for_tma(elem_dtype), shape_divisible_by=(1, 1), base_addr_divisible_by=16, ) -def _relaxed_vector_constraint(elem_dtype, *, tma_friendly: bool = False): - base_div = 16 if tma_friendly else 1 +def _cuvs_vector_constraint(elem_dtype): + """1-D device vectors: contiguous, 16-byte base. Length need not be divisible by 16.""" return ArrayConstraint( elem_dtype, ndim=1, - index_dtype=ct.int64, + index_dtype=ct.int32, stride_lower_bound_incl=(None,), alias_groups=(), may_alias_internally=False, stride_constant=(1,), stride_divisible_by=(1,), shape_divisible_by=(1,), - base_addr_divisible_by=base_div, + base_addr_divisible_by=16, ) +def _relaxed_matrix_constraint(elem_dtype): + """Deprecated alias; use _cuvs_matrix_constraint.""" + return _cuvs_matrix_constraint(elem_dtype) + + +def _relaxed_vector_constraint(elem_dtype, *, tma_friendly: bool = False): + """Deprecated alias; use _cuvs_vector_constraint.""" + del tma_friendly + return _cuvs_vector_constraint(elem_dtype) + + def _kernel_signature( data_type: str, metric: str, + index_type: str, tile_m: int, tile_n: int, tile_k: int, ) -> KernelSignature: elem = _dtype_for(data_type) - matrix = _relaxed_matrix_constraint(elem) - norm_array = _relaxed_vector_constraint(elem, tma_friendly=True) - idx_array = _relaxed_vector_constraint(ct.int64) - dist_array = _relaxed_vector_constraint(ct.float32) + matrix = _cuvs_matrix_constraint(elem) + norm_array = _cuvs_vector_constraint(elem) + idx_elem = ct.int32 if index_type == "int32" else ct.int64 + idx_array = _cuvs_vector_constraint(idx_elem) + dist_array = _cuvs_vector_constraint(elem) abbrev = _data_abbrev(data_type) - symbol = kernel_symbol(abbrev, metric_abbrev(metric)) + symbol = kernel_symbol( + abbrev, metric_abbrev(metric), index_abbrev(index_type) + ) return KernelSignature( parameters=[ @@ -98,6 +134,8 @@ def _kernel_signature( ScalarConstraint(ct.int64), ScalarConstraint(ct.int64), ScalarConstraint(ct.int64), + ScalarConstraint(ct.int64), + ScalarConstraint(ct.int64), ConstantConstraint(tile_m), ConstantConstraint(tile_n), ConstantConstraint(tile_k), @@ -112,14 +150,19 @@ def export_binary( output_format: Literal["cubin", "tileir_bytecode"], data_type: str, metric: str, + index_type: str, tile_m: int, tile_n: int, tile_k: int, gpu_code: str, bytecode_version: str | None = None, ) -> str: - kernel = make_kernel(data_type, metric, tile_m, tile_n, tile_k) - signature = _kernel_signature(data_type, metric, tile_m, tile_n, tile_k) + kernel = make_kernel( + data_type, metric, tile_m, tile_n, tile_k, index_type=index_type + ) + signature = _kernel_signature( + data_type, metric, index_type, tile_m, tile_n, tile_k + ) export_kwargs = { "kernel": kernel, @@ -148,6 +191,7 @@ def main() -> int: "--data-type", choices=("half", "float"), required=True ) parser.add_argument("--metric", choices=METRICS, required=True) + parser.add_argument("--index-type", choices=INDEX_TYPES, required=True) parser.add_argument("--tile-m", type=int, required=True) parser.add_argument("--tile-n", type=int, required=True) parser.add_argument("--tile-k", type=int, required=True) @@ -167,6 +211,7 @@ def main() -> int: output_format=args.format, data_type=args.data_type, metric=args.metric, + index_type=args.index_type, tile_m=args.tile_m, tile_n=args.tile_n, tile_k=args.tile_k, diff --git a/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_cutile_matrix.json b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_cutile_matrix.json index 3aa9dffd8a..7d9b723b39 100644 --- a/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_cutile_matrix.json +++ b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_cutile_matrix.json @@ -24,18 +24,28 @@ "metric_abbrev": "cos" } ], + "_index": [ + { + "index_type": "int32", + "index_abbrev": "i32" + }, + { + "index_type": "int64", + "index_abbrev": "i64" + } + ], "_tile": [ { - "tile_m": 128, - "tile_n": 128, - "tile_k": 64 + "tile_m": 256, + "tile_n": 64, + "tile_k": 32 } ], "_export": [ { "output_format": "cubin", "artifact_ext": "cubin", - "artifact_basename": "@data_type@_@metric_abbrev@_@gpu_code@", + "artifact_basename": "@data_type@_@metric_abbrev@_@index_abbrev@_@gpu_code@", "register": "cubin", "gpu_code": "sm_80", "cc_major": 8, @@ -45,7 +55,7 @@ { "output_format": "cubin", "artifact_ext": "cubin", - "artifact_basename": "@data_type@_@metric_abbrev@_@gpu_code@", + "artifact_basename": "@data_type@_@metric_abbrev@_@index_abbrev@_@gpu_code@", "register": "cubin", "gpu_code": "sm_86", "cc_major": 8, @@ -55,7 +65,7 @@ { "output_format": "cubin", "artifact_ext": "cubin", - "artifact_basename": "@data_type@_@metric_abbrev@_@gpu_code@", + "artifact_basename": "@data_type@_@metric_abbrev@_@index_abbrev@_@gpu_code@", "register": "cubin", "gpu_code": "sm_90", "cc_major": 9, @@ -65,7 +75,7 @@ { "output_format": "cubin", "artifact_ext": "cubin", - "artifact_basename": "@data_type@_@metric_abbrev@_@gpu_code@", + "artifact_basename": "@data_type@_@metric_abbrev@_@index_abbrev@_@gpu_code@", "register": "cubin", "gpu_code": "sm_120", "cc_major": 12, @@ -75,7 +85,7 @@ { "output_format": "tileir_bytecode", "artifact_ext": "tilebc", - "artifact_basename": "@data_type@_@metric_abbrev@", + "artifact_basename": "@data_type@_@metric_abbrev@_@index_abbrev@", "register": "tileir", "gpu_code": "sm_80", "bytecode_version": "13.1" diff --git a/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_kernel.py b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_kernel.py index 7d78525869..b2ff25555b 100644 --- a/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_kernel.py +++ b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_kernel.py @@ -9,11 +9,20 @@ ConstInt = ct.Constant[int] # Default tile geometry; overridden per export via make_kernel(..., tile_m, tile_n, tile_k). -DEFAULT_TILE_M = 128 -DEFAULT_TILE_N = 128 -DEFAULT_TILE_K = 64 +DEFAULT_TILE_M = 256 +DEFAULT_TILE_N = 64 +DEFAULT_TILE_K = 32 METRICS = ("inner_product", "l2_expanded", "cosine_expanded") +INDEX_TYPES = ("int32", "int64") + + +def _idx_dtype(index_type: str): + if index_type == "int32": + return ct.int32 + if index_type == "int64": + return ct.int64 + raise ValueError(f"Unsupported index_type {index_type!r}") def make_kernel( @@ -22,14 +31,20 @@ def make_kernel( tile_m: int = DEFAULT_TILE_M, tile_n: int = DEFAULT_TILE_N, tile_k: int = DEFAULT_TILE_K, + *, + index_type: str = "int32", ): - """Build a cuTile kernel with metric and tile sizes baked in at compile time.""" + """Build a cuTile kernel with metric, index width, and tile sizes baked in at compile time.""" if data_type not in ("half", "float"): raise ValueError(f"Unsupported data_type {data_type!r}") if metric not in METRICS: raise ValueError(f"Unsupported metric {metric!r}") + if index_type not in INDEX_TYPES: + raise ValueError(f"Unsupported index_type {index_type!r}") acc_dtype = ct.float32 + idx_dtype = _idx_dtype(index_type) + out_dist_dtype = ct.float16 if data_type == "half" else ct.float32 is_ip = metric == "inner_product" is_l2 = metric == "l2_expanded" is_cos = metric == "cosine_expanded" @@ -45,6 +60,8 @@ def fused_1nn_kernel( M, N, K, + apply_sqrt, + store_idx, tm: ConstInt, tn: ConstInt, tk: ConstInt, @@ -55,7 +72,7 @@ def fused_1nn_kernel( best_dist = ct.full((tm,), -3.4e38, acc_dtype) else: best_dist = ct.full((tm,), 3.4e38, acc_dtype) - best_idx = ct.zeros((tm,), ct.int64) + best_idx = ct.zeros((tm,), idx_dtype) num_tiles_k = ct.num_tiles(A, axis=1, shape=(tm, tk)) num_tiles_n = ct.num_tiles(B, axis=0, shape=(tn, tk)) @@ -65,12 +82,15 @@ def fused_1nn_kernel( accumulator = ct.full((tm, tn), 0, dtype=acc_dtype) for k in range(num_tiles_k): + dtype = ct.tfloat32 if A.dtype == ct.float32 else A.dtype + a = ct.load( A, index=(bidm, k), shape=(tm, tk), padding_mode=zero_pad - ) + ).astype(dtype) b_T = ct.load( B, index=(n, k), shape=(tn, tk), padding_mode=zero_pad - ) + ).astype(dtype) + accumulator = ct.mma(a, ct.transpose(b_T), accumulator) if is_ip: @@ -89,14 +109,13 @@ def fused_1nn_kernel( ) elif is_cos: # Cosine expanded distance: 1 - dot / (||x|| * ||y||); norms are L2 (not squared). - # No sqrt during the reduction — only arithmetic on stored distance if needed. denom = a_norm[:, None] * b_norm[None, :] score = 1.0 - (accumulator / denom) # Only the final N-tile can include zero-padded centroid columns. if n == num_tiles_n - 1: - col = ct.arange(tn, dtype=ct.int64) - global_col = n * tn + col + col = ct.arange(tn, dtype=idx_dtype) + global_col = (n * tn + col).astype(idx_dtype) valid = global_col < N if is_ip: score = ct.where(valid[None, :], score, -3.4e38) @@ -114,17 +133,25 @@ def fused_1nn_kernel( update = curr_best < best_dist best_dist = ct.where(update, curr_best, best_dist) - best_idx = ct.where(update, n * tn + curr_idx, best_idx) + best_idx = ct.where( + update, (n * tn + curr_idx).astype(idx_dtype), best_idx + ) - ct.store(OutIdx, index=(bidm,), tile=best_idx) - ct.store(OutDist, index=(bidm,), tile=best_dist) + out_dist = best_dist + if is_l2: + out_dist = ct.where(apply_sqrt != 0, ct.sqrt(best_dist), best_dist) + if store_idx != 0: + ct.store(OutIdx, index=(bidm,), tile=best_idx) + ct.store(OutDist, index=(bidm,), tile=out_dist.astype(out_dist_dtype)) return fused_1nn_kernel -def kernel_symbol(data_abbrev: str, metric_abbrev: str) -> str: +def kernel_symbol( + data_abbrev: str, metric_abbrev: str, index_abbrev: str +) -> str: """Must stay in sync with fused_1nn_kernel_entrypoint() in fused_1nn_planner.hpp.""" - return f"fused_1nn_{data_abbrev}_{metric_abbrev}" + return f"fused_1nn_{data_abbrev}_{metric_abbrev}_{index_abbrev}" def metric_abbrev(metric: str) -> str: @@ -133,3 +160,7 @@ def metric_abbrev(metric: str) -> str: "l2_expanded": "l2", "cosine_expanded": "cos", }[metric] + + +def index_abbrev(index_type: str) -> str: + return {"int32": "i32", "int64": "i64"}[index_type] diff --git a/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_planner.hpp b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_planner.hpp index c70ab3f87b..017fb72d48 100644 --- a/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_planner.hpp +++ b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_planner.hpp @@ -17,42 +17,48 @@ namespace cuvs::distance::detail { /** Must match kernel_symbol() in fused_1nn_kernel.py (export uses with_symbol). */ -template +template inline const char* fused_1nn_kernel_entrypoint() { + constexpr bool is_i32 = std::is_same_v; + constexpr bool is_i64 = std::is_same_v; + static_assert(is_i32 || is_i64, "unsupported fused 1-NN cuTile index width"); + if constexpr (std::is_same_v && std::is_same_v) { - return "fused_1nn_f_ip"; + return is_i32 ? "fused_1nn_f_ip_i32" : "fused_1nn_f_ip_i64"; } else if constexpr (std::is_same_v && std::is_same_v) { - return "fused_1nn_f_l2"; + return is_i32 ? "fused_1nn_f_l2_i32" : "fused_1nn_f_l2_i64"; } else if constexpr (std::is_same_v && std::is_same_v) { - return "fused_1nn_f_cos"; + return is_i32 ? "fused_1nn_f_cos_i32" : "fused_1nn_f_cos_i64"; } else if constexpr (std::is_same_v && std::is_same_v) { - return "fused_1nn_h_ip"; + return is_i32 ? "fused_1nn_h_ip_i32" : "fused_1nn_h_ip_i64"; } else if constexpr (std::is_same_v && std::is_same_v) { - return "fused_1nn_h_l2"; + return is_i32 ? "fused_1nn_h_l2_i32" : "fused_1nn_h_l2_i64"; } else if constexpr (std::is_same_v && std::is_same_v) { - return "fused_1nn_h_cos"; + return is_i32 ? "fused_1nn_h_cos_i32" : "fused_1nn_h_cos_i64"; } else { static_assert(sizeof(DataTag) == 0, "unsupported fused 1-NN cuTile data/metric combination"); return ""; } } -template +template struct Fused1nnTilePlanner : TileAlgorithmPlanner { using DataTag = fused_1nn_data_tag_t; using MetricTag = fused_1nn_metric_tag_t; + using IndexTag = fused_1nn_index_tag_t; inline static LauncherJitCache launcher_jit_cache{}; Fused1nnTilePlanner() - : TileAlgorithmPlanner(fused_1nn_kernel_entrypoint(), launcher_jit_cache) + : TileAlgorithmPlanner(fused_1nn_kernel_entrypoint(), + launcher_jit_cache) { } @@ -65,20 +71,32 @@ struct Fused1nnTilePlanner : TileAlgorithmPlanner { using cuvs::detail::jit_lto::cutile_arch_8_6; using cuvs::detail::jit_lto::cutile_arch_9_0; - this->add_static_fragment< - fragment_tag_fused_1nn_cubin>(); - this->add_static_fragment< - fragment_tag_fused_1nn_cubin>(); - this->add_static_fragment< - fragment_tag_fused_1nn_cubin>(); - this->add_static_fragment< - fragment_tag_fused_1nn_cubin>(); + this->add_static_fragment>(); + this->add_static_fragment>(); + this->add_static_fragment>(); + this->add_static_fragment>(); } void add_tileir_fallback() { this->add_static_tileir_fragment< - fragment_tag_fused_1nn_tileir>(); + fragment_tag_fused_1nn_tileir>(); } }; diff --git a/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_tile.cu b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_tile.cu index d292f0522b..cf01c12ce1 100644 --- a/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_tile.cu +++ b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_tile.cu @@ -8,6 +8,7 @@ #include "fused_1nn_planner.hpp" #include +#include #include namespace cuvs { @@ -16,25 +17,13 @@ namespace detail { namespace { -template -__global__ void pack_fused_1nn_kvp( - OutT* out, const int64_t* idx, const float* dist, IdxT len, bool apply_sqrt) -{ - IdxT i = blockIdx.x * blockDim.x + threadIdx.x; - if (i < len) { - out[i].key = static_cast(idx[i]); - float value = dist[i]; - if (apply_sqrt) { value = sqrtf(value); } - out[i].value = static_cast(value); - } -} - -template -bool launch_fused_1nn_tile(const DataT* x, +template +bool launch_fused_1nn_tile(IdxT* nearest_idx, + DataT* nearest_dist, + const DataT* x, const DataT* y, const DataT* xn, const DataT* yn, - OutT* out, IdxT m, IdxT n, IdxT k, @@ -43,7 +32,9 @@ bool launch_fused_1nn_tile(const DataT* x, { if constexpr (!std::is_same_v && !std::is_same_v) { return false; } - Fused1nnTilePlanner planner; + if (nearest_dist == nullptr) { return false; } + + Fused1nnTilePlanner planner; planner.add_entrypoint(); planner.add_tileir_fallback(); const CutileTileConfig tile_cfg = planner.tile_config(); @@ -52,11 +43,6 @@ bool launch_fused_1nn_tile(const DataT* x, const bool apply_sqrt = fused_1nn_apply_sqrt_at_pack(is_sqrt); - int64_t* d_idx = nullptr; - float* d_dist = nullptr; - RAFT_CUDA_TRY(cudaMallocAsync(&d_idx, m * sizeof(int64_t), stream)); - RAFT_CUDA_TRY(cudaMallocAsync(&d_dist, m * sizeof(float), stream)); - int64_t shape_x[2] = {m, k}; int64_t stride_x[2] = {k, 1}; int64_t shape_y[2] = {n, k}; @@ -72,19 +58,21 @@ bool launch_fused_1nn_tile(const DataT* x, int64_t M = m, N = n, K = k; - void* x_ptr = const_cast(x); - void* y_ptr = const_cast(y); - void* xn_ptr = const_cast(xn); - void* yn_ptr = const_cast(yn); - void* idx_ptr = d_idx; - void* dist_ptr = d_dist; + void* x_ptr = const_cast(x); + void* y_ptr = const_cast(y); + void* xn_ptr = const_cast(xn); + void* yn_ptr = const_cast(yn); + // OutIdx must be a valid device pointer for the launch ABI; when store_idx is 0 the kernel + // does not write it (dist-only callers pass nearest_dist as a stand-in). + const int64_t store_idx = nearest_idx != nullptr ? 1 : 0; + void* idx_ptr = + nearest_idx != nullptr ? static_cast(nearest_idx) : static_cast(nearest_dist); + void* dist_ptr = nearest_dist; const int64_t tile_m = tile_cfg.tile_m; dim3 grid((m + tile_m - 1) / tile_m, 1, 1); dim3 block(1, 1, 1); - // cutile_python_v1: 2D array (ptr, shape0, shape1, stride0, stride1); - // 1D array (ptr, shape, stride); tile sizes are embedded constants. using fused_1nn_cutile_kernel_t = void(void*, int64_t, int64_t, @@ -109,8 +97,9 @@ bool launch_fused_1nn_tile(const DataT* x, int64_t, int64_t, int64_t, + int64_t, + int64_t, int64_t); - std::cout << "Launching cuTile kernel" << std::endl; launcher->template dispatch(stream, grid, block, @@ -139,18 +128,16 @@ bool launch_fused_1nn_tile(const DataT* x, stride_dist, M, N, - K); - - pack_fused_1nn_kvp - <<<(m + 255) / 256, 256, 0, stream>>>(out, d_idx, d_dist, m, apply_sqrt); + K, + static_cast(apply_sqrt), + store_idx); RAFT_CUDA_TRY(cudaGetLastError()); - RAFT_CUDA_TRY(cudaFreeAsync(d_idx, stream)); - RAFT_CUDA_TRY(cudaFreeAsync(d_dist, stream)); return true; } -template -bool try_fused_1nn_tile_dispatch(OutT* min, +template +bool try_fused_1nn_tile_dispatch(IdxT* nearest_idx, + DataT* nearest_dist, const DataT* x, const DataT* y, const DataT* xn, @@ -164,26 +151,27 @@ bool try_fused_1nn_tile_dispatch(OutT* min, { switch (metric) { case cuvs::distance::DistanceType::InnerProduct: - return launch_fused_1nn_tile( - x, y, xn, yn, min, m, n, k, is_sqrt, stream); + return launch_fused_1nn_tile( + nearest_idx, nearest_dist, x, y, xn, yn, m, n, k, is_sqrt, stream); case cuvs::distance::DistanceType::L2Expanded: - return launch_fused_1nn_tile( - x, y, xn, yn, min, m, n, k, is_sqrt, stream); + return launch_fused_1nn_tile( + nearest_idx, nearest_dist, x, y, xn, yn, m, n, k, is_sqrt, stream); case cuvs::distance::DistanceType::L2SqrtExpanded: - return launch_fused_1nn_tile( - x, y, xn, yn, min, m, n, k, is_sqrt, stream); + return launch_fused_1nn_tile( + nearest_idx, nearest_dist, x, y, xn, yn, m, n, k, is_sqrt, stream); case cuvs::distance::DistanceType::CosineExpanded: - return launch_fused_1nn_tile( - x, y, xn, yn, min, m, n, k, is_sqrt, stream); + return launch_fused_1nn_tile( + nearest_idx, nearest_dist, x, y, xn, yn, m, n, k, is_sqrt, stream); default: return false; } } } // namespace -template - requires Fused1nnKvpOutput -bool try_fused_1nn_tile(OutT* min, +template + requires is_fused_1nn_cutile_data_v +bool try_fused_1nn_tile(IdxT* nearest_idx, + DataT* nearest_dist, const DataT* x, const DataT* y, const DataT* xn, @@ -196,35 +184,28 @@ bool try_fused_1nn_tile(OutT* min, cudaStream_t stream) { if (!cuvs::detail::jit_lto::cutile_launch_available_on_current_device()) { return false; } - return try_fused_1nn_tile_dispatch( - min, x, y, xn, yn, m, n, k, metric, is_sqrt, stream); + return try_fused_1nn_tile_dispatch( + nearest_idx, nearest_dist, x, y, xn, yn, m, n, k, metric, is_sqrt, stream); } -using kvp_i_f = raft::KeyValuePair; -using kvp_i64_f = raft::KeyValuePair; -using kvp_i_h = raft::KeyValuePair; -using kvp_i64_h = raft::KeyValuePair; - -#define CUVS_INST_TRY_FUSED_1NN_TILE(DataT, OutT, IdxT) \ - template CUVS_EXPORT bool try_fused_1nn_tile(OutT*, \ - const DataT*, \ - const DataT*, \ - const DataT*, \ - const DataT*, \ - IdxT, \ - IdxT, \ - IdxT, \ - cuvs::distance::DistanceType, \ - bool, \ - cudaStream_t) - -// int and int64_t are the same on LP64; one instantiation covers both. -CUVS_INST_TRY_FUSED_1NN_TILE(float, kvp_i_f, int); -CUVS_INST_TRY_FUSED_1NN_TILE(float, kvp_i64_f, int64_t); -CUVS_INST_TRY_FUSED_1NN_TILE(half, kvp_i_f, int); -CUVS_INST_TRY_FUSED_1NN_TILE(half, kvp_i64_f, int64_t); -CUVS_INST_TRY_FUSED_1NN_TILE(half, kvp_i_h, int); -CUVS_INST_TRY_FUSED_1NN_TILE(half, kvp_i64_h, int64_t); +#define CUVS_INST_TRY_FUSED_1NN_TILE(DataT, IdxT) \ + template CUVS_EXPORT bool try_fused_1nn_tile(IdxT*, \ + DataT*, \ + const DataT*, \ + const DataT*, \ + const DataT*, \ + const DataT*, \ + IdxT, \ + IdxT, \ + IdxT, \ + cuvs::distance::DistanceType, \ + bool, \ + cudaStream_t) + +CUVS_INST_TRY_FUSED_1NN_TILE(float, int); +CUVS_INST_TRY_FUSED_1NN_TILE(float, int64_t); +CUVS_INST_TRY_FUSED_1NN_TILE(half, int); +CUVS_INST_TRY_FUSED_1NN_TILE(half, int64_t); #undef CUVS_INST_TRY_FUSED_1NN_TILE diff --git a/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_tile.hpp b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_tile.hpp index 563d2583d8..4c0964631c 100644 --- a/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_tile.hpp +++ b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_tile.hpp @@ -5,11 +5,9 @@ #pragma once -#include #include #include -#include #include #include @@ -26,18 +24,11 @@ template inline constexpr bool is_fused_1nn_cutile_data_v = std::is_same_v || std::is_same_v; -template -inline constexpr bool is_fused_1nn_kvp_output_v = - is_fused_1nn_cutile_data_v && (std::is_same_v> || - std::is_same_v>); - -template -concept Fused1nnKvpOutput = is_fused_1nn_kvp_output_v; - #if CUVS_CUTILE_ENABLED -template - requires Fused1nnKvpOutput -bool try_fused_1nn_tile(OutT* min, +template + requires is_fused_1nn_cutile_data_v +bool try_fused_1nn_tile(IdxT* nearest_idx, + DataT* nearest_dist, const DataT* x, const DataT* y, const DataT* xn, @@ -49,9 +40,9 @@ bool try_fused_1nn_tile(OutT* min, bool is_sqrt, cudaStream_t stream); #else -template - requires Fused1nnKvpOutput -bool try_fused_1nn_tile(OutT*, +template +bool try_fused_1nn_tile(IdxT*, + DataT*, const DataT*, const DataT*, const DataT*, @@ -67,23 +58,6 @@ bool try_fused_1nn_tile(OutT*, } #endif -template - requires(!Fused1nnKvpOutput) -bool try_fused_1nn_tile(OutT*, - const DataT*, - const DataT*, - const DataT*, - const DataT*, - IdxT, - IdxT, - IdxT, - cuvs::distance::DistanceType, - bool, - cudaStream_t) -{ - return false; -} - } // namespace detail } // namespace distance } // namespace cuvs diff --git a/cpp/src/distance/detail/fused_distance_nn/fused_cosine_nn.cuh b/cpp/src/distance/detail/fused_distance_nn/fused_cosine_nn.cuh index 12f4f17cac..cc16d8a2e1 100644 --- a/cpp/src/distance/detail/fused_distance_nn/fused_cosine_nn.cuh +++ b/cpp/src/distance/detail/fused_distance_nn/fused_cosine_nn.cuh @@ -1,11 +1,11 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2024, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ #pragma once -#include "../distance_ops/cosine.cuh" // ops::l2_exp_distance_op +#include "../distance_ops/cosine.cuh" // ops::cosine_distance_op #include "../pairwise_distance_base.cuh" // PairwiseDistances #include "cutlass_base.cuh" #include "helper_structs.cuh" @@ -24,13 +24,9 @@ namespace distance { namespace detail { -template -void fusedCosineNN(OutT* min, +template +void fusedCosineNN(IdxT* nearest_idx, + DataT* nearest_dist, const DataT* x, const DataT* y, const DataT* xn, @@ -42,15 +38,20 @@ void fusedCosineNN(OutT* min, ReduceOpT redOp, KVPReduceOpT pairRedOp, bool sqrt, + raft::KeyValuePair* cutlass_out, cudaStream_t stream) { - // The kernel policy is determined by fusedL2NN. typedef Policy P; dim3 blk(P::Nthreads); constexpr auto maxVal = std::numeric_limits::max(); typedef raft::KeyValuePair KVPair; + if (cutlass_out == nullptr) { + initFused1nnOutput(nearest_idx, nearest_dist, m, maxVal, stream); + RAFT_CUDA_TRY(cudaGetLastError()); + } + namespace arch = raft::util::arch; using AccT = DataT; ops::cosine_distance_op distance_op{}; @@ -58,7 +59,7 @@ void fusedCosineNN(OutT* min, raft::identity_op fin_op{}; auto kernel = fusedDistanceNNkernel; - // Get pointer to fp32 SIMT kernel to determine the runtime architecture of the - // current system. Other methods to determine the architecture (that do not - // require a pointer) can be error prone. See: - // https://github.com/NVIDIA/cub/issues/545 void* kernel_ptr = reinterpret_cast(kernel); auto runtime_arch = arch::kernel_virtual_arch(kernel_ptr); auto cutlass_range = arch::SM_range(arch::SM_80(), arch::SM_future()); if (cutlass_range.contains(runtime_arch)) { - // If device is SM_80 or later, use CUTLASS-based kernel. using cosineOp = cuvs::distance::detail::ops::cosine_cutlass_op; - using kvp_cg_min_reduce_op_ = kvp_cg_min_reduce_op; + using kvp_cg_min_reduce_op_ = kvp_cg_min_reduce_op; kvp_cg_min_reduce_op_ cg_reduce_op; cosineOp cosine_dist_op; @@ -86,7 +82,7 @@ void fusedCosineNN(OutT* min, cutlassFusedDistanceNN(m, n, shmemSize, kernel); kernel<<>>( - min, x, y, xn, yn, m, n, k, maxVal, workspace, redOp, pairRedOp, distance_op, fin_op); + cutlass_out, x, y, xn, yn, m, n, k, maxVal, workspace, redOp, pairRedOp, distance_op, fin_op); RAFT_CUDA_TRY(cudaGetLastError()); } } diff --git a/cpp/src/distance/detail/fused_distance_nn/fused_l2_nn.cuh b/cpp/src/distance/detail/fused_distance_nn/fused_l2_nn.cuh index f1aad72110..8c532e2932 100644 --- a/cpp/src/distance/detail/fused_distance_nn/fused_l2_nn.cuh +++ b/cpp/src/distance/detail/fused_distance_nn/fused_l2_nn.cuh @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2021-2025, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2021-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ @@ -24,13 +24,9 @@ namespace distance { namespace detail { -template -void fusedL2NNImpl(OutT* min, +template +void fusedL2NNImpl(IdxT* nearest_idx, + DataT* nearest_dist, const DataT* x, const DataT* y, const DataT* xn, @@ -43,19 +39,17 @@ void fusedL2NNImpl(OutT* min, KVPReduceOpT pairRedOp, bool sqrt, bool initOutBuffer, + raft::KeyValuePair* cutlass_out, cudaStream_t stream) { - // The kernel policy is determined by fusedL2NN. typedef Policy P; dim3 blk(P::Nthreads); - auto nblks = raft::ceildiv(m, P::Nthreads); constexpr auto maxVal = std::numeric_limits::max(); typedef raft::KeyValuePair KVPair; - if (initOutBuffer) { - initKernel - <<>>(min, m, maxVal, redOp); + if (initOutBuffer && cutlass_out == nullptr) { + initFused1nnOutput(nearest_idx, nearest_dist, m, maxVal, stream); RAFT_CUDA_TRY(cudaGetLastError()); } @@ -66,7 +60,7 @@ void fusedL2NNImpl(OutT* min, raft::identity_op fin_op{}; auto kernel = fusedDistanceNNkernel; - // Get pointer to fp32 SIMT kernel to determine the best compute architecture - // out of all for which the kernel was compiled for that matches closely - // to the current device. Other methods to determine the architecture (that do not - // require a pointer) can be error prone. See: - // https://github.com/NVIDIA/cub/issues/545 void* kernel_ptr = reinterpret_cast(kernel); auto runtime_arch = arch::kernel_virtual_arch(kernel_ptr); auto cutlass_range = arch::SM_range(arch::SM_80(), arch::SM_future()); if (cutlass_range.contains(runtime_arch)) { - // If device is SM_80 or later, use CUTLASS-based kernel. using L2Op = cuvs::distance::detail::ops::l2_exp_cutlass_op; - using kvp_cg_min_reduce_op_ = kvp_cg_min_reduce_op; + using kvp_cg_min_reduce_op_ = kvp_cg_min_reduce_op; kvp_cg_min_reduce_op_ cg_reduce_op; L2Op L2_dist_op(sqrt); @@ -95,7 +83,7 @@ void fusedL2NNImpl(OutT* min, cutlassFusedDistanceNN(m, n, shmemSize, kernel); kernel<<>>( - min, x, y, xn, yn, m, n, k, maxVal, workspace, redOp, pairRedOp, distance_op, fin_op); + cutlass_out, x, y, xn, yn, m, n, k, maxVal, workspace, redOp, pairRedOp, distance_op, fin_op); RAFT_CUDA_TRY(cudaGetLastError()); } } diff --git a/cpp/src/distance/detail/fused_distance_nn/helper_structs.cuh b/cpp/src/distance/detail/fused_distance_nn/helper_structs.cuh index 762c720568..3bd78ba5ab 100644 --- a/cpp/src/distance/detail/fused_distance_nn/helper_structs.cuh +++ b/cpp/src/distance/detail/fused_distance_nn/helper_structs.cuh @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2021-2024, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2021-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ @@ -32,20 +32,43 @@ struct KVPMinReduceImpl { }; // KVPMinReduce +/** Writes fused 1-NN results to separate idx/dist arrays (dist may be null). */ template struct MinAndDistanceReduceOpImpl { typedef typename raft::KeyValuePair KVP; + LabelT* out_idx{nullptr}; + DataT* out_dist{nullptr}; + /** When set, CUTLASS/SIMT global merge writes here instead of SoA (caller unpacks). */ + KVP* out_kvp{nullptr}; + + DI void merge(LabelT rid, const KVP& other) const + { + if (out_kvp != nullptr) { + if (other.value < out_kvp[rid].value) { out_kvp[rid] = other; } + } else if (out_dist != nullptr) { + if (other.value < out_dist[rid]) { + out_dist[rid] = other.value; + if (out_idx != nullptr) { out_idx[rid] = other.key; } + } + } else if (out_idx != nullptr) { + // Idx-only output: dist must still be tracked for multi-tile merge; caller must provide + // out_dist or use a single-pass backend (cuTile). KMeans always passes both buffers. + out_idx[rid] = other.key; + } + } + DI void operator()(LabelT rid, KVP* out, const KVP& other) const { - if (other.value < out->value) { + if (out != nullptr && other.value < out->value) { out->key = other.key; out->value = other.value; } } + DI void operator()(LabelT rid, volatile KVP* out, const KVP& other) const { - if (other.value < out->value) { + if (out != nullptr && other.value < out->value) { out->key = other.key; out->value = other.value; } @@ -53,35 +76,41 @@ struct MinAndDistanceReduceOpImpl { DI void operator()(LabelT rid, DataT* out, const KVP& other) const { - if (other.value < *out) { *out = other.value; } + if (out != nullptr && other.value < *out) { *out = other.value; } } DI void operator()(LabelT rid, volatile DataT* out, const KVP& other) const { - if (other.value < *out) { *out = other.value; } + if (out != nullptr && other.value < *out) { *out = other.value; } } DI void operator()(LabelT rid, DataT* out, const DataT& other) const { - if (other < *out) { *out = other; } + if (out != nullptr && other < *out) { *out = other; } } DI void operator()(LabelT rid, volatile DataT* out, const DataT& other) const { - if (other < *out) { *out = other; } + if (out != nullptr && other < *out) { *out = other; } + } + + DI void init(DataT* out, DataT maxVal) const + { + if (out != nullptr) { *out = maxVal; } } - DI void init(DataT* out, DataT maxVal) const { *out = maxVal; } DI void init(KVP* out, DataT maxVal) const { out->value = maxVal; - out->key = 0xfffffff0; + out->key = LabelT(0); } - DI void init_key(DataT& out, LabelT idx) const { return; } + DI void init_key(DataT& /*out*/, LabelT /*idx*/) const {} + DI void init_key(KVP& out, LabelT idx) const { out.key = idx; } DI DataT get_value(KVP& out) const { return out.value; } + DI DataT get_value(DataT& out) const { return out; } }; @@ -96,6 +125,53 @@ struct MinReduceOpImpl { DI void init(DataT* out, DataT maxVal) { *out = maxVal; } }; +template +RAFT_KERNEL initFused1nnOutputKernel(IdxT* nearest_idx, DataT* nearest_dist, IdxT m, DataT maxVal) +{ + IdxT tid = IdxT(blockIdx.x) * blockDim.x + threadIdx.x; + if (tid < m) { + if (nearest_idx != nullptr) { nearest_idx[tid] = IdxT(0); } + if (nearest_dist != nullptr) { nearest_dist[tid] = maxVal; } + } +} + +template +void initFused1nnOutput( + IdxT* nearest_idx, DataT* nearest_dist, IdxT m, DataT maxVal, cudaStream_t stream) +{ + if (nearest_idx == nullptr && nearest_dist == nullptr) { return; } + auto blks = raft::ceildiv(m, 256); + initFused1nnOutputKernel + <<>>(nearest_idx, nearest_dist, m, maxVal); +} + +template +RAFT_KERNEL unpackFused1nnKvpToSoaKernel(IdxT* nearest_idx, + DataT* nearest_dist, + const raft::KeyValuePair* kvp, + IdxT n) +{ + IdxT i = IdxT(blockIdx.x) * blockDim.x + threadIdx.x; + if (i < n) { + if (nearest_idx != nullptr) { nearest_idx[i] = kvp[i].key; } + if (nearest_dist != nullptr) { nearest_dist[i] = kvp[i].value; } + } +} + +template +void unpackFused1nnKvpToSoa(IdxT* nearest_idx, + DataT* nearest_dist, + const raft::KeyValuePair* kvp, + IdxT m, + cudaStream_t stream) +{ + if (nearest_idx == nullptr && nearest_dist == nullptr) { return; } + auto blks = raft::ceildiv(m, 256); + unpackFused1nnKvpToSoaKernel + <<>>(nearest_idx, nearest_dist, kvp, m); + RAFT_CUDA_TRY(cudaGetLastError()); +} + template RAFT_KERNEL initKernel(OutT* min, IdxT m, DataT maxVal, ReduceOpT redOp) { @@ -106,15 +182,13 @@ RAFT_KERNEL initKernel(OutT* min, IdxT m, DataT maxVal, ReduceOpT redOp) template void initialize(OutT* min, IdxT m, DataT maxVal, ReduceOpT redOp, cudaStream_t stream) { - auto blks = raft::ceildiv(m, 256); - initKernel<<>>(min, m, maxVal, redOp); + auto blks = raft::ceildiv(m, 256); + initKernel<<>>(min, m, maxVal, redOp); } // cg::reduce functor for FusedDistanceNN used in its cutlass version // to output the min distance value & key(loc id). -// This is used in fused_distance_nn/predicated_tile_iterator_reduced_vec.h -// store_with_byte_offset() passed to cg::reduce() & select_reduce. -template +template struct kvp_cg_min_reduce_op { typedef typename raft::KeyValuePair KVP; @@ -122,7 +196,6 @@ struct kvp_cg_min_reduce_op { using AccTypeT = AccType; using IndexT = Index; - // functor signature. __host__ __device__ KVP operator()(KVP a, KVP b) const { return a.value < b.value ? a : b; } __host__ __device__ AccType operator()(AccType a, AccType b) const { return min(a, b); } diff --git a/cpp/src/distance/detail/fused_distance_nn/predicated_tile_iterator_reduced_vec.h b/cpp/src/distance/detail/fused_distance_nn/predicated_tile_iterator_reduced_vec.h index caa6a36d53..0d9f5333af 100644 --- a/cpp/src/distance/detail/fused_distance_nn/predicated_tile_iterator_reduced_vec.h +++ b/cpp/src/distance/detail/fused_distance_nn/predicated_tile_iterator_reduced_vec.h @@ -437,10 +437,8 @@ class PredicatedTileIteratorReducedVec { __syncthreads(); if (row < total_rows) { - volatile Element* gmem_ptr = reinterpret_cast(first_tile_byte_pointer_); - if ((block_start_row_first_tile_ + row) < extent_row_) { - user_params.red_op_(block_start_row_first_tile_ + row, (gmem_ptr + row), row_local_min); + user_params.red_op_.merge(block_start_row_first_tile_ + row, row_local_min); } } diff --git a/cpp/src/distance/fused_distance_nn-inl.cuh b/cpp/src/distance/fused_distance_nn-inl.cuh index 3fa80a9b60..13c4faa472 100644 --- a/cpp/src/distance/fused_distance_nn-inl.cuh +++ b/cpp/src/distance/fused_distance_nn-inl.cuh @@ -28,48 +28,10 @@ namespace distance { * \ingroup fused_l2_nn * @{ */ -/** - * @brief Fused L2 distance and 1-nearest-neighbor computation in a single call. - * - * The benefits of such a call are 2-fold: 1) eliminate the need for an - * intermediate buffer to store the output of gemm 2) reduce the memory read - * traffic on this intermediate buffer, otherwise needed during the reduction - * phase for 1-NN. - * - * @tparam DataT data type - * @tparam OutT output type to either store 1-NN indices and their minimum - * distances or store only the min distances. Accordingly, one - * has to pass an appropriate `ReduceOpT` - * @tparam IdxT indexing arithmetic type - * @tparam ReduceOpT A struct to perform the final needed reduction operation - * and also to initialize the output array elements with the - * appropriate initial value needed for reduction. - * @tparam KVPReduceOpT A struct providing functions for key-value pair comparison. - * - * @param[out] min will contain the reduced output (Length = `m`) - * (on device) - * @param[in] x first matrix. Row major. Dim = `m x k`. - * (on device). - * @param[in] y second matrix. Row major. Dim = `n x k`. - * (on device). - * @param[in] xn L2 squared norm of `x`. Length = `m`. (on device). - * @param[in] yn L2 squared norm of `y`. Length = `n`. (on device) - * @param[in] m gemm m - * @param[in] n gemm n - * @param[in] k gemm k - * @param[in] workspace temp workspace. Size = sizeof(int)*m. (on device) - * @param[in] redOp reduction operator in the epilogue - * @param[in] pairRedOp reduction operation on key value pairs - * @param[in] sqrt Whether the output `minDist` should contain L2-sqrt - * @param[in] initOutBuffer whether to initialize the output buffer before the - * main kernel launch - * @param[in] isRowMajor whether the input/output is row or column major. - * @param[in] metric Distance metric to be used (supports L2, cosine) - * @param[in] metric_arg power argument for distances like Minkowski (not supported for now) - * @param[in] stream cuda stream - */ -template -void fusedDistanceNN(OutT* min, + +template +void fusedDistanceNN(IdxT* nearest_idx, + DataT* nearest_dist, const DataT* x, const DataT* y, const DataT* xn, @@ -85,12 +47,10 @@ void fusedDistanceNN(OutT* min, bool isRowMajor, cuvs::distance::DistanceType metric, float metric_arg, + raft::KeyValuePair* cutlass_kvp_scratch, cudaStream_t stream) { ASSERT(isRowMajor, "fusedDistanceNN only supports row major inputs"); - // When k is smaller than 32, the Policy4x4 results in redundant calculations - // as it uses tiles that have k=32. Therefore, use a "skinny" policy instead - // that uses tiles with a smaller value of k. bool is_skinny = k < 32; size_t bytes = sizeof(DataT) * k; @@ -100,10 +60,10 @@ void fusedDistanceNN(OutT* min, if (is_skinny) { detail::fusedDistanceNNImpl< DataT, - OutT, IdxT, typename raft::linalg::Policy4x4Skinny::Policy, - ReduceOpT>(min, + ReduceOpT>(nearest_idx, + nearest_dist, x, y, xn, @@ -119,14 +79,15 @@ void fusedDistanceNN(OutT* min, isRowMajor, metric, metric_arg, + cutlass_kvp_scratch, stream); } else { detail::fusedDistanceNNImpl< DataT, - OutT, IdxT, typename raft::linalg::Policy4x4::Policy, - ReduceOpT>(min, + ReduceOpT>(nearest_idx, + nearest_dist, x, y, xn, @@ -142,16 +103,17 @@ void fusedDistanceNN(OutT* min, isRowMajor, metric, metric_arg, + cutlass_kvp_scratch, stream); } } else if (8 % sizeof(DataT) == 0 && bytes % 8 == 0 && px % 8 == 0 && py % 8 == 0) { if (is_skinny) { detail::fusedDistanceNNImpl< DataT, - OutT, IdxT, typename raft::linalg::Policy4x4Skinny::Policy, - ReduceOpT>(min, + ReduceOpT>(nearest_idx, + nearest_dist, x, y, xn, @@ -167,14 +129,15 @@ void fusedDistanceNN(OutT* min, isRowMajor, metric, metric_arg, + cutlass_kvp_scratch, stream); } else { detail::fusedDistanceNNImpl< DataT, - OutT, IdxT, typename raft::linalg::Policy4x4::Policy, - ReduceOpT>(min, + ReduceOpT>(nearest_idx, + nearest_dist, x, y, xn, @@ -190,15 +153,16 @@ void fusedDistanceNN(OutT* min, isRowMajor, metric, metric_arg, + cutlass_kvp_scratch, stream); } } else { if (is_skinny) { detail::fusedDistanceNNImpl::Policy, - ReduceOpT>(min, + ReduceOpT>(nearest_idx, + nearest_dist, x, y, xn, @@ -214,13 +178,14 @@ void fusedDistanceNN(OutT* min, isRowMajor, metric, metric_arg, + cutlass_kvp_scratch, stream); } else { detail::fusedDistanceNNImpl::Policy, - ReduceOpT>(min, + ReduceOpT>(nearest_idx, + nearest_dist, x, y, xn, @@ -236,44 +201,23 @@ void fusedDistanceNN(OutT* min, isRowMajor, metric, metric_arg, + cutlass_kvp_scratch, stream); } } } /** - * @brief Wrapper around fusedDistanceNN with minimum reduction operators. - * - * fusedDistanceNN cannot be compiled in the distance library due to the lambda - * operators, so this wrapper covers the most common case (minimum). + * @brief Fused GEMM + 1-NN minimum reduction. * - * @tparam DataT data type - * @tparam OutT output type to either store 1-NN indices and their minimum - * distances (e.g. raft::KeyValuePair) or store only the min - * distances. - * @tparam IdxT indexing arithmetic type - * @param[out] min will contain the reduced output (Length = `m`) - * (on device) - * @param[in] x first matrix. Row major. Dim = `m x k`. - * (on device). - * @param[in] y second matrix. Row major. Dim = `n x k`. - * (on device). - * @param[in] xn L2 squared norm of `x`. Length = `m`. (on device). - * @param[in] yn L2 squared norm of `y`. Length = `n`. (on device) - * @param[in] m gemm m - * @param[in] n gemm n - * @param[in] k gemm k - * @param[in] workspace temp workspace. Size = sizeof(int)*m. (on device) - * @param[in] sqrt Whether the output `minDist` should contain L2-sqrt - * @param[in] initOutBuffer whether to initialize the output buffer before the - * main kernel launch - * @param[in] isRowMajor whether the input/output is row or column major. - * @param[in] metric Distance metric to be used (supports L2, cosine) - * @param[in] metric_arg power argument for distances like Minkowski (not supported for now) - * @param[in] stream cuda stream + * @param[out] nearest_idx Nearest neighbor index per row, length `m` (required). + * @param[out] nearest_dist Minimum distance per row, length `m` (optional, may be null). + * @param[in] cutlass_kvp_scratch Temp KVP buffer, length `m`; required when CUTLASS/SIMT runs. + * Unused when cuTile handles the launch. */ -template -void fusedDistanceNNMinReduce(OutT* min, +template +void fusedDistanceNNMinReduce(IdxT* nearest_idx, + DataT* nearest_dist, const DataT* x, const DataT* y, const DataT* xn, @@ -287,28 +231,33 @@ void fusedDistanceNNMinReduce(OutT* min, bool isRowMajor, cuvs::distance::DistanceType metric, float metric_arg, + raft::KeyValuePair* cutlass_kvp_scratch, cudaStream_t stream) { MinAndDistanceReduceOp redOp; + redOp.out_idx = nearest_idx; + redOp.out_dist = nearest_dist; KVPMinReduce pairRedOp; - fusedDistanceNN(min, - x, - y, - xn, - yn, - m, - n, - k, - workspace, - redOp, - pairRedOp, - sqrt, - initOutBuffer, - isRowMajor, - metric, - metric_arg, - stream); + fusedDistanceNN(nearest_idx, + nearest_dist, + x, + y, + xn, + yn, + m, + n, + k, + workspace, + redOp, + pairRedOp, + sqrt, + initOutBuffer, + isRowMajor, + metric, + metric_arg, + cutlass_kvp_scratch, + stream); } /** @} */ diff --git a/cpp/tests/CMakeLists.txt b/cpp/tests/CMakeLists.txt index 9b96f94bf0..d4e0099035 100644 --- a/cpp/tests/CMakeLists.txt +++ b/cpp/tests/CMakeLists.txt @@ -120,7 +120,7 @@ ConfigureTest( ConfigureTest( NAME CLUSTER_TEST PATH cluster/kmeans.cu cluster/kmeans_balanced.cu cluster/kmeans_find_k.cu cluster/linkage.cu - cluster/connect_knn.cu cluster/spectral.cu + cluster/connect_knn.cu cluster/spectral.cu cluster/soa_unpack_trace.cu GPUS 1 PERCENT 100 ) diff --git a/cpp/tests/neighbors/distance_nn.cu b/cpp/tests/neighbors/distance_nn.cu index f5efaa5bec..6b17fc646b 100644 --- a/cpp/tests/neighbors/distance_nn.cu +++ b/cpp/tests/neighbors/distance_nn.cu @@ -42,7 +42,7 @@ __global__ void fill_int8(int8_t* buff, int len, int seed_offset) template class NNTest : public ::testing::TestWithParam> { public: - using OutT = raft::KeyValuePair; + using RefOutT = raft::KeyValuePair; NNTest() : params_{::testing::TestWithParam>::GetParam()}, m{params_.m}, @@ -55,8 +55,10 @@ class NNTest : public ::testing::TestWithParam> { y{raft::make_device_matrix(handle, n, k)}, x_norm{raft::make_device_vector(handle, m)}, y_norm{raft::make_device_vector(handle, n)}, - out{raft::make_device_vector(handle, m)}, - ref_out{raft::make_device_vector(handle, m)} + out_idx{raft::make_device_vector(handle, m)}, + out_dist{raft::make_device_vector(handle, m)}, + out_kvp{raft::make_device_vector(handle, m)}, + ref_out{raft::make_device_vector(handle, m)} { } @@ -92,15 +94,11 @@ class NNTest : public ::testing::TestWithParam> { workspace_size = m * n * sizeof(AccT); } - // Reset buffer - if constexpr (std::is_same_v>) { - // OutT is a RAFT KeyValuePair - raft::matrix::fill( - handle, raft::make_device_matrix_view(out.data_handle(), m, 1), OutT{0, 0}); - } else { - // OutT is a scalar type - raft::matrix::fill(handle, raft::make_device_matrix_view(out.data_handle(), m, 1), OutT{0}); - } + raft::matrix::fill(handle, raft::make_device_matrix_view(out_idx.data_handle(), m, 1), IdxT{0}); + raft::matrix::fill( + handle, raft::make_device_matrix_view(out_dist.data_handle(), m, 1), AccT{0}); + raft::matrix::fill( + handle, raft::make_device_matrix_view(ref_out.data_handle(), m, 1), RefOutT{0, 0}); raft::resource::sync_stream(handle, stream); } @@ -109,34 +107,36 @@ class NNTest : public ::testing::TestWithParam> { raft::device_vector workspace = raft::make_device_vector(handle, workspace_size); - ref_nn( + ref_nn( ref_out.data_handle(), x.data_handle(), y.data_handle(), m, n, k, sqrt, metric, stream); if constexpr (impl == ImplType::fused) { if constexpr (std::is_same_v) { - cuvs::distance::fusedDistanceNNMinReduce(out.data_handle(), - x.data_handle(), - y.data_handle(), - x_norm.data_handle(), - y_norm.data_handle(), - m, - n, - k, - (void*)workspace.data_handle(), - sqrt, - true, - true, - metric, - 0.0, - stream); + cuvs::distance::fusedDistanceNNMinReduce(out_idx.data_handle(), + out_dist.data_handle(), + x.data_handle(), + y.data_handle(), + x_norm.data_handle(), + y_norm.data_handle(), + m, + n, + k, + (void*)workspace.data_handle(), + sqrt, + true, + true, + metric, + 0.0, + out_kvp.data_handle(), + stream); } else { static_assert(sizeof(DataT) == 0, "fusedDistanceNNMinReduce is not implemented for datatype other than float"); } } else if constexpr (impl == ImplType::unfused) { - cuvs::distance::unfusedDistanceNNMinReduce( + cuvs::distance::unfusedDistanceNNMinReduce( handle, - out.data_handle(), + out_kvp.data_handle(), x.data_handle(), y.data_handle(), x_norm.data_handle(), @@ -156,7 +156,12 @@ class NNTest : public ::testing::TestWithParam> { void compare() { - vector_compare(handle, ref_out.data_handle(), out.data_handle(), m, summary); + if constexpr (impl == ImplType::fused) { + vector_compare_soa( + handle, ref_out.data_handle(), out_idx.data_handle(), out_dist.data_handle(), m, summary); + } else { + vector_compare(handle, ref_out.data_handle(), out_kvp.data_handle(), m, summary); + } ASSERT_TRUE(summary.max_diff < params_.tol) << summary; } @@ -174,8 +179,10 @@ class NNTest : public ::testing::TestWithParam> { raft::device_matrix y; raft::device_vector x_norm; raft::device_vector y_norm; - raft::device_vector out; - raft::device_vector ref_out; + raft::device_vector out_idx; + raft::device_vector out_dist; + raft::device_vector out_kvp; + raft::device_vector ref_out; size_t workspace_size; }; diff --git a/cpp/tests/neighbors/distance_nn_helper.cuh b/cpp/tests/neighbors/distance_nn_helper.cuh index ea440387b4..dfa71f71a4 100644 --- a/cpp/tests/neighbors/distance_nn_helper.cuh +++ b/cpp/tests/neighbors/distance_nn_helper.cuh @@ -209,6 +209,34 @@ class ComparisonSummary { } }; +template +void vector_compare_soa(raft::resources const& handle, + const raft::KeyValuePair* ref, + const IdxT* out_idx, + const AccT* out_dist, + IdxT n, + ComparisonSummary& summary) +{ + auto ref_h = raft::make_host_vector, IdxT>(n); + auto idx_h = raft::make_host_vector(n); + auto dist_h = raft::make_host_vector(n); + + raft::copy(ref_h.data_handle(), ref, n, raft::resource::get_cuda_stream(handle)); + raft::copy(idx_h.data_handle(), out_idx, n, raft::resource::get_cuda_stream(handle)); + raft::copy(dist_h.data_handle(), out_dist, n, raft::resource::get_cuda_stream(handle)); + raft::resource::sync_stream(handle, raft::resource::get_cuda_stream(handle)); + + summary.init(); + + for (IdxT i = 0; i < n; i++) { + const double a_val = double(dist_h(i)); + const double b_val = double(ref_h(i).value); + const bool missed = idx_h(i) != ref_h(i).key; + const double diff = std::abs(a_val - b_val); + summary.update(diff, i, a_val, b_val, missed); + } +} + template void vector_compare( raft::resources const& handle, const OutT* a, const OutT* b, IdxT n, ComparisonSummary& summary) From 0a3da06d706a05a676f288288a78e9d3070fcbf7 Mon Sep 17 00:00:00 2001 From: divyegala Date: Tue, 30 Jun 2026 04:38:36 +0000 Subject: [PATCH 10/82] add reproducible benchmark scripts --- benchmark_kmeans.py | 410 ++++++++++++++++++++++++++++++++++++++++ run_benchmark_kmeans.sh | 47 +++++ 2 files changed, 457 insertions(+) create mode 100644 benchmark_kmeans.py create mode 100755 run_benchmark_kmeans.sh diff --git a/benchmark_kmeans.py b/benchmark_kmeans.py new file mode 100644 index 0000000000..cfabe7570a --- /dev/null +++ b/benchmark_kmeans.py @@ -0,0 +1,410 @@ +#!/usr/bin/env python3 +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. +# SPDX-License-Identifier: Apache-2.0 +r"""KMeans fit+predict benchmark: baseline / cuTile / flash-kmeans. + +Single impl (activate the target conda env first): + python benchmark_kmeans.py --impl baseline|cutile|flash --n N --d D --k K \\ + --max-iter 5 --tol 1e-4 --seed 42 \\ + --warmup-fit 1 --iters-fit 3 --warmup-pred 1 --iters-pred 3 + +Compare (subprocess per impl; export env vars, then --compare): + export BENCH_CONDA=/path/to/miniforge3 + export BENCH_ENV_BASE=cuvs_2608_base + export BENCH_ENV_CUTILE=cuvs_2608 + export BENCH_ENV_FLASH=cuvs_2608_base + python benchmark_kmeans.py --compare --n 33554432 --d 32 --k 64 \\ + --max-iter 5 --tol 1e-4 --seed 42 \\ + --warmup-fit 1 --iters-fit 3 --warmup-pred 1 --iters-pred 3 + +Smoke test (small shape, single impl): + conda activate cuvs_2608 + python benchmark_kmeans.py --impl cutile --n 10000 --d 32 --k 8 \\ + --max-iter 2 --tol 1e-4 --seed 42 \\ + --warmup-fit 0 --iters-fit 1 --warmup-pred 0 --iters-pred 1 + +Required for --compare (no defaults): + BENCH_CONDA path to miniforge/conda root + BENCH_ENV_BASE conda env name for baseline libcuvs + BENCH_ENV_CUTILE conda env name for cuTile libcuvs + BENCH_ENV_FLASH conda env name for flash-kmeans +""" + +from __future__ import annotations + +import argparse +import os +import re +import subprocess +import sys +import time +from dataclasses import dataclass +from pathlib import Path + +ROOT = Path(__file__).resolve().parent +IMPLS = ("baseline", "cutile", "flash") + + +def _require_env(name: str) -> str: + val = os.environ.get(name) + if not val: + raise SystemExit(f"required environment variable {name} is not set") + return val + + +def _impl_config() -> dict[str, dict]: + conda = Path(_require_env("BENCH_CONDA")) + return { + "baseline": { + "bench_mode": "cuvs_base", + "conda": conda, + "conda_env": _require_env("BENCH_ENV_BASE"), + }, + "cutile": { + "bench_mode": "cuvs_cutile", + "conda": conda, + "conda_env": _require_env("BENCH_ENV_CUTILE"), + }, + "flash": { + "bench_mode": "flash", + "conda": conda, + "conda_env": _require_env("BENCH_ENV_FLASH"), + }, + } + + +@dataclass +class BenchResult: + impl: str + fit_median_ms: float | None = None + predict_median_ms: float | None = None + n_iter: int | None = None + inertia: float | None = None + error: str | None = None + + +def median(xs: list[float]) -> float: + import numpy as np + + return float(np.median(xs)) + + +def run_benchmark( + bench_mode: str, + n: int, + d: int, + k: int, + *, + max_iter: int, + tol: float, + seed: int, + warmup_fit: int, + iters_fit: int, + warmup_pred: int, + iters_pred: int, +) -> BenchResult: + import numpy as np + + rng = np.random.default_rng(seed) + init_centroids_host = rng.standard_normal((k, d), dtype=np.float32) + x_host = rng.standard_normal((n, d), dtype=np.float32) + input_gib = n * d * 4 / (1024**3) + + label = { + "cuvs_base": "baseline", + "cuvs_cutile": "cutile", + "flash": "flash", + }[bench_mode] + print( + f"=== N={n:,} D={d} K={k:,} iters={max_iter} input={input_gib:.2f} GiB ===", + flush=True, + ) + + if bench_mode in ("cuvs_base", "cuvs_cutile"): + from cuda.bindings import runtime as cudart + from pylibraft.common import device_ndarray + + from cuvs.cluster.kmeans import KMeansParams, fit, predict + + def sync(): + cudart.cudaDeviceSynchronize() + + x = device_ndarray(x_host) + params = KMeansParams( + n_clusters=k, + max_iter=max_iter, + tol=tol, + metric="sqeuclidean", + hierarchical=False, + init_method="Array", + n_init=1, + ) + + for _ in range(warmup_fit): + fit( + params, x, centroids=device_ndarray(init_centroids_host.copy()) + ) + sync() + + fit_times: list[float] = [] + n_iter = 0 + inertia = 0.0 + for _ in range(iters_fit): + t0 = time.perf_counter() + _, inertia, n_iter = fit( + params, x, centroids=device_ndarray(init_centroids_host.copy()) + ) + sync() + fit_times.append((time.perf_counter() - t0) * 1e3) + + centroids, _, _ = fit( + params, x, centroids=device_ndarray(init_centroids_host.copy()) + ) + sync() + + for _ in range(warmup_pred): + predict(params, x, centroids) + sync() + + pred_times: list[float] = [] + for _ in range(iters_pred): + t0 = time.perf_counter() + predict(params, x, centroids) + sync() + pred_times.append((time.perf_counter() - t0) * 1e3) + + print(f"impl={label} init=Array", flush=True) + print(f"fit_median_ms={median(fit_times):.2f}", flush=True) + print(f"predict_median_ms={median(pred_times):.2f}", flush=True) + print(f"n_iter={n_iter} inertia={inertia:.6g}", flush=True) + return BenchResult( + impl=label, + fit_median_ms=median(fit_times), + predict_median_ms=median(pred_times), + n_iter=n_iter, + inertia=inertia, + ) + + if bench_mode == "flash": + import torch + from flash_kmeans.assign_euclid_triton import euclid_assign_triton + from flash_kmeans.kmeans_triton_impl import batch_kmeans_Euclid + + def sync(): + torch.cuda.synchronize() + + x = torch.from_numpy(x_host).cuda() + init_c = ( + torch.from_numpy(init_centroids_host.copy()).cuda().unsqueeze(0) + ) + + def run_fit(init): + x_b = x.unsqueeze(0) + _, centroids_b, _ = batch_kmeans_Euclid( + x_b, + k, + max_iters=max_iter, + tol=tol, + init_centroids=init, + verbose=False, + ) + return centroids_b + + for _ in range(warmup_fit): + run_fit(init_c.clone()) + sync() + + fit_times = [] + for _ in range(iters_fit): + t0 = time.perf_counter() + run_fit(init_c.clone()) + sync() + fit_times.append((time.perf_counter() - t0) * 1e3) + + centroids_b = run_fit(init_c.clone()) + sync() + + x_b = x.unsqueeze(0) + x_sq = (x_b**2).sum(dim=-1) + + for _ in range(warmup_pred): + euclid_assign_triton(x_b, centroids_b, x_sq) + sync() + + pred_times = [] + for _ in range(iters_pred): + t0 = time.perf_counter() + euclid_assign_triton(x_b, centroids_b, x_sq) + sync() + pred_times.append((time.perf_counter() - t0) * 1e3) + + print("impl=flash-kmeans init=Array", flush=True) + print(f"fit_median_ms={median(fit_times):.2f}", flush=True) + print(f"predict_median_ms={median(pred_times):.2f}", flush=True) + return BenchResult( + impl="flash", + fit_median_ms=median(fit_times), + predict_median_ms=median(pred_times), + ) + + raise ValueError(f"unknown bench_mode={bench_mode!r}") + + +def _parse_output(text: str, impl: str) -> BenchResult: + fit_m = re.search(r"^fit_median_ms=([0-9.]+)", text, re.M) + pred_m = re.search(r"^predict_median_ms=([0-9.]+)", text, re.M) + if not fit_m or not pred_m: + return BenchResult(impl=impl, error=text.strip() or "no output") + n_iter_m = re.search(r"^n_iter=([0-9]+)", text, re.M) + inertia_m = re.search(r"^inertia=([0-9.eE+-]+)", text, re.M) + return BenchResult( + impl=impl, + fit_median_ms=float(fit_m.group(1)), + predict_median_ms=float(pred_m.group(1)), + n_iter=int(n_iter_m.group(1)) if n_iter_m else None, + inertia=float(inertia_m.group(1)) if inertia_m else None, + ) + + +def _run_subprocess( + impl: str, + n: int, + d: int, + k: int, + args: argparse.Namespace, +) -> BenchResult: + cfg = _impl_config()[impl] + conda = cfg["conda"] + env_exports = " ".join( + f'export {key}="{val}"' + for key, val in ( + ( + "CUDA_VISIBLE_DEVICES", + os.environ.get("CUDA_VISIBLE_DEVICES", ""), + ), + ("MAX_ITER", args.max_iter), + ("TOL", args.tol), + ("SEED", args.seed), + ("WARMUP_FIT", args.warmup_fit), + ("ITERS_FIT", args.iters_fit), + ("WARMUP_PRED", args.warmup_pred), + ("ITERS_PRED", args.iters_pred), + ) + if val != "" + ) + cmd = f""" +set -eo pipefail +source "{conda}/etc/profile.d/conda.sh" +conda activate "{cfg["conda_env"]}" +{env_exports} +python3 "{ROOT / "benchmark_kmeans.py"}" --impl {impl} --n {n} --d {d} --k {k} \\ + --max-iter {args.max_iter} --tol {args.tol} --seed {args.seed} \\ + --warmup-fit {args.warmup_fit} --iters-fit {args.iters_fit} \\ + --warmup-pred {args.warmup_pred} --iters-pred {args.iters_pred} +""" + proc = subprocess.run(["bash", "-lc", cmd], capture_output=True, text=True) + out = proc.stdout + proc.stderr + if proc.returncode != 0: + return BenchResult( + impl=impl, error=out.strip() or f"exit {proc.returncode}" + ) + return _parse_output(out, impl) + + +def _speedup(base: float, other: float) -> str: + if other <= 0: + return "n/a" + return f"{base / other:.2f}x" + + +def print_compare_table( + results: list[BenchResult], n: int, d: int, k: int +) -> None: + print(f"\n######## compare N={n} D={d} K={k} ########") + print(f"{'impl':<10} {'fit_ms':>10} {'pred_ms':>10} {'notes'}") + print("-" * 50) + by_impl = {r.impl: r for r in results} + for impl in IMPLS: + r = by_impl.get(impl) + if r is None: + print(f"{impl:<10} {'—':>10} {'—':>10} missing") + continue + if r.error: + print( + f"{impl:<10} {'FAIL':>10} {'FAIL':>10} {r.error.splitlines()[-1][:40]}" + ) + continue + print( + f"{impl:<10} {r.fit_median_ms:10.2f} {r.predict_median_ms:10.2f}" + ) + + flash = by_impl.get("flash") + cutile = by_impl.get("cutile") + if flash and cutile and flash.fit_median_ms and cutile.fit_median_ms: + if not flash.error and not cutile.error: + print( + f"\nflash vs cutile fit: {_speedup(cutile.fit_median_ms, flash.fit_median_ms)}" + f" predict: {_speedup(cutile.predict_median_ms, flash.predict_median_ms)}" + ) + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--compare", action="store_true", help="run baseline, cutile, flash" + ) + parser.add_argument("--impl", choices=IMPLS, help="single impl") + parser.add_argument("--n", type=int, required=True) + parser.add_argument("--d", type=int, required=True) + parser.add_argument("--k", type=int, required=True) + parser.add_argument("--max-iter", type=int, required=True) + parser.add_argument("--tol", type=float, required=True) + parser.add_argument("--seed", type=int, required=True) + parser.add_argument("--warmup-fit", type=int, required=True) + parser.add_argument("--iters-fit", type=int, required=True) + parser.add_argument("--warmup-pred", type=int, required=True) + parser.add_argument("--iters-pred", type=int, required=True) + args = parser.parse_args() + + if args.compare: + if args.impl: + parser.error("--compare and --impl are mutually exclusive") + _impl_config() # validate required env before launching subprocesses + results = [ + _run_subprocess(impl, args.n, args.d, args.k, args) + for impl in IMPLS + ] + print_compare_table(results, args.n, args.d, args.k) + return 0 if all(r.error is None for r in results) else 1 + + if not args.impl: + parser.error("set --impl for single-run mode, or use --compare") + + bench_mode = { + "baseline": "cuvs_base", + "cutile": "cuvs_cutile", + "flash": "flash", + }[args.impl] + + try: + run_benchmark( + bench_mode, + args.n, + args.d, + args.k, + max_iter=args.max_iter, + tol=args.tol, + seed=args.seed, + warmup_fit=args.warmup_fit, + iters_fit=args.iters_fit, + warmup_pred=args.warmup_pred, + iters_pred=args.iters_pred, + ) + except Exception as exc: + print(f"ERROR: {exc}", file=sys.stderr) + return 1 + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/run_benchmark_kmeans.sh b/run_benchmark_kmeans.sh new file mode 100755 index 0000000000..37291b5518 --- /dev/null +++ b/run_benchmark_kmeans.sh @@ -0,0 +1,47 @@ +#!/usr/bin/env bash +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +# Compare baseline / cuTile / flash-kmeans for one shape. +# +# Usage: +# export BENCH_CONDA=/path/to/miniforge3 +# export BENCH_ENV_BASE=... +# export BENCH_ENV_CUTILE=... +# export BENCH_ENV_FLASH=... +# export MAX_ITER=5 TOL=1e-4 SEED=42 +# export WARMUP_FIT=1 ITERS_FIT=3 WARMUP_PRED=1 ITERS_PRED=3 +# ./run_benchmark_kmeans.sh N D K +# + +set -u + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" + +if [[ $# -ne 3 ]]; then + echo "usage: $0 N D K" >&2 + echo "See script header for required env vars and examples." >&2 + exit 2 +fi + +: "${BENCH_CONDA:?set BENCH_CONDA to conda/miniforge root}" +: "${BENCH_ENV_BASE:?set BENCH_ENV_BASE}" +: "${BENCH_ENV_CUTILE:?set BENCH_ENV_CUTILE}" +: "${BENCH_ENV_FLASH:?set BENCH_ENV_FLASH}" +: "${MAX_ITER:?set MAX_ITER}" +: "${SEED:?set SEED}" +: "${WARMUP_FIT:?set WARMUP_FIT}" +: "${ITERS_FIT:?set ITERS_FIT}" +: "${WARMUP_PRED:?set WARMUP_PRED}" +: "${ITERS_PRED:?set ITERS_PRED}" +: "${TOL:?set TOL}" + +N=$1 +D=$2 +K=$3 + +exec python3 "$SCRIPT_DIR/benchmark_kmeans.py" --compare \ + --n "$N" --d "$D" --k "$K" \ + --max-iter "$MAX_ITER" --tol "$TOL" --seed "$SEED" \ + --warmup-fit "$WARMUP_FIT" --iters-fit "$ITERS_FIT" \ + --warmup-pred "$WARMUP_PRED" --iters-pred "$ITERS_PRED" From 9da31ae40ec280e2aa591b477aa5498c24a4a86d Mon Sep 17 00:00:00 2001 From: divyegala Date: Tue, 30 Jun 2026 18:43:19 +0000 Subject: [PATCH 11/82] benchmark sweep --- benchmark_kmeans.py | 266 ++++++++++++------ .../cutile/fused_1nn_cutile_matrix.json | 4 +- cpp/tests/CMakeLists.txt | 2 +- run_benchmark_kmeans.sh | 65 +++-- 4 files changed, 233 insertions(+), 104 deletions(-) diff --git a/benchmark_kmeans.py b/benchmark_kmeans.py index cfabe7570a..3ff965527e 100644 --- a/benchmark_kmeans.py +++ b/benchmark_kmeans.py @@ -3,8 +3,18 @@ # SPDX-License-Identifier: Apache-2.0 r"""KMeans fit+predict benchmark: baseline / cuTile / flash-kmeans. +Dimension glossary (same for cuVS and flash; matches fused GEMM A[M,D] @ B[K,D]^T): + M (--n) n_samples rows of X + D (--d) n_features inner / contraction dimension + K (--k) n_clusters centroids count (GEMM N) + +Shapes: + cuVS: X (M, D), centroids (K, D), KMeansParams(n_clusters=K) + flash: x (1, M, D), init_centroids (1, K, D), batch_kmeans_Euclid(..., K) + Single impl (activate the target conda env first): - python benchmark_kmeans.py --impl baseline|cutile|flash --n N --d D --k K \\ + python benchmark_kmeans.py --impl baseline|cutile|flash --n M --d D --k K \\ + --phase fit|predict|both \\ --max-iter 5 --tol 1e-4 --seed 42 \\ --warmup-fit 1 --iters-fit 3 --warmup-pred 1 --iters-pred 3 @@ -13,7 +23,7 @@ export BENCH_ENV_BASE=cuvs_2608_base export BENCH_ENV_CUTILE=cuvs_2608 export BENCH_ENV_FLASH=cuvs_2608_base - python benchmark_kmeans.py --compare --n 33554432 --d 32 --k 64 \\ + python benchmark_kmeans.py --compare --n 1000000 --d 128 --k 256 \\ --max-iter 5 --tol 1e-4 --seed 42 \\ --warmup-fit 1 --iters-fit 3 --warmup-pred 1 --iters-pred 3 @@ -95,6 +105,7 @@ def run_benchmark( d: int, k: int, *, + phase: str, max_iter: int, tol: float, seed: int, @@ -106,6 +117,7 @@ def run_benchmark( import numpy as np rng = np.random.default_rng(seed) + # Shared host data: X (M, D), centroids (K, D) — same layout for cuVS and flash. init_centroids_host = rng.standard_normal((k, d), dtype=np.float32) x_host = rng.standard_normal((n, d), dtype=np.float32) input_gib = n * d * 4 / (1024**3) @@ -115,8 +127,12 @@ def run_benchmark( "cuvs_cutile": "cutile", "flash": "flash", }[bench_mode] + run_fit = phase in ("fit", "both") + run_predict = phase in ("predict", "both") + print( - f"=== N={n:,} D={d} K={k:,} iters={max_iter} input={input_gib:.2f} GiB ===", + f"=== M={n:,} D={d} K={k:,} phase={phase} iters={max_iter} " + f"input={input_gib:.2f} GiB ===", flush=True, ) @@ -129,7 +145,7 @@ def run_benchmark( def sync(): cudart.cudaDeviceSynchronize() - x = device_ndarray(x_host) + x = device_ndarray(x_host) # (M, D) params = KMeansParams( n_clusters=k, max_iter=max_iter, @@ -140,7 +156,7 @@ def sync(): n_init=1, ) - for _ in range(warmup_fit): + for _ in range(warmup_fit if run_fit else 0): fit( params, x, centroids=device_ndarray(init_centroids_host.copy()) ) @@ -149,40 +165,46 @@ def sync(): fit_times: list[float] = [] n_iter = 0 inertia = 0.0 - for _ in range(iters_fit): - t0 = time.perf_counter() - _, inertia, n_iter = fit( + if run_fit: + for _ in range(iters_fit): + t0 = time.perf_counter() + _, inertia, n_iter = fit( + params, + x, + centroids=device_ndarray(init_centroids_host.copy()), + ) + sync() + fit_times.append((time.perf_counter() - t0) * 1e3) + + pred_times: list[float] = [] + if run_predict: + centroids, _, _ = fit( params, x, centroids=device_ndarray(init_centroids_host.copy()) ) sync() - fit_times.append((time.perf_counter() - t0) * 1e3) - centroids, _, _ = fit( - params, x, centroids=device_ndarray(init_centroids_host.copy()) - ) - sync() - - for _ in range(warmup_pred): - predict(params, x, centroids) - sync() + for _ in range(warmup_pred): + predict(params, x, centroids) + sync() - pred_times: list[float] = [] - for _ in range(iters_pred): - t0 = time.perf_counter() - predict(params, x, centroids) - sync() - pred_times.append((time.perf_counter() - t0) * 1e3) + for _ in range(iters_pred): + t0 = time.perf_counter() + predict(params, x, centroids) + sync() + pred_times.append((time.perf_counter() - t0) * 1e3) print(f"impl={label} init=Array", flush=True) - print(f"fit_median_ms={median(fit_times):.2f}", flush=True) - print(f"predict_median_ms={median(pred_times):.2f}", flush=True) - print(f"n_iter={n_iter} inertia={inertia:.6g}", flush=True) + if run_fit: + print(f"fit_median_ms={median(fit_times):.2f}", flush=True) + print(f"n_iter={n_iter} inertia={inertia:.6g}", flush=True) + if run_predict: + print(f"predict_median_ms={median(pred_times):.2f}", flush=True) return BenchResult( impl=label, - fit_median_ms=median(fit_times), - predict_median_ms=median(pred_times), - n_iter=n_iter, - inertia=inertia, + fit_median_ms=median(fit_times) if run_fit else None, + predict_median_ms=median(pred_times) if run_predict else None, + n_iter=n_iter if run_fit else None, + inertia=inertia if run_fit else None, ) if bench_mode == "flash": @@ -193,16 +215,16 @@ def sync(): def sync(): torch.cuda.synchronize() - x = torch.from_numpy(x_host).cuda() + x = torch.from_numpy(x_host).cuda() # (M, D) init_c = ( torch.from_numpy(init_centroids_host.copy()).cuda().unsqueeze(0) - ) + ) # (1, K, D) def run_fit(init): - x_b = x.unsqueeze(0) + x_b = x.unsqueeze(0) # (1, M, D) _, centroids_b, _ = batch_kmeans_Euclid( x_b, - k, + k, # n_clusters max_iters=max_iter, tol=tol, init_centroids=init, @@ -210,62 +232,84 @@ def run_fit(init): ) return centroids_b - for _ in range(warmup_fit): + for _ in range(warmup_fit if run_fit else 0): run_fit(init_c.clone()) sync() - fit_times = [] - for _ in range(iters_fit): - t0 = time.perf_counter() - run_fit(init_c.clone()) + fit_times: list[float] = [] + if run_fit: + for _ in range(iters_fit): + t0 = time.perf_counter() + run_fit(init_c.clone()) + sync() + fit_times.append((time.perf_counter() - t0) * 1e3) + + pred_times: list[float] = [] + if run_predict: + centroids_b = run_fit(init_c.clone()) sync() - fit_times.append((time.perf_counter() - t0) * 1e3) - centroids_b = run_fit(init_c.clone()) - sync() + x_b = x.unsqueeze(0) + x_sq = (x_b**2).sum(dim=-1) - x_b = x.unsqueeze(0) - x_sq = (x_b**2).sum(dim=-1) + for _ in range(warmup_pred): + euclid_assign_triton(x_b, centroids_b, x_sq) + sync() - for _ in range(warmup_pred): - euclid_assign_triton(x_b, centroids_b, x_sq) - sync() - - pred_times = [] - for _ in range(iters_pred): - t0 = time.perf_counter() - euclid_assign_triton(x_b, centroids_b, x_sq) - sync() - pred_times.append((time.perf_counter() - t0) * 1e3) + for _ in range(iters_pred): + t0 = time.perf_counter() + euclid_assign_triton(x_b, centroids_b, x_sq) + sync() + pred_times.append((time.perf_counter() - t0) * 1e3) print("impl=flash-kmeans init=Array", flush=True) - print(f"fit_median_ms={median(fit_times):.2f}", flush=True) - print(f"predict_median_ms={median(pred_times):.2f}", flush=True) + if run_fit: + print(f"fit_median_ms={median(fit_times):.2f}", flush=True) + if run_predict: + print(f"predict_median_ms={median(pred_times):.2f}", flush=True) return BenchResult( impl="flash", - fit_median_ms=median(fit_times), - predict_median_ms=median(pred_times), + fit_median_ms=median(fit_times) if run_fit else None, + predict_median_ms=median(pred_times) if run_predict else None, ) raise ValueError(f"unknown bench_mode={bench_mode!r}") -def _parse_output(text: str, impl: str) -> BenchResult: +def _parse_output(text: str, impl: str, phase: str) -> BenchResult: fit_m = re.search(r"^fit_median_ms=([0-9.]+)", text, re.M) pred_m = re.search(r"^predict_median_ms=([0-9.]+)", text, re.M) - if not fit_m or not pred_m: - return BenchResult(impl=impl, error=text.strip() or "no output") + if phase in ("fit", "both") and not fit_m: + return BenchResult(impl=impl, error=text.strip() or "no fit output") + if phase in ("predict", "both") and not pred_m: + return BenchResult( + impl=impl, error=text.strip() or "no predict output" + ) n_iter_m = re.search(r"^n_iter=([0-9]+)", text, re.M) inertia_m = re.search(r"^inertia=([0-9.eE+-]+)", text, re.M) return BenchResult( impl=impl, - fit_median_ms=float(fit_m.group(1)), - predict_median_ms=float(pred_m.group(1)), + fit_median_ms=float(fit_m.group(1)) if fit_m else None, + predict_median_ms=float(pred_m.group(1)) if pred_m else None, n_iter=int(n_iter_m.group(1)) if n_iter_m else None, inertia=float(inertia_m.group(1)) if inertia_m else None, ) +def _result_ok(result: BenchResult, phase: str) -> bool: + if result.error: + return False + if phase in ("fit", "both") and result.fit_median_ms is None: + return False + if phase in ("predict", "both") and result.predict_median_ms is None: + return False + return True + + +def _fmt_ms(value: float | None) -> str: + return f"{value:10.2f}" if value is not None else f"{'—':>10}" + + def _run_subprocess( impl: str, n: int, @@ -298,6 +342,7 @@ def _run_subprocess( conda activate "{cfg["conda_env"]}" {env_exports} python3 "{ROOT / "benchmark_kmeans.py"}" --impl {impl} --n {n} --d {d} --k {k} \\ + --phase {args.phase} \\ --max-iter {args.max_iter} --tol {args.tol} --seed {args.seed} \\ --warmup-fit {args.warmup_fit} --iters-fit {args.iters_fit} \\ --warmup-pred {args.warmup_pred} --iters-pred {args.iters_pred} @@ -308,7 +353,7 @@ def _run_subprocess( return BenchResult( impl=impl, error=out.strip() or f"exit {proc.returncode}" ) - return _parse_output(out, impl) + return _parse_output(out, impl, args.phase) def _speedup(base: float, other: float) -> str: @@ -318,34 +363,60 @@ def _speedup(base: float, other: float) -> str: def print_compare_table( - results: list[BenchResult], n: int, d: int, k: int + results: list[BenchResult], n: int, d: int, k: int, phase: str ) -> None: - print(f"\n######## compare N={n} D={d} K={k} ########") - print(f"{'impl':<10} {'fit_ms':>10} {'pred_ms':>10} {'notes'}") - print("-" * 50) + print(f"\n######## compare M={n} D={d} K={k} phase={phase} ########") + show_fit = phase in ("fit", "both") + show_pred = phase in ("predict", "both") + header = f"{'impl':<10}" + if show_fit: + header += f" {'fit_ms':>10}" + if show_pred: + header += f" {'pred_ms':>10}" + header += " notes" + print(header) + print("-" * len(header)) by_impl = {r.impl: r for r in results} for impl in IMPLS: r = by_impl.get(impl) if r is None: - print(f"{impl:<10} {'—':>10} {'—':>10} missing") + row = f"{impl:<10}" + if show_fit: + row += f" {'—':>10}" + if show_pred: + row += f" {'—':>10}" + print(f"{row} missing") continue if r.error: - print( - f"{impl:<10} {'FAIL':>10} {'FAIL':>10} {r.error.splitlines()[-1][:40]}" - ) + row = f"{impl:<10}" + if show_fit: + row += f" {'FAIL':>10}" + if show_pred: + row += f" {'FAIL':>10}" + print(f"{row} {r.error.splitlines()[-1][:40]}") continue - print( - f"{impl:<10} {r.fit_median_ms:10.2f} {r.predict_median_ms:10.2f}" - ) + row = f"{impl:<10}" + if show_fit: + row += _fmt_ms(r.fit_median_ms) + if show_pred: + row += _fmt_ms(r.predict_median_ms) + print(row) flash = by_impl.get("flash") cutile = by_impl.get("cutile") - if flash and cutile and flash.fit_median_ms and cutile.fit_median_ms: - if not flash.error and not cutile.error: - print( - f"\nflash vs cutile fit: {_speedup(cutile.fit_median_ms, flash.fit_median_ms)}" - f" predict: {_speedup(cutile.predict_median_ms, flash.predict_median_ms)}" + if flash and cutile and not flash.error and not cutile.error: + parts: list[str] = [] + if show_fit and flash.fit_median_ms and cutile.fit_median_ms: + parts.append( + f"fit: {_speedup(cutile.fit_median_ms, flash.fit_median_ms)}" ) + if show_pred and flash.predict_median_ms and cutile.predict_median_ms: + parts.append( + "predict: " + f"{_speedup(cutile.predict_median_ms, flash.predict_median_ms)}" + ) + if parts: + print(f"\nflash vs cutile {' '.join(parts)}") def main() -> int: @@ -354,9 +425,33 @@ def main() -> int: "--compare", action="store_true", help="run baseline, cutile, flash" ) parser.add_argument("--impl", choices=IMPLS, help="single impl") - parser.add_argument("--n", type=int, required=True) - parser.add_argument("--d", type=int, required=True) - parser.add_argument("--k", type=int, required=True) + parser.add_argument( + "--n", + type=int, + required=True, + metavar="M", + help="n_samples (GEMM M)", + ) + parser.add_argument( + "--d", + type=int, + required=True, + metavar="D", + help="n_features (GEMM inner dimension)", + ) + parser.add_argument( + "--k", + type=int, + required=True, + metavar="K", + help="n_clusters (GEMM N)", + ) + parser.add_argument( + "--phase", + choices=("fit", "predict", "both"), + default="both", + help="benchmark fit only, predict only, or both (default: both)", + ) parser.add_argument("--max-iter", type=int, required=True) parser.add_argument("--tol", type=float, required=True) parser.add_argument("--seed", type=int, required=True) @@ -374,8 +469,8 @@ def main() -> int: _run_subprocess(impl, args.n, args.d, args.k, args) for impl in IMPLS ] - print_compare_table(results, args.n, args.d, args.k) - return 0 if all(r.error is None for r in results) else 1 + print_compare_table(results, args.n, args.d, args.k, args.phase) + return 0 if all(_result_ok(r, args.phase) for r in results) else 1 if not args.impl: parser.error("set --impl for single-run mode, or use --compare") @@ -392,6 +487,7 @@ def main() -> int: args.n, args.d, args.k, + phase=args.phase, max_iter=args.max_iter, tol=args.tol, seed=args.seed, diff --git a/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_cutile_matrix.json b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_cutile_matrix.json index 7d9b723b39..7dcf05796b 100644 --- a/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_cutile_matrix.json +++ b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_cutile_matrix.json @@ -36,8 +36,8 @@ ], "_tile": [ { - "tile_m": 256, - "tile_n": 64, + "tile_m": 128, + "tile_n": 128, "tile_k": 32 } ], diff --git a/cpp/tests/CMakeLists.txt b/cpp/tests/CMakeLists.txt index d4e0099035..9b96f94bf0 100644 --- a/cpp/tests/CMakeLists.txt +++ b/cpp/tests/CMakeLists.txt @@ -120,7 +120,7 @@ ConfigureTest( ConfigureTest( NAME CLUSTER_TEST PATH cluster/kmeans.cu cluster/kmeans_balanced.cu cluster/kmeans_find_k.cu cluster/linkage.cu - cluster/connect_knn.cu cluster/spectral.cu cluster/soa_unpack_trace.cu + cluster/connect_knn.cu cluster/spectral.cu GPUS 1 PERCENT 100 ) diff --git a/run_benchmark_kmeans.sh b/run_benchmark_kmeans.sh index 37291b5518..46b05420b6 100755 --- a/run_benchmark_kmeans.sh +++ b/run_benchmark_kmeans.sh @@ -2,7 +2,7 @@ # SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -# Compare baseline / cuTile / flash-kmeans for one shape. +# Compare baseline / cuTile / flash-kmeans for one shape or the default sweep. # # Usage: # export BENCH_CONDA=/path/to/miniforge3 @@ -11,19 +11,22 @@ # export BENCH_ENV_FLASH=... # export MAX_ITER=5 TOL=1e-4 SEED=42 # export WARMUP_FIT=1 ITERS_FIT=3 WARMUP_PRED=1 ITERS_PRED=3 -# ./run_benchmark_kmeans.sh N D K +# export BENCH_PHASE=both # fit | predict | both (default: both) +# ./run_benchmark_kmeans.sh # default sweep (M=1M, all D and K below) +# ./run_benchmark_kmeans.sh N D K # single shape +# ./run_benchmark_kmeans.sh fit # sweep, fit only +# ./run_benchmark_kmeans.sh predict N D K # single shape, predict only +# +# Default sweep grid: +# M (n_samples) = 1_000_000 +# D (n_features) = 16 64 128 384 768 1024 1536 # GEMM inner (K) dimension +# K (n_clusters) = 10 100 1000 10000 100000 # GEMM N dimension # set -u SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" -if [[ $# -ne 3 ]]; then - echo "usage: $0 N D K" >&2 - echo "See script header for required env vars and examples." >&2 - exit 2 -fi - : "${BENCH_CONDA:?set BENCH_CONDA to conda/miniforge root}" : "${BENCH_ENV_BASE:?set BENCH_ENV_BASE}" : "${BENCH_ENV_CUTILE:?set BENCH_ENV_CUTILE}" @@ -36,12 +39,42 @@ fi : "${ITERS_PRED:?set ITERS_PRED}" : "${TOL:?set TOL}" -N=$1 -D=$2 -K=$3 +PHASE="${BENCH_PHASE:-both}" +if [[ $# -ge 1 && $1 =~ ^(fit|predict|both)$ ]]; then + PHASE=$1 + shift +fi + +run_shape() { + local n=$1 d=$2 k=$3 + echo "=== benchmark M=${n} D=${d} K=${k} phase=${PHASE} ===" + python3 "$SCRIPT_DIR/benchmark_kmeans.py" --compare \ + --n "$n" --d "$d" --k "$k" \ + --phase "$PHASE" \ + --max-iter "$MAX_ITER" --tol "$TOL" --seed "$SEED" \ + --warmup-fit "$WARMUP_FIT" --iters-fit "$ITERS_FIT" \ + --warmup-pred "$WARMUP_PRED" --iters-pred "$ITERS_PRED" +} + +if [[ $# -eq 3 ]]; then + run_shape "$1" "$2" "$3" + exit $? +fi + +if [[ $# -ne 0 ]]; then + echo "usage: $0 [fit|predict|both] [N D K]" >&2 + echo " no args — sweep all shapes, phase=\${BENCH_PHASE:-both}" >&2 + echo " fit|predict|both — optional phase override, then sweep" >&2 + echo " [phase] N D K — run one shape" >&2 + exit 2 +fi + +M=1000000 +D_VALUES=(16 64 128 384 768 1024 1536) +K_VALUES=(10 100 1000 10000 100000) -exec python3 "$SCRIPT_DIR/benchmark_kmeans.py" --compare \ - --n "$N" --d "$D" --k "$K" \ - --max-iter "$MAX_ITER" --tol "$TOL" --seed "$SEED" \ - --warmup-fit "$WARMUP_FIT" --iters-fit "$ITERS_FIT" \ - --warmup-pred "$WARMUP_PRED" --iters-pred "$ITERS_PRED" +for d in "${D_VALUES[@]}"; do + for k in "${K_VALUES[@]}"; do + run_shape "$M" "$d" "$k" || true + done +done From 35291ecec596630bcdc7e3e067fed7f2625ca9bd Mon Sep 17 00:00:00 2001 From: divyegala Date: Thu, 2 Jul 2026 17:46:35 +0000 Subject: [PATCH 12/82] index dtype --- c/src/cluster/kmeans.cpp | 194 +++++++++++++++--- .../cutile/export_fused_1nn.py | 29 +-- .../cutile/fused_1nn_kernel.py | 6 +- .../cutile/fused_1nn_tile.cu | 80 ++++---- run_benchmark_kmeans.sh | 7 + 5 files changed, 232 insertions(+), 84 deletions(-) diff --git a/c/src/cluster/kmeans.cpp b/c/src/cluster/kmeans.cpp index 8e46764ce4..f58a10b902 100644 --- a/c/src/cluster/kmeans.cpp +++ b/c/src/cluster/kmeans.cpp @@ -4,6 +4,7 @@ */ #include +#include #include @@ -44,7 +45,57 @@ cuvs::cluster::kmeans::balanced_params convert_balanced_params(const ParamsT& pa return kmeans_params; } -template +constexpr int64_t kKMeansInt32IndexMax = std::numeric_limits::max(); + +bool dlpack_shape_exceeds_int32_index(const DLTensor& tensor) +{ + for (int i = 0; i < tensor.ndim; ++i) { + if (tensor.shape[i] > kKMeansInt32IndexMax) { return true; } + } + return false; +} + +bool kmeans_tensor_shapes_use_int64_index(DLManagedTensor* X, DLManagedTensor* centroids) +{ + if (dlpack_shape_exceeds_int32_index(X->dl_tensor)) { return true; } + if (dlpack_shape_exceeds_int32_index(centroids->dl_tensor)) { return true; } + return false; +} + +bool kmeans_fit_uses_int64_index(DLManagedTensor* X, + DLManagedTensor* centroids, + int n_clusters) +{ + if (static_cast(n_clusters) > kKMeansInt32IndexMax) { return true; } + return kmeans_tensor_shapes_use_int64_index(X, centroids); +} + +bool kmeans_labels_use_int64_index(const DLTensor& labels) +{ + return labels.dtype.code == kDLInt && labels.dtype.bits == 64; +} + +void validate_kmeans_labels_dtype(const DLTensor& labels) +{ + if (labels.dtype.code == kDLInt && (labels.dtype.bits == 32 || labels.dtype.bits == 64)) { + return; + } + RAFT_FAIL("Unsupported labels DLtensor dtype: %d and bits: %d", + labels.dtype.code, + labels.dtype.bits); +} + +bool kmeans_predict_uses_int64_index(DLManagedTensor* X, + DLManagedTensor* centroids, + DLManagedTensor* labels, + int n_clusters) +{ + validate_kmeans_labels_dtype(labels->dl_tensor); + if (kmeans_labels_use_int64_index(labels->dl_tensor)) { return true; } + return kmeans_fit_uses_int64_index(X, centroids, n_clusters); +} + +template void _fit(cuvsResources_t res, const ParamsT& params, DLManagedTensor* X_tensor, @@ -57,8 +108,10 @@ void _fit(cuvsResources_t res, auto res_ptr = reinterpret_cast(res); if (!cuvs::core::is_dlpack_device_compatible(X)) { - auto n_samples = static_cast(X.shape[0]); - auto n_features = static_cast(X.shape[1]); + // Host fit overloads are only exposed with int64_t index types. + using HostIdxT = int64_t; + auto n_samples = static_cast(X.shape[0]); + auto n_features = static_cast(X.shape[1]); if (params.hierarchical) { RAFT_FAIL("hierarchical kmeans is not supported with host data"); @@ -69,24 +122,24 @@ void _fit(cuvsResources_t res, RAFT_FAIL("centroids must be on device memory"); } - auto X_view = raft::make_host_matrix_view( + auto X_view = raft::make_host_matrix_view( reinterpret_cast(X.data), n_samples, n_features); auto centroids_view = - cuvs::core::from_dlpack>( + cuvs::core::from_dlpack>( centroids_tensor); - std::optional> sample_weight; + std::optional> sample_weight; if (sample_weight_tensor != NULL) { auto sw = sample_weight_tensor->dl_tensor; if (!cuvs::core::is_dlpack_host_compatible(sw)) { RAFT_FAIL("sample_weight must be host accessible when X is on host"); } - sample_weight = raft::make_host_vector_view( + sample_weight = raft::make_host_vector_view( reinterpret_cast(sw.data), n_samples); } T inertia_temp; - IdxT n_iter_temp; + HostIdxT n_iter_temp; auto kmeans_params = convert_params(params); cuvs::cluster::kmeans::fit(*res_ptr, @@ -95,7 +148,7 @@ void _fit(cuvsResources_t res, sample_weight, centroids_view, raft::make_host_scalar_view(&inertia_temp), - raft::make_host_scalar_view(&n_iter_temp)); + raft::make_host_scalar_view(&n_iter_temp)); *inertia = inertia_temp; *n_iter = n_iter_temp; @@ -112,10 +165,20 @@ void _fit(cuvsResources_t res, if constexpr (std::is_same_v) { RAFT_FAIL("float64 is an unsupported dtype for hierarchical kmeans"); } else { - auto kmeans_params = convert_balanced_params(params); + // Balanced fit overloads are only exposed with int64_t index types. + using BalancedIdxT = int64_t; + using balanced_const_mdspan_type = + raft::device_matrix_view; + using balanced_mdspan_type = raft::device_matrix_view; + auto kmeans_params = convert_balanced_params(params); T inertia_temp; auto inertia_view = raft::make_host_scalar_view(&inertia_temp); - cuvs::cluster::kmeans::fit(*res_ptr, kmeans_params, cuvs::core::from_dlpack(X_tensor), cuvs::core::from_dlpack(centroids_tensor), std::make_optional(inertia_view)); + cuvs::cluster::kmeans::fit( + *res_ptr, + kmeans_params, + cuvs::core::from_dlpack(X_tensor), + cuvs::core::from_dlpack(centroids_tensor), + std::make_optional(inertia_view)); *inertia = inertia_temp; *n_iter = params.hierarchical_n_iters; } @@ -143,7 +206,7 @@ void _fit(cuvsResources_t res, } } -template +template void _predict(cuvsResources_t res, const ParamsT& params, DLManagedTensor* X_tensor, @@ -167,13 +230,22 @@ void _predict(cuvsResources_t res, if constexpr (std::is_same_v) { RAFT_FAIL("float64 is an unsupported dtype for hierarchical kmeans"); + } else if constexpr (!std::is_same_v) { + RAFT_FAIL("int64 labels are unsupported for hierarchical kmeans"); } else { + // Balanced predict overloads are only exposed with int64_t index types and int32 labels. + using BalancedIdxT = int64_t; + using balanced_const_mdspan_type = + raft::device_matrix_view; + using balanced_labels_mdspan_type = + raft::device_vector_view; auto kmeans_params = convert_balanced_params(params); - cuvs::cluster::kmeans::predict(*res_ptr, - kmeans_params, - cuvs::core::from_dlpack(X_tensor), - cuvs::core::from_dlpack(centroids_tensor), - cuvs::core::from_dlpack(labels_tensor)); + cuvs::cluster::kmeans::predict( + *res_ptr, + kmeans_params, + cuvs::core::from_dlpack(X_tensor), + cuvs::core::from_dlpack(centroids_tensor), + cuvs::core::from_dlpack(labels_tensor)); *inertia = 0; } } else { @@ -199,7 +271,7 @@ void _predict(cuvsResources_t res, } } -template +template void _cluster_cost(cuvsResources_t res, DLManagedTensor* X_tensor, DLManagedTensor* centroids_tensor, @@ -263,10 +335,24 @@ extern "C" cuvsError_t cuvsKMeansFit(cuvsResources_t res, { return cuvs::core::translate_exceptions([=] { auto dataset = X->dl_tensor; + const bool use_int64_index = + kmeans_fit_uses_int64_index(X, centroids, params->n_clusters); if (dataset.dtype.code == kDLFloat && dataset.dtype.bits == 32) { - _fit(res, *params, X, sample_weight, centroids, inertia, n_iter); + if (use_int64_index) { + _fit( + res, *params, X, sample_weight, centroids, inertia, n_iter); + } else { + _fit( + res, *params, X, sample_weight, centroids, inertia, n_iter); + } } else if (dataset.dtype.code == kDLFloat && dataset.dtype.bits == 64) { - _fit(res, *params, X, sample_weight, centroids, inertia, n_iter); + if (use_int64_index) { + _fit( + res, *params, X, sample_weight, centroids, inertia, n_iter); + } else { + _fit( + res, *params, X, sample_weight, centroids, inertia, n_iter); + } } else { RAFT_FAIL("Unsupported dataset DLtensor dtype: %d and bits: %d", dataset.dtype.code, @@ -286,10 +372,24 @@ extern "C" cuvsError_t cuvsKMeansPredict(cuvsResources_t res, { return cuvs::core::translate_exceptions([=] { auto dataset = X->dl_tensor; + const bool use_int64_index = + kmeans_predict_uses_int64_index(X, centroids, labels, params->n_clusters); if (dataset.dtype.code == kDLFloat && dataset.dtype.bits == 32) { - _predict(res, *params, X, sample_weight, centroids, labels, normalize_weight, inertia); + if (use_int64_index) { + _predict( + res, *params, X, sample_weight, centroids, labels, normalize_weight, inertia); + } else { + _predict( + res, *params, X, sample_weight, centroids, labels, normalize_weight, inertia); + } } else if (dataset.dtype.code == kDLFloat && dataset.dtype.bits == 64) { - _predict(res, *params, X, sample_weight, centroids, labels, normalize_weight, inertia); + if (use_int64_index) { + _predict( + res, *params, X, sample_weight, centroids, labels, normalize_weight, inertia); + } else { + _predict( + res, *params, X, sample_weight, centroids, labels, normalize_weight, inertia); + } } else { RAFT_FAIL("Unsupported dataset DLtensor dtype: %d and bits: %d", dataset.dtype.code, @@ -335,10 +435,24 @@ extern "C" cuvsError_t cuvsKMeansFit_v2(cuvsResources_t res, { return cuvs::core::translate_exceptions([=] { auto dataset = X->dl_tensor; + const bool use_int64_index = + kmeans_fit_uses_int64_index(X, centroids, params->n_clusters); if (dataset.dtype.code == kDLFloat && dataset.dtype.bits == 32) { - _fit(res, *params, X, sample_weight, centroids, inertia, n_iter); + if (use_int64_index) { + _fit( + res, *params, X, sample_weight, centroids, inertia, n_iter); + } else { + _fit( + res, *params, X, sample_weight, centroids, inertia, n_iter); + } } else if (dataset.dtype.code == kDLFloat && dataset.dtype.bits == 64) { - _fit(res, *params, X, sample_weight, centroids, inertia, n_iter); + if (use_int64_index) { + _fit( + res, *params, X, sample_weight, centroids, inertia, n_iter); + } else { + _fit( + res, *params, X, sample_weight, centroids, inertia, n_iter); + } } else { RAFT_FAIL("Unsupported dataset DLtensor dtype: %d and bits: %d", dataset.dtype.code, @@ -358,11 +472,24 @@ extern "C" cuvsError_t cuvsKMeansPredict_v2(cuvsResources_t res, { return cuvs::core::translate_exceptions([=] { auto dataset = X->dl_tensor; + const bool use_int64_index = + kmeans_predict_uses_int64_index(X, centroids, labels, params->n_clusters); if (dataset.dtype.code == kDLFloat && dataset.dtype.bits == 32) { - _predict(res, *params, X, sample_weight, centroids, labels, normalize_weight, inertia); + if (use_int64_index) { + _predict( + res, *params, X, sample_weight, centroids, labels, normalize_weight, inertia); + } else { + _predict( + res, *params, X, sample_weight, centroids, labels, normalize_weight, inertia); + } } else if (dataset.dtype.code == kDLFloat && dataset.dtype.bits == 64) { - _predict( - res, *params, X, sample_weight, centroids, labels, normalize_weight, inertia); + if (use_int64_index) { + _predict( + res, *params, X, sample_weight, centroids, labels, normalize_weight, inertia); + } else { + _predict( + res, *params, X, sample_weight, centroids, labels, normalize_weight, inertia); + } } else { RAFT_FAIL("Unsupported dataset DLtensor dtype: %d and bits: %d", dataset.dtype.code, @@ -378,10 +505,19 @@ extern "C" cuvsError_t cuvsKMeansClusterCost(cuvsResources_t res, { return cuvs::core::translate_exceptions([=] { auto dataset = X->dl_tensor; + const bool use_int64_index = kmeans_tensor_shapes_use_int64_index(X, centroids); if (dataset.dtype.code == kDLFloat && dataset.dtype.bits == 32) { - _cluster_cost(res, X, centroids, cost); + if (use_int64_index) { + _cluster_cost(res, X, centroids, cost); + } else { + _cluster_cost(res, X, centroids, cost); + } } else if (dataset.dtype.code == kDLFloat && dataset.dtype.bits == 64) { - _cluster_cost(res, X, centroids, cost); + if (use_int64_index) { + _cluster_cost(res, X, centroids, cost); + } else { + _cluster_cost(res, X, centroids, cost); + } } else { RAFT_FAIL("Unsupported dataset DLtensor dtype: %d and bits: %d", dataset.dtype.code, diff --git a/cpp/src/distance/detail/fused_distance_nn/cutile/export_fused_1nn.py b/cpp/src/distance/detail/fused_distance_nn/cutile/export_fused_1nn.py index 1211456c9e..582ad5b8a4 100644 --- a/cpp/src/distance/detail/fused_distance_nn/cutile/export_fused_1nn.py +++ b/cpp/src/distance/detail/fused_distance_nn/cutile/export_fused_1nn.py @@ -22,6 +22,7 @@ from fused_1nn_kernel import ( INDEX_TYPES, METRICS, + _idx_dtype, index_abbrev, kernel_symbol, make_kernel, @@ -52,7 +53,7 @@ def _elem_stride_divisible_for_tma(elem_dtype) -> tuple[int, int]: return (16 // bytes_per_elem, 1) -def _cuvs_matrix_constraint(elem_dtype): +def _cuvs_matrix_constraint(elem_dtype, *, index_dtype=ct.int32): """Row-major device matrices for cuVS KMeans benchmarks. Assumes raft/cupy-style contiguous layout: stride[-1]==1, stride[0]==D, @@ -65,7 +66,7 @@ def _cuvs_matrix_constraint(elem_dtype): return ArrayConstraint( elem_dtype, ndim=2, - index_dtype=ct.int32, + index_dtype=index_dtype, stride_lower_bound_incl=(0, None), alias_groups=(), may_alias_internally=False, @@ -76,12 +77,12 @@ def _cuvs_matrix_constraint(elem_dtype): ) -def _cuvs_vector_constraint(elem_dtype): +def _cuvs_vector_constraint(elem_dtype, *, index_dtype=ct.int32): """1-D device vectors: contiguous, 16-byte base. Length need not be divisible by 16.""" return ArrayConstraint( elem_dtype, ndim=1, - index_dtype=ct.int32, + index_dtype=index_dtype, stride_lower_bound_incl=(None,), alias_groups=(), may_alias_internally=False, @@ -112,11 +113,11 @@ def _kernel_signature( tile_k: int, ) -> KernelSignature: elem = _dtype_for(data_type) - matrix = _cuvs_matrix_constraint(elem) - norm_array = _cuvs_vector_constraint(elem) - idx_elem = ct.int32 if index_type == "int32" else ct.int64 - idx_array = _cuvs_vector_constraint(idx_elem) - dist_array = _cuvs_vector_constraint(elem) + idx_dtype = _idx_dtype(index_type) + matrix = _cuvs_matrix_constraint(elem, index_dtype=idx_dtype) + norm_array = _cuvs_vector_constraint(elem, index_dtype=idx_dtype) + idx_array = _cuvs_vector_constraint(idx_dtype, index_dtype=idx_dtype) + dist_array = _cuvs_vector_constraint(elem, index_dtype=idx_dtype) abbrev = _data_abbrev(data_type) symbol = kernel_symbol( @@ -131,11 +132,11 @@ def _kernel_signature( norm_array, idx_array, dist_array, - ScalarConstraint(ct.int64), - ScalarConstraint(ct.int64), - ScalarConstraint(ct.int64), - ScalarConstraint(ct.int64), - ScalarConstraint(ct.int64), + ScalarConstraint(idx_dtype), + ScalarConstraint(idx_dtype), + ScalarConstraint(idx_dtype), + ScalarConstraint(idx_dtype), + ScalarConstraint(idx_dtype), ConstantConstraint(tile_m), ConstantConstraint(tile_n), ConstantConstraint(tile_k), diff --git a/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_kernel.py b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_kernel.py index b2ff25555b..ea0c83164b 100644 --- a/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_kernel.py +++ b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_kernel.py @@ -9,9 +9,9 @@ ConstInt = ct.Constant[int] # Default tile geometry; overridden per export via make_kernel(..., tile_m, tile_n, tile_k). -DEFAULT_TILE_M = 256 -DEFAULT_TILE_N = 64 -DEFAULT_TILE_K = 32 +DEFAULT_TILE_M = 128 +DEFAULT_TILE_N = 128 +DEFAULT_TILE_K = 64 METRICS = ("inner_product", "l2_expanded", "cosine_expanded") INDEX_TYPES = ("int32", "int64") diff --git a/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_tile.cu b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_tile.cu index cf01c12ce1..275ae33403 100644 --- a/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_tile.cu +++ b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_tile.cu @@ -11,6 +11,8 @@ #include #include +#include + namespace cuvs { namespace distance { namespace detail { @@ -43,20 +45,22 @@ bool launch_fused_1nn_tile(IdxT* nearest_idx, const bool apply_sqrt = fused_1nn_apply_sqrt_at_pack(is_sqrt); - int64_t shape_x[2] = {m, k}; - int64_t stride_x[2] = {k, 1}; - int64_t shape_y[2] = {n, k}; - int64_t stride_y[2] = {k, 1}; - int64_t shape_xn = m; - int64_t stride_xn = 1; - int64_t shape_yn = n; - int64_t stride_yn = 1; - int64_t shape_idx = m; - int64_t stride_idx = 1; - int64_t shape_dist = m; - int64_t stride_dist = 1; - - int64_t M = m, N = n, K = k; + IdxT shape_x[2] = {m, k}; + IdxT stride_x[2] = {k, IdxT{1}}; + IdxT shape_y[2] = {n, k}; + IdxT stride_y[2] = {k, IdxT{1}}; + IdxT shape_xn = m; + IdxT stride_xn = IdxT{1}; + IdxT shape_yn = n; + IdxT stride_yn = IdxT{1}; + IdxT shape_idx = m; + IdxT stride_idx = IdxT{1}; + IdxT shape_dist = m; + IdxT stride_dist = IdxT{1}; + + IdxT M = m; + IdxT N = n; + IdxT K = k; void* x_ptr = const_cast(x); void* y_ptr = const_cast(y); @@ -64,42 +68,42 @@ bool launch_fused_1nn_tile(IdxT* nearest_idx, void* yn_ptr = const_cast(yn); // OutIdx must be a valid device pointer for the launch ABI; when store_idx is 0 the kernel // does not write it (dist-only callers pass nearest_dist as a stand-in). - const int64_t store_idx = nearest_idx != nullptr ? 1 : 0; + const IdxT store_idx = nearest_idx != nullptr ? IdxT{1} : IdxT{0}; void* idx_ptr = nearest_idx != nullptr ? static_cast(nearest_idx) : static_cast(nearest_dist); void* dist_ptr = nearest_dist; - const int64_t tile_m = tile_cfg.tile_m; + const int tile_m = tile_cfg.tile_m; dim3 grid((m + tile_m - 1) / tile_m, 1, 1); dim3 block(1, 1, 1); using fused_1nn_cutile_kernel_t = void(void*, - int64_t, - int64_t, - int64_t, - int64_t, + IdxT, + IdxT, + IdxT, + IdxT, void*, - int64_t, - int64_t, - int64_t, - int64_t, + IdxT, + IdxT, + IdxT, + IdxT, void*, - int64_t, - int64_t, + IdxT, + IdxT, void*, - int64_t, - int64_t, + IdxT, + IdxT, void*, - int64_t, - int64_t, + IdxT, + IdxT, void*, - int64_t, - int64_t, - int64_t, - int64_t, - int64_t, - int64_t, - int64_t); + IdxT, + IdxT, + IdxT, + IdxT, + IdxT, + IdxT, + IdxT); launcher->template dispatch(stream, grid, block, @@ -129,7 +133,7 @@ bool launch_fused_1nn_tile(IdxT* nearest_idx, M, N, K, - static_cast(apply_sqrt), + static_cast(apply_sqrt), store_idx); RAFT_CUDA_TRY(cudaGetLastError()); return true; diff --git a/run_benchmark_kmeans.sh b/run_benchmark_kmeans.sh index 46b05420b6..8109bd1f81 100755 --- a/run_benchmark_kmeans.sh +++ b/run_benchmark_kmeans.sh @@ -12,6 +12,8 @@ # export MAX_ITER=5 TOL=1e-4 SEED=42 # export WARMUP_FIT=1 ITERS_FIT=3 WARMUP_PRED=1 ITERS_PRED=3 # export BENCH_PHASE=both # fit | predict | both (default: both) +# export BENCH_GPU_NAME=rtx_pro_6000 # sweep log tag (e.g. h200 on Hopper) +# export BENCH_LOG=path/to.log # optional sweep log override # ./run_benchmark_kmeans.sh # default sweep (M=1M, all D and K below) # ./run_benchmark_kmeans.sh N D K # single shape # ./run_benchmark_kmeans.sh fit # sweep, fit only @@ -69,6 +71,11 @@ if [[ $# -ne 0 ]]; then exit 2 fi +: "${BENCH_GPU_NAME:?set BENCH_GPU_NAME e.g. rtx_pro_6000}" +BENCH_LOG="${BENCH_LOG:-${SCRIPT_DIR}/benchmark_kmeans_sweep_${BENCH_GPU_NAME}_$(date +%Y%m%d_%H%M%S).log}" +echo "Logging to ${BENCH_LOG}" +exec > >(tee "$BENCH_LOG") 2>&1 + M=1000000 D_VALUES=(16 64 128 384 768 1024 1536) K_VALUES=(10 100 1000 10000 100000) From f48ec5fb299c52ed9aa57cc512a5c9523c68793c Mon Sep 17 00:00:00 2001 From: divyegala Date: Thu, 9 Jul 2026 17:25:54 +0000 Subject: [PATCH 13/82] new reduce --- .../cutile/export_fused_1nn.py | 8 +- .../cutile/fused_1nn_kernel.py | 84 +++++++++++++++++-- 2 files changed, 82 insertions(+), 10 deletions(-) diff --git a/cpp/src/distance/detail/fused_distance_nn/cutile/export_fused_1nn.py b/cpp/src/distance/detail/fused_distance_nn/cutile/export_fused_1nn.py index 582ad5b8a4..c176d4b9ad 100644 --- a/cpp/src/distance/detail/fused_distance_nn/cutile/export_fused_1nn.py +++ b/cpp/src/distance/detail/fused_distance_nn/cutile/export_fused_1nn.py @@ -159,7 +159,13 @@ def export_binary( bytecode_version: str | None = None, ) -> str: kernel = make_kernel( - data_type, metric, tile_m, tile_n, tile_k, index_type=index_type + data_type, + metric, + tile_m, + tile_n, + tile_k, + index_type=index_type, + gpu_code=gpu_code, ) signature = _kernel_signature( data_type, metric, index_type, tile_m, tile_n, tile_k diff --git a/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_kernel.py b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_kernel.py index ea0c83164b..e1da35a9a4 100644 --- a/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_kernel.py +++ b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_kernel.py @@ -33,6 +33,7 @@ def make_kernel( tile_k: int = DEFAULT_TILE_K, *, index_type: str = "int32", + gpu_code: str = "sm_80", ): """Build a cuTile kernel with metric, index width, and tile sizes baked in at compile time.""" if data_type not in ("half", "float"): @@ -49,6 +50,22 @@ def make_kernel( is_l2 = metric == "l2_expanded" is_cos = metric == "cosine_expanded" + if gpu_code == "sm_100": + # Blackwell tcgen05 maps one logical row of the output tile to the reduction owner. + core_shape = (tile_m, tile_n) + best_shape = (tile_m, 1) + inner_reduction_axes = (1,) + outer_reduction_axes = (1,) + else: + # SM80/86/90/110/120 distribute eight consecutive N values across four threads, + # two per thread. + # Reducing in this physical layout keeps the score/index pair reduction local and avoids + # the slower generic argmin/argmax lowering over the full N dimension. + core_shape = (tile_m, tile_n // 8, 4, 2) + best_shape = (tile_m, 1, 4, 1) + inner_reduction_axes = (1, 3) + outer_reduction_axes = (2,) + @ct.kernel def fused_1nn_kernel( A, @@ -69,15 +86,48 @@ def fused_1nn_kernel( bidm = ct.bid(0) if is_ip: - best_dist = ct.full((tm,), -3.4e38, acc_dtype) + best_dist = ct.full(best_shape, -3.4e38, acc_dtype) + neutral_dist = -3.4e38 else: - best_dist = ct.full((tm,), 3.4e38, acc_dtype) - best_idx = ct.zeros((tm,), idx_dtype) + best_dist = ct.full(best_shape, 3.4e38, acc_dtype) + neutral_dist = 3.4e38 + best_idx = ct.zeros(best_shape, idx_dtype) num_tiles_k = ct.num_tiles(A, axis=1, shape=(tm, tk)) num_tiles_n = ct.num_tiles(B, axis=0, shape=(tn, tk)) zero_pad = ct.PaddingMode.ZERO + def reduce_scores(best, best_idx, axes): + def red_op(a_score, a_idx, b_score, b_idx): + if is_ip: + cond = a_score > b_score + else: + cond = a_score < b_score + + return ( + ct.where(cond, a_score, b_score), + ct.where(cond, a_idx, b_idx), + ) + + if len(axes) >= 1: + best, best_idx = ct.reduce( + (best, best_idx), + axes[0], + red_op, + (neutral_dist, -1), + keepdims=True, + ) + if len(axes) >= 2: + best, best_idx = ct.reduce( + (best, best_idx), + axes[1], + red_op, + (neutral_dist, -1), + keepdims=True, + ) + + return best, best_idx + for n in range(num_tiles_n): accumulator = ct.full((tm, tn), 0, dtype=acc_dtype) @@ -123,13 +173,21 @@ def fused_1nn_kernel( score = ct.where(valid[None, :], score, 3.4e38) if is_ip: - curr_best = ct.max(score, axis=1) - curr_idx = ct.argmax(score, axis=1) + curr_idx = ct.arange(tn, dtype=idx_dtype).reshape( + core_shape[1:] + )[None, ...] + curr_best, curr_idx = reduce_scores( + score.reshape(core_shape), curr_idx, inner_reduction_axes + ) update = curr_best > best_dist best_dist = ct.where(update, curr_best, best_dist) else: - curr_best = ct.min(score, axis=1) - curr_idx = ct.argmin(score, axis=1) + curr_idx = ct.arange(tn, dtype=idx_dtype).reshape( + core_shape[1:] + )[None, ...] + curr_best, curr_idx = reduce_scores( + score.reshape(core_shape), curr_idx, inner_reduction_axes + ) update = curr_best < best_dist best_dist = ct.where(update, curr_best, best_dist) @@ -137,12 +195,20 @@ def fused_1nn_kernel( update, (n * tn + curr_idx).astype(idx_dtype), best_idx ) + best_dist, best_idx = reduce_scores( + best_dist, best_idx, outer_reduction_axes + ) + out_dist = best_dist if is_l2: out_dist = ct.where(apply_sqrt != 0, ct.sqrt(best_dist), best_dist) if store_idx != 0: - ct.store(OutIdx, index=(bidm,), tile=best_idx) - ct.store(OutDist, index=(bidm,), tile=out_dist.astype(out_dist_dtype)) + ct.store(OutIdx, index=(bidm,), tile=best_idx.reshape((tm,))) + ct.store( + OutDist, + index=(bidm,), + tile=out_dist.reshape((tm,)).astype(out_dist_dtype), + ) return fused_1nn_kernel From 842863edf2d99a0821cf20c51323d24e3c08582f Mon Sep 17 00:00:00 2001 From: divyegala Date: Thu, 9 Jul 2026 19:22:18 +0000 Subject: [PATCH 14/82] try int64 vector constant --- .../cutile/export_fused_1nn.py | 9 +++-- .../cutile/fused_1nn_tile.cu | 40 +++++++++---------- 2 files changed, 26 insertions(+), 23 deletions(-) diff --git a/cpp/src/distance/detail/fused_distance_nn/cutile/export_fused_1nn.py b/cpp/src/distance/detail/fused_distance_nn/cutile/export_fused_1nn.py index c176d4b9ad..b680468b8d 100644 --- a/cpp/src/distance/detail/fused_distance_nn/cutile/export_fused_1nn.py +++ b/cpp/src/distance/detail/fused_distance_nn/cutile/export_fused_1nn.py @@ -114,10 +114,13 @@ def _kernel_signature( ) -> KernelSignature: elem = _dtype_for(data_type) idx_dtype = _idx_dtype(index_type) + vector_index_dtype = ct.int64 matrix = _cuvs_matrix_constraint(elem, index_dtype=idx_dtype) - norm_array = _cuvs_vector_constraint(elem, index_dtype=idx_dtype) - idx_array = _cuvs_vector_constraint(idx_dtype, index_dtype=idx_dtype) - dist_array = _cuvs_vector_constraint(elem, index_dtype=idx_dtype) + norm_array = _cuvs_vector_constraint(elem, index_dtype=vector_index_dtype) + idx_array = _cuvs_vector_constraint( + idx_dtype, index_dtype=vector_index_dtype + ) + dist_array = _cuvs_vector_constraint(elem, index_dtype=vector_index_dtype) abbrev = _data_abbrev(data_type) symbol = kernel_symbol( diff --git a/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_tile.cu b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_tile.cu index 275ae33403..4651a4b7d3 100644 --- a/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_tile.cu +++ b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_tile.cu @@ -45,18 +45,18 @@ bool launch_fused_1nn_tile(IdxT* nearest_idx, const bool apply_sqrt = fused_1nn_apply_sqrt_at_pack(is_sqrt); - IdxT shape_x[2] = {m, k}; - IdxT stride_x[2] = {k, IdxT{1}}; - IdxT shape_y[2] = {n, k}; - IdxT stride_y[2] = {k, IdxT{1}}; - IdxT shape_xn = m; - IdxT stride_xn = IdxT{1}; - IdxT shape_yn = n; - IdxT stride_yn = IdxT{1}; - IdxT shape_idx = m; - IdxT stride_idx = IdxT{1}; - IdxT shape_dist = m; - IdxT stride_dist = IdxT{1}; + IdxT shape_x[2] = {m, k}; + IdxT stride_x[2] = {k, IdxT{1}}; + IdxT shape_y[2] = {n, k}; + IdxT stride_y[2] = {k, IdxT{1}}; + int64_t shape_xn = m; + int64_t stride_xn = 1; + int64_t shape_yn = n; + int64_t stride_yn = 1; + int64_t shape_idx = m; + int64_t stride_idx = 1; + int64_t shape_dist = m; + int64_t stride_dist = 1; IdxT M = m; IdxT N = n; @@ -88,17 +88,17 @@ bool launch_fused_1nn_tile(IdxT* nearest_idx, IdxT, IdxT, void*, - IdxT, - IdxT, + int64_t, + int64_t, void*, - IdxT, - IdxT, + int64_t, + int64_t, void*, - IdxT, - IdxT, + int64_t, + int64_t, void*, - IdxT, - IdxT, + int64_t, + int64_t, IdxT, IdxT, IdxT, From 9fb768df7dde69cc343ca1646399ad7eda322f0f Mon Sep 17 00:00:00 2001 From: divyegala Date: Thu, 9 Jul 2026 20:14:34 +0000 Subject: [PATCH 15/82] new tile shape for sm90 --- .../modules/generate_cutile_kernels.cmake | 27 ++++++++ .../cutile/export_fused_1nn.py | 9 +-- .../cutile/fused_1nn_cutile_matrix.json | 65 ++++++++++++++++--- .../cutile/fused_1nn_planner.hpp | 8 +-- .../cutile/fused_1nn_tile.cu | 40 ++++++------ 5 files changed, 109 insertions(+), 40 deletions(-) diff --git a/cpp/cmake/modules/generate_cutile_kernels.cmake b/cpp/cmake/modules/generate_cutile_kernels.cmake index 9cd8a207c8..97e06e79aa 100644 --- a/cpp/cmake/modules/generate_cutile_kernels.cmake +++ b/cpp/cmake/modules/generate_cutile_kernels.cmake @@ -98,6 +98,32 @@ function(_cutile_generate_matrix_tiles_header header_path matrix_json_file) string(JSON _tile_m GET "${_tile0}" "tile_m") string(JSON _tile_n GET "${_tile0}" "tile_n") string(JSON _tile_k GET "${_tile0}" "tile_k") + set(_arch_tile_aliases "") + string(JSON _entry_len LENGTH "${_matrix_json}") + set(_entry_idx 0) + while(_entry_idx LESS _entry_len) + string(JSON _entry GET "${_matrix_json}" "${_entry_idx}") + string(JSON _entry_tile GET "${_entry}" "_tile" 0) + string(JSON _entry_tile_m GET "${_entry_tile}" "tile_m") + string(JSON _entry_tile_n GET "${_entry_tile}" "tile_n") + string(JSON _entry_tile_k GET "${_entry_tile}" "tile_k") + string(JSON _export_len LENGTH "${_entry}" "_export") + set(_export_idx 0) + while(_export_idx LESS _export_len) + string(JSON _export_entry GET "${_entry}" "_export" "${_export_idx}") + string(JSON _register GET "${_export_entry}" "register") + if(_register STREQUAL "cubin") + string(JSON _arch_tag GET "${_export_entry}" "arch_tag") + string( + APPEND + _arch_tile_aliases + "using fused_1nn_matrix_tile_${_arch_tag} = cutile_tile_config<${_entry_tile_m}, ${_entry_tile_n}, ${_entry_tile_k}>;\n" + ) + endif() + math(EXPR _export_idx "${_export_idx} + 1") + endwhile() + math(EXPR _entry_idx "${_entry_idx} + 1") + endwhile() file( WRITE "${header_path}" "/* @@ -110,6 +136,7 @@ function(_cutile_generate_matrix_tiles_header header_path matrix_json_file) namespace cuvs::distance::detail { using fused_1nn_matrix_tile = cutile_tile_config<${_tile_m}, ${_tile_n}, ${_tile_k}>; +${_arch_tile_aliases} } // namespace cuvs::distance::detail " diff --git a/cpp/src/distance/detail/fused_distance_nn/cutile/export_fused_1nn.py b/cpp/src/distance/detail/fused_distance_nn/cutile/export_fused_1nn.py index b680468b8d..c176d4b9ad 100644 --- a/cpp/src/distance/detail/fused_distance_nn/cutile/export_fused_1nn.py +++ b/cpp/src/distance/detail/fused_distance_nn/cutile/export_fused_1nn.py @@ -114,13 +114,10 @@ def _kernel_signature( ) -> KernelSignature: elem = _dtype_for(data_type) idx_dtype = _idx_dtype(index_type) - vector_index_dtype = ct.int64 matrix = _cuvs_matrix_constraint(elem, index_dtype=idx_dtype) - norm_array = _cuvs_vector_constraint(elem, index_dtype=vector_index_dtype) - idx_array = _cuvs_vector_constraint( - idx_dtype, index_dtype=vector_index_dtype - ) - dist_array = _cuvs_vector_constraint(elem, index_dtype=vector_index_dtype) + norm_array = _cuvs_vector_constraint(elem, index_dtype=idx_dtype) + idx_array = _cuvs_vector_constraint(idx_dtype, index_dtype=idx_dtype) + dist_array = _cuvs_vector_constraint(elem, index_dtype=idx_dtype) abbrev = _data_abbrev(data_type) symbol = kernel_symbol( diff --git a/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_cutile_matrix.json b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_cutile_matrix.json index 7dcf05796b..0aca5bdf14 100644 --- a/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_cutile_matrix.json +++ b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_cutile_matrix.json @@ -62,16 +62,6 @@ "cc_minor": 6, "arch_tag": "cutile_arch_8_6" }, - { - "output_format": "cubin", - "artifact_ext": "cubin", - "artifact_basename": "@data_type@_@metric_abbrev@_@index_abbrev@_@gpu_code@", - "register": "cubin", - "gpu_code": "sm_90", - "cc_major": 9, - "cc_minor": 0, - "arch_tag": "cutile_arch_9_0" - }, { "output_format": "cubin", "artifact_ext": "cubin", @@ -91,5 +81,60 @@ "bytecode_version": "13.1" } ] + }, + { + "_data": [ + { + "data_type": "half", + "data_abbrev": "h" + }, + { + "data_type": "float", + "data_abbrev": "f" + } + ], + "_metric": [ + { + "metric": "inner_product", + "metric_abbrev": "ip" + }, + { + "metric": "l2_expanded", + "metric_abbrev": "l2" + }, + { + "metric": "cosine_expanded", + "metric_abbrev": "cos" + } + ], + "_index": [ + { + "index_type": "int32", + "index_abbrev": "i32" + }, + { + "index_type": "int64", + "index_abbrev": "i64" + } + ], + "_tile": [ + { + "tile_m": 64, + "tile_n": 128, + "tile_k": 32 + } + ], + "_export": [ + { + "output_format": "cubin", + "artifact_ext": "cubin", + "artifact_basename": "@data_type@_@metric_abbrev@_@index_abbrev@_@gpu_code@", + "register": "cubin", + "gpu_code": "sm_90", + "cc_major": 9, + "cc_minor": 0, + "arch_tag": "cutile_arch_9_0" + } + ] } ] diff --git a/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_planner.hpp b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_planner.hpp index 017fb72d48..41185e5a61 100644 --- a/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_planner.hpp +++ b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_planner.hpp @@ -74,22 +74,22 @@ struct Fused1nnTilePlanner : TileAlgorithmPlanner { this->add_static_fragment>(); this->add_static_fragment>(); this->add_static_fragment>(); this->add_static_fragment>(); } diff --git a/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_tile.cu b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_tile.cu index 4651a4b7d3..275ae33403 100644 --- a/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_tile.cu +++ b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_tile.cu @@ -45,18 +45,18 @@ bool launch_fused_1nn_tile(IdxT* nearest_idx, const bool apply_sqrt = fused_1nn_apply_sqrt_at_pack(is_sqrt); - IdxT shape_x[2] = {m, k}; - IdxT stride_x[2] = {k, IdxT{1}}; - IdxT shape_y[2] = {n, k}; - IdxT stride_y[2] = {k, IdxT{1}}; - int64_t shape_xn = m; - int64_t stride_xn = 1; - int64_t shape_yn = n; - int64_t stride_yn = 1; - int64_t shape_idx = m; - int64_t stride_idx = 1; - int64_t shape_dist = m; - int64_t stride_dist = 1; + IdxT shape_x[2] = {m, k}; + IdxT stride_x[2] = {k, IdxT{1}}; + IdxT shape_y[2] = {n, k}; + IdxT stride_y[2] = {k, IdxT{1}}; + IdxT shape_xn = m; + IdxT stride_xn = IdxT{1}; + IdxT shape_yn = n; + IdxT stride_yn = IdxT{1}; + IdxT shape_idx = m; + IdxT stride_idx = IdxT{1}; + IdxT shape_dist = m; + IdxT stride_dist = IdxT{1}; IdxT M = m; IdxT N = n; @@ -88,17 +88,17 @@ bool launch_fused_1nn_tile(IdxT* nearest_idx, IdxT, IdxT, void*, - int64_t, - int64_t, + IdxT, + IdxT, void*, - int64_t, - int64_t, + IdxT, + IdxT, void*, - int64_t, - int64_t, + IdxT, + IdxT, void*, - int64_t, - int64_t, + IdxT, + IdxT, IdxT, IdxT, IdxT, From 88ddda8e6929b49423817c02a4cfc388ff20225c Mon Sep 17 00:00:00 2001 From: divyegala Date: Sat, 11 Jul 2026 00:12:11 +0000 Subject: [PATCH 16/82] using prebuilt cubins --- .../modules/generate_cutile_kernels.cmake | 54 ++++-- .../cutile/fused_1nn_cutile_matrix.json | 183 ++++++++++++++++-- 2 files changed, 207 insertions(+), 30 deletions(-) diff --git a/cpp/cmake/modules/generate_cutile_kernels.cmake b/cpp/cmake/modules/generate_cutile_kernels.cmake index 97e06e79aa..307ffec463 100644 --- a/cpp/cmake/modules/generate_cutile_kernels.cmake +++ b/cpp/cmake/modules/generate_cutile_kernels.cmake @@ -114,11 +114,21 @@ function(_cutile_generate_matrix_tiles_header header_path matrix_json_file) string(JSON _register GET "${_export_entry}" "register") if(_register STREQUAL "cubin") string(JSON _arch_tag GET "${_export_entry}" "arch_tag") - string( - APPEND - _arch_tile_aliases - "using fused_1nn_matrix_tile_${_arch_tag} = cutile_tile_config<${_entry_tile_m}, ${_entry_tile_n}, ${_entry_tile_k}>;\n" - ) + set(_arch_tile_value "${_entry_tile_m},${_entry_tile_n},${_entry_tile_k}") + if(DEFINED _arch_tile_alias_value_${_arch_tag}) + if(NOT "${_arch_tile_alias_value_${_arch_tag}}" STREQUAL "${_arch_tile_value}") + message(FATAL_ERROR "Conflicting cuTile tile geometry for ${_arch_tag}: " + "${_arch_tile_alias_value_${_arch_tag}} vs ${_arch_tile_value}" + ) + endif() + else() + set(_arch_tile_alias_value_${_arch_tag} "${_arch_tile_value}") + string( + APPEND + _arch_tile_aliases + "using fused_1nn_matrix_tile_${_arch_tag} = cutile_tile_config<${_entry_tile_m}, ${_entry_tile_n}, ${_entry_tile_k}>;\n" + ) + endif() endif() math(EXPR _export_idx "${_export_idx} + 1") endwhile() @@ -198,16 +208,30 @@ function(process_cutile_matrix_entry source_list_var) list(APPEND _python_args --bytecode-version "${bytecode_version}") endif() - add_custom_command( - OUTPUT "${_artifact_file}" - COMMAND "${Python3_EXECUTABLE}" "${_CUTILE_KERNEL_DIR}/${_CUTILE_EXPORT_SCRIPT}" - "${_artifact_file}" ${_python_args} - WORKING_DIRECTORY "${_CUTILE_KERNEL_DIR}" - DEPENDS "${_CUTILE_KERNEL_DIR}/${_CUTILE_EXPORT_SCRIPT}" - "${_CUTILE_KERNEL_DIR}/${_CUTILE_KERNEL_PYTHON}" - COMMENT "Exporting cuTile ${_CUTILE_KERNEL_BASENAME} ${output_format} ${data_type}" - VERBATIM - ) + if(DEFINED prebuilt_artifact AND NOT "${prebuilt_artifact}" STREQUAL "") + string(CONFIGURE "${prebuilt_artifact}" _prebuilt_artifact @ONLY) + if(NOT IS_ABSOLUTE "${_prebuilt_artifact}") + set(_prebuilt_artifact "${_CUTILE_KERNEL_DIR}/${_prebuilt_artifact}") + endif() + add_custom_command( + OUTPUT "${_artifact_file}" + COMMAND "${CMAKE_COMMAND}" -E copy_if_different "${_prebuilt_artifact}" "${_artifact_file}" + DEPENDS "${_prebuilt_artifact}" + COMMENT "Copying prebuilt cuTile ${_CUTILE_KERNEL_BASENAME} ${output_format} ${data_type}" + VERBATIM + ) + else() + add_custom_command( + OUTPUT "${_artifact_file}" + COMMAND "${Python3_EXECUTABLE}" "${_CUTILE_KERNEL_DIR}/${_CUTILE_EXPORT_SCRIPT}" + "${_artifact_file}" ${_python_args} + WORKING_DIRECTORY "${_CUTILE_KERNEL_DIR}" + DEPENDS "${_CUTILE_KERNEL_DIR}/${_CUTILE_EXPORT_SCRIPT}" + "${_CUTILE_KERNEL_DIR}/${_CUTILE_KERNEL_PYTHON}" + COMMENT "Exporting cuTile ${_CUTILE_KERNEL_BASENAME} ${output_format} ${data_type}" + VERBATIM + ) + endif() add_custom_command( OUTPUT "${_embedded_header}" diff --git a/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_cutile_matrix.json b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_cutile_matrix.json index 0aca5bdf14..27ca834921 100644 --- a/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_cutile_matrix.json +++ b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_cutile_matrix.json @@ -62,16 +62,6 @@ "cc_minor": 6, "arch_tag": "cutile_arch_8_6" }, - { - "output_format": "cubin", - "artifact_ext": "cubin", - "artifact_basename": "@data_type@_@metric_abbrev@_@index_abbrev@_@gpu_code@", - "register": "cubin", - "gpu_code": "sm_120", - "cc_major": 12, - "cc_minor": 0, - "arch_tag": "cutile_arch_12_0" - }, { "output_format": "tileir_bytecode", "artifact_ext": "tilebc", @@ -87,10 +77,6 @@ { "data_type": "half", "data_abbrev": "h" - }, - { - "data_type": "float", - "data_abbrev": "f" } ], "_metric": [ @@ -119,7 +105,113 @@ ], "_tile": [ { - "tile_m": 64, + "tile_m": 128, + "tile_n": 128, + "tile_k": 32 + } + ], + "_export": [ + { + "output_format": "cubin", + "artifact_ext": "cubin", + "artifact_basename": "@data_type@_@metric_abbrev@_@index_abbrev@_@gpu_code@", + "register": "cubin", + "gpu_code": "sm_90", + "cc_major": 9, + "cc_minor": 0, + "arch_tag": "cutile_arch_9_0" + }, + { + "output_format": "cubin", + "artifact_ext": "cubin", + "artifact_basename": "@data_type@_@metric_abbrev@_@index_abbrev@_@gpu_code@", + "register": "cubin", + "gpu_code": "sm_120", + "cc_major": 12, + "cc_minor": 0, + "arch_tag": "cutile_arch_12_0" + } + ] + }, + { + "_data": [ + { + "data_type": "float", + "data_abbrev": "f" + } + ], + "_metric": [ + { + "metric": "inner_product", + "metric_abbrev": "ip" + }, + { + "metric": "cosine_expanded", + "metric_abbrev": "cos" + } + ], + "_index": [ + { + "index_type": "int32", + "index_abbrev": "i32" + }, + { + "index_type": "int64", + "index_abbrev": "i64" + } + ], + "_tile": [ + { + "tile_m": 128, + "tile_n": 128, + "tile_k": 32 + } + ], + "_export": [ + { + "output_format": "cubin", + "artifact_ext": "cubin", + "artifact_basename": "@data_type@_@metric_abbrev@_@index_abbrev@_@gpu_code@", + "register": "cubin", + "gpu_code": "sm_90", + "cc_major": 9, + "cc_minor": 0, + "arch_tag": "cutile_arch_9_0" + }, + { + "output_format": "cubin", + "artifact_ext": "cubin", + "artifact_basename": "@data_type@_@metric_abbrev@_@index_abbrev@_@gpu_code@", + "register": "cubin", + "gpu_code": "sm_120", + "cc_major": 12, + "cc_minor": 0, + "arch_tag": "cutile_arch_12_0" + } + ] + }, + { + "_data": [ + { + "data_type": "float", + "data_abbrev": "f" + } + ], + "_metric": [ + { + "metric": "l2_expanded", + "metric_abbrev": "l2" + } + ], + "_index": [ + { + "index_type": "int64", + "index_abbrev": "i64" + } + ], + "_tile": [ + { + "tile_m": 128, "tile_n": 128, "tile_k": 32 } @@ -134,6 +226,67 @@ "cc_major": 9, "cc_minor": 0, "arch_tag": "cutile_arch_9_0" + }, + { + "output_format": "cubin", + "artifact_ext": "cubin", + "artifact_basename": "@data_type@_@metric_abbrev@_@index_abbrev@_@gpu_code@", + "register": "cubin", + "gpu_code": "sm_120", + "cc_major": 12, + "cc_minor": 0, + "arch_tag": "cutile_arch_12_0" + } + ] + }, + { + "_data": [ + { + "data_type": "float", + "data_abbrev": "f" + } + ], + "_metric": [ + { + "metric": "l2_expanded", + "metric_abbrev": "l2" + } + ], + "_index": [ + { + "index_type": "int32", + "index_abbrev": "i32" + } + ], + "_tile": [ + { + "tile_m": 128, + "tile_n": 128, + "tile_k": 32 + } + ], + "_export": [ + { + "output_format": "cubin", + "artifact_ext": "cubin", + "artifact_basename": "@data_type@_@metric_abbrev@_@index_abbrev@_@gpu_code@", + "register": "cubin", + "gpu_code": "sm_90", + "cc_major": 9, + "cc_minor": 0, + "arch_tag": "cutile_arch_9_0", + "prebuilt_artifact": "../../../../../../tileir13.4_cubin/export_cubins/fused_1nn_float_l2_i32_sm_90_128x128x32.cubin" + }, + { + "output_format": "cubin", + "artifact_ext": "cubin", + "artifact_basename": "@data_type@_@metric_abbrev@_@index_abbrev@_@gpu_code@", + "register": "cubin", + "gpu_code": "sm_120", + "cc_major": 12, + "cc_minor": 0, + "arch_tag": "cutile_arch_12_0", + "prebuilt_artifact": "../../../../../../tileir13.4_cubin/export_cubins/fused_1nn_float_l2_i32_sm_120_128x128x32.cubin" } ] } From 099324abcc5e86db565fc914843ffe89bf0a21e4 Mon Sep 17 00:00:00 2001 From: divyegala Date: Wed, 15 Jul 2026 04:24:11 +0000 Subject: [PATCH 17/82] use 13.4 tile --- .../modules/generate_cutile_kernels.cmake | 7 +- .../cutile/fused_1nn_cutile_matrix.json | 183 ++---------------- .../cutile/fused_1nn_kernel.py | 147 +++++++++++--- 3 files changed, 137 insertions(+), 200 deletions(-) diff --git a/cpp/cmake/modules/generate_cutile_kernels.cmake b/cpp/cmake/modules/generate_cutile_kernels.cmake index 307ffec463..8a626e2b6e 100644 --- a/cpp/cmake/modules/generate_cutile_kernels.cmake +++ b/cpp/cmake/modules/generate_cutile_kernels.cmake @@ -208,6 +208,11 @@ function(process_cutile_matrix_entry source_list_var) list(APPEND _python_args --bytecode-version "${bytecode_version}") endif() + set(_export_python_executable "${Python3_EXECUTABLE}") + if(DEFINED python_executable AND NOT "${python_executable}" STREQUAL "") + string(CONFIGURE "${python_executable}" _export_python_executable @ONLY) + endif() + if(DEFINED prebuilt_artifact AND NOT "${prebuilt_artifact}" STREQUAL "") string(CONFIGURE "${prebuilt_artifact}" _prebuilt_artifact @ONLY) if(NOT IS_ABSOLUTE "${_prebuilt_artifact}") @@ -223,7 +228,7 @@ function(process_cutile_matrix_entry source_list_var) else() add_custom_command( OUTPUT "${_artifact_file}" - COMMAND "${Python3_EXECUTABLE}" "${_CUTILE_KERNEL_DIR}/${_CUTILE_EXPORT_SCRIPT}" + COMMAND "${_export_python_executable}" "${_CUTILE_KERNEL_DIR}/${_CUTILE_EXPORT_SCRIPT}" "${_artifact_file}" ${_python_args} WORKING_DIRECTORY "${_CUTILE_KERNEL_DIR}" DEPENDS "${_CUTILE_KERNEL_DIR}/${_CUTILE_EXPORT_SCRIPT}" diff --git a/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_cutile_matrix.json b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_cutile_matrix.json index 27ca834921..a53b0b7efa 100644 --- a/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_cutile_matrix.json +++ b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_cutile_matrix.json @@ -1,5 +1,6 @@ [ { + "python_executable": "../../../../../../tileiras13.4/.venv/bin/python", "_data": [ { "data_type": "half", @@ -62,55 +63,6 @@ "cc_minor": 6, "arch_tag": "cutile_arch_8_6" }, - { - "output_format": "tileir_bytecode", - "artifact_ext": "tilebc", - "artifact_basename": "@data_type@_@metric_abbrev@_@index_abbrev@", - "register": "tileir", - "gpu_code": "sm_80", - "bytecode_version": "13.1" - } - ] - }, - { - "_data": [ - { - "data_type": "half", - "data_abbrev": "h" - } - ], - "_metric": [ - { - "metric": "inner_product", - "metric_abbrev": "ip" - }, - { - "metric": "l2_expanded", - "metric_abbrev": "l2" - }, - { - "metric": "cosine_expanded", - "metric_abbrev": "cos" - } - ], - "_index": [ - { - "index_type": "int32", - "index_abbrev": "i32" - }, - { - "index_type": "int64", - "index_abbrev": "i64" - } - ], - "_tile": [ - { - "tile_m": 128, - "tile_n": 128, - "tile_k": 32 - } - ], - "_export": [ { "output_format": "cubin", "artifact_ext": "cubin", @@ -122,19 +74,22 @@ "arch_tag": "cutile_arch_9_0" }, { - "output_format": "cubin", - "artifact_ext": "cubin", - "artifact_basename": "@data_type@_@metric_abbrev@_@index_abbrev@_@gpu_code@", - "register": "cubin", - "gpu_code": "sm_120", - "cc_major": 12, - "cc_minor": 0, - "arch_tag": "cutile_arch_12_0" + "output_format": "tileir_bytecode", + "artifact_ext": "tilebc", + "artifact_basename": "@data_type@_@metric_abbrev@_@index_abbrev@", + "register": "tileir", + "gpu_code": "sm_80", + "bytecode_version": "13.1" } ] }, { + "python_executable": "../../../../../../tileiras13.4/.venv/bin/python", "_data": [ + { + "data_type": "half", + "data_abbrev": "h" + }, { "data_type": "float", "data_abbrev": "f" @@ -145,6 +100,10 @@ "metric": "inner_product", "metric_abbrev": "ip" }, + { + "metric": "l2_expanded", + "metric_abbrev": "l2" + }, { "metric": "cosine_expanded", "metric_abbrev": "cos" @@ -168,16 +127,6 @@ } ], "_export": [ - { - "output_format": "cubin", - "artifact_ext": "cubin", - "artifact_basename": "@data_type@_@metric_abbrev@_@index_abbrev@_@gpu_code@", - "register": "cubin", - "gpu_code": "sm_90", - "cc_major": 9, - "cc_minor": 0, - "arch_tag": "cutile_arch_9_0" - }, { "output_format": "cubin", "artifact_ext": "cubin", @@ -189,105 +138,5 @@ "arch_tag": "cutile_arch_12_0" } ] - }, - { - "_data": [ - { - "data_type": "float", - "data_abbrev": "f" - } - ], - "_metric": [ - { - "metric": "l2_expanded", - "metric_abbrev": "l2" - } - ], - "_index": [ - { - "index_type": "int64", - "index_abbrev": "i64" - } - ], - "_tile": [ - { - "tile_m": 128, - "tile_n": 128, - "tile_k": 32 - } - ], - "_export": [ - { - "output_format": "cubin", - "artifact_ext": "cubin", - "artifact_basename": "@data_type@_@metric_abbrev@_@index_abbrev@_@gpu_code@", - "register": "cubin", - "gpu_code": "sm_90", - "cc_major": 9, - "cc_minor": 0, - "arch_tag": "cutile_arch_9_0" - }, - { - "output_format": "cubin", - "artifact_ext": "cubin", - "artifact_basename": "@data_type@_@metric_abbrev@_@index_abbrev@_@gpu_code@", - "register": "cubin", - "gpu_code": "sm_120", - "cc_major": 12, - "cc_minor": 0, - "arch_tag": "cutile_arch_12_0" - } - ] - }, - { - "_data": [ - { - "data_type": "float", - "data_abbrev": "f" - } - ], - "_metric": [ - { - "metric": "l2_expanded", - "metric_abbrev": "l2" - } - ], - "_index": [ - { - "index_type": "int32", - "index_abbrev": "i32" - } - ], - "_tile": [ - { - "tile_m": 128, - "tile_n": 128, - "tile_k": 32 - } - ], - "_export": [ - { - "output_format": "cubin", - "artifact_ext": "cubin", - "artifact_basename": "@data_type@_@metric_abbrev@_@index_abbrev@_@gpu_code@", - "register": "cubin", - "gpu_code": "sm_90", - "cc_major": 9, - "cc_minor": 0, - "arch_tag": "cutile_arch_9_0", - "prebuilt_artifact": "../../../../../../tileir13.4_cubin/export_cubins/fused_1nn_float_l2_i32_sm_90_128x128x32.cubin" - }, - { - "output_format": "cubin", - "artifact_ext": "cubin", - "artifact_basename": "@data_type@_@metric_abbrev@_@index_abbrev@_@gpu_code@", - "register": "cubin", - "gpu_code": "sm_120", - "cc_major": 12, - "cc_minor": 0, - "arch_tag": "cutile_arch_12_0", - "prebuilt_artifact": "../../../../../../tileir13.4_cubin/export_cubins/fused_1nn_float_l2_i32_sm_120_128x128x32.cubin" - } - ] } ] diff --git a/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_kernel.py b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_kernel.py index e1da35a9a4..97037fe455 100644 --- a/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_kernel.py +++ b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_kernel.py @@ -50,24 +50,105 @@ def make_kernel( is_l2 = metric == "l2_expanded" is_cos = metric == "cosine_expanded" - if gpu_code == "sm_100": - # Blackwell tcgen05 maps one logical row of the output tile to the reduction owner. - core_shape = (tile_m, tile_n) - best_shape = (tile_m, 1) - inner_reduction_axes = (1,) - outer_reduction_axes = (1,) - else: - # SM80/86/90/110/120 distribute eight consecutive N values across four threads, - # two per thread. - # Reducing in this physical layout keeps the score/index pair reduction local and avoids - # the slower generic argmin/argmax lowering over the full N dimension. - core_shape = (tile_m, tile_n // 8, 4, 2) - best_shape = (tile_m, 1, 4, 1) - inner_reduction_axes = (1, 3) - outer_reduction_axes = (2,) + @ct.kernel + def fused_1nn_argmin_kernel( + A, + B, + A_norm, + B_norm, + OutIdx, + OutDist, + M, + N, + K, + apply_sqrt, + store_idx, + tm: ConstInt, + tn: ConstInt, + tk: ConstInt, + ): + bidm = ct.bid(0) + + if is_ip: + best_dist = ct.full((tm,), -3.4e38, acc_dtype) + else: + best_dist = ct.full((tm,), 3.4e38, acc_dtype) + best_idx = ct.zeros((tm,), idx_dtype) + + num_tiles_k = ct.num_tiles(A, axis=1, shape=(tm, tk)) + num_tiles_n = ct.num_tiles(B, axis=0, shape=(tn, tk)) + zero_pad = ct.PaddingMode.ZERO + + for n in range(num_tiles_n): + accumulator = ct.full((tm, tn), 0, dtype=acc_dtype) + + for k in range(num_tiles_k): + dtype = ct.tfloat32 if A.dtype == ct.float32 else A.dtype + + a = ct.load( + A, index=(bidm, k), shape=(tm, tk), padding_mode=zero_pad + ).astype(dtype) + b_T = ct.load( + B, index=(n, k), shape=(tn, tk), padding_mode=zero_pad + ).astype(dtype) + + accumulator = ct.mma(a, ct.transpose(b_T), accumulator) + + if is_ip: + score = accumulator + elif is_l2 or is_cos: + a_norm = ct.load( + A_norm, index=(bidm,), shape=(tm,), padding_mode=zero_pad + ) + b_norm = ct.load( + B_norm, index=(n,), shape=(tn,), padding_mode=zero_pad + ) + if is_l2: + # L2 expanded: ||x||^2 + ||y||^2 - 2 * dot(x, y); norms are squared. + score = ( + a_norm[:, None] + b_norm[None, :] - (2.0 * accumulator) + ) + elif is_cos: + # Cosine expanded distance: 1 - dot / (||x|| * ||y||); norms are L2 (not squared). + denom = a_norm[:, None] * b_norm[None, :] + score = 1.0 - (accumulator / denom) + + # Only the final N-tile can include zero-padded centroid columns. + if n == num_tiles_n - 1: + col = ct.arange(tn, dtype=idx_dtype) + global_col = (n * tn + col).astype(idx_dtype) + valid = global_col < N + if is_ip: + score = ct.where(valid[None, :], score, -3.4e38) + else: + score = ct.where(valid[None, :], score, 3.4e38) + + if is_ip: + curr_best = ct.max(score, axis=1) + curr_idx = ct.argmax(score, axis=1) + update = curr_best > best_dist + else: + curr_best = ct.min(score, axis=1) + curr_idx = ct.argmin(score, axis=1) + update = curr_best < best_dist + best_dist = ct.where(update, curr_best, best_dist) + best_idx = ct.where( + update, (n * tn + curr_idx).astype(idx_dtype), best_idx + ) + + out_dist = best_dist + if is_l2: + out_dist = ct.where(apply_sqrt != 0, ct.sqrt(best_dist), best_dist) + if store_idx != 0: + ct.store(OutIdx, index=(bidm,), tile=best_idx) + ct.store( + OutDist, + index=(bidm,), + tile=out_dist.astype(out_dist_dtype), + ) @ct.kernel - def fused_1nn_kernel( + def fused_1nn_reduce_kernel( A, B, A_norm, @@ -85,6 +166,14 @@ def fused_1nn_kernel( ): bidm = ct.bid(0) + # SM80/86/90/110 distribute eight consecutive N values across four threads, + # two per thread. Reducing in this physical layout keeps score/index selection + # local and avoids slower generic argmin/argmax lowering on these targets. + core_shape = (tile_m, tile_n // 8, 4, 2) + best_shape = (tile_m, 1, 4, 1) + inner_reduction_axes = (1, 3) + outer_reduction_axes = (2,) + if is_ip: best_dist = ct.full(best_shape, -3.4e38, acc_dtype) neutral_dist = -3.4e38 @@ -172,25 +261,17 @@ def red_op(a_score, a_idx, b_score, b_idx): else: score = ct.where(valid[None, :], score, 3.4e38) + curr_idx = ct.arange(tn, dtype=idx_dtype).reshape(core_shape[1:])[ + None, ... + ] + curr_best, curr_idx = reduce_scores( + score.reshape(core_shape), curr_idx, inner_reduction_axes + ) if is_ip: - curr_idx = ct.arange(tn, dtype=idx_dtype).reshape( - core_shape[1:] - )[None, ...] - curr_best, curr_idx = reduce_scores( - score.reshape(core_shape), curr_idx, inner_reduction_axes - ) update = curr_best > best_dist - best_dist = ct.where(update, curr_best, best_dist) else: - curr_idx = ct.arange(tn, dtype=idx_dtype).reshape( - core_shape[1:] - )[None, ...] - curr_best, curr_idx = reduce_scores( - score.reshape(core_shape), curr_idx, inner_reduction_axes - ) update = curr_best < best_dist - best_dist = ct.where(update, curr_best, best_dist) - + best_dist = ct.where(update, curr_best, best_dist) best_idx = ct.where( update, (n * tn + curr_idx).astype(idx_dtype), best_idx ) @@ -210,7 +291,9 @@ def red_op(a_score, a_idx, b_score, b_idx): tile=out_dist.reshape((tm,)).astype(out_dist_dtype), ) - return fused_1nn_kernel + if gpu_code in ("sm_100", "sm_120"): + return fused_1nn_argmin_kernel + return fused_1nn_reduce_kernel def kernel_symbol( From fe77c616f5f51370d27b00f00b0560e4dc3129b7 Mon Sep 17 00:00:00 2001 From: divyegala Date: Sat, 18 Jul 2026 04:25:15 +0000 Subject: [PATCH 18/82] relaxed cubin, sm100 dispatch, new reduction strat --- cpp/CMakeLists.txt | 4 +- .../modules/generate_cutile_kernels.cmake | 3 + .../cuvs/detail/jit_lto/cutile_arch_tags.hpp | 6 + .../fused_distance_nn/fused_1nn_fragments.hpp | 9 +- cpp/src/cluster/detail/kmeans_balanced.cuh | 4 + .../cutile/export_fused_1nn.py | 46 +++++- .../cutile/fused_1nn_cutile_matrix.json | 44 +++++- .../cutile/fused_1nn_kernel.py | 132 ++++-------------- .../cutile/fused_1nn_planner.hpp | 44 ++++-- .../cutile/fused_1nn_tile.cu | 27 ++-- 10 files changed, 176 insertions(+), 143 deletions(-) diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt index 84979d05c1..ef0967d1a5 100644 --- a/cpp/CMakeLists.txt +++ b/cpp/CMakeLists.txt @@ -980,9 +980,9 @@ if(NOT BUILD_CPU_ONLY) MATRIX_JSON_FILE "${fused_1nn_cutile_dir}/fused_1nn_cutile_matrix.json" FRAGMENT_TAG_FORMAT_CUBIN - "cuvs::distance::detail::fragment_tag_fused_1nn_cubin, cuvs::detail::jit_lto::@arch_tag@>" + "cuvs::distance::detail::fragment_tag_fused_1nn_cubin, cuvs::distance::detail::@abi_tag@, cuvs::detail::jit_lto::@arch_tag@>" FRAGMENT_TAG_FORMAT_TILEIR - "cuvs::distance::detail::fragment_tag_fused_1nn_tileir>" + "cuvs::distance::detail::fragment_tag_fused_1nn_tileir, cuvs::distance::detail::@abi_tag@>" FRAGMENT_TAG_HEADER_FILES "" "" diff --git a/cpp/cmake/modules/generate_cutile_kernels.cmake b/cpp/cmake/modules/generate_cutile_kernels.cmake index 8a626e2b6e..bd1e3efa13 100644 --- a/cpp/cmake/modules/generate_cutile_kernels.cmake +++ b/cpp/cmake/modules/generate_cutile_kernels.cmake @@ -207,6 +207,9 @@ function(process_cutile_matrix_entry source_list_var) if(DEFINED bytecode_version AND NOT "${bytecode_version}" STREQUAL "") list(APPEND _python_args --bytecode-version "${bytecode_version}") endif() + if(DEFINED matrix_layout AND NOT "${matrix_layout}" STREQUAL "") + list(APPEND _python_args --matrix-layout "${matrix_layout}") + endif() set(_export_python_executable "${Python3_EXECUTABLE}") if(DEFINED python_executable AND NOT "${python_executable}" STREQUAL "") diff --git a/cpp/include/cuvs/detail/jit_lto/cutile_arch_tags.hpp b/cpp/include/cuvs/detail/jit_lto/cutile_arch_tags.hpp index 2c915a278b..1b9f58837c 100644 --- a/cpp/include/cuvs/detail/jit_lto/cutile_arch_tags.hpp +++ b/cpp/include/cuvs/detail/jit_lto/cutile_arch_tags.hpp @@ -29,6 +29,11 @@ struct cutile_arch_9_0 { static constexpr int cc_minor = 0; }; +struct cutile_arch_10_0 { + static constexpr int cc_major = 10; + static constexpr int cc_minor = 0; +}; + struct cutile_arch_12_0 { static constexpr int cc_major = 12; static constexpr int cc_minor = 0; @@ -39,6 +44,7 @@ inline bool is_embedded_cubin_arch(int cc_major, int cc_minor) if (cc_major == 8 && cc_minor == 0) { return true; } if (cc_major == 8 && cc_minor == 6) { return true; } if (cc_major == 9 && cc_minor == 0) { return true; } + if (cc_major == 10 && cc_minor == 0) { return true; } if (cc_major == 12 && cc_minor == 0) { return true; } return false; } diff --git a/cpp/include/cuvs/detail/jit_lto/fused_distance_nn/fused_1nn_fragments.hpp b/cpp/include/cuvs/detail/jit_lto/fused_distance_nn/fused_1nn_fragments.hpp index c6afe16b5c..fe27441b0f 100644 --- a/cpp/include/cuvs/detail/jit_lto/fused_distance_nn/fused_1nn_fragments.hpp +++ b/cpp/include/cuvs/detail/jit_lto/fused_distance_nn/fused_1nn_fragments.hpp @@ -15,6 +15,8 @@ namespace cuvs::distance::detail { struct metric_tag_ip {}; struct metric_tag_l2 {}; struct metric_tag_cos {}; +struct cutile_abi_strict {}; +struct cutile_abi_relaxed {}; template struct cutile_tile_config { @@ -97,13 +99,18 @@ template struct fragment_tag_fused_1nn_cubin { static constexpr int cc_major = ArchTag::cc_major; static constexpr int cc_minor = ArchTag::cc_minor; }; -template +template struct fragment_tag_fused_1nn_tileir {}; } // namespace cuvs::distance::detail diff --git a/cpp/src/cluster/detail/kmeans_balanced.cuh b/cpp/src/cluster/detail/kmeans_balanced.cuh index 007c462247..e4b4ae2026 100644 --- a/cpp/src/cluster/detail/kmeans_balanced.cuh +++ b/cpp/src/cluster/detail/kmeans_balanced.cuh @@ -279,6 +279,10 @@ auto calc_minibatch_size(const raft::resources& handle, const auto available_ws_size = std::min((free_ws_size * size_t{8}) / size_t{10}, size_t{1} << 29); + // A fused implementation may require no per-row temporary workspace. In that case, + // process the complete input rather than dividing the available workspace by zero. + if (mem_per_row == 0) { return std::make_tuple(n_rows, mem_per_row); } + IdxT minibatch_size = std::max(IdxT{1}, static_cast(available_ws_size / mem_per_row)); minibatch_size = raft::round_down_safe(minibatch_size, IdxT{64}); diff --git a/cpp/src/distance/detail/fused_distance_nn/cutile/export_fused_1nn.py b/cpp/src/distance/detail/fused_distance_nn/cutile/export_fused_1nn.py index c176d4b9ad..2b1e37b5fd 100644 --- a/cpp/src/distance/detail/fused_distance_nn/cutile/export_fused_1nn.py +++ b/cpp/src/distance/detail/fused_distance_nn/cutile/export_fused_1nn.py @@ -53,7 +53,12 @@ def _elem_stride_divisible_for_tma(elem_dtype) -> tuple[int, int]: return (16 // bytes_per_elem, 1) -def _cuvs_matrix_constraint(elem_dtype, *, index_dtype=ct.int32): +def _cuvs_matrix_constraint( + elem_dtype, + *, + index_dtype=ct.int32, + require_tma_friendly_pitch: bool = True, +): """Row-major device matrices for cuVS KMeans benchmarks. Assumes raft/cupy-style contiguous layout: stride[-1]==1, stride[0]==D, @@ -71,7 +76,11 @@ def _cuvs_matrix_constraint(elem_dtype, *, index_dtype=ct.int32): alias_groups=(), may_alias_internally=False, stride_constant=(None, 1), - stride_divisible_by=_elem_stride_divisible_for_tma(elem_dtype), + stride_divisible_by=( + _elem_stride_divisible_for_tma(elem_dtype) + if require_tma_friendly_pitch + else (1, 1) + ), shape_divisible_by=(1, 1), base_addr_divisible_by=16, ) @@ -94,8 +103,10 @@ def _cuvs_vector_constraint(elem_dtype, *, index_dtype=ct.int32): def _relaxed_matrix_constraint(elem_dtype): - """Deprecated alias; use _cuvs_matrix_constraint.""" - return _cuvs_matrix_constraint(elem_dtype) + """Deprecated alias for the arbitrary-row-pitch matrix constraint.""" + return _cuvs_matrix_constraint( + elem_dtype, require_tma_friendly_pitch=False + ) def _relaxed_vector_constraint(elem_dtype, *, tma_friendly: bool = False): @@ -111,17 +122,25 @@ def _kernel_signature( tile_m: int, tile_n: int, tile_k: int, + matrix_layout: str, ) -> KernelSignature: elem = _dtype_for(data_type) idx_dtype = _idx_dtype(index_type) - matrix = _cuvs_matrix_constraint(elem, index_dtype=idx_dtype) + matrix = _cuvs_matrix_constraint( + elem, + index_dtype=idx_dtype, + require_tma_friendly_pitch=matrix_layout == "strict", + ) norm_array = _cuvs_vector_constraint(elem, index_dtype=idx_dtype) idx_array = _cuvs_vector_constraint(idx_dtype, index_dtype=idx_dtype) dist_array = _cuvs_vector_constraint(elem, index_dtype=idx_dtype) abbrev = _data_abbrev(data_type) symbol = kernel_symbol( - abbrev, metric_abbrev(metric), index_abbrev(index_type) + abbrev, + metric_abbrev(metric), + index_abbrev(index_type), + matrix_layout, ) return KernelSignature( @@ -156,6 +175,7 @@ def export_binary( tile_n: int, tile_k: int, gpu_code: str, + matrix_layout: str = "strict", bytecode_version: str | None = None, ) -> str: kernel = make_kernel( @@ -168,7 +188,13 @@ def export_binary( gpu_code=gpu_code, ) signature = _kernel_signature( - data_type, metric, index_type, tile_m, tile_n, tile_k + data_type, + metric, + index_type, + tile_m, + tile_n, + tile_k, + matrix_layout, ) export_kwargs = { @@ -207,6 +233,11 @@ def main() -> int: default=DEFAULT_TILEIR_EXPORT_GPU_CODE, help="Target SM for cubin export, or compile hint for TileIR bytecode export", ) + parser.add_argument( + "--matrix-layout", + choices=("strict", "relaxed"), + default="strict", + ) parser.add_argument( "--bytecode-version", default=DEFAULT_TILEIR_BYTECODE_VERSION ) @@ -223,6 +254,7 @@ def main() -> int: tile_n=args.tile_n, tile_k=args.tile_k, gpu_code=args.gpu_code, + matrix_layout=args.matrix_layout, bytecode_version=args.bytecode_version, ) ) diff --git a/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_cutile_matrix.json b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_cutile_matrix.json index a53b0b7efa..1c5da9983b 100644 --- a/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_cutile_matrix.json +++ b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_cutile_matrix.json @@ -1,6 +1,18 @@ [ { "python_executable": "../../../../../../tileiras13.4/.venv/bin/python", + "_abi": [ + { + "matrix_layout": "strict", + "abi_abbrev": "strict", + "abi_tag": "cutile_abi_strict" + }, + { + "matrix_layout": "relaxed", + "abi_abbrev": "relaxed", + "abi_tag": "cutile_abi_relaxed" + } + ], "_data": [ { "data_type": "half", @@ -46,7 +58,17 @@ { "output_format": "cubin", "artifact_ext": "cubin", - "artifact_basename": "@data_type@_@metric_abbrev@_@index_abbrev@_@gpu_code@", + "artifact_basename": "@data_type@_@metric_abbrev@_@index_abbrev@_@abi_abbrev@_@gpu_code@", + "register": "cubin", + "gpu_code": "sm_100", + "cc_major": 10, + "cc_minor": 0, + "arch_tag": "cutile_arch_10_0" + }, + { + "output_format": "cubin", + "artifact_ext": "cubin", + "artifact_basename": "@data_type@_@metric_abbrev@_@index_abbrev@_@abi_abbrev@_@gpu_code@", "register": "cubin", "gpu_code": "sm_80", "cc_major": 8, @@ -56,7 +78,7 @@ { "output_format": "cubin", "artifact_ext": "cubin", - "artifact_basename": "@data_type@_@metric_abbrev@_@index_abbrev@_@gpu_code@", + "artifact_basename": "@data_type@_@metric_abbrev@_@index_abbrev@_@abi_abbrev@_@gpu_code@", "register": "cubin", "gpu_code": "sm_86", "cc_major": 8, @@ -66,7 +88,7 @@ { "output_format": "cubin", "artifact_ext": "cubin", - "artifact_basename": "@data_type@_@metric_abbrev@_@index_abbrev@_@gpu_code@", + "artifact_basename": "@data_type@_@metric_abbrev@_@index_abbrev@_@abi_abbrev@_@gpu_code@", "register": "cubin", "gpu_code": "sm_90", "cc_major": 9, @@ -76,7 +98,7 @@ { "output_format": "tileir_bytecode", "artifact_ext": "tilebc", - "artifact_basename": "@data_type@_@metric_abbrev@_@index_abbrev@", + "artifact_basename": "@data_type@_@metric_abbrev@_@index_abbrev@_@abi_abbrev@", "register": "tileir", "gpu_code": "sm_80", "bytecode_version": "13.1" @@ -85,6 +107,18 @@ }, { "python_executable": "../../../../../../tileiras13.4/.venv/bin/python", + "_abi": [ + { + "matrix_layout": "strict", + "abi_abbrev": "strict", + "abi_tag": "cutile_abi_strict" + }, + { + "matrix_layout": "relaxed", + "abi_abbrev": "relaxed", + "abi_tag": "cutile_abi_relaxed" + } + ], "_data": [ { "data_type": "half", @@ -130,7 +164,7 @@ { "output_format": "cubin", "artifact_ext": "cubin", - "artifact_basename": "@data_type@_@metric_abbrev@_@index_abbrev@_@gpu_code@", + "artifact_basename": "@data_type@_@metric_abbrev@_@index_abbrev@_@abi_abbrev@_@gpu_code@", "register": "cubin", "gpu_code": "sm_120", "cc_major": 12, diff --git a/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_kernel.py b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_kernel.py index 97037fe455..2c139831cb 100644 --- a/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_kernel.py +++ b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_kernel.py @@ -49,103 +49,16 @@ def make_kernel( is_ip = metric == "inner_product" is_l2 = metric == "l2_expanded" is_cos = metric == "cosine_expanded" - - @ct.kernel - def fused_1nn_argmin_kernel( - A, - B, - A_norm, - B_norm, - OutIdx, - OutDist, - M, - N, - K, - apply_sqrt, - store_idx, - tm: ConstInt, - tn: ConstInt, - tk: ConstInt, - ): - bidm = ct.bid(0) - - if is_ip: - best_dist = ct.full((tm,), -3.4e38, acc_dtype) - else: - best_dist = ct.full((tm,), 3.4e38, acc_dtype) - best_idx = ct.zeros((tm,), idx_dtype) - - num_tiles_k = ct.num_tiles(A, axis=1, shape=(tm, tk)) - num_tiles_n = ct.num_tiles(B, axis=0, shape=(tn, tk)) - zero_pad = ct.PaddingMode.ZERO - - for n in range(num_tiles_n): - accumulator = ct.full((tm, tn), 0, dtype=acc_dtype) - - for k in range(num_tiles_k): - dtype = ct.tfloat32 if A.dtype == ct.float32 else A.dtype - - a = ct.load( - A, index=(bidm, k), shape=(tm, tk), padding_mode=zero_pad - ).astype(dtype) - b_T = ct.load( - B, index=(n, k), shape=(tn, tk), padding_mode=zero_pad - ).astype(dtype) - - accumulator = ct.mma(a, ct.transpose(b_T), accumulator) - - if is_ip: - score = accumulator - elif is_l2 or is_cos: - a_norm = ct.load( - A_norm, index=(bidm,), shape=(tm,), padding_mode=zero_pad - ) - b_norm = ct.load( - B_norm, index=(n,), shape=(tn,), padding_mode=zero_pad - ) - if is_l2: - # L2 expanded: ||x||^2 + ||y||^2 - 2 * dot(x, y); norms are squared. - score = ( - a_norm[:, None] + b_norm[None, :] - (2.0 * accumulator) - ) - elif is_cos: - # Cosine expanded distance: 1 - dot / (||x|| * ||y||); norms are L2 (not squared). - denom = a_norm[:, None] * b_norm[None, :] - score = 1.0 - (accumulator / denom) - - # Only the final N-tile can include zero-padded centroid columns. - if n == num_tiles_n - 1: - col = ct.arange(tn, dtype=idx_dtype) - global_col = (n * tn + col).astype(idx_dtype) - valid = global_col < N - if is_ip: - score = ct.where(valid[None, :], score, -3.4e38) - else: - score = ct.where(valid[None, :], score, 3.4e38) - - if is_ip: - curr_best = ct.max(score, axis=1) - curr_idx = ct.argmax(score, axis=1) - update = curr_best > best_dist - else: - curr_best = ct.min(score, axis=1) - curr_idx = ct.argmin(score, axis=1) - update = curr_best < best_dist - best_dist = ct.where(update, curr_best, best_dist) - best_idx = ct.where( - update, (n * tn + curr_idx).astype(idx_dtype), best_idx - ) - - out_dist = best_dist - if is_l2: - out_dist = ct.where(apply_sqrt != 0, ct.sqrt(best_dist), best_dist) - if store_idx != 0: - ct.store(OutIdx, index=(bidm,), tile=best_idx) - ct.store( - OutDist, - index=(bidm,), - tile=out_dist.astype(out_dist_dtype), - ) + items_per_thread = 4 if gpu_code in ("sm_100", "sm_120") else 2 + core_shape = ( + tile_m, + tile_n // (4 * items_per_thread), + 4, + items_per_thread, + ) + best_shape = (tile_m, 1, 4, 1) + inner_reduction_axes = (1, 3) + outer_reduction_axes = (2,) @ct.kernel def fused_1nn_reduce_kernel( @@ -166,13 +79,10 @@ def fused_1nn_reduce_kernel( ): bidm = ct.bid(0) - # SM80/86/90/110 distribute eight consecutive N values across four threads, - # two per thread. Reducing in this physical layout keeps score/index selection - # local and avoids slower generic argmin/argmax lowering on these targets. - core_shape = (tile_m, tile_n // 8, 4, 2) - best_shape = (tile_m, 1, 4, 1) - inner_reduction_axes = (1, 3) - outer_reduction_axes = (2,) + # Reduce groups and per-thread items inside each N tile, carry four + # partial winners across N tiles, then reduce those winners once. + # Blackwell uses four items per logical thread slot; earlier targets + # retain the existing two-item grouping. if is_ip: best_dist = ct.full(best_shape, -3.4e38, acc_dtype) @@ -291,16 +201,22 @@ def red_op(a_score, a_idx, b_score, b_idx): tile=out_dist.reshape((tm,)).astype(out_dist_dtype), ) - if gpu_code in ("sm_100", "sm_120"): - return fused_1nn_argmin_kernel return fused_1nn_reduce_kernel def kernel_symbol( - data_abbrev: str, metric_abbrev: str, index_abbrev: str + data_abbrev: str, + metric_abbrev: str, + index_abbrev: str, + matrix_layout: str = "strict", ) -> str: """Must stay in sync with fused_1nn_kernel_entrypoint() in fused_1nn_planner.hpp.""" - return f"fused_1nn_{data_abbrev}_{metric_abbrev}_{index_abbrev}" + base = f"fused_1nn_{data_abbrev}_{metric_abbrev}_{index_abbrev}" + if matrix_layout == "strict": + return base + if matrix_layout == "relaxed": + return f"{base}_relaxed" + raise ValueError(f"Unsupported matrix layout {matrix_layout!r}") def metric_abbrev(metric: str) -> str: diff --git a/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_planner.hpp b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_planner.hpp index 41185e5a61..a80f4a4eb3 100644 --- a/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_planner.hpp +++ b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_planner.hpp @@ -17,38 +17,47 @@ namespace cuvs::distance::detail { /** Must match kernel_symbol() in fused_1nn_kernel.py (export uses with_symbol). */ -template +template inline const char* fused_1nn_kernel_entrypoint() { - constexpr bool is_i32 = std::is_same_v; - constexpr bool is_i64 = std::is_same_v; + constexpr bool is_i32 = std::is_same_v; + constexpr bool is_i64 = std::is_same_v; + constexpr bool is_relaxed = std::is_same_v; static_assert(is_i32 || is_i64, "unsupported fused 1-NN cuTile index width"); + static_assert(is_relaxed || std::is_same_v, + "unsupported fused 1-NN cuTile ABI"); if constexpr (std::is_same_v && std::is_same_v) { - return is_i32 ? "fused_1nn_f_ip_i32" : "fused_1nn_f_ip_i64"; + return is_i32 ? (is_relaxed ? "fused_1nn_f_ip_i32_relaxed" : "fused_1nn_f_ip_i32") + : (is_relaxed ? "fused_1nn_f_ip_i64_relaxed" : "fused_1nn_f_ip_i64"); } else if constexpr (std::is_same_v && std::is_same_v) { - return is_i32 ? "fused_1nn_f_l2_i32" : "fused_1nn_f_l2_i64"; + return is_i32 ? (is_relaxed ? "fused_1nn_f_l2_i32_relaxed" : "fused_1nn_f_l2_i32") + : (is_relaxed ? "fused_1nn_f_l2_i64_relaxed" : "fused_1nn_f_l2_i64"); } else if constexpr (std::is_same_v && std::is_same_v) { - return is_i32 ? "fused_1nn_f_cos_i32" : "fused_1nn_f_cos_i64"; + return is_i32 ? (is_relaxed ? "fused_1nn_f_cos_i32_relaxed" : "fused_1nn_f_cos_i32") + : (is_relaxed ? "fused_1nn_f_cos_i64_relaxed" : "fused_1nn_f_cos_i64"); } else if constexpr (std::is_same_v && std::is_same_v) { - return is_i32 ? "fused_1nn_h_ip_i32" : "fused_1nn_h_ip_i64"; + return is_i32 ? (is_relaxed ? "fused_1nn_h_ip_i32_relaxed" : "fused_1nn_h_ip_i32") + : (is_relaxed ? "fused_1nn_h_ip_i64_relaxed" : "fused_1nn_h_ip_i64"); } else if constexpr (std::is_same_v && std::is_same_v) { - return is_i32 ? "fused_1nn_h_l2_i32" : "fused_1nn_h_l2_i64"; + return is_i32 ? (is_relaxed ? "fused_1nn_h_l2_i32_relaxed" : "fused_1nn_h_l2_i32") + : (is_relaxed ? "fused_1nn_h_l2_i64_relaxed" : "fused_1nn_h_l2_i64"); } else if constexpr (std::is_same_v && std::is_same_v) { - return is_i32 ? "fused_1nn_h_cos_i32" : "fused_1nn_h_cos_i64"; + return is_i32 ? (is_relaxed ? "fused_1nn_h_cos_i32_relaxed" : "fused_1nn_h_cos_i32") + : (is_relaxed ? "fused_1nn_h_cos_i64_relaxed" : "fused_1nn_h_cos_i64"); } else { static_assert(sizeof(DataTag) == 0, "unsupported fused 1-NN cuTile data/metric combination"); return ""; } } -template +template struct Fused1nnTilePlanner : TileAlgorithmPlanner { using DataTag = fused_1nn_data_tag_t; using MetricTag = fused_1nn_metric_tag_t; @@ -57,7 +66,7 @@ struct Fused1nnTilePlanner : TileAlgorithmPlanner { inline static LauncherJitCache launcher_jit_cache{}; Fused1nnTilePlanner() - : TileAlgorithmPlanner(fused_1nn_kernel_entrypoint(), + : TileAlgorithmPlanner(fused_1nn_kernel_entrypoint(), launcher_jit_cache) { } @@ -66,6 +75,7 @@ struct Fused1nnTilePlanner : TileAlgorithmPlanner { */ void add_entrypoint() { + using cuvs::detail::jit_lto::cutile_arch_10_0; using cuvs::detail::jit_lto::cutile_arch_12_0; using cuvs::detail::jit_lto::cutile_arch_8_0; using cuvs::detail::jit_lto::cutile_arch_8_6; @@ -75,28 +85,38 @@ struct Fused1nnTilePlanner : TileAlgorithmPlanner { MetricTag, IndexTag, fused_1nn_matrix_tile_cutile_arch_8_0, + AbiTag, cutile_arch_8_0>>(); this->add_static_fragment>(); this->add_static_fragment>(); + this->add_static_fragment>(); this->add_static_fragment>(); } void add_tileir_fallback() { this->add_static_tileir_fragment< - fragment_tag_fused_1nn_tileir>(); + fragment_tag_fused_1nn_tileir>(); } }; diff --git a/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_tile.cu b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_tile.cu index 275ae33403..5c5edf35ce 100644 --- a/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_tile.cu +++ b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_tile.cu @@ -19,7 +19,7 @@ namespace detail { namespace { -template +template bool launch_fused_1nn_tile(IdxT* nearest_idx, DataT* nearest_dist, const DataT* x, @@ -36,7 +36,7 @@ bool launch_fused_1nn_tile(IdxT* nearest_idx, if (nearest_dist == nullptr) { return false; } - Fused1nnTilePlanner planner; + Fused1nnTilePlanner planner; planner.add_entrypoint(); planner.add_tileir_fallback(); const CutileTileConfig tile_cfg = planner.tile_config(); @@ -139,7 +139,7 @@ bool launch_fused_1nn_tile(IdxT* nearest_idx, return true; } -template +template bool try_fused_1nn_tile_dispatch(IdxT* nearest_idx, DataT* nearest_dist, const DataT* x, @@ -155,16 +155,22 @@ bool try_fused_1nn_tile_dispatch(IdxT* nearest_idx, { switch (metric) { case cuvs::distance::DistanceType::InnerProduct: - return launch_fused_1nn_tile( + return launch_fused_1nn_tile( nearest_idx, nearest_dist, x, y, xn, yn, m, n, k, is_sqrt, stream); case cuvs::distance::DistanceType::L2Expanded: - return launch_fused_1nn_tile( + return launch_fused_1nn_tile( nearest_idx, nearest_dist, x, y, xn, yn, m, n, k, is_sqrt, stream); case cuvs::distance::DistanceType::L2SqrtExpanded: - return launch_fused_1nn_tile( + return launch_fused_1nn_tile( nearest_idx, nearest_dist, x, y, xn, yn, m, n, k, is_sqrt, stream); case cuvs::distance::DistanceType::CosineExpanded: - return launch_fused_1nn_tile( + return launch_fused_1nn_tile( nearest_idx, nearest_dist, x, y, xn, yn, m, n, k, is_sqrt, stream); default: return false; } @@ -188,7 +194,12 @@ bool try_fused_1nn_tile(IdxT* nearest_idx, cudaStream_t stream) { if (!cuvs::detail::jit_lto::cutile_launch_available_on_current_device()) { return false; } - return try_fused_1nn_tile_dispatch( + constexpr IdxT tma_pitch_elements = IdxT{16 / sizeof(DataT)}; + if (k % tma_pitch_elements == 0) { + return try_fused_1nn_tile_dispatch( + nearest_idx, nearest_dist, x, y, xn, yn, m, n, k, metric, is_sqrt, stream); + } + return try_fused_1nn_tile_dispatch( nearest_idx, nearest_dist, x, y, xn, yn, m, n, k, metric, is_sqrt, stream); } From 7e086b16b0f6d95dcdb008278ef27755db1d6fd6 Mon Sep 17 00:00:00 2001 From: divyegala Date: Wed, 22 Jul 2026 01:57:13 +0000 Subject: [PATCH 19/82] split tile for relaxed and strict sm120 abi --- .../cutile/fused_1nn_cutile_matrix.json | 17 ++++++++--------- 1 file changed, 8 insertions(+), 9 deletions(-) diff --git a/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_cutile_matrix.json b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_cutile_matrix.json index 1c5da9983b..7fcb786833 100644 --- a/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_cutile_matrix.json +++ b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_cutile_matrix.json @@ -5,12 +5,18 @@ { "matrix_layout": "strict", "abi_abbrev": "strict", - "abi_tag": "cutile_abi_strict" + "abi_tag": "cutile_abi_strict", + "tile_m": 128, + "tile_n": 64, + "tile_k": 64 }, { "matrix_layout": "relaxed", "abi_abbrev": "relaxed", - "abi_tag": "cutile_abi_relaxed" + "abi_tag": "cutile_abi_relaxed", + "tile_m": 128, + "tile_n": 128, + "tile_k": 32 } ], "_data": [ @@ -153,13 +159,6 @@ "index_abbrev": "i64" } ], - "_tile": [ - { - "tile_m": 128, - "tile_n": 128, - "tile_k": 32 - } - ], "_export": [ { "output_format": "cubin", From 82d14a6fe1863e2649c075ebc132fb9265be6db6 Mon Sep 17 00:00:00 2001 From: divyegala Date: Wed, 22 Jul 2026 03:33:25 +0000 Subject: [PATCH 20/82] remove i64, batch i32 to save binary size --- cpp/src/cluster/detail/kmeans_balanced.cuh | 7 +- cpp/src/cluster/detail/kmeans_common.cuh | 5 +- .../detail/minClusterDistanceCompute.cu | 7 +- cpp/src/distance/detail/fused_distance_nn.cuh | 2 +- .../cutile/fused_1nn_cutile_matrix.json | 8 --- .../cutile/fused_1nn_tile.cu | 70 +++++++++++++++++-- .../cutile/fused_1nn_tile.hpp | 2 + 7 files changed, 83 insertions(+), 18 deletions(-) diff --git a/cpp/src/cluster/detail/kmeans_balanced.cuh b/cpp/src/cluster/detail/kmeans_balanced.cuh index e4b4ae2026..7e98a26351 100644 --- a/cpp/src/cluster/detail/kmeans_balanced.cuh +++ b/cpp/src/cluster/detail/kmeans_balanced.cuh @@ -251,7 +251,12 @@ auto calc_minibatch_size(const raft::resources& handle, case distance::DistanceType::L2SqrtExpanded: case distance::DistanceType::InnerProduct: { switch (use_fused(handle, n_rows, n_clusters, dim, metric)) { - case FusedDistancePath::FusedCutile: break; + case FusedDistancePath::FusedCutile: + if constexpr (std::is_same_v) { + // cuTile computes labels with i32 and widens them after each launch. + mem_per_row += sizeof(int); + } + break; case FusedDistancePath::FusedCutlass: // fusedDistanceNNMinReduce CUTLASS fallback: mutex workspace + scratch KVP per row. mem_per_row += sizeof(int); diff --git a/cpp/src/cluster/detail/kmeans_common.cuh b/cpp/src/cluster/detail/kmeans_common.cuh index f0d9bde801..d720df00b3 100644 --- a/cpp/src/cluster/detail/kmeans_common.cuh +++ b/cpp/src/cluster/detail/kmeans_common.cuh @@ -105,7 +105,10 @@ FusedDistancePath use_fused( if constexpr (is_cutile_fused_data_type_v) { if constexpr (cuvs::detail::jit_lto::library_built_with_cutile()) { - if (cuvs::detail::jit_lto::cutile_launch_available_on_current_device()) { + const bool dimensions_fit_i32 = n <= static_cast(std::numeric_limits::max()) && + k <= static_cast(std::numeric_limits::max()); + if (dimensions_fit_i32 && + cuvs::detail::jit_lto::cutile_launch_available_on_current_device()) { return FusedDistancePath::FusedCutile; } } diff --git a/cpp/src/cluster/detail/minClusterDistanceCompute.cu b/cpp/src/cluster/detail/minClusterDistanceCompute.cu index d01f48fc0a..aac3fae307 100644 --- a/cpp/src/cluster/detail/minClusterDistanceCompute.cu +++ b/cpp/src/cluster/detail/minClusterDistanceCompute.cu @@ -89,6 +89,9 @@ void minClusterAndDistanceCompute(raft::resources const& handle, temp_kvp.resize(n_samples, stream); cutlass_kvp_scratch = temp_kvp.data(); workspace.resize(sizeof(int) * n_samples, stream); + } else if constexpr (std::is_same_v) { + // The cuTile kernel uses i32 internally and widens labels after the launch. + workspace.resize(sizeof(int) * static_cast(n_samples), stream); } cuvs::distance::fusedDistanceNNMinReduce( @@ -101,7 +104,9 @@ void minClusterAndDistanceCompute(raft::resources const& handle, n_samples, n_clusters, n_features, - needs_fused_mutex_workspace(fused_path) ? (void*)workspace.data() : nullptr, + needs_fused_mutex_workspace(fused_path) || std::is_same_v + ? (void*)workspace.data() + : nullptr, metric != cuvs::distance::DistanceType::L2Expanded, true, true, diff --git a/cpp/src/distance/detail/fused_distance_nn.cuh b/cpp/src/distance/detail/fused_distance_nn.cuh index a2ec5422dd..ec662fcbfa 100644 --- a/cpp/src/distance/detail/fused_distance_nn.cuh +++ b/cpp/src/distance/detail/fused_distance_nn.cuh @@ -57,7 +57,7 @@ void fusedDistanceNNImpl(IdxT* nearest_idx, if constexpr (is_fused_1nn_cutile_data_v) { if constexpr (cuvs::detail::jit_lto::library_built_with_cutile()) { if (try_fused_1nn_tile( - nearest_idx, nearest_dist, x, y, xn, yn, m, n, k, metric, sqrt, stream)) { + nearest_idx, nearest_dist, x, y, xn, yn, m, n, k, metric, sqrt, workspace, stream)) { return; } } diff --git a/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_cutile_matrix.json b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_cutile_matrix.json index 7fcb786833..c6c13f17b9 100644 --- a/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_cutile_matrix.json +++ b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_cutile_matrix.json @@ -47,10 +47,6 @@ { "index_type": "int32", "index_abbrev": "i32" - }, - { - "index_type": "int64", - "index_abbrev": "i64" } ], "_tile": [ @@ -153,10 +149,6 @@ { "index_type": "int32", "index_abbrev": "i32" - }, - { - "index_type": "int64", - "index_abbrev": "i64" } ], "_export": [ diff --git a/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_tile.cu b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_tile.cu index 5c5edf35ce..1a4dd24130 100644 --- a/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_tile.cu +++ b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_tile.cu @@ -9,8 +9,13 @@ #include #include +#include +#include #include +#include +#include +#include #include namespace cuvs { @@ -74,7 +79,7 @@ bool launch_fused_1nn_tile(IdxT* nearest_idx, void* dist_ptr = nearest_dist; const int tile_m = tile_cfg.tile_m; - dim3 grid((m + tile_m - 1) / tile_m, 1, 1); + dim3 grid((static_cast(m) + tile_m - 1) / tile_m, 1, 1); dim3 block(1, 1, 1); using fused_1nn_cutile_kernel_t = void(void*, @@ -191,16 +196,68 @@ bool try_fused_1nn_tile(IdxT* nearest_idx, IdxT k, cuvs::distance::DistanceType metric, bool is_sqrt, + void* index_workspace, cudaStream_t stream) { if (!cuvs::detail::jit_lto::cutile_launch_available_on_current_device()) { return false; } - constexpr IdxT tma_pitch_elements = IdxT{16 / sizeof(DataT)}; - if (k % tma_pitch_elements == 0) { - return try_fused_1nn_tile_dispatch( + static_assert(std::is_same_v || std::is_same_v); + + if constexpr (std::is_same_v) { + constexpr int tma_pitch_elements = 16 / sizeof(DataT); + if (k % tma_pitch_elements == 0) { + return try_fused_1nn_tile_dispatch( + nearest_idx, nearest_dist, x, y, xn, yn, m, n, k, metric, is_sqrt, stream); + } + return try_fused_1nn_tile_dispatch( nearest_idx, nearest_dist, x, y, xn, yn, m, n, k, metric, is_sqrt, stream); + } else { + constexpr int64_t max_i32 = std::numeric_limits::max(); + if (n > max_i32 || k > max_i32) { return false; } + if (nearest_idx != nullptr && index_workspace == nullptr) { return false; } + + auto* tmp_idx = static_cast(index_workspace); + for (int64_t offset = 0; offset < m; offset += max_i32) { + const int batch_m = static_cast(std::min(max_i32, m - offset)); + const auto* batch_x = x + static_cast(offset) * static_cast(k); + const auto* batch_xn = xn == nullptr ? nullptr : xn + offset; + auto* batch_dist = nearest_dist == nullptr ? nullptr : nearest_dist + offset; + + constexpr int tma_pitch_elements = 16 / sizeof(DataT); + const bool launched = + k % tma_pitch_elements == 0 + ? try_fused_1nn_tile_dispatch(tmp_idx, + batch_dist, + batch_x, + y, + batch_xn, + yn, + batch_m, + static_cast(n), + static_cast(k), + metric, + is_sqrt, + stream) + : try_fused_1nn_tile_dispatch(tmp_idx, + batch_dist, + batch_x, + y, + batch_xn, + yn, + batch_m, + static_cast(n), + static_cast(k), + metric, + is_sqrt, + stream); + if (!launched) { return false; } + + if (nearest_idx != nullptr) { + raft::linalg::unaryOp( + nearest_idx + offset, tmp_idx, batch_m, raft::cast_op{}, stream); + } + } + return true; } - return try_fused_1nn_tile_dispatch( - nearest_idx, nearest_dist, x, y, xn, yn, m, n, k, metric, is_sqrt, stream); } #define CUVS_INST_TRY_FUSED_1NN_TILE(DataT, IdxT) \ @@ -215,6 +272,7 @@ bool try_fused_1nn_tile(IdxT* nearest_idx, IdxT, \ cuvs::distance::DistanceType, \ bool, \ + void*, \ cudaStream_t) CUVS_INST_TRY_FUSED_1NN_TILE(float, int); diff --git a/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_tile.hpp b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_tile.hpp index 4c0964631c..953f78f63e 100644 --- a/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_tile.hpp +++ b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_tile.hpp @@ -38,6 +38,7 @@ bool try_fused_1nn_tile(IdxT* nearest_idx, IdxT k, cuvs::distance::DistanceType metric, bool is_sqrt, + void* index_workspace, cudaStream_t stream); #else template @@ -52,6 +53,7 @@ bool try_fused_1nn_tile(IdxT*, IdxT, cuvs::distance::DistanceType, bool, + void*, cudaStream_t) { return false; From b0efee67a0ad1af8d2046932e125d9460f387273 Mon Sep 17 00:00:00 2001 From: divyegala Date: Tue, 28 Jul 2026 16:00:23 +0000 Subject: [PATCH 21/82] working through --- cpp/CMakeLists.txt | 4 +- .../modules/generate_cutile_kernels.cmake | 98 ++++++++++------ .../fused_distance_nn/fused_1nn_fragments.hpp | 58 +--------- .../cutile/export_fused_1nn.py | 3 +- .../cutile/fused_1nn_cutile_matrix.json | 54 +++------ .../cutile/fused_1nn_kernel.py | 66 ++++------- .../cutile/fused_1nn_planner.hpp | 108 +++++++----------- .../cutile/fused_1nn_tile.cu | 51 ++++----- 8 files changed, 181 insertions(+), 261 deletions(-) diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt index ef0967d1a5..3617e64e9c 100644 --- a/cpp/CMakeLists.txt +++ b/cpp/CMakeLists.txt @@ -980,9 +980,9 @@ if(NOT BUILD_CPU_ONLY) MATRIX_JSON_FILE "${fused_1nn_cutile_dir}/fused_1nn_cutile_matrix.json" FRAGMENT_TAG_FORMAT_CUBIN - "cuvs::distance::detail::fragment_tag_fused_1nn_cubin, cuvs::distance::detail::@abi_tag@, cuvs::detail::jit_lto::@arch_tag@>" + "cuvs::distance::detail::fragment_tag_fused_1nn_cubin, cuvs::distance::detail::@abi_tag@, cuvs::detail::jit_lto::@arch_tag@>" FRAGMENT_TAG_FORMAT_TILEIR - "cuvs::distance::detail::fragment_tag_fused_1nn_tileir, cuvs::distance::detail::@abi_tag@>" + "cuvs::distance::detail::fragment_tag_fused_1nn_tileir, cuvs::distance::detail::@abi_tag@>" FRAGMENT_TAG_HEADER_FILES "" "" diff --git a/cpp/cmake/modules/generate_cutile_kernels.cmake b/cpp/cmake/modules/generate_cutile_kernels.cmake index bd1e3efa13..f06ad3801f 100644 --- a/cpp/cmake/modules/generate_cutile_kernels.cmake +++ b/cpp/cmake/modules/generate_cutile_kernels.cmake @@ -92,45 +92,78 @@ function(_cutile_kernels_setup) ) endfunction() +macro(_cutile_append_matrix_tile_aliases entry abi_abbrev tile_m tile_n tile_k) + string(JSON _cutile_export_len LENGTH "${entry}" "_export") + set(_cutile_export_idx 0) + while(_cutile_export_idx LESS _cutile_export_len) + string(JSON _cutile_export_entry GET "${entry}" "_export" "${_cutile_export_idx}") + string(JSON _cutile_register GET "${_cutile_export_entry}" "register") + if(_cutile_register STREQUAL "cubin") + string(JSON _cutile_arch_tag GET "${_cutile_export_entry}" "arch_tag") + set(_cutile_alias_suffix "${_cutile_arch_tag}_${abi_abbrev}") + elseif(_cutile_register STREQUAL "tileir") + set(_cutile_alias_suffix "tileir_${abi_abbrev}") + else() + message(FATAL_ERROR "Unknown cuTile register kind '${_cutile_register}'") + endif() + + set(_cutile_tile_value "${tile_m},${tile_n},${tile_k}") + if(DEFINED _tile_alias_value_${_cutile_alias_suffix}) + if(NOT "${_tile_alias_value_${_cutile_alias_suffix}}" STREQUAL "${_cutile_tile_value}") + message(FATAL_ERROR "Conflicting cuTile tile geometry for ${_cutile_alias_suffix}: " + "${_tile_alias_value_${_cutile_alias_suffix}} vs ${_cutile_tile_value}" + ) + endif() + else() + set(_tile_alias_value_${_cutile_alias_suffix} "${_cutile_tile_value}") + string( + APPEND + _tile_aliases + "using fused_1nn_matrix_tile_${_cutile_alias_suffix} = cutile_tile_config<${tile_m}, ${tile_n}, ${tile_k}>;\n" + ) + endif() + math(EXPR _cutile_export_idx "${_cutile_export_idx} + 1") + endwhile() +endmacro() + function(_cutile_generate_matrix_tiles_header header_path matrix_json_file) file(READ "${matrix_json_file}" _matrix_json) - string(JSON _tile0 GET "${_matrix_json}" 0 "_tile" 0) - string(JSON _tile_m GET "${_tile0}" "tile_m") - string(JSON _tile_n GET "${_tile0}" "tile_n") - string(JSON _tile_k GET "${_tile0}" "tile_k") - set(_arch_tile_aliases "") + set(_tile_aliases "") string(JSON _entry_len LENGTH "${_matrix_json}") set(_entry_idx 0) while(_entry_idx LESS _entry_len) string(JSON _entry GET "${_matrix_json}" "${_entry_idx}") - string(JSON _entry_tile GET "${_entry}" "_tile" 0) - string(JSON _entry_tile_m GET "${_entry_tile}" "tile_m") - string(JSON _entry_tile_n GET "${_entry_tile}" "tile_n") - string(JSON _entry_tile_k GET "${_entry_tile}" "tile_k") - string(JSON _export_len LENGTH "${_entry}" "_export") - set(_export_idx 0) - while(_export_idx LESS _export_len) - string(JSON _export_entry GET "${_entry}" "_export" "${_export_idx}") - string(JSON _register GET "${_export_entry}" "register") - if(_register STREQUAL "cubin") - string(JSON _arch_tag GET "${_export_entry}" "arch_tag") - set(_arch_tile_value "${_entry_tile_m},${_entry_tile_n},${_entry_tile_k}") - if(DEFINED _arch_tile_alias_value_${_arch_tag}) - if(NOT "${_arch_tile_alias_value_${_arch_tag}}" STREQUAL "${_arch_tile_value}") - message(FATAL_ERROR "Conflicting cuTile tile geometry for ${_arch_tag}: " - "${_arch_tile_alias_value_${_arch_tag}} vs ${_arch_tile_value}" - ) - endif() - else() - set(_arch_tile_alias_value_${_arch_tag} "${_arch_tile_value}") - string( - APPEND - _arch_tile_aliases - "using fused_1nn_matrix_tile_${_arch_tag} = cutile_tile_config<${_entry_tile_m}, ${_entry_tile_n}, ${_entry_tile_k}>;\n" - ) + string(JSON _entry_tile ERROR_VARIABLE _entry_tile_error GET "${_entry}" "_tile" 0) + if(NOT _entry_tile_error) + string(JSON _default_tile_m GET "${_entry_tile}" "tile_m") + string(JSON _default_tile_n GET "${_entry_tile}" "tile_n") + string(JSON _default_tile_k GET "${_entry_tile}" "tile_k") + endif() + + string(JSON _abi_len LENGTH "${_entry}" "_abi") + set(_abi_idx 0) + while(_abi_idx LESS _abi_len) + string(JSON _abi_entry GET "${_entry}" "_abi" "${_abi_idx}") + string(JSON _abi_abbrev GET "${_abi_entry}" "abi_abbrev") + string(JSON _tile_m ERROR_VARIABLE _tile_m_error GET "${_abi_entry}" "tile_m") + string(JSON _tile_n ERROR_VARIABLE _tile_n_error GET "${_abi_entry}" "tile_n") + string(JSON _tile_k ERROR_VARIABLE _tile_k_error GET "${_abi_entry}" "tile_k") + if(_tile_m_error + OR _tile_n_error + OR _tile_k_error + ) + if(_entry_tile_error) + message(FATAL_ERROR "Missing cuTile geometry for ABI ${_abi_abbrev}") endif() + set(_tile_m "${_default_tile_m}") + set(_tile_n "${_default_tile_n}") + set(_tile_k "${_default_tile_k}") endif() - math(EXPR _export_idx "${_export_idx} + 1") + + _cutile_append_matrix_tile_aliases( + "${_entry}" "${_abi_abbrev}" "${_tile_m}" "${_tile_n}" "${_tile_k}" + ) + math(EXPR _abi_idx "${_abi_idx} + 1") endwhile() math(EXPR _entry_idx "${_entry_idx} + 1") endwhile() @@ -145,8 +178,7 @@ function(_cutile_generate_matrix_tiles_header header_path matrix_json_file) namespace cuvs::distance::detail { -using fused_1nn_matrix_tile = cutile_tile_config<${_tile_m}, ${_tile_n}, ${_tile_k}>; -${_arch_tile_aliases} +${_tile_aliases} } // namespace cuvs::distance::detail " diff --git a/cpp/include/cuvs/detail/jit_lto/fused_distance_nn/fused_1nn_fragments.hpp b/cpp/include/cuvs/detail/jit_lto/fused_distance_nn/fused_1nn_fragments.hpp index fe27441b0f..f4b64233ad 100644 --- a/cpp/include/cuvs/detail/jit_lto/fused_distance_nn/fused_1nn_fragments.hpp +++ b/cpp/include/cuvs/detail/jit_lto/fused_distance_nn/fused_1nn_fragments.hpp @@ -5,16 +5,13 @@ #pragma once +#include + #include #include -#include - namespace cuvs::distance::detail { -struct metric_tag_ip {}; -struct metric_tag_l2 {}; -struct metric_tag_cos {}; struct cutile_abi_strict {}; struct cutile_abi_relaxed {}; @@ -25,44 +22,6 @@ struct cutile_tile_config { static constexpr int tile_k = TileK; }; -template -struct fused_1nn_metric_tag; - -template <> -struct fused_1nn_metric_tag { - using type = metric_tag_ip; -}; - -template <> -struct fused_1nn_metric_tag { - using type = metric_tag_l2; -}; - -template <> -struct fused_1nn_metric_tag { - using type = metric_tag_l2; -}; - -template <> -struct fused_1nn_metric_tag { - using type = metric_tag_cos; -}; - -/** Whether sqrt is applied when packing distance into KVP output. */ -template -constexpr bool fused_1nn_apply_sqrt_at_pack(bool is_sqrt) -{ - if constexpr (Metric == cuvs::distance::DistanceType::L2Expanded || - Metric == cuvs::distance::DistanceType::L2SqrtExpanded) { - return is_sqrt; - } else { - return false; - } -} - -template -using fused_1nn_metric_tag_t = typename fused_1nn_metric_tag::type; - template struct fused_1nn_data_tag; @@ -95,22 +54,13 @@ struct fused_1nn_index_tag { template using fused_1nn_index_tag_t = typename fused_1nn_index_tag::type; -template +template struct fragment_tag_fused_1nn_cubin { static constexpr int cc_major = ArchTag::cc_major; static constexpr int cc_minor = ArchTag::cc_minor; }; -template +template struct fragment_tag_fused_1nn_tileir {}; } // namespace cuvs::distance::detail diff --git a/cpp/src/distance/detail/fused_distance_nn/cutile/export_fused_1nn.py b/cpp/src/distance/detail/fused_distance_nn/cutile/export_fused_1nn.py index 2b1e37b5fd..929dd028dd 100644 --- a/cpp/src/distance/detail/fused_distance_nn/cutile/export_fused_1nn.py +++ b/cpp/src/distance/detail/fused_distance_nn/cutile/export_fused_1nn.py @@ -26,7 +26,6 @@ index_abbrev, kernel_symbol, make_kernel, - metric_abbrev, ) DEFAULT_TILEIR_BYTECODE_VERSION = "13.1" @@ -138,7 +137,6 @@ def _kernel_signature( abbrev = _data_abbrev(data_type) symbol = kernel_symbol( abbrev, - metric_abbrev(metric), index_abbrev(index_type), matrix_layout, ) @@ -156,6 +154,7 @@ def _kernel_signature( ScalarConstraint(idx_dtype), ScalarConstraint(idx_dtype), ScalarConstraint(idx_dtype), + ScalarConstraint(ct.int32), ConstantConstraint(tile_m), ConstantConstraint(tile_n), ConstantConstraint(tile_k), diff --git a/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_cutile_matrix.json b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_cutile_matrix.json index c6c13f17b9..e578f6e19f 100644 --- a/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_cutile_matrix.json +++ b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_cutile_matrix.json @@ -5,18 +5,12 @@ { "matrix_layout": "strict", "abi_abbrev": "strict", - "abi_tag": "cutile_abi_strict", - "tile_m": 128, - "tile_n": 64, - "tile_k": 64 + "abi_tag": "cutile_abi_strict" }, { "matrix_layout": "relaxed", "abi_abbrev": "relaxed", - "abi_tag": "cutile_abi_relaxed", - "tile_m": 128, - "tile_n": 128, - "tile_k": 32 + "abi_tag": "cutile_abi_relaxed" } ], "_data": [ @@ -31,16 +25,7 @@ ], "_metric": [ { - "metric": "inner_product", - "metric_abbrev": "ip" - }, - { - "metric": "l2_expanded", - "metric_abbrev": "l2" - }, - { - "metric": "cosine_expanded", - "metric_abbrev": "cos" + "metric": "runtime" } ], "_index": [ @@ -60,7 +45,7 @@ { "output_format": "cubin", "artifact_ext": "cubin", - "artifact_basename": "@data_type@_@metric_abbrev@_@index_abbrev@_@abi_abbrev@_@gpu_code@", + "artifact_basename": "@data_type@_@index_abbrev@_@abi_abbrev@_@gpu_code@", "register": "cubin", "gpu_code": "sm_100", "cc_major": 10, @@ -70,7 +55,7 @@ { "output_format": "cubin", "artifact_ext": "cubin", - "artifact_basename": "@data_type@_@metric_abbrev@_@index_abbrev@_@abi_abbrev@_@gpu_code@", + "artifact_basename": "@data_type@_@index_abbrev@_@abi_abbrev@_@gpu_code@", "register": "cubin", "gpu_code": "sm_80", "cc_major": 8, @@ -80,7 +65,7 @@ { "output_format": "cubin", "artifact_ext": "cubin", - "artifact_basename": "@data_type@_@metric_abbrev@_@index_abbrev@_@abi_abbrev@_@gpu_code@", + "artifact_basename": "@data_type@_@index_abbrev@_@abi_abbrev@_@gpu_code@", "register": "cubin", "gpu_code": "sm_86", "cc_major": 8, @@ -90,7 +75,7 @@ { "output_format": "cubin", "artifact_ext": "cubin", - "artifact_basename": "@data_type@_@metric_abbrev@_@index_abbrev@_@abi_abbrev@_@gpu_code@", + "artifact_basename": "@data_type@_@index_abbrev@_@abi_abbrev@_@gpu_code@", "register": "cubin", "gpu_code": "sm_90", "cc_major": 9, @@ -100,7 +85,7 @@ { "output_format": "tileir_bytecode", "artifact_ext": "tilebc", - "artifact_basename": "@data_type@_@metric_abbrev@_@index_abbrev@_@abi_abbrev@", + "artifact_basename": "@data_type@_@index_abbrev@_@abi_abbrev@", "register": "tileir", "gpu_code": "sm_80", "bytecode_version": "13.1" @@ -113,12 +98,18 @@ { "matrix_layout": "strict", "abi_abbrev": "strict", - "abi_tag": "cutile_abi_strict" + "abi_tag": "cutile_abi_strict", + "tile_m": 128, + "tile_n": 64, + "tile_k": 64 }, { "matrix_layout": "relaxed", "abi_abbrev": "relaxed", - "abi_tag": "cutile_abi_relaxed" + "abi_tag": "cutile_abi_relaxed", + "tile_m": 128, + "tile_n": 128, + "tile_k": 32 } ], "_data": [ @@ -133,16 +124,7 @@ ], "_metric": [ { - "metric": "inner_product", - "metric_abbrev": "ip" - }, - { - "metric": "l2_expanded", - "metric_abbrev": "l2" - }, - { - "metric": "cosine_expanded", - "metric_abbrev": "cos" + "metric": "runtime" } ], "_index": [ @@ -155,7 +137,7 @@ { "output_format": "cubin", "artifact_ext": "cubin", - "artifact_basename": "@data_type@_@metric_abbrev@_@index_abbrev@_@abi_abbrev@_@gpu_code@", + "artifact_basename": "@data_type@_@index_abbrev@_@abi_abbrev@_@gpu_code@", "register": "cubin", "gpu_code": "sm_120", "cc_major": 12, diff --git a/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_kernel.py b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_kernel.py index 2c139831cb..5130739f47 100644 --- a/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_kernel.py +++ b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_kernel.py @@ -1,6 +1,6 @@ # SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. # SPDX-License-Identifier: Apache-2.0 -"""cuTile fused GEMM + 1-NN kernels (InnerProduct, L2Expanded, CosineExpanded).""" +"""cuTile fused GEMM + 1-NN kernel with runtime metric selection.""" from __future__ import annotations @@ -13,8 +13,11 @@ DEFAULT_TILE_N = 128 DEFAULT_TILE_K = 64 -METRICS = ("inner_product", "l2_expanded", "cosine_expanded") +METRICS = ("runtime",) INDEX_TYPES = ("int32", "int64") +METRIC_L2_EXPANDED = 0 +METRIC_COSINE_EXPANDED = 2 +METRIC_INNER_PRODUCT = 6 def _idx_dtype(index_type: str): @@ -35,7 +38,7 @@ def make_kernel( index_type: str = "int32", gpu_code: str = "sm_80", ): - """Build a cuTile kernel with metric, index width, and tile sizes baked in at compile time.""" + """Build a cuTile kernel with index width and tile sizes baked in.""" if data_type not in ("half", "float"): raise ValueError(f"Unsupported data_type {data_type!r}") if metric not in METRICS: @@ -46,9 +49,6 @@ def make_kernel( acc_dtype = ct.float32 idx_dtype = _idx_dtype(index_type) out_dist_dtype = ct.float16 if data_type == "half" else ct.float32 - is_ip = metric == "inner_product" - is_l2 = metric == "l2_expanded" - is_cos = metric == "cosine_expanded" items_per_thread = 4 if gpu_code in ("sm_100", "sm_120") else 2 core_shape = ( tile_m, @@ -73,6 +73,7 @@ def fused_1nn_reduce_kernel( K, apply_sqrt, store_idx, + metric_code, tm: ConstInt, tn: ConstInt, tk: ConstInt, @@ -84,12 +85,7 @@ def fused_1nn_reduce_kernel( # Blackwell uses four items per logical thread slot; earlier targets # retain the existing two-item grouping. - if is_ip: - best_dist = ct.full(best_shape, -3.4e38, acc_dtype) - neutral_dist = -3.4e38 - else: - best_dist = ct.full(best_shape, 3.4e38, acc_dtype) - neutral_dist = 3.4e38 + best_dist = ct.full(best_shape, 3.4e38, acc_dtype) best_idx = ct.zeros(best_shape, idx_dtype) num_tiles_k = ct.num_tiles(A, axis=1, shape=(tm, tk)) @@ -98,10 +94,7 @@ def fused_1nn_reduce_kernel( def reduce_scores(best, best_idx, axes): def red_op(a_score, a_idx, b_score, b_idx): - if is_ip: - cond = a_score > b_score - else: - cond = a_score < b_score + cond = a_score < b_score return ( ct.where(cond, a_score, b_score), @@ -113,7 +106,7 @@ def red_op(a_score, a_idx, b_score, b_idx): (best, best_idx), axes[0], red_op, - (neutral_dist, -1), + (3.4e38, -1), keepdims=True, ) if len(axes) >= 2: @@ -121,7 +114,7 @@ def red_op(a_score, a_idx, b_score, b_idx): (best, best_idx), axes[1], red_op, - (neutral_dist, -1), + (3.4e38, -1), keepdims=True, ) @@ -142,21 +135,23 @@ def red_op(a_score, a_idx, b_score, b_idx): accumulator = ct.mma(a, ct.transpose(b_T), accumulator) - if is_ip: - score = accumulator - elif is_l2 or is_cos: + if metric_code == METRIC_INNER_PRODUCT: + # Keep one min reduction for every metric, then restore the + # inner-product sign before writing the result. + score = -accumulator + else: a_norm = ct.load( A_norm, index=(bidm,), shape=(tm,), padding_mode=zero_pad ) b_norm = ct.load( B_norm, index=(n,), shape=(tn,), padding_mode=zero_pad ) - if is_l2: + if metric_code == METRIC_L2_EXPANDED: # L2 expanded: ||x||^2 + ||y||^2 - 2 * dot(x, y); norms are squared. score = ( a_norm[:, None] + b_norm[None, :] - (2.0 * accumulator) ) - elif is_cos: + else: # Cosine expanded distance: 1 - dot / (||x|| * ||y||); norms are L2 (not squared). denom = a_norm[:, None] * b_norm[None, :] score = 1.0 - (accumulator / denom) @@ -166,10 +161,7 @@ def red_op(a_score, a_idx, b_score, b_idx): col = ct.arange(tn, dtype=idx_dtype) global_col = (n * tn + col).astype(idx_dtype) valid = global_col < N - if is_ip: - score = ct.where(valid[None, :], score, -3.4e38) - else: - score = ct.where(valid[None, :], score, 3.4e38) + score = ct.where(valid[None, :], score, 3.4e38) curr_idx = ct.arange(tn, dtype=idx_dtype).reshape(core_shape[1:])[ None, ... @@ -177,10 +169,7 @@ def red_op(a_score, a_idx, b_score, b_idx): curr_best, curr_idx = reduce_scores( score.reshape(core_shape), curr_idx, inner_reduction_axes ) - if is_ip: - update = curr_best > best_dist - else: - update = curr_best < best_dist + update = curr_best < best_dist best_dist = ct.where(update, curr_best, best_dist) best_idx = ct.where( update, (n * tn + curr_idx).astype(idx_dtype), best_idx @@ -191,7 +180,9 @@ def red_op(a_score, a_idx, b_score, b_idx): ) out_dist = best_dist - if is_l2: + if metric_code == METRIC_INNER_PRODUCT: + out_dist = -best_dist + elif metric_code == METRIC_L2_EXPANDED: out_dist = ct.where(apply_sqrt != 0, ct.sqrt(best_dist), best_dist) if store_idx != 0: ct.store(OutIdx, index=(bidm,), tile=best_idx.reshape((tm,))) @@ -206,12 +197,11 @@ def red_op(a_score, a_idx, b_score, b_idx): def kernel_symbol( data_abbrev: str, - metric_abbrev: str, index_abbrev: str, matrix_layout: str = "strict", ) -> str: """Must stay in sync with fused_1nn_kernel_entrypoint() in fused_1nn_planner.hpp.""" - base = f"fused_1nn_{data_abbrev}_{metric_abbrev}_{index_abbrev}" + base = f"fused_1nn_{data_abbrev}_{index_abbrev}" if matrix_layout == "strict": return base if matrix_layout == "relaxed": @@ -219,13 +209,5 @@ def kernel_symbol( raise ValueError(f"Unsupported matrix layout {matrix_layout!r}") -def metric_abbrev(metric: str) -> str: - return { - "inner_product": "ip", - "l2_expanded": "l2", - "cosine_expanded": "cos", - }[metric] - - def index_abbrev(index_type: str) -> str: return {"int32": "i32", "int64": "i64"}[index_type] diff --git a/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_planner.hpp b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_planner.hpp index a80f4a4eb3..995f4c787e 100644 --- a/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_planner.hpp +++ b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_planner.hpp @@ -17,57 +17,32 @@ namespace cuvs::distance::detail { /** Must match kernel_symbol() in fused_1nn_kernel.py (export uses with_symbol). */ -template +template inline const char* fused_1nn_kernel_entrypoint() { - constexpr bool is_i32 = std::is_same_v; - constexpr bool is_i64 = std::is_same_v; constexpr bool is_relaxed = std::is_same_v; - static_assert(is_i32 || is_i64, "unsupported fused 1-NN cuTile index width"); static_assert(is_relaxed || std::is_same_v, "unsupported fused 1-NN cuTile ABI"); - if constexpr (std::is_same_v && - std::is_same_v) { - return is_i32 ? (is_relaxed ? "fused_1nn_f_ip_i32_relaxed" : "fused_1nn_f_ip_i32") - : (is_relaxed ? "fused_1nn_f_ip_i64_relaxed" : "fused_1nn_f_ip_i64"); - } else if constexpr (std::is_same_v && - std::is_same_v) { - return is_i32 ? (is_relaxed ? "fused_1nn_f_l2_i32_relaxed" : "fused_1nn_f_l2_i32") - : (is_relaxed ? "fused_1nn_f_l2_i64_relaxed" : "fused_1nn_f_l2_i64"); - } else if constexpr (std::is_same_v && - std::is_same_v) { - return is_i32 ? (is_relaxed ? "fused_1nn_f_cos_i32_relaxed" : "fused_1nn_f_cos_i32") - : (is_relaxed ? "fused_1nn_f_cos_i64_relaxed" : "fused_1nn_f_cos_i64"); - } else if constexpr (std::is_same_v && - std::is_same_v) { - return is_i32 ? (is_relaxed ? "fused_1nn_h_ip_i32_relaxed" : "fused_1nn_h_ip_i32") - : (is_relaxed ? "fused_1nn_h_ip_i64_relaxed" : "fused_1nn_h_ip_i64"); - } else if constexpr (std::is_same_v && - std::is_same_v) { - return is_i32 ? (is_relaxed ? "fused_1nn_h_l2_i32_relaxed" : "fused_1nn_h_l2_i32") - : (is_relaxed ? "fused_1nn_h_l2_i64_relaxed" : "fused_1nn_h_l2_i64"); - } else if constexpr (std::is_same_v && - std::is_same_v) { - return is_i32 ? (is_relaxed ? "fused_1nn_h_cos_i32_relaxed" : "fused_1nn_h_cos_i32") - : (is_relaxed ? "fused_1nn_h_cos_i64_relaxed" : "fused_1nn_h_cos_i64"); + if constexpr (std::is_same_v) { + return is_relaxed ? "fused_1nn_f_i32_relaxed" : "fused_1nn_f_i32"; + } else if constexpr (std::is_same_v) { + return is_relaxed ? "fused_1nn_h_i32_relaxed" : "fused_1nn_h_i32"; } else { - static_assert(sizeof(DataTag) == 0, "unsupported fused 1-NN cuTile data/metric combination"); + static_assert(sizeof(DataTag) == 0, "unsupported fused 1-NN cuTile data type"); return ""; } } -template +template struct Fused1nnTilePlanner : TileAlgorithmPlanner { - using DataTag = fused_1nn_data_tag_t; - using MetricTag = fused_1nn_metric_tag_t; - using IndexTag = fused_1nn_index_tag_t; + using DataTag = fused_1nn_data_tag_t; + using IndexTag = cuvs::neighbors::detail::tag_index_i32; inline static LauncherJitCache launcher_jit_cache{}; Fused1nnTilePlanner() - : TileAlgorithmPlanner(fused_1nn_kernel_entrypoint(), - launcher_jit_cache) + : TileAlgorithmPlanner(fused_1nn_kernel_entrypoint(), launcher_jit_cache) { } @@ -81,42 +56,43 @@ struct Fused1nnTilePlanner : TileAlgorithmPlanner { using cuvs::detail::jit_lto::cutile_arch_8_6; using cuvs::detail::jit_lto::cutile_arch_9_0; - this->add_static_fragment>(); - this->add_static_fragment>(); - this->add_static_fragment>(); - this->add_static_fragment>(); - this->add_static_fragment>(); + constexpr bool is_relaxed = std::is_same_v; + using Tile80 = std::conditional_t; + using Tile86 = std::conditional_t; + using Tile90 = std::conditional_t; + using Tile100 = std::conditional_t; + using Tile120 = std::conditional_t; + + this->add_static_fragment< + fragment_tag_fused_1nn_cubin>(); + this->add_static_fragment< + fragment_tag_fused_1nn_cubin>(); + this->add_static_fragment< + fragment_tag_fused_1nn_cubin>(); + this->add_static_fragment< + fragment_tag_fused_1nn_cubin>(); + this->add_static_fragment< + fragment_tag_fused_1nn_cubin>(); } void add_tileir_fallback() { + constexpr bool is_relaxed = std::is_same_v; + using TileIr = std::conditional_t; this->add_static_tileir_fragment< - fragment_tag_fused_1nn_tileir>(); + fragment_tag_fused_1nn_tileir>(); } }; diff --git a/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_tile.cu b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_tile.cu index 1a4dd24130..0ab4e1ef84 100644 --- a/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_tile.cu +++ b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_tile.cu @@ -24,7 +24,7 @@ namespace detail { namespace { -template +template bool launch_fused_1nn_tile(IdxT* nearest_idx, DataT* nearest_dist, const DataT* x, @@ -34,6 +34,7 @@ bool launch_fused_1nn_tile(IdxT* nearest_idx, IdxT m, IdxT n, IdxT k, + cuvs::distance::DistanceType metric, bool is_sqrt, cudaStream_t stream) { @@ -41,14 +42,29 @@ bool launch_fused_1nn_tile(IdxT* nearest_idx, if (nearest_dist == nullptr) { return false; } - Fused1nnTilePlanner planner; + Fused1nnTilePlanner planner; planner.add_entrypoint(); planner.add_tileir_fallback(); const CutileTileConfig tile_cfg = planner.tile_config(); auto launcher = planner.try_get_launcher(); if (!launcher) { return false; } - const bool apply_sqrt = fused_1nn_apply_sqrt_at_pack(is_sqrt); + int metric_code; + bool apply_sqrt = false; + switch (metric) { + case cuvs::distance::DistanceType::InnerProduct: + metric_code = static_cast(cuvs::distance::DistanceType::InnerProduct); + break; + case cuvs::distance::DistanceType::L2Expanded: + case cuvs::distance::DistanceType::L2SqrtExpanded: + metric_code = static_cast(cuvs::distance::DistanceType::L2Expanded); + apply_sqrt = is_sqrt; + break; + case cuvs::distance::DistanceType::CosineExpanded: + metric_code = static_cast(cuvs::distance::DistanceType::CosineExpanded); + break; + default: return false; + } IdxT shape_x[2] = {m, k}; IdxT stride_x[2] = {k, IdxT{1}}; @@ -108,7 +124,8 @@ bool launch_fused_1nn_tile(IdxT* nearest_idx, IdxT, IdxT, IdxT, - IdxT); + IdxT, + int); launcher->template dispatch(stream, grid, block, @@ -139,7 +156,8 @@ bool launch_fused_1nn_tile(IdxT* nearest_idx, N, K, static_cast(apply_sqrt), - store_idx); + store_idx, + metric_code); RAFT_CUDA_TRY(cudaGetLastError()); return true; } @@ -158,27 +176,8 @@ bool try_fused_1nn_tile_dispatch(IdxT* nearest_idx, bool is_sqrt, cudaStream_t stream) { - switch (metric) { - case cuvs::distance::DistanceType::InnerProduct: - return launch_fused_1nn_tile( - nearest_idx, nearest_dist, x, y, xn, yn, m, n, k, is_sqrt, stream); - case cuvs::distance::DistanceType::L2Expanded: - return launch_fused_1nn_tile( - nearest_idx, nearest_dist, x, y, xn, yn, m, n, k, is_sqrt, stream); - case cuvs::distance::DistanceType::L2SqrtExpanded: - return launch_fused_1nn_tile( - nearest_idx, nearest_dist, x, y, xn, yn, m, n, k, is_sqrt, stream); - case cuvs::distance::DistanceType::CosineExpanded: - return launch_fused_1nn_tile( - nearest_idx, nearest_dist, x, y, xn, yn, m, n, k, is_sqrt, stream); - default: return false; - } + return launch_fused_1nn_tile( + nearest_idx, nearest_dist, x, y, xn, yn, m, n, k, metric, is_sqrt, stream); } } // namespace From 9b2d2cbb02842cc8e7f7b2a1f86b609bfd8d077e Mon Sep 17 00:00:00 2001 From: divyegala Date: Wed, 29 Jul 2026 21:45:48 +0000 Subject: [PATCH 22/82] new tiles and reduction shapes for sm120 --- .../cutile/export_fused_1nn.py | 1 + .../cutile/fused_1nn_cutile_matrix.json | 2 +- .../cutile/fused_1nn_kernel.py | 28 ++++++++++++------- 3 files changed, 20 insertions(+), 11 deletions(-) diff --git a/cpp/src/distance/detail/fused_distance_nn/cutile/export_fused_1nn.py b/cpp/src/distance/detail/fused_distance_nn/cutile/export_fused_1nn.py index 929dd028dd..7d212c698e 100644 --- a/cpp/src/distance/detail/fused_distance_nn/cutile/export_fused_1nn.py +++ b/cpp/src/distance/detail/fused_distance_nn/cutile/export_fused_1nn.py @@ -185,6 +185,7 @@ def export_binary( tile_k, index_type=index_type, gpu_code=gpu_code, + matrix_layout=matrix_layout, ) signature = _kernel_signature( data_type, diff --git a/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_cutile_matrix.json b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_cutile_matrix.json index e578f6e19f..4b675c7c11 100644 --- a/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_cutile_matrix.json +++ b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_cutile_matrix.json @@ -109,7 +109,7 @@ "abi_tag": "cutile_abi_relaxed", "tile_m": 128, "tile_n": 128, - "tile_k": 32 + "tile_k": 16 } ], "_data": [ diff --git a/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_kernel.py b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_kernel.py index 5130739f47..9e4fd5e91d 100644 --- a/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_kernel.py +++ b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_kernel.py @@ -37,26 +37,36 @@ def make_kernel( *, index_type: str = "int32", gpu_code: str = "sm_80", + matrix_layout: str = "strict", ): - """Build a cuTile kernel with index width and tile sizes baked in.""" + """Build a cuTile kernel with index width, tile, and reduction baked in.""" if data_type not in ("half", "float"): raise ValueError(f"Unsupported data_type {data_type!r}") if metric not in METRICS: raise ValueError(f"Unsupported metric {metric!r}") if index_type not in INDEX_TYPES: raise ValueError(f"Unsupported index_type {index_type!r}") + if matrix_layout not in ("strict", "relaxed"): + raise ValueError(f"Unsupported matrix_layout {matrix_layout!r}") acc_dtype = ct.float32 idx_dtype = _idx_dtype(index_type) out_dist_dtype = ct.float16 if data_type == "half" else ct.float32 - items_per_thread = 4 if gpu_code in ("sm_100", "sm_120") else 2 + if gpu_code == "sm_120": + reduction_groups, reduction_items = ( + (8, 2) if matrix_layout == "strict" else (2, 8) + ) + elif gpu_code == "sm_100": + reduction_groups, reduction_items = 4, 4 + else: + reduction_groups, reduction_items = 4, 2 core_shape = ( tile_m, - tile_n // (4 * items_per_thread), - 4, - items_per_thread, + tile_n // (reduction_groups * reduction_items), + reduction_groups, + reduction_items, ) - best_shape = (tile_m, 1, 4, 1) + best_shape = (tile_m, 1, reduction_groups, 1) inner_reduction_axes = (1, 3) outer_reduction_axes = (2,) @@ -80,10 +90,8 @@ def fused_1nn_reduce_kernel( ): bidm = ct.bid(0) - # Reduce groups and per-thread items inside each N tile, carry four - # partial winners across N tiles, then reduce those winners once. - # Blackwell uses four items per logical thread slot; earlier targets - # retain the existing two-item grouping. + # Reduce groups and per-thread items inside each N tile, carry the + # remaining group winners across N tiles, then reduce them once. best_dist = ct.full(best_shape, 3.4e38, acc_dtype) best_idx = ct.zeros(best_shape, idx_dtype) From 54125c035a908f821546b9e01ea22b18833f0bae Mon Sep 17 00:00:00 2001 From: divyegala Date: Tue, 4 Aug 2026 17:38:33 +0000 Subject: [PATCH 23/82] new kernel from huy --- .../cutile/fused_1nn_cutile_matrix.json | 6 +- .../cutile/fused_1nn_kernel.py | 131 ++++++------------ 2 files changed, 49 insertions(+), 88 deletions(-) diff --git a/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_cutile_matrix.json b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_cutile_matrix.json index 4b675c7c11..c57e8a2658 100644 --- a/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_cutile_matrix.json +++ b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_cutile_matrix.json @@ -100,8 +100,8 @@ "abi_abbrev": "strict", "abi_tag": "cutile_abi_strict", "tile_m": 128, - "tile_n": 64, - "tile_k": 64 + "tile_n": 128, + "tile_k": 32 }, { "matrix_layout": "relaxed", @@ -109,7 +109,7 @@ "abi_tag": "cutile_abi_relaxed", "tile_m": 128, "tile_n": 128, - "tile_k": 16 + "tile_k": 32 } ], "_data": [ diff --git a/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_kernel.py b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_kernel.py index 9e4fd5e91d..465bbcbb4e 100644 --- a/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_kernel.py +++ b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_kernel.py @@ -11,7 +11,7 @@ # Default tile geometry; overridden per export via make_kernel(..., tile_m, tile_n, tile_k). DEFAULT_TILE_M = 128 DEFAULT_TILE_N = 128 -DEFAULT_TILE_K = 64 +DEFAULT_TILE_K = 32 METRICS = ("runtime",) INDEX_TYPES = ("int32", "int64") @@ -39,7 +39,7 @@ def make_kernel( gpu_code: str = "sm_80", matrix_layout: str = "strict", ): - """Build a cuTile kernel with index width, tile, and reduction baked in.""" + """Build the flat-reduction runtime-metric cuTile kernel.""" if data_type not in ("half", "float"): raise ValueError(f"Unsupported data_type {data_type!r}") if metric not in METRICS: @@ -52,26 +52,11 @@ def make_kernel( acc_dtype = ct.float32 idx_dtype = _idx_dtype(index_type) out_dist_dtype = ct.float16 if data_type == "half" else ct.float32 - if gpu_code == "sm_120": - reduction_groups, reduction_items = ( - (8, 2) if matrix_layout == "strict" else (2, 8) - ) - elif gpu_code == "sm_100": - reduction_groups, reduction_items = 4, 4 - else: - reduction_groups, reduction_items = 4, 2 - core_shape = ( - tile_m, - tile_n // (reduction_groups * reduction_items), - reduction_groups, - reduction_items, - ) - best_shape = (tile_m, 1, reduction_groups, 1) - inner_reduction_axes = (1, 3) - outer_reduction_axes = (2,) - - @ct.kernel - def fused_1nn_reduce_kernel( + core_shape = (tile_m, tile_n) + best_shape = (tile_m, 1) + + @ct.kernel(occupancy=ct.ByTarget(sm_100=2, sm_120=2)) + def fused_1nn_kernel( A, B, A_norm, @@ -89,109 +74,85 @@ def fused_1nn_reduce_kernel( tk: ConstInt, ): bidm = ct.bid(0) - - # Reduce groups and per-thread items inside each N tile, carry the - # remaining group winners across N tiles, then reduce them once. - best_dist = ct.full(best_shape, 3.4e38, acc_dtype) best_idx = ct.zeros(best_shape, idx_dtype) - num_tiles_k = ct.num_tiles(A, axis=1, shape=(tm, tk)) num_tiles_n = ct.num_tiles(B, axis=0, shape=(tn, tk)) zero_pad = ct.PaddingMode.ZERO - def reduce_scores(best, best_idx, axes): + def reduce_scores(dists, indices): def red_op(a_score, a_idx, b_score, b_idx): cond = a_score < b_score - return ( ct.where(cond, a_score, b_score), ct.where(cond, a_idx, b_idx), ) - if len(axes) >= 1: - best, best_idx = ct.reduce( - (best, best_idx), - axes[0], - red_op, - (3.4e38, -1), - keepdims=True, - ) - if len(axes) >= 2: - best, best_idx = ct.reduce( - (best, best_idx), - axes[1], - red_op, - (3.4e38, -1), - keepdims=True, - ) - - return best, best_idx + return ct.reduce( + (dists, indices), + 1, + red_op, + (3.4e38, -1), + keepdims=True, + ) + local_indices = ct.arange(tn, dtype=ct.int16)[None, :] for n in range(num_tiles_n): accumulator = ct.full((tm, tn), 0, dtype=acc_dtype) - for k in range(num_tiles_k): dtype = ct.tfloat32 if A.dtype == ct.float32 else A.dtype - a = ct.load( A, index=(bidm, k), shape=(tm, tk), padding_mode=zero_pad ).astype(dtype) b_T = ct.load( - B, index=(n, k), shape=(tn, tk), padding_mode=zero_pad + B, + index=(k, n), + shape=(tk, tn), + padding_mode=zero_pad, + order=(1, 0), ).astype(dtype) - - accumulator = ct.mma(a, ct.transpose(b_T), accumulator) + accumulator = ct.mma(a, b_T, accumulator) if metric_code == METRIC_INNER_PRODUCT: - # Keep one min reduction for every metric, then restore the - # inner-product sign before writing the result. score = -accumulator else: - a_norm = ct.load( - A_norm, index=(bidm,), shape=(tm,), padding_mode=zero_pad - ) b_norm = ct.load( B_norm, index=(n,), shape=(tn,), padding_mode=zero_pad ) if metric_code == METRIC_L2_EXPANDED: - # L2 expanded: ||x||^2 + ||y||^2 - 2 * dot(x, y); norms are squared. - score = ( - a_norm[:, None] + b_norm[None, :] - (2.0 * accumulator) - ) + # The A norm is constant across centroids. Reduce + # 0.5 * ||y||^2 - dot(x, y), then recover full L2 once. + score = (0.5 * b_norm)[None, :] - accumulator else: - # Cosine expanded distance: 1 - dot / (||x|| * ||y||); norms are L2 (not squared). - denom = a_norm[:, None] * b_norm[None, :] - score = 1.0 - (accumulator / denom) + # Defer the A-norm division until after selecting the + # winning centroid. + score = -(accumulator / b_norm[None, :]) - # Only the final N-tile can include zero-padded centroid columns. if n == num_tiles_n - 1: - col = ct.arange(tn, dtype=idx_dtype) - global_col = (n * tn + col).astype(idx_dtype) - valid = global_col < N - score = ct.where(valid[None, :], score, 3.4e38) - - curr_idx = ct.arange(tn, dtype=idx_dtype).reshape(core_shape[1:])[ - None, ... - ] + col = ct.arange(tn, dtype=ct.int16) + score = ct.where((n * tn + col)[None, :] < N, score, 3.4e38) + curr_best, curr_idx = reduce_scores( - score.reshape(core_shape), curr_idx, inner_reduction_axes + score.reshape(core_shape), local_indices ) update = curr_best < best_dist best_dist = ct.where(update, curr_best, best_dist) - best_idx = ct.where( - update, (n * tn + curr_idx).astype(idx_dtype), best_idx - ) + best_idx = ct.where(update, n * tn + curr_idx, best_idx) - best_dist, best_idx = reduce_scores( - best_dist, best_idx, outer_reduction_axes - ) - - out_dist = best_dist if metric_code == METRIC_INNER_PRODUCT: out_dist = -best_dist - elif metric_code == METRIC_L2_EXPANDED: - out_dist = ct.where(apply_sqrt != 0, ct.sqrt(best_dist), best_dist) + else: + a_norm = ct.load( + A_norm, index=(bidm,), shape=(tm,), padding_mode=zero_pad + )[:, None] + if metric_code == METRIC_L2_EXPANDED: + out_dist = a_norm + 2.0 * best_dist + out_dist = ct.where( + apply_sqrt != 0, ct.sqrt(out_dist), out_dist + ) + else: + out_dist = 1.0 + best_dist / a_norm + if store_idx != 0: ct.store(OutIdx, index=(bidm,), tile=best_idx.reshape((tm,))) ct.store( @@ -200,7 +161,7 @@ def red_op(a_score, a_idx, b_score, b_idx): tile=out_dist.reshape((tm,)).astype(out_dist_dtype), ) - return fused_1nn_reduce_kernel + return fused_1nn_kernel def kernel_symbol( From 27f20699a4b2a6222bc1d14fd5e06cd562ff7a63 Mon Sep 17 00:00:00 2001 From: divyegala Date: Thu, 6 Aug 2026 22:11:11 +0000 Subject: [PATCH 24/82] cutile pip at build time --- .../all_cuda-133_arch-aarch64.yaml | 3 ++ .../all_cuda-133_arch-x86_64.yaml | 3 ++ .../bench_ann_cuda-133_arch-aarch64.yaml | 3 ++ .../bench_ann_cuda-133_arch-x86_64.yaml | 3 ++ conda/recipes/libcuvs/recipe.yaml | 15 ++++--- .../cutile/fused_1nn_cutile_matrix.json | 2 - dependencies.yaml | 39 +++++++++++++++++-- python/libcuvs/pyproject.toml | 2 +- 8 files changed, 59 insertions(+), 11 deletions(-) diff --git a/conda/environments/all_cuda-133_arch-aarch64.yaml b/conda/environments/all_cuda-133_arch-aarch64.yaml index 2f10891439..c43a75fbd2 100644 --- a/conda/environments/all_cuda-133_arch-aarch64.yaml +++ b/conda/environments/all_cuda-133_arch-aarch64.yaml @@ -36,6 +36,7 @@ dependencies: - nodejs>=22 - numpy>=1.23,<3.0 - openblas +- pip - pre-commit - pylibraft==26.8.*,>=0.0.0a0 - pytest @@ -45,4 +46,6 @@ dependencies: - scikit-build-core>=0.11.0 - scikit-learn>=1.5 - sysroot_linux-aarch64==2.28 +- pip: + - cuda-tile[tileiras] name: all_cuda-133_arch-aarch64 diff --git a/conda/environments/all_cuda-133_arch-x86_64.yaml b/conda/environments/all_cuda-133_arch-x86_64.yaml index 0d98002ba2..a21ffc42d7 100644 --- a/conda/environments/all_cuda-133_arch-x86_64.yaml +++ b/conda/environments/all_cuda-133_arch-x86_64.yaml @@ -35,6 +35,7 @@ dependencies: - nodejs>=22 - numpy>=1.23,<3.0 - openblas +- pip - pre-commit - pylibraft==26.8.*,>=0.0.0a0 - pytest @@ -44,4 +45,6 @@ dependencies: - scikit-build-core>=0.11.0 - scikit-learn>=1.5 - sysroot_linux-64==2.28 +- pip: + - cuda-tile[tileiras] name: all_cuda-133_arch-x86_64 diff --git a/conda/environments/bench_ann_cuda-133_arch-aarch64.yaml b/conda/environments/bench_ann_cuda-133_arch-aarch64.yaml index c874298d7d..889567eb43 100644 --- a/conda/environments/bench_ann_cuda-133_arch-aarch64.yaml +++ b/conda/environments/bench_ann_cuda-133_arch-aarch64.yaml @@ -39,6 +39,7 @@ dependencies: - openblas - opensearch-py>=2.4.0 - pandas +- pip - pylibraft==26.8.*,>=0.0.0a0 - pyyaml - rapids-build-backend>=0.4.0,<0.5.0 @@ -47,4 +48,6 @@ dependencies: - setuptools>=77.0.0 - sysroot_linux-aarch64==2.28 - wheel +- pip: + - cuda-tile[tileiras] name: bench_ann_cuda-133_arch-aarch64 diff --git a/conda/environments/bench_ann_cuda-133_arch-x86_64.yaml b/conda/environments/bench_ann_cuda-133_arch-x86_64.yaml index 43b2dbce46..d952d7721f 100644 --- a/conda/environments/bench_ann_cuda-133_arch-x86_64.yaml +++ b/conda/environments/bench_ann_cuda-133_arch-x86_64.yaml @@ -42,6 +42,7 @@ dependencies: - openblas - opensearch-py>=2.4.0 - pandas +- pip - pylibraft==26.8.*,>=0.0.0a0 - pyyaml - rapids-build-backend>=0.4.0,<0.5.0 @@ -50,4 +51,6 @@ dependencies: - setuptools>=77.0.0 - sysroot_linux-64==2.28 - wheel +- pip: + - cuda-tile[tileiras] name: bench_ann_cuda-133_arch-x86_64 diff --git a/conda/recipes/libcuvs/recipe.yaml b/conda/recipes/libcuvs/recipe.yaml index 93f31f8cf2..b48011812f 100644 --- a/conda/recipes/libcuvs/recipe.yaml +++ b/conda/recipes/libcuvs/recipe.yaml @@ -30,6 +30,12 @@ cache: export CXXFLAGS=$(echo $CXXFLAGS | sed -E 's@\-fdebug\-prefix\-map[^ ]*@@g') set +x + # cuTile and TileIRAS are PyPI-only build tools. Their CUDA components + # are installed under the Python package namespace, alongside the conda toolchain. + if [[ "${{ cuda_major }}" == "13" ]]; then + python -m pip install --no-cache-dir "cuda-tile[tileiras]" + fi + ./build.sh libcuvs bench-ann tests --allgpuarch --build-metrics=compile_lib --incl-cache-stats --no-nvtx -n secrets: @@ -80,11 +86,14 @@ cache: - cuda-cudart-dev - cuda-nvrtc-dev - cuda-profiler-api - - cutile-python - libcublas-dev - libcurand-dev - libcusolver-dev - libcusparse-dev + - if: cuda_major == "13" + then: + - pip + - python # These are used for bench-ann - openblas - if: linux64 @@ -118,7 +127,6 @@ outputs: - cuda-cudart-dev - cuda-nvrtc-dev - cuda-profiler-api - - cutile-python - libcublas-dev - libcurand-dev - libcusolver-dev @@ -181,7 +189,6 @@ outputs: - cuda-cudart-dev - cuda-nvrtc-dev - cuda-profiler-api - - cutile-python - libcublas-dev - libcurand-dev - libcusolver-dev @@ -243,7 +250,6 @@ outputs: - cuda-cudart-dev - cuda-nvrtc-dev - cuda-profiler-api - - cutile-python - libcublas-dev - libcurand-dev - libcusolver-dev @@ -303,7 +309,6 @@ outputs: - openblas # required by some CPU algos in benchmarks - cuda-cudart-dev - cuda-profiler-api - - cutile-python - libcublas-dev - libcurand-dev - libcusolver-dev diff --git a/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_cutile_matrix.json b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_cutile_matrix.json index c57e8a2658..c94aec352d 100644 --- a/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_cutile_matrix.json +++ b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_cutile_matrix.json @@ -1,6 +1,5 @@ [ { - "python_executable": "../../../../../../tileiras13.4/.venv/bin/python", "_abi": [ { "matrix_layout": "strict", @@ -93,7 +92,6 @@ ] }, { - "python_executable": "../../../../../../tileiras13.4/.venv/bin/python", "_abi": [ { "matrix_layout": "strict", diff --git a/dependencies.yaml b/dependencies.yaml index 756041e60c..dae4e2c25e 100644 --- a/dependencies.yaml +++ b/dependencies.yaml @@ -13,6 +13,7 @@ files: - checks - clang - cuda + - cutile_python - cuda_version - depends_on_cuda_python - depends_on_cupy @@ -40,6 +41,7 @@ files: - build_py_cuvs - clang - cuda + - cutile_python - cuda_version - depends_on_cuda_python - depends_on_cupy @@ -77,6 +79,7 @@ files: includes: - clang - cuda + - cutile_python - cuda_version - depends_on_cupy - docs @@ -138,6 +141,7 @@ files: table: tool.rapids-build-backend key: requires includes: + - cutile_python - depends_on_libraft - depends_on_librmm - depends_on_nccl @@ -395,12 +399,43 @@ dependencies: - cuda-nvrtc-dev - cuda-nvtx-dev - cuda-profiler-api - - cutile-python - libcublas-dev - libcurand-dev - libcusolver-dev - libcusparse-dev - libnvjitlink-dev + cutile_python: + specific: + - output_types: conda + matrices: + - matrix: + cuda: "12.*" + packages: + - matrix: + cuda: "13.*" + packages: + - pip + - pip: + - cuda-tile[tileiras] + - matrix: + packages: + - pip + - pip: + - cuda-tile[tileiras] + - output_types: [requirements, pyproject] + matrices: + - matrix: + cuda: "12.*" + packages: + - matrix: + cuda: "13.*" + packages: + - &cutile_python_cu13 cuda-tile[tileiras] + # if no matching matrix selectors passed, list the CUDA 13 requirement + # (as a source of documentation in the generated pyproject.toml) + - matrix: + packages: + - *cutile_python_cu13 cuda_wheels: specific: # cuVS needs 'nvJitLink>={whatever-cuvs-was-built-against}' at runtime, and mixing @@ -431,14 +466,12 @@ dependencies: packages: - &ctk_cu13 cuda-toolkit[cublas,curand,cusolver,cusparse,nvrtc]==13.* - &nvjitlink_cu13 nvidia-nvjitlink>=13.0,<14 - - &cutile_cu13 cuda-tile[tileiras] # if no matching matrix selectors passed, list the CUDA 13 requirement # (just as a source of documentation, as this populates pyproject.toml in source control) - matrix: packages: - *ctk_cu13 - *nvjitlink_cu13 - - *cutile_cu13 depends_on_cudart: common: - output_types: conda diff --git a/python/libcuvs/pyproject.toml b/python/libcuvs/pyproject.toml index b4e848304f..3c69264dfd 100644 --- a/python/libcuvs/pyproject.toml +++ b/python/libcuvs/pyproject.toml @@ -19,7 +19,6 @@ authors = [ license = "Apache-2.0" requires-python = ">=3.11" dependencies = [ - "cuda-tile[tileiras]", "cuda-toolkit[cublas,curand,cusolver,cusparse,nvrtc]==13.*", "libraft==26.8.*,>=0.0.0a0", "librmm==26.8.*,>=0.0.0a0", @@ -82,6 +81,7 @@ regex = "(?P.*)" build-backend = "scikit_build_core.build" requires = [ "cmake>=4.0", + "cuda-tile[tileiras]", "libraft==26.8.*,>=0.0.0a0", "librmm==26.8.*,>=0.0.0a0", "ninja", From a69cb0116d0324cf5c5dd07bb3edaea6392605d7 Mon Sep 17 00:00:00 2001 From: Divye Gala Date: Thu, 6 Aug 2026 18:39:30 -0400 Subject: [PATCH 25/82] Delete cpp/src/cluster/detail/.nfs000000001cb8b01800001a98 --- .../detail/.nfs000000001cb8b01800001a98 | 414 ------------------ 1 file changed, 414 deletions(-) delete mode 100644 cpp/src/cluster/detail/.nfs000000001cb8b01800001a98 diff --git a/cpp/src/cluster/detail/.nfs000000001cb8b01800001a98 b/cpp/src/cluster/detail/.nfs000000001cb8b01800001a98 deleted file mode 100644 index 63f5d20ba6..0000000000 --- a/cpp/src/cluster/detail/.nfs000000001cb8b01800001a98 +++ /dev/null @@ -1,414 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: Apache-2.0 - */ - -#include "../../distance/fused_distance_nn.cuh" -#include "../../distance/unfused_distance_nn.cuh" -#include "kmeans_common.cuh" - -#include - -namespace cuvs::cluster::kmeans::detail { - -namespace { - -template -__global__ void unpack_kvp_to_soa(IndexT* nearest_idx, - DataT* nearest_dist, - const raft::KeyValuePair* kvp, - IndexT n) -{ - IndexT i = blockIdx.x * blockDim.x + threadIdx.x; - if (i < n) { - if (nearest_idx != nullptr) { nearest_idx[i] = kvp[i].key; } - if (nearest_dist != nullptr) { nearest_dist[i] = kvp[i].value; } - } -} - -template -void unpack_kvp(raft::resources const& handle, - raft::device_vector_view nearest_idx, - raft::device_vector_view nearest_dist, - raft::device_vector_view, IndexT> kvp) -{ - auto stream = raft::resource::get_cuda_stream(handle); - auto n = static_cast(kvp.extent(0)); - int blks = static_cast((n + 255) / 256); - unpack_kvp_to_soa<<>>( - nearest_idx.data_handle(), nearest_dist.data_handle(), kvp.data_handle(), n); - RAFT_CUDA_TRY(cudaGetLastError()); -} - -} // namespace - -template -void minClusterAndDistanceCompute(raft::resources const& handle, - raft::device_matrix_view X, - raft::device_matrix_view centroids, - raft::device_vector_view nearest_idx, - raft::device_vector_view nearest_dist, - raft::device_vector_view L2NormX, - rmm::device_uvector& L2NormBuf_OR_DistBuf, - cuvs::distance::DistanceType metric, - int batch_samples, - int batch_centroids, - rmm::device_uvector& workspace) -{ - cudaStream_t stream = raft::resource::get_cuda_stream(handle); - auto n_samples = X.extent(0); - auto n_features = X.extent(1); - auto n_clusters = centroids.extent(0); - const bool is_l2_cos = metric == cuvs::distance::DistanceType::L2Expanded || - metric == cuvs::distance::DistanceType::L2SqrtExpanded || - metric == cuvs::distance::DistanceType::CosineExpanded; - const FusedDistancePath fused_path = - use_fused(handle, n_samples, n_clusters, n_features, metric); - - if (uses_fused_distance_nn(fused_path)) { - L2NormBuf_OR_DistBuf.resize(n_clusters, stream); - auto centroidsNorm = - raft::make_device_vector_view(L2NormBuf_OR_DistBuf.data(), n_clusters); - - if (is_l2_cos) { - if (metric == cuvs::distance::DistanceType::CosineExpanded) { - raft::linalg::norm( - handle, centroids, centroidsNorm, raft::sqrt_op{}); - } else { - raft::linalg::norm( - handle, centroids, centroidsNorm); - } - } - - auto centroidsNormConst = - raft::make_device_vector_view(L2NormBuf_OR_DistBuf.data(), n_clusters); - - raft::KeyValuePair* cutlass_kvp_scratch = nullptr; - rmm::device_uvector> temp_kvp(0, stream); - if (needs_cutlass_kvp_scratch(fused_path)) { - temp_kvp.resize(n_samples, stream); - cutlass_kvp_scratch = temp_kvp.data(); - workspace.resize(sizeof(int) * n_samples, stream); - } else if constexpr (std::is_same_v) { - // The cuTile kernel uses i32 internally and widens labels after the launch. - workspace.resize(sizeof(int) * static_cast(n_samples), stream); - } - - cuvs::distance::fusedDistanceNNMinReduce( - nearest_idx.data_handle(), - nearest_dist.data_handle(), - X.data_handle(), - centroids.data_handle(), - L2NormX.data_handle(), - centroidsNormConst.data_handle(), - n_samples, - n_clusters, - n_features, - needs_fused_mutex_workspace(fused_path) || std::is_same_v - ? (void*)workspace.data() - : nullptr, - metric != cuvs::distance::DistanceType::L2Expanded, - true, - true, - metric, - 0.0f, - cutlass_kvp_scratch, - stream); - } else if (is_l2_cos) { - L2NormBuf_OR_DistBuf.resize(n_clusters, stream); - auto centroidsNorm = - raft::make_device_vector_view(L2NormBuf_OR_DistBuf.data(), n_clusters); - - if (metric == cuvs::distance::DistanceType::CosineExpanded) { - raft::linalg::norm( - handle, centroids, centroidsNorm, raft::sqrt_op{}); - } else { - raft::linalg::norm( - handle, centroids, centroidsNorm); - } - - auto centroidsNormConst = - raft::make_device_vector_view(L2NormBuf_OR_DistBuf.data(), n_clusters); - - auto dataBatchSize = getDataBatchSize(batch_samples, n_samples); - auto centroidsBatchSize = getCentroidsBatchSize(batch_centroids, n_clusters); - - // The unfused reduction indexes its distance matrix with IndexT. - dataBatchSize = - std::min(dataBatchSize, std::numeric_limits::max() / centroidsBatchSize); - workspace.resize(sizeof(DataT) * dataBatchSize * centroidsBatchSize, stream); - - using KeyValueT = raft::KeyValuePair; - auto temp_kvp = raft::make_device_vector(handle, n_samples); - KeyValueT initial_value(0, std::numeric_limits::max()); - raft::matrix::fill(handle, temp_kvp.view(), initial_value); - - const bool tileCentroids = centroidsBatchSize < n_clusters; - rmm::device_uvector batchMinClusterAndDistance(tileCentroids ? dataBatchSize : 0, - stream); - - for (IndexT dIdx = 0; dIdx < n_samples;) { - auto ns = std::min(dataBatchSize, n_samples - dIdx); - auto minClusterAndDistanceView = raft::make_device_vector_view( - temp_kvp.data_handle() + dIdx, ns); - - for (IndexT cIdx = 0; cIdx < n_clusters;) { - auto nc = std::min(centroidsBatchSize, n_clusters - cIdx); - auto batchMin = tileCentroids ? batchMinClusterAndDistance.data() - : minClusterAndDistanceView.data_handle(); - - cuvs::distance::unfusedDistanceNNMinReduce( - handle, - batchMin, - X.data_handle() + dIdx * n_features, - centroids.data_handle() + cIdx * n_features, - L2NormX.data_handle() + dIdx, - centroidsNormConst.data_handle() + cIdx, - ns, - nc, - n_features, - (void*)workspace.data(), - metric != cuvs::distance::DistanceType::L2Expanded, - tileCentroids, - true, - metric, - 0.0f, - stream); - - if (tileCentroids) { - // Convert tile-local centroid indices and merge the tile minima. - auto batchMinView = raft::make_device_vector_view(batchMin, ns); - raft::linalg::map( - handle, - minClusterAndDistanceView, - [cIdx] __device__(KeyValueT current, KeyValueT batch) { - batch.key += cIdx; - return batch.value < current.value ? batch : current; - }, - raft::make_const_mdspan(minClusterAndDistanceView), - batchMinView); - } - cIdx += nc; - } - dIdx += ns; - } - - unpack_kvp(handle, nearest_idx, nearest_dist, raft::make_const_mdspan(temp_kvp.view())); - } else { - auto dataBatchSize = getDataBatchSize(batch_samples, n_samples); - auto centroidsBatchSize = getCentroidsBatchSize(batch_centroids, n_clusters); - - L2NormBuf_OR_DistBuf.resize(dataBatchSize * centroidsBatchSize, stream); - - auto pairwiseDistance = raft::make_device_matrix_view( - L2NormBuf_OR_DistBuf.data(), dataBatchSize, centroidsBatchSize); - - auto temp_kvp = - raft::make_device_vector, IndexT>(handle, n_samples); - raft::KeyValuePair initial_value(0, std::numeric_limits::max()); - raft::matrix::fill(handle, temp_kvp.view(), initial_value); - - for (IndexT dIdx = 0; dIdx < n_samples; dIdx += dataBatchSize) { - auto ns = std::min((IndexT)dataBatchSize, n_samples - dIdx); - - auto datasetView = raft::make_device_matrix_view( - X.data_handle() + (dIdx * n_features), ns, n_features); - - auto temp_kvp_view = raft::make_device_vector_view, IndexT>( - temp_kvp.data_handle() + dIdx, ns); - - for (IndexT cIdx = 0; cIdx < n_clusters; cIdx += centroidsBatchSize) { - auto nc = std::min((IndexT)centroidsBatchSize, n_clusters - cIdx); - - auto centroidsView = raft::make_device_matrix_view( - centroids.data_handle() + (cIdx * n_features), nc, n_features); - - auto pairwiseDistanceView = - raft::make_device_matrix_view(pairwiseDistance.data_handle(), ns, nc); - - pairwise_distance_kmeans( - handle, datasetView, centroidsView, pairwiseDistanceView, metric); - - raft::linalg::coalescedReduction( - temp_kvp_view.data_handle(), - pairwiseDistanceView.data_handle(), - pairwiseDistanceView.extent(1), - pairwiseDistanceView.extent(0), - initial_value, - stream, - true, - [=] __device__(const DataT val, const IndexT i) { - raft::KeyValuePair pair; - pair.key = cIdx + i; - pair.value = val; - return pair; - }, - raft::argmin_op{}, - raft::identity_op{}); - } - } - - unpack_kvp(handle, nearest_idx, nearest_dist, raft::make_const_mdspan(temp_kvp.view())); - } -} - -#define INSTANTIATE_MIN_CLUSTER_AND_DISTANCE(DataT, IndexT) \ - template void minClusterAndDistanceCompute( \ - raft::resources const& handle, \ - raft::device_matrix_view X, \ - raft::device_matrix_view centroids, \ - raft::device_vector_view nearest_idx, \ - raft::device_vector_view nearest_dist, \ - raft::device_vector_view L2NormX, \ - rmm::device_uvector& L2NormBuf_OR_DistBuf, \ - cuvs::distance::DistanceType metric, \ - int batch_samples, \ - int batch_centroids, \ - rmm::device_uvector& workspace); - -INSTANTIATE_MIN_CLUSTER_AND_DISTANCE(float, int64_t) -INSTANTIATE_MIN_CLUSTER_AND_DISTANCE(double, int64_t) -INSTANTIATE_MIN_CLUSTER_AND_DISTANCE(float, int) -INSTANTIATE_MIN_CLUSTER_AND_DISTANCE(double, int) - -#undef INSTANTIATE_MIN_CLUSTER_AND_DISTANCE - -template -void minClusterDistanceCompute(raft::resources const& handle, - raft::device_matrix_view X, - raft::device_matrix_view centroids, - raft::device_vector_view minClusterDistance, - raft::device_vector_view L2NormX, - rmm::device_uvector& L2NormBuf_OR_DistBuf, - cuvs::distance::DistanceType metric, - int batch_samples, - int batch_centroids, - rmm::device_uvector& workspace) -{ - cudaStream_t stream = raft::resource::get_cuda_stream(handle); - auto n_samples = X.extent(0); - auto n_features = X.extent(1); - auto n_clusters = centroids.extent(0); - - const bool is_l2_cos = metric == cuvs::distance::DistanceType::L2Expanded || - metric == cuvs::distance::DistanceType::L2SqrtExpanded || - metric == cuvs::distance::DistanceType::CosineExpanded; - - raft::matrix::fill(handle, minClusterDistance, std::numeric_limits::max()); - - const FusedDistancePath fused_path = - is_l2_cos ? use_fused(handle, n_samples, n_clusters, n_features, metric) - : FusedDistancePath::Unfused; - - if (uses_fused_distance_nn(fused_path)) { - L2NormBuf_OR_DistBuf.resize(n_clusters, stream); - auto centroidsNorm = - raft::make_device_vector_view(L2NormBuf_OR_DistBuf.data(), n_clusters); - - if (metric == cuvs::distance::DistanceType::CosineExpanded) { - raft::linalg::norm( - handle, - raft::make_device_matrix_view( - centroids.data_handle(), centroids.extent(0), centroids.extent(1)), - centroidsNorm, - raft::sqrt_op{}); - } else { - raft::linalg::norm( - handle, - raft::make_device_matrix_view( - centroids.data_handle(), centroids.extent(0), centroids.extent(1)), - centroidsNorm); - } - - raft::KeyValuePair* cutlass_kvp_scratch = nullptr; - rmm::device_uvector> temp_kvp(0, stream); - if (needs_cutlass_kvp_scratch(fused_path)) { - temp_kvp.resize(n_samples, stream); - cutlass_kvp_scratch = temp_kvp.data(); - workspace.resize(sizeof(int) * n_samples, stream); - } - - cuvs::distance::fusedDistanceNNMinReduce( - nullptr, - minClusterDistance.data_handle(), - X.data_handle(), - centroids.data_handle(), - L2NormX.data_handle(), - centroidsNorm.data_handle(), - n_samples, - n_clusters, - n_features, - needs_fused_mutex_workspace(fused_path) ? (void*)workspace.data() : nullptr, - metric != cuvs::distance::DistanceType::L2Expanded, - true, - true, - metric, - 0.0f, - cutlass_kvp_scratch, - stream); - } else { - auto dataBatchSize = getDataBatchSize(batch_samples, n_samples); - auto centroidsBatchSize = getCentroidsBatchSize(batch_centroids, n_clusters); - - L2NormBuf_OR_DistBuf.resize(dataBatchSize * centroidsBatchSize, stream); - - auto pairwiseDistance = raft::make_device_matrix_view( - L2NormBuf_OR_DistBuf.data(), dataBatchSize, centroidsBatchSize); - - for (IndexT dIdx = 0; dIdx < n_samples; dIdx += dataBatchSize) { - auto ns = std::min((IndexT)dataBatchSize, n_samples - dIdx); - - auto datasetView = raft::make_device_matrix_view( - X.data_handle() + dIdx * n_features, ns, n_features); - - auto minClusterDistanceView = - raft::make_device_vector_view(minClusterDistance.data_handle() + dIdx, ns); - - for (IndexT cIdx = 0; cIdx < n_clusters; cIdx += centroidsBatchSize) { - auto nc = std::min((IndexT)centroidsBatchSize, n_clusters - cIdx); - - auto centroidsView = raft::make_device_matrix_view( - centroids.data_handle() + cIdx * n_features, nc, n_features); - - auto pairwiseDistanceView = - raft::make_device_matrix_view(pairwiseDistance.data_handle(), ns, nc); - - pairwise_distance_kmeans( - handle, datasetView, centroidsView, pairwiseDistanceView, metric); - - raft::linalg::coalescedReduction(minClusterDistanceView.data_handle(), - pairwiseDistanceView.data_handle(), - pairwiseDistanceView.extent(1), - pairwiseDistanceView.extent(0), - std::numeric_limits::max(), - stream, - true, - raft::identity_op{}, - raft::min_op{}, - raft::identity_op{}); - } - } - } -} - -#define INSTANTIATE_MIN_CLUSTER_DISTANCE(DataT, IndexT) \ - template void minClusterDistanceCompute( \ - raft::resources const& handle, \ - raft::device_matrix_view X, \ - raft::device_matrix_view centroids, \ - raft::device_vector_view minClusterDistance, \ - raft::device_vector_view L2NormX, \ - rmm::device_uvector& L2NormBuf_OR_DistBuf, \ - cuvs::distance::DistanceType metric, \ - int batch_samples, \ - int batch_centroids, \ - rmm::device_uvector& workspace); - -INSTANTIATE_MIN_CLUSTER_DISTANCE(float, int64_t) -INSTANTIATE_MIN_CLUSTER_DISTANCE(double, int64_t) -INSTANTIATE_MIN_CLUSTER_DISTANCE(float, int) -INSTANTIATE_MIN_CLUSTER_DISTANCE(double, int) - -#undef INSTANTIATE_MIN_CLUSTER_DISTANCE - -} // namespace cuvs::cluster::kmeans::detail From 2185f1c675a730ad93e2c5b5d6a4dd3738063dd4 Mon Sep 17 00:00:00 2001 From: Divye Gala Date: Thu, 6 Aug 2026 18:40:57 -0400 Subject: [PATCH 26/82] Delete benchmark_kmeans.py --- benchmark_kmeans.py | 506 -------------------------------------------- 1 file changed, 506 deletions(-) delete mode 100644 benchmark_kmeans.py diff --git a/benchmark_kmeans.py b/benchmark_kmeans.py deleted file mode 100644 index 3ff965527e..0000000000 --- a/benchmark_kmeans.py +++ /dev/null @@ -1,506 +0,0 @@ -#!/usr/bin/env python3 -# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. -# SPDX-License-Identifier: Apache-2.0 -r"""KMeans fit+predict benchmark: baseline / cuTile / flash-kmeans. - -Dimension glossary (same for cuVS and flash; matches fused GEMM A[M,D] @ B[K,D]^T): - M (--n) n_samples rows of X - D (--d) n_features inner / contraction dimension - K (--k) n_clusters centroids count (GEMM N) - -Shapes: - cuVS: X (M, D), centroids (K, D), KMeansParams(n_clusters=K) - flash: x (1, M, D), init_centroids (1, K, D), batch_kmeans_Euclid(..., K) - -Single impl (activate the target conda env first): - python benchmark_kmeans.py --impl baseline|cutile|flash --n M --d D --k K \\ - --phase fit|predict|both \\ - --max-iter 5 --tol 1e-4 --seed 42 \\ - --warmup-fit 1 --iters-fit 3 --warmup-pred 1 --iters-pred 3 - -Compare (subprocess per impl; export env vars, then --compare): - export BENCH_CONDA=/path/to/miniforge3 - export BENCH_ENV_BASE=cuvs_2608_base - export BENCH_ENV_CUTILE=cuvs_2608 - export BENCH_ENV_FLASH=cuvs_2608_base - python benchmark_kmeans.py --compare --n 1000000 --d 128 --k 256 \\ - --max-iter 5 --tol 1e-4 --seed 42 \\ - --warmup-fit 1 --iters-fit 3 --warmup-pred 1 --iters-pred 3 - -Smoke test (small shape, single impl): - conda activate cuvs_2608 - python benchmark_kmeans.py --impl cutile --n 10000 --d 32 --k 8 \\ - --max-iter 2 --tol 1e-4 --seed 42 \\ - --warmup-fit 0 --iters-fit 1 --warmup-pred 0 --iters-pred 1 - -Required for --compare (no defaults): - BENCH_CONDA path to miniforge/conda root - BENCH_ENV_BASE conda env name for baseline libcuvs - BENCH_ENV_CUTILE conda env name for cuTile libcuvs - BENCH_ENV_FLASH conda env name for flash-kmeans -""" - -from __future__ import annotations - -import argparse -import os -import re -import subprocess -import sys -import time -from dataclasses import dataclass -from pathlib import Path - -ROOT = Path(__file__).resolve().parent -IMPLS = ("baseline", "cutile", "flash") - - -def _require_env(name: str) -> str: - val = os.environ.get(name) - if not val: - raise SystemExit(f"required environment variable {name} is not set") - return val - - -def _impl_config() -> dict[str, dict]: - conda = Path(_require_env("BENCH_CONDA")) - return { - "baseline": { - "bench_mode": "cuvs_base", - "conda": conda, - "conda_env": _require_env("BENCH_ENV_BASE"), - }, - "cutile": { - "bench_mode": "cuvs_cutile", - "conda": conda, - "conda_env": _require_env("BENCH_ENV_CUTILE"), - }, - "flash": { - "bench_mode": "flash", - "conda": conda, - "conda_env": _require_env("BENCH_ENV_FLASH"), - }, - } - - -@dataclass -class BenchResult: - impl: str - fit_median_ms: float | None = None - predict_median_ms: float | None = None - n_iter: int | None = None - inertia: float | None = None - error: str | None = None - - -def median(xs: list[float]) -> float: - import numpy as np - - return float(np.median(xs)) - - -def run_benchmark( - bench_mode: str, - n: int, - d: int, - k: int, - *, - phase: str, - max_iter: int, - tol: float, - seed: int, - warmup_fit: int, - iters_fit: int, - warmup_pred: int, - iters_pred: int, -) -> BenchResult: - import numpy as np - - rng = np.random.default_rng(seed) - # Shared host data: X (M, D), centroids (K, D) — same layout for cuVS and flash. - init_centroids_host = rng.standard_normal((k, d), dtype=np.float32) - x_host = rng.standard_normal((n, d), dtype=np.float32) - input_gib = n * d * 4 / (1024**3) - - label = { - "cuvs_base": "baseline", - "cuvs_cutile": "cutile", - "flash": "flash", - }[bench_mode] - run_fit = phase in ("fit", "both") - run_predict = phase in ("predict", "both") - - print( - f"=== M={n:,} D={d} K={k:,} phase={phase} iters={max_iter} " - f"input={input_gib:.2f} GiB ===", - flush=True, - ) - - if bench_mode in ("cuvs_base", "cuvs_cutile"): - from cuda.bindings import runtime as cudart - from pylibraft.common import device_ndarray - - from cuvs.cluster.kmeans import KMeansParams, fit, predict - - def sync(): - cudart.cudaDeviceSynchronize() - - x = device_ndarray(x_host) # (M, D) - params = KMeansParams( - n_clusters=k, - max_iter=max_iter, - tol=tol, - metric="sqeuclidean", - hierarchical=False, - init_method="Array", - n_init=1, - ) - - for _ in range(warmup_fit if run_fit else 0): - fit( - params, x, centroids=device_ndarray(init_centroids_host.copy()) - ) - sync() - - fit_times: list[float] = [] - n_iter = 0 - inertia = 0.0 - if run_fit: - for _ in range(iters_fit): - t0 = time.perf_counter() - _, inertia, n_iter = fit( - params, - x, - centroids=device_ndarray(init_centroids_host.copy()), - ) - sync() - fit_times.append((time.perf_counter() - t0) * 1e3) - - pred_times: list[float] = [] - if run_predict: - centroids, _, _ = fit( - params, x, centroids=device_ndarray(init_centroids_host.copy()) - ) - sync() - - for _ in range(warmup_pred): - predict(params, x, centroids) - sync() - - for _ in range(iters_pred): - t0 = time.perf_counter() - predict(params, x, centroids) - sync() - pred_times.append((time.perf_counter() - t0) * 1e3) - - print(f"impl={label} init=Array", flush=True) - if run_fit: - print(f"fit_median_ms={median(fit_times):.2f}", flush=True) - print(f"n_iter={n_iter} inertia={inertia:.6g}", flush=True) - if run_predict: - print(f"predict_median_ms={median(pred_times):.2f}", flush=True) - return BenchResult( - impl=label, - fit_median_ms=median(fit_times) if run_fit else None, - predict_median_ms=median(pred_times) if run_predict else None, - n_iter=n_iter if run_fit else None, - inertia=inertia if run_fit else None, - ) - - if bench_mode == "flash": - import torch - from flash_kmeans.assign_euclid_triton import euclid_assign_triton - from flash_kmeans.kmeans_triton_impl import batch_kmeans_Euclid - - def sync(): - torch.cuda.synchronize() - - x = torch.from_numpy(x_host).cuda() # (M, D) - init_c = ( - torch.from_numpy(init_centroids_host.copy()).cuda().unsqueeze(0) - ) # (1, K, D) - - def run_fit(init): - x_b = x.unsqueeze(0) # (1, M, D) - _, centroids_b, _ = batch_kmeans_Euclid( - x_b, - k, # n_clusters - max_iters=max_iter, - tol=tol, - init_centroids=init, - verbose=False, - ) - return centroids_b - - for _ in range(warmup_fit if run_fit else 0): - run_fit(init_c.clone()) - sync() - - fit_times: list[float] = [] - if run_fit: - for _ in range(iters_fit): - t0 = time.perf_counter() - run_fit(init_c.clone()) - sync() - fit_times.append((time.perf_counter() - t0) * 1e3) - - pred_times: list[float] = [] - if run_predict: - centroids_b = run_fit(init_c.clone()) - sync() - - x_b = x.unsqueeze(0) - x_sq = (x_b**2).sum(dim=-1) - - for _ in range(warmup_pred): - euclid_assign_triton(x_b, centroids_b, x_sq) - sync() - - for _ in range(iters_pred): - t0 = time.perf_counter() - euclid_assign_triton(x_b, centroids_b, x_sq) - sync() - pred_times.append((time.perf_counter() - t0) * 1e3) - - print("impl=flash-kmeans init=Array", flush=True) - if run_fit: - print(f"fit_median_ms={median(fit_times):.2f}", flush=True) - if run_predict: - print(f"predict_median_ms={median(pred_times):.2f}", flush=True) - return BenchResult( - impl="flash", - fit_median_ms=median(fit_times) if run_fit else None, - predict_median_ms=median(pred_times) if run_predict else None, - ) - - raise ValueError(f"unknown bench_mode={bench_mode!r}") - - -def _parse_output(text: str, impl: str, phase: str) -> BenchResult: - fit_m = re.search(r"^fit_median_ms=([0-9.]+)", text, re.M) - pred_m = re.search(r"^predict_median_ms=([0-9.]+)", text, re.M) - if phase in ("fit", "both") and not fit_m: - return BenchResult(impl=impl, error=text.strip() or "no fit output") - if phase in ("predict", "both") and not pred_m: - return BenchResult( - impl=impl, error=text.strip() or "no predict output" - ) - n_iter_m = re.search(r"^n_iter=([0-9]+)", text, re.M) - inertia_m = re.search(r"^inertia=([0-9.eE+-]+)", text, re.M) - return BenchResult( - impl=impl, - fit_median_ms=float(fit_m.group(1)) if fit_m else None, - predict_median_ms=float(pred_m.group(1)) if pred_m else None, - n_iter=int(n_iter_m.group(1)) if n_iter_m else None, - inertia=float(inertia_m.group(1)) if inertia_m else None, - ) - - -def _result_ok(result: BenchResult, phase: str) -> bool: - if result.error: - return False - if phase in ("fit", "both") and result.fit_median_ms is None: - return False - if phase in ("predict", "both") and result.predict_median_ms is None: - return False - return True - - -def _fmt_ms(value: float | None) -> str: - return f"{value:10.2f}" if value is not None else f"{'—':>10}" - - -def _run_subprocess( - impl: str, - n: int, - d: int, - k: int, - args: argparse.Namespace, -) -> BenchResult: - cfg = _impl_config()[impl] - conda = cfg["conda"] - env_exports = " ".join( - f'export {key}="{val}"' - for key, val in ( - ( - "CUDA_VISIBLE_DEVICES", - os.environ.get("CUDA_VISIBLE_DEVICES", ""), - ), - ("MAX_ITER", args.max_iter), - ("TOL", args.tol), - ("SEED", args.seed), - ("WARMUP_FIT", args.warmup_fit), - ("ITERS_FIT", args.iters_fit), - ("WARMUP_PRED", args.warmup_pred), - ("ITERS_PRED", args.iters_pred), - ) - if val != "" - ) - cmd = f""" -set -eo pipefail -source "{conda}/etc/profile.d/conda.sh" -conda activate "{cfg["conda_env"]}" -{env_exports} -python3 "{ROOT / "benchmark_kmeans.py"}" --impl {impl} --n {n} --d {d} --k {k} \\ - --phase {args.phase} \\ - --max-iter {args.max_iter} --tol {args.tol} --seed {args.seed} \\ - --warmup-fit {args.warmup_fit} --iters-fit {args.iters_fit} \\ - --warmup-pred {args.warmup_pred} --iters-pred {args.iters_pred} -""" - proc = subprocess.run(["bash", "-lc", cmd], capture_output=True, text=True) - out = proc.stdout + proc.stderr - if proc.returncode != 0: - return BenchResult( - impl=impl, error=out.strip() or f"exit {proc.returncode}" - ) - return _parse_output(out, impl, args.phase) - - -def _speedup(base: float, other: float) -> str: - if other <= 0: - return "n/a" - return f"{base / other:.2f}x" - - -def print_compare_table( - results: list[BenchResult], n: int, d: int, k: int, phase: str -) -> None: - print(f"\n######## compare M={n} D={d} K={k} phase={phase} ########") - show_fit = phase in ("fit", "both") - show_pred = phase in ("predict", "both") - header = f"{'impl':<10}" - if show_fit: - header += f" {'fit_ms':>10}" - if show_pred: - header += f" {'pred_ms':>10}" - header += " notes" - print(header) - print("-" * len(header)) - by_impl = {r.impl: r for r in results} - for impl in IMPLS: - r = by_impl.get(impl) - if r is None: - row = f"{impl:<10}" - if show_fit: - row += f" {'—':>10}" - if show_pred: - row += f" {'—':>10}" - print(f"{row} missing") - continue - if r.error: - row = f"{impl:<10}" - if show_fit: - row += f" {'FAIL':>10}" - if show_pred: - row += f" {'FAIL':>10}" - print(f"{row} {r.error.splitlines()[-1][:40]}") - continue - row = f"{impl:<10}" - if show_fit: - row += _fmt_ms(r.fit_median_ms) - if show_pred: - row += _fmt_ms(r.predict_median_ms) - print(row) - - flash = by_impl.get("flash") - cutile = by_impl.get("cutile") - if flash and cutile and not flash.error and not cutile.error: - parts: list[str] = [] - if show_fit and flash.fit_median_ms and cutile.fit_median_ms: - parts.append( - f"fit: {_speedup(cutile.fit_median_ms, flash.fit_median_ms)}" - ) - if show_pred and flash.predict_median_ms and cutile.predict_median_ms: - parts.append( - "predict: " - f"{_speedup(cutile.predict_median_ms, flash.predict_median_ms)}" - ) - if parts: - print(f"\nflash vs cutile {' '.join(parts)}") - - -def main() -> int: - parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument( - "--compare", action="store_true", help="run baseline, cutile, flash" - ) - parser.add_argument("--impl", choices=IMPLS, help="single impl") - parser.add_argument( - "--n", - type=int, - required=True, - metavar="M", - help="n_samples (GEMM M)", - ) - parser.add_argument( - "--d", - type=int, - required=True, - metavar="D", - help="n_features (GEMM inner dimension)", - ) - parser.add_argument( - "--k", - type=int, - required=True, - metavar="K", - help="n_clusters (GEMM N)", - ) - parser.add_argument( - "--phase", - choices=("fit", "predict", "both"), - default="both", - help="benchmark fit only, predict only, or both (default: both)", - ) - parser.add_argument("--max-iter", type=int, required=True) - parser.add_argument("--tol", type=float, required=True) - parser.add_argument("--seed", type=int, required=True) - parser.add_argument("--warmup-fit", type=int, required=True) - parser.add_argument("--iters-fit", type=int, required=True) - parser.add_argument("--warmup-pred", type=int, required=True) - parser.add_argument("--iters-pred", type=int, required=True) - args = parser.parse_args() - - if args.compare: - if args.impl: - parser.error("--compare and --impl are mutually exclusive") - _impl_config() # validate required env before launching subprocesses - results = [ - _run_subprocess(impl, args.n, args.d, args.k, args) - for impl in IMPLS - ] - print_compare_table(results, args.n, args.d, args.k, args.phase) - return 0 if all(_result_ok(r, args.phase) for r in results) else 1 - - if not args.impl: - parser.error("set --impl for single-run mode, or use --compare") - - bench_mode = { - "baseline": "cuvs_base", - "cutile": "cuvs_cutile", - "flash": "flash", - }[args.impl] - - try: - run_benchmark( - bench_mode, - args.n, - args.d, - args.k, - phase=args.phase, - max_iter=args.max_iter, - tol=args.tol, - seed=args.seed, - warmup_fit=args.warmup_fit, - iters_fit=args.iters_fit, - warmup_pred=args.warmup_pred, - iters_pred=args.iters_pred, - ) - except Exception as exc: - print(f"ERROR: {exc}", file=sys.stderr) - return 1 - return 0 - - -if __name__ == "__main__": - raise SystemExit(main()) From 2fdd7f721068f4b2ede84665f4c7b83e54f50535 Mon Sep 17 00:00:00 2001 From: Divye Gala Date: Thu, 6 Aug 2026 18:43:33 -0400 Subject: [PATCH 27/82] Delete run_benchmark_kmeans.sh --- run_benchmark_kmeans.sh | 87 ----------------------------------------- 1 file changed, 87 deletions(-) delete mode 100755 run_benchmark_kmeans.sh diff --git a/run_benchmark_kmeans.sh b/run_benchmark_kmeans.sh deleted file mode 100755 index 8109bd1f81..0000000000 --- a/run_benchmark_kmeans.sh +++ /dev/null @@ -1,87 +0,0 @@ -#!/usr/bin/env bash -# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -# Compare baseline / cuTile / flash-kmeans for one shape or the default sweep. -# -# Usage: -# export BENCH_CONDA=/path/to/miniforge3 -# export BENCH_ENV_BASE=... -# export BENCH_ENV_CUTILE=... -# export BENCH_ENV_FLASH=... -# export MAX_ITER=5 TOL=1e-4 SEED=42 -# export WARMUP_FIT=1 ITERS_FIT=3 WARMUP_PRED=1 ITERS_PRED=3 -# export BENCH_PHASE=both # fit | predict | both (default: both) -# export BENCH_GPU_NAME=rtx_pro_6000 # sweep log tag (e.g. h200 on Hopper) -# export BENCH_LOG=path/to.log # optional sweep log override -# ./run_benchmark_kmeans.sh # default sweep (M=1M, all D and K below) -# ./run_benchmark_kmeans.sh N D K # single shape -# ./run_benchmark_kmeans.sh fit # sweep, fit only -# ./run_benchmark_kmeans.sh predict N D K # single shape, predict only -# -# Default sweep grid: -# M (n_samples) = 1_000_000 -# D (n_features) = 16 64 128 384 768 1024 1536 # GEMM inner (K) dimension -# K (n_clusters) = 10 100 1000 10000 100000 # GEMM N dimension -# - -set -u - -SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" - -: "${BENCH_CONDA:?set BENCH_CONDA to conda/miniforge root}" -: "${BENCH_ENV_BASE:?set BENCH_ENV_BASE}" -: "${BENCH_ENV_CUTILE:?set BENCH_ENV_CUTILE}" -: "${BENCH_ENV_FLASH:?set BENCH_ENV_FLASH}" -: "${MAX_ITER:?set MAX_ITER}" -: "${SEED:?set SEED}" -: "${WARMUP_FIT:?set WARMUP_FIT}" -: "${ITERS_FIT:?set ITERS_FIT}" -: "${WARMUP_PRED:?set WARMUP_PRED}" -: "${ITERS_PRED:?set ITERS_PRED}" -: "${TOL:?set TOL}" - -PHASE="${BENCH_PHASE:-both}" -if [[ $# -ge 1 && $1 =~ ^(fit|predict|both)$ ]]; then - PHASE=$1 - shift -fi - -run_shape() { - local n=$1 d=$2 k=$3 - echo "=== benchmark M=${n} D=${d} K=${k} phase=${PHASE} ===" - python3 "$SCRIPT_DIR/benchmark_kmeans.py" --compare \ - --n "$n" --d "$d" --k "$k" \ - --phase "$PHASE" \ - --max-iter "$MAX_ITER" --tol "$TOL" --seed "$SEED" \ - --warmup-fit "$WARMUP_FIT" --iters-fit "$ITERS_FIT" \ - --warmup-pred "$WARMUP_PRED" --iters-pred "$ITERS_PRED" -} - -if [[ $# -eq 3 ]]; then - run_shape "$1" "$2" "$3" - exit $? -fi - -if [[ $# -ne 0 ]]; then - echo "usage: $0 [fit|predict|both] [N D K]" >&2 - echo " no args — sweep all shapes, phase=\${BENCH_PHASE:-both}" >&2 - echo " fit|predict|both — optional phase override, then sweep" >&2 - echo " [phase] N D K — run one shape" >&2 - exit 2 -fi - -: "${BENCH_GPU_NAME:?set BENCH_GPU_NAME e.g. rtx_pro_6000}" -BENCH_LOG="${BENCH_LOG:-${SCRIPT_DIR}/benchmark_kmeans_sweep_${BENCH_GPU_NAME}_$(date +%Y%m%d_%H%M%S).log}" -echo "Logging to ${BENCH_LOG}" -exec > >(tee "$BENCH_LOG") 2>&1 - -M=1000000 -D_VALUES=(16 64 128 384 768 1024 1536) -K_VALUES=(10 100 1000 10000 100000) - -for d in "${D_VALUES[@]}"; do - for k in "${K_VALUES[@]}"; do - run_shape "$M" "$d" "$k" || true - done -done From 2fb62539cf2aa3a66cc200be73410c9697af64e7 Mon Sep 17 00:00:00 2001 From: divyegala Date: Thu, 6 Aug 2026 23:15:30 +0000 Subject: [PATCH 28/82] correct dependencies --- .../all_cuda-133_arch-aarch64.yaml | 3 ++- .../all_cuda-133_arch-x86_64.yaml | 3 ++- .../bench_ann_cuda-133_arch-aarch64.yaml | 3 ++- .../bench_ann_cuda-133_arch-x86_64.yaml | 3 ++- conda/recipes/libcuvs/recipe.yaml | 4 +++- dependencies.yaml | 22 ++++++++++++++++--- python/libcuvs/pyproject.toml | 3 ++- 7 files changed, 32 insertions(+), 9 deletions(-) diff --git a/conda/environments/all_cuda-133_arch-aarch64.yaml b/conda/environments/all_cuda-133_arch-aarch64.yaml index 83b154436d..bc20593702 100644 --- a/conda/environments/all_cuda-133_arch-aarch64.yaml +++ b/conda/environments/all_cuda-133_arch-aarch64.yaml @@ -48,5 +48,6 @@ dependencies: - scikit-learn>=1.5 - sysroot_linux-aarch64==2.28 - pip: - - cuda-tile[tileiras] + - cuda-tile + - cuda-toolkit[tileiras]==13.3.* name: all_cuda-133_arch-aarch64 diff --git a/conda/environments/all_cuda-133_arch-x86_64.yaml b/conda/environments/all_cuda-133_arch-x86_64.yaml index 97edea8dd2..c1347e6c67 100644 --- a/conda/environments/all_cuda-133_arch-x86_64.yaml +++ b/conda/environments/all_cuda-133_arch-x86_64.yaml @@ -47,5 +47,6 @@ dependencies: - scikit-learn>=1.5 - sysroot_linux-64==2.28 - pip: - - cuda-tile[tileiras] + - cuda-tile + - cuda-toolkit[tileiras]==13.3.* name: all_cuda-133_arch-x86_64 diff --git a/conda/environments/bench_ann_cuda-133_arch-aarch64.yaml b/conda/environments/bench_ann_cuda-133_arch-aarch64.yaml index 53a5c0032c..87e75a3e09 100644 --- a/conda/environments/bench_ann_cuda-133_arch-aarch64.yaml +++ b/conda/environments/bench_ann_cuda-133_arch-aarch64.yaml @@ -48,5 +48,6 @@ dependencies: - sysroot_linux-aarch64==2.28 - wheel - pip: - - cuda-tile[tileiras] + - cuda-tile + - cuda-toolkit[tileiras]==13.3.* name: bench_ann_cuda-133_arch-aarch64 diff --git a/conda/environments/bench_ann_cuda-133_arch-x86_64.yaml b/conda/environments/bench_ann_cuda-133_arch-x86_64.yaml index 7dd72df378..36cba2b0f5 100644 --- a/conda/environments/bench_ann_cuda-133_arch-x86_64.yaml +++ b/conda/environments/bench_ann_cuda-133_arch-x86_64.yaml @@ -51,5 +51,6 @@ dependencies: - sysroot_linux-64==2.28 - wheel - pip: - - cuda-tile[tileiras] + - cuda-tile + - cuda-toolkit[tileiras]==13.3.* name: bench_ann_cuda-133_arch-x86_64 diff --git a/conda/recipes/libcuvs/recipe.yaml b/conda/recipes/libcuvs/recipe.yaml index 4d3ca69a36..27bfd5c14b 100644 --- a/conda/recipes/libcuvs/recipe.yaml +++ b/conda/recipes/libcuvs/recipe.yaml @@ -33,7 +33,9 @@ cache: # cuTile and TileIRAS are PyPI-only build tools. Their CUDA components # are installed under the Python package namespace, alongside the conda toolchain. if [[ "${{ cuda_major }}" == "13" ]]; then - python -m pip install --no-cache-dir "cuda-tile[tileiras]" + python -m pip install --no-cache-dir \ + "cuda-tile" \ + "cuda-toolkit[tileiras]==${{ cuda_version }}.*" fi ./build.sh libcuvs bench-ann tests --allgpuarch --mnmg-tests --build-metrics=compile_lib --incl-cache-stats --no-nvtx -n diff --git a/dependencies.yaml b/dependencies.yaml index e1ff92f7bf..258b7edfa8 100644 --- a/dependencies.yaml +++ b/dependencies.yaml @@ -424,31 +424,47 @@ dependencies: - matrix: cuda: "12.*" packages: + - matrix: + cuda: "13.3" + packages: + - pip + - pip: + - cuda-tile + - cuda-toolkit[tileiras]==13.3.* - matrix: cuda: "13.*" packages: - pip - pip: - - cuda-tile[tileiras] + - cuda-tile + - cuda-toolkit[tileiras]==13.* - matrix: packages: - pip - pip: - - cuda-tile[tileiras] + - cuda-tile + - cuda-toolkit[tileiras]==13.* - output_types: [requirements, pyproject] matrices: - matrix: cuda: "12.*" packages: + - matrix: + cuda: "13.3" + packages: + - cuda-tile + - cuda-toolkit[tileiras]==13.3.* - matrix: cuda: "13.*" packages: - - &cutile_python_cu13 cuda-tile[tileiras] + - &cutile_python_cu13 cuda-tile + - &cutile_toolkit_cu13 cuda-toolkit[tileiras]==13.* # if no matching matrix selectors passed, list the CUDA 13 requirement # (as a source of documentation in the generated pyproject.toml) - matrix: packages: - *cutile_python_cu13 + - *cutile_toolkit_cu13 cuda_wheels: specific: # cuVS needs 'nvJitLink>={whatever-cuvs-was-built-against}' at runtime, and mixing diff --git a/python/libcuvs/pyproject.toml b/python/libcuvs/pyproject.toml index fbcee95c0f..323679975c 100644 --- a/python/libcuvs/pyproject.toml +++ b/python/libcuvs/pyproject.toml @@ -81,7 +81,8 @@ regex = "(?P.*)" build-backend = "scikit_build_core.build" requires = [ "cmake>=4.0", - "cuda-tile[tileiras]", + "cuda-tile", + "cuda-toolkit[tileiras]==13.*", "libraft==26.10.*,>=0.0.0a0", "librmm==26.10.*,>=0.0.0a0", "ninja", From afc3cb5567d6562597af760d4c368f6aa9dbd1c9 Mon Sep 17 00:00:00 2001 From: divyegala Date: Thu, 6 Aug 2026 23:52:06 +0000 Subject: [PATCH 29/82] package and style check --- .../modules/generate_cutile_kernels.cmake | 2 +- .../modules/register_cutile_fragment.cpp.in | 2 +- .../cuvs/detail/jit_lto/AlgorithmPlanner.hpp | 2 +- .../cuvs/detail/jit_lto/FragmentEntry.hpp | 2 +- .../cuvs/detail/jit_lto/cutile_arch_tags.hpp | 2 +- .../cuvs/detail/jit_lto/cutile_module.hpp | 2 +- .../fused_distance_nn/fused_1nn_fragments.hpp | 2 +- .../cuvs/detail/jit_lto/tileir_compat.hpp | 2 +- .../detail/.nfs000000001cb8b01800001a98 | 414 ++++++++++++++++++ cpp/src/detail/jit_lto/AlgorithmPlanner.cpp | 2 +- .../detail/jit_lto/LTOAlgorithmPlanner.cpp | 2 +- .../detail/jit_lto/TileAlgorithmPlanner.cpp | 2 +- cpp/src/distance/detail/fused_distance_nn.cuh | 2 +- .../cutile/export_fused_1nn.py | 10 +- .../cutile/fused_1nn_kernel.py | 2 +- .../cutile/fused_1nn_planner.hpp | 2 +- .../cutile/fused_1nn_tile.cu | 2 +- .../cutile/fused_1nn_tile.hpp | 2 +- .../fused_distance_nn/fused_cosine_nn.cuh | 2 +- .../detail/fused_distance_nn/fused_l2_nn.cuh | 2 +- .../fused_distance_nn/helper_structs.cuh | 2 +- .../predicated_tile_iterator_reduced_vec.h | 2 +- .../pairwise_matrix_planner.hpp | 2 +- cpp/src/distance/fused_distance_nn-inl.cuh | 2 +- .../interleaved_scan_planner.hpp | 2 +- .../compute_similarity_planner.hpp | 2 +- .../detail/jit_lto_kernels/scan_planner.hpp | 2 +- cpp/tests/neighbors/distance_nn.cu | 2 +- cpp/tests/neighbors/distance_nn_helper.cuh | 2 +- 29 files changed, 449 insertions(+), 29 deletions(-) create mode 100644 cpp/src/cluster/detail/.nfs000000001cb8b01800001a98 diff --git a/cpp/cmake/modules/generate_cutile_kernels.cmake b/cpp/cmake/modules/generate_cutile_kernels.cmake index f06ad3801f..d37513d817 100644 --- a/cpp/cmake/modules/generate_cutile_kernels.cmake +++ b/cpp/cmake/modules/generate_cutile_kernels.cmake @@ -1,6 +1,6 @@ # ============================================================================= # cmake-format: off -# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # cmake-format: on # ============================================================================= diff --git a/cpp/cmake/modules/register_cutile_fragment.cpp.in b/cpp/cmake/modules/register_cutile_fragment.cpp.in index de0472a779..79d4d00008 100644 --- a/cpp/cmake/modules/register_cutile_fragment.cpp.in +++ b/cpp/cmake/modules/register_cutile_fragment.cpp.in @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/include/cuvs/detail/jit_lto/AlgorithmPlanner.hpp b/cpp/include/cuvs/detail/jit_lto/AlgorithmPlanner.hpp index 7ff8487d20..0cb2b9a876 100644 --- a/cpp/include/cuvs/detail/jit_lto/AlgorithmPlanner.hpp +++ b/cpp/include/cuvs/detail/jit_lto/AlgorithmPlanner.hpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/include/cuvs/detail/jit_lto/FragmentEntry.hpp b/cpp/include/cuvs/detail/jit_lto/FragmentEntry.hpp index 0961595f8d..ba477faf7c 100644 --- a/cpp/include/cuvs/detail/jit_lto/FragmentEntry.hpp +++ b/cpp/include/cuvs/detail/jit_lto/FragmentEntry.hpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/include/cuvs/detail/jit_lto/cutile_arch_tags.hpp b/cpp/include/cuvs/detail/jit_lto/cutile_arch_tags.hpp index 1b9f58837c..95f7c81ecd 100644 --- a/cpp/include/cuvs/detail/jit_lto/cutile_arch_tags.hpp +++ b/cpp/include/cuvs/detail/jit_lto/cutile_arch_tags.hpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/include/cuvs/detail/jit_lto/cutile_module.hpp b/cpp/include/cuvs/detail/jit_lto/cutile_module.hpp index dff0f472a7..29f30ae15a 100644 --- a/cpp/include/cuvs/detail/jit_lto/cutile_module.hpp +++ b/cpp/include/cuvs/detail/jit_lto/cutile_module.hpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/include/cuvs/detail/jit_lto/fused_distance_nn/fused_1nn_fragments.hpp b/cpp/include/cuvs/detail/jit_lto/fused_distance_nn/fused_1nn_fragments.hpp index f4b64233ad..807dc50e24 100644 --- a/cpp/include/cuvs/detail/jit_lto/fused_distance_nn/fused_1nn_fragments.hpp +++ b/cpp/include/cuvs/detail/jit_lto/fused_distance_nn/fused_1nn_fragments.hpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/include/cuvs/detail/jit_lto/tileir_compat.hpp b/cpp/include/cuvs/detail/jit_lto/tileir_compat.hpp index f114233179..029d53bfa0 100644 --- a/cpp/include/cuvs/detail/jit_lto/tileir_compat.hpp +++ b/cpp/include/cuvs/detail/jit_lto/tileir_compat.hpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/src/cluster/detail/.nfs000000001cb8b01800001a98 b/cpp/src/cluster/detail/.nfs000000001cb8b01800001a98 new file mode 100644 index 0000000000..63f5d20ba6 --- /dev/null +++ b/cpp/src/cluster/detail/.nfs000000001cb8b01800001a98 @@ -0,0 +1,414 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include "../../distance/fused_distance_nn.cuh" +#include "../../distance/unfused_distance_nn.cuh" +#include "kmeans_common.cuh" + +#include + +namespace cuvs::cluster::kmeans::detail { + +namespace { + +template +__global__ void unpack_kvp_to_soa(IndexT* nearest_idx, + DataT* nearest_dist, + const raft::KeyValuePair* kvp, + IndexT n) +{ + IndexT i = blockIdx.x * blockDim.x + threadIdx.x; + if (i < n) { + if (nearest_idx != nullptr) { nearest_idx[i] = kvp[i].key; } + if (nearest_dist != nullptr) { nearest_dist[i] = kvp[i].value; } + } +} + +template +void unpack_kvp(raft::resources const& handle, + raft::device_vector_view nearest_idx, + raft::device_vector_view nearest_dist, + raft::device_vector_view, IndexT> kvp) +{ + auto stream = raft::resource::get_cuda_stream(handle); + auto n = static_cast(kvp.extent(0)); + int blks = static_cast((n + 255) / 256); + unpack_kvp_to_soa<<>>( + nearest_idx.data_handle(), nearest_dist.data_handle(), kvp.data_handle(), n); + RAFT_CUDA_TRY(cudaGetLastError()); +} + +} // namespace + +template +void minClusterAndDistanceCompute(raft::resources const& handle, + raft::device_matrix_view X, + raft::device_matrix_view centroids, + raft::device_vector_view nearest_idx, + raft::device_vector_view nearest_dist, + raft::device_vector_view L2NormX, + rmm::device_uvector& L2NormBuf_OR_DistBuf, + cuvs::distance::DistanceType metric, + int batch_samples, + int batch_centroids, + rmm::device_uvector& workspace) +{ + cudaStream_t stream = raft::resource::get_cuda_stream(handle); + auto n_samples = X.extent(0); + auto n_features = X.extent(1); + auto n_clusters = centroids.extent(0); + const bool is_l2_cos = metric == cuvs::distance::DistanceType::L2Expanded || + metric == cuvs::distance::DistanceType::L2SqrtExpanded || + metric == cuvs::distance::DistanceType::CosineExpanded; + const FusedDistancePath fused_path = + use_fused(handle, n_samples, n_clusters, n_features, metric); + + if (uses_fused_distance_nn(fused_path)) { + L2NormBuf_OR_DistBuf.resize(n_clusters, stream); + auto centroidsNorm = + raft::make_device_vector_view(L2NormBuf_OR_DistBuf.data(), n_clusters); + + if (is_l2_cos) { + if (metric == cuvs::distance::DistanceType::CosineExpanded) { + raft::linalg::norm( + handle, centroids, centroidsNorm, raft::sqrt_op{}); + } else { + raft::linalg::norm( + handle, centroids, centroidsNorm); + } + } + + auto centroidsNormConst = + raft::make_device_vector_view(L2NormBuf_OR_DistBuf.data(), n_clusters); + + raft::KeyValuePair* cutlass_kvp_scratch = nullptr; + rmm::device_uvector> temp_kvp(0, stream); + if (needs_cutlass_kvp_scratch(fused_path)) { + temp_kvp.resize(n_samples, stream); + cutlass_kvp_scratch = temp_kvp.data(); + workspace.resize(sizeof(int) * n_samples, stream); + } else if constexpr (std::is_same_v) { + // The cuTile kernel uses i32 internally and widens labels after the launch. + workspace.resize(sizeof(int) * static_cast(n_samples), stream); + } + + cuvs::distance::fusedDistanceNNMinReduce( + nearest_idx.data_handle(), + nearest_dist.data_handle(), + X.data_handle(), + centroids.data_handle(), + L2NormX.data_handle(), + centroidsNormConst.data_handle(), + n_samples, + n_clusters, + n_features, + needs_fused_mutex_workspace(fused_path) || std::is_same_v + ? (void*)workspace.data() + : nullptr, + metric != cuvs::distance::DistanceType::L2Expanded, + true, + true, + metric, + 0.0f, + cutlass_kvp_scratch, + stream); + } else if (is_l2_cos) { + L2NormBuf_OR_DistBuf.resize(n_clusters, stream); + auto centroidsNorm = + raft::make_device_vector_view(L2NormBuf_OR_DistBuf.data(), n_clusters); + + if (metric == cuvs::distance::DistanceType::CosineExpanded) { + raft::linalg::norm( + handle, centroids, centroidsNorm, raft::sqrt_op{}); + } else { + raft::linalg::norm( + handle, centroids, centroidsNorm); + } + + auto centroidsNormConst = + raft::make_device_vector_view(L2NormBuf_OR_DistBuf.data(), n_clusters); + + auto dataBatchSize = getDataBatchSize(batch_samples, n_samples); + auto centroidsBatchSize = getCentroidsBatchSize(batch_centroids, n_clusters); + + // The unfused reduction indexes its distance matrix with IndexT. + dataBatchSize = + std::min(dataBatchSize, std::numeric_limits::max() / centroidsBatchSize); + workspace.resize(sizeof(DataT) * dataBatchSize * centroidsBatchSize, stream); + + using KeyValueT = raft::KeyValuePair; + auto temp_kvp = raft::make_device_vector(handle, n_samples); + KeyValueT initial_value(0, std::numeric_limits::max()); + raft::matrix::fill(handle, temp_kvp.view(), initial_value); + + const bool tileCentroids = centroidsBatchSize < n_clusters; + rmm::device_uvector batchMinClusterAndDistance(tileCentroids ? dataBatchSize : 0, + stream); + + for (IndexT dIdx = 0; dIdx < n_samples;) { + auto ns = std::min(dataBatchSize, n_samples - dIdx); + auto minClusterAndDistanceView = raft::make_device_vector_view( + temp_kvp.data_handle() + dIdx, ns); + + for (IndexT cIdx = 0; cIdx < n_clusters;) { + auto nc = std::min(centroidsBatchSize, n_clusters - cIdx); + auto batchMin = tileCentroids ? batchMinClusterAndDistance.data() + : minClusterAndDistanceView.data_handle(); + + cuvs::distance::unfusedDistanceNNMinReduce( + handle, + batchMin, + X.data_handle() + dIdx * n_features, + centroids.data_handle() + cIdx * n_features, + L2NormX.data_handle() + dIdx, + centroidsNormConst.data_handle() + cIdx, + ns, + nc, + n_features, + (void*)workspace.data(), + metric != cuvs::distance::DistanceType::L2Expanded, + tileCentroids, + true, + metric, + 0.0f, + stream); + + if (tileCentroids) { + // Convert tile-local centroid indices and merge the tile minima. + auto batchMinView = raft::make_device_vector_view(batchMin, ns); + raft::linalg::map( + handle, + minClusterAndDistanceView, + [cIdx] __device__(KeyValueT current, KeyValueT batch) { + batch.key += cIdx; + return batch.value < current.value ? batch : current; + }, + raft::make_const_mdspan(minClusterAndDistanceView), + batchMinView); + } + cIdx += nc; + } + dIdx += ns; + } + + unpack_kvp(handle, nearest_idx, nearest_dist, raft::make_const_mdspan(temp_kvp.view())); + } else { + auto dataBatchSize = getDataBatchSize(batch_samples, n_samples); + auto centroidsBatchSize = getCentroidsBatchSize(batch_centroids, n_clusters); + + L2NormBuf_OR_DistBuf.resize(dataBatchSize * centroidsBatchSize, stream); + + auto pairwiseDistance = raft::make_device_matrix_view( + L2NormBuf_OR_DistBuf.data(), dataBatchSize, centroidsBatchSize); + + auto temp_kvp = + raft::make_device_vector, IndexT>(handle, n_samples); + raft::KeyValuePair initial_value(0, std::numeric_limits::max()); + raft::matrix::fill(handle, temp_kvp.view(), initial_value); + + for (IndexT dIdx = 0; dIdx < n_samples; dIdx += dataBatchSize) { + auto ns = std::min((IndexT)dataBatchSize, n_samples - dIdx); + + auto datasetView = raft::make_device_matrix_view( + X.data_handle() + (dIdx * n_features), ns, n_features); + + auto temp_kvp_view = raft::make_device_vector_view, IndexT>( + temp_kvp.data_handle() + dIdx, ns); + + for (IndexT cIdx = 0; cIdx < n_clusters; cIdx += centroidsBatchSize) { + auto nc = std::min((IndexT)centroidsBatchSize, n_clusters - cIdx); + + auto centroidsView = raft::make_device_matrix_view( + centroids.data_handle() + (cIdx * n_features), nc, n_features); + + auto pairwiseDistanceView = + raft::make_device_matrix_view(pairwiseDistance.data_handle(), ns, nc); + + pairwise_distance_kmeans( + handle, datasetView, centroidsView, pairwiseDistanceView, metric); + + raft::linalg::coalescedReduction( + temp_kvp_view.data_handle(), + pairwiseDistanceView.data_handle(), + pairwiseDistanceView.extent(1), + pairwiseDistanceView.extent(0), + initial_value, + stream, + true, + [=] __device__(const DataT val, const IndexT i) { + raft::KeyValuePair pair; + pair.key = cIdx + i; + pair.value = val; + return pair; + }, + raft::argmin_op{}, + raft::identity_op{}); + } + } + + unpack_kvp(handle, nearest_idx, nearest_dist, raft::make_const_mdspan(temp_kvp.view())); + } +} + +#define INSTANTIATE_MIN_CLUSTER_AND_DISTANCE(DataT, IndexT) \ + template void minClusterAndDistanceCompute( \ + raft::resources const& handle, \ + raft::device_matrix_view X, \ + raft::device_matrix_view centroids, \ + raft::device_vector_view nearest_idx, \ + raft::device_vector_view nearest_dist, \ + raft::device_vector_view L2NormX, \ + rmm::device_uvector& L2NormBuf_OR_DistBuf, \ + cuvs::distance::DistanceType metric, \ + int batch_samples, \ + int batch_centroids, \ + rmm::device_uvector& workspace); + +INSTANTIATE_MIN_CLUSTER_AND_DISTANCE(float, int64_t) +INSTANTIATE_MIN_CLUSTER_AND_DISTANCE(double, int64_t) +INSTANTIATE_MIN_CLUSTER_AND_DISTANCE(float, int) +INSTANTIATE_MIN_CLUSTER_AND_DISTANCE(double, int) + +#undef INSTANTIATE_MIN_CLUSTER_AND_DISTANCE + +template +void minClusterDistanceCompute(raft::resources const& handle, + raft::device_matrix_view X, + raft::device_matrix_view centroids, + raft::device_vector_view minClusterDistance, + raft::device_vector_view L2NormX, + rmm::device_uvector& L2NormBuf_OR_DistBuf, + cuvs::distance::DistanceType metric, + int batch_samples, + int batch_centroids, + rmm::device_uvector& workspace) +{ + cudaStream_t stream = raft::resource::get_cuda_stream(handle); + auto n_samples = X.extent(0); + auto n_features = X.extent(1); + auto n_clusters = centroids.extent(0); + + const bool is_l2_cos = metric == cuvs::distance::DistanceType::L2Expanded || + metric == cuvs::distance::DistanceType::L2SqrtExpanded || + metric == cuvs::distance::DistanceType::CosineExpanded; + + raft::matrix::fill(handle, minClusterDistance, std::numeric_limits::max()); + + const FusedDistancePath fused_path = + is_l2_cos ? use_fused(handle, n_samples, n_clusters, n_features, metric) + : FusedDistancePath::Unfused; + + if (uses_fused_distance_nn(fused_path)) { + L2NormBuf_OR_DistBuf.resize(n_clusters, stream); + auto centroidsNorm = + raft::make_device_vector_view(L2NormBuf_OR_DistBuf.data(), n_clusters); + + if (metric == cuvs::distance::DistanceType::CosineExpanded) { + raft::linalg::norm( + handle, + raft::make_device_matrix_view( + centroids.data_handle(), centroids.extent(0), centroids.extent(1)), + centroidsNorm, + raft::sqrt_op{}); + } else { + raft::linalg::norm( + handle, + raft::make_device_matrix_view( + centroids.data_handle(), centroids.extent(0), centroids.extent(1)), + centroidsNorm); + } + + raft::KeyValuePair* cutlass_kvp_scratch = nullptr; + rmm::device_uvector> temp_kvp(0, stream); + if (needs_cutlass_kvp_scratch(fused_path)) { + temp_kvp.resize(n_samples, stream); + cutlass_kvp_scratch = temp_kvp.data(); + workspace.resize(sizeof(int) * n_samples, stream); + } + + cuvs::distance::fusedDistanceNNMinReduce( + nullptr, + minClusterDistance.data_handle(), + X.data_handle(), + centroids.data_handle(), + L2NormX.data_handle(), + centroidsNorm.data_handle(), + n_samples, + n_clusters, + n_features, + needs_fused_mutex_workspace(fused_path) ? (void*)workspace.data() : nullptr, + metric != cuvs::distance::DistanceType::L2Expanded, + true, + true, + metric, + 0.0f, + cutlass_kvp_scratch, + stream); + } else { + auto dataBatchSize = getDataBatchSize(batch_samples, n_samples); + auto centroidsBatchSize = getCentroidsBatchSize(batch_centroids, n_clusters); + + L2NormBuf_OR_DistBuf.resize(dataBatchSize * centroidsBatchSize, stream); + + auto pairwiseDistance = raft::make_device_matrix_view( + L2NormBuf_OR_DistBuf.data(), dataBatchSize, centroidsBatchSize); + + for (IndexT dIdx = 0; dIdx < n_samples; dIdx += dataBatchSize) { + auto ns = std::min((IndexT)dataBatchSize, n_samples - dIdx); + + auto datasetView = raft::make_device_matrix_view( + X.data_handle() + dIdx * n_features, ns, n_features); + + auto minClusterDistanceView = + raft::make_device_vector_view(minClusterDistance.data_handle() + dIdx, ns); + + for (IndexT cIdx = 0; cIdx < n_clusters; cIdx += centroidsBatchSize) { + auto nc = std::min((IndexT)centroidsBatchSize, n_clusters - cIdx); + + auto centroidsView = raft::make_device_matrix_view( + centroids.data_handle() + cIdx * n_features, nc, n_features); + + auto pairwiseDistanceView = + raft::make_device_matrix_view(pairwiseDistance.data_handle(), ns, nc); + + pairwise_distance_kmeans( + handle, datasetView, centroidsView, pairwiseDistanceView, metric); + + raft::linalg::coalescedReduction(minClusterDistanceView.data_handle(), + pairwiseDistanceView.data_handle(), + pairwiseDistanceView.extent(1), + pairwiseDistanceView.extent(0), + std::numeric_limits::max(), + stream, + true, + raft::identity_op{}, + raft::min_op{}, + raft::identity_op{}); + } + } + } +} + +#define INSTANTIATE_MIN_CLUSTER_DISTANCE(DataT, IndexT) \ + template void minClusterDistanceCompute( \ + raft::resources const& handle, \ + raft::device_matrix_view X, \ + raft::device_matrix_view centroids, \ + raft::device_vector_view minClusterDistance, \ + raft::device_vector_view L2NormX, \ + rmm::device_uvector& L2NormBuf_OR_DistBuf, \ + cuvs::distance::DistanceType metric, \ + int batch_samples, \ + int batch_centroids, \ + rmm::device_uvector& workspace); + +INSTANTIATE_MIN_CLUSTER_DISTANCE(float, int64_t) +INSTANTIATE_MIN_CLUSTER_DISTANCE(double, int64_t) +INSTANTIATE_MIN_CLUSTER_DISTANCE(float, int) +INSTANTIATE_MIN_CLUSTER_DISTANCE(double, int) + +#undef INSTANTIATE_MIN_CLUSTER_DISTANCE + +} // namespace cuvs::cluster::kmeans::detail diff --git a/cpp/src/detail/jit_lto/AlgorithmPlanner.cpp b/cpp/src/detail/jit_lto/AlgorithmPlanner.cpp index 486d6f1aa5..a9734ed568 100644 --- a/cpp/src/detail/jit_lto/AlgorithmPlanner.cpp +++ b/cpp/src/detail/jit_lto/AlgorithmPlanner.cpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/src/detail/jit_lto/LTOAlgorithmPlanner.cpp b/cpp/src/detail/jit_lto/LTOAlgorithmPlanner.cpp index da7c0408b4..4d2b12bc51 100644 --- a/cpp/src/detail/jit_lto/LTOAlgorithmPlanner.cpp +++ b/cpp/src/detail/jit_lto/LTOAlgorithmPlanner.cpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/src/detail/jit_lto/TileAlgorithmPlanner.cpp b/cpp/src/detail/jit_lto/TileAlgorithmPlanner.cpp index 1487abb239..907b9fbc47 100644 --- a/cpp/src/detail/jit_lto/TileAlgorithmPlanner.cpp +++ b/cpp/src/detail/jit_lto/TileAlgorithmPlanner.cpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/src/distance/detail/fused_distance_nn.cuh b/cpp/src/distance/detail/fused_distance_nn.cuh index ec662fcbfa..dbc87f468d 100644 --- a/cpp/src/distance/detail/fused_distance_nn.cuh +++ b/cpp/src/distance/detail/fused_distance_nn.cuh @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/src/distance/detail/fused_distance_nn/cutile/export_fused_1nn.py b/cpp/src/distance/detail/fused_distance_nn/cutile/export_fused_1nn.py index 7d212c698e..2ddc717b60 100644 --- a/cpp/src/distance/detail/fused_distance_nn/cutile/export_fused_1nn.py +++ b/cpp/src/distance/detail/fused_distance_nn/cutile/export_fused_1nn.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 """Export fused 1-NN cuTile kernels to cubin or TileIR bytecode.""" @@ -19,7 +19,13 @@ export_kernel, ) -from fused_1nn_kernel import ( +# CI enables Python safe-path mode, so the script directory is not guaranteed +# to be importable even when this file is executed directly. +SCRIPT_DIR = Path(__file__).resolve().parent +if str(SCRIPT_DIR) not in sys.path: + sys.path.insert(0, str(SCRIPT_DIR)) + +from fused_1nn_kernel import ( # noqa: E402 INDEX_TYPES, METRICS, _idx_dtype, diff --git a/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_kernel.py b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_kernel.py index 465bbcbb4e..0937dc8116 100644 --- a/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_kernel.py +++ b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_kernel.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 """cuTile fused GEMM + 1-NN kernel with runtime metric selection.""" diff --git a/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_planner.hpp b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_planner.hpp index 995f4c787e..e9e66a7a26 100644 --- a/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_planner.hpp +++ b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_planner.hpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_tile.cu b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_tile.cu index 0ab4e1ef84..622af314a7 100644 --- a/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_tile.cu +++ b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_tile.cu @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_tile.hpp b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_tile.hpp index 953f78f63e..7cdbabd411 100644 --- a/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_tile.hpp +++ b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_tile.hpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/src/distance/detail/fused_distance_nn/fused_cosine_nn.cuh b/cpp/src/distance/detail/fused_distance_nn/fused_cosine_nn.cuh index cc16d8a2e1..43059a681c 100644 --- a/cpp/src/distance/detail/fused_distance_nn/fused_cosine_nn.cuh +++ b/cpp/src/distance/detail/fused_distance_nn/fused_cosine_nn.cuh @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/src/distance/detail/fused_distance_nn/fused_l2_nn.cuh b/cpp/src/distance/detail/fused_distance_nn/fused_l2_nn.cuh index 8c532e2932..49948951fe 100644 --- a/cpp/src/distance/detail/fused_distance_nn/fused_l2_nn.cuh +++ b/cpp/src/distance/detail/fused_distance_nn/fused_l2_nn.cuh @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2021-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2021-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/src/distance/detail/fused_distance_nn/helper_structs.cuh b/cpp/src/distance/detail/fused_distance_nn/helper_structs.cuh index 3bd78ba5ab..3f4e839f35 100644 --- a/cpp/src/distance/detail/fused_distance_nn/helper_structs.cuh +++ b/cpp/src/distance/detail/fused_distance_nn/helper_structs.cuh @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2021-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2021-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/src/distance/detail/fused_distance_nn/predicated_tile_iterator_reduced_vec.h b/cpp/src/distance/detail/fused_distance_nn/predicated_tile_iterator_reduced_vec.h index 0d9f5333af..8d16c72c04 100644 --- a/cpp/src/distance/detail/fused_distance_nn/predicated_tile_iterator_reduced_vec.h +++ b/cpp/src/distance/detail/fused_distance_nn/predicated_tile_iterator_reduced_vec.h @@ -1,7 +1,7 @@ // clang-format off /* * SPDX-FileCopyrightText: Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 AND BSD-3-Clause */ // clang-format on diff --git a/cpp/src/distance/detail/pairwise_matrix/jit_lto_kernels/pairwise_matrix_planner.hpp b/cpp/src/distance/detail/pairwise_matrix/jit_lto_kernels/pairwise_matrix_planner.hpp index f89a383596..3b92b63f1b 100644 --- a/cpp/src/distance/detail/pairwise_matrix/jit_lto_kernels/pairwise_matrix_planner.hpp +++ b/cpp/src/distance/detail/pairwise_matrix/jit_lto_kernels/pairwise_matrix_planner.hpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/src/distance/fused_distance_nn-inl.cuh b/cpp/src/distance/fused_distance_nn-inl.cuh index 13c4faa472..ccaef64319 100644 --- a/cpp/src/distance/fused_distance_nn-inl.cuh +++ b/cpp/src/distance/fused_distance_nn-inl.cuh @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/src/neighbors/ivf_flat/detail/jit_lto_kernels/interleaved_scan_planner.hpp b/cpp/src/neighbors/ivf_flat/detail/jit_lto_kernels/interleaved_scan_planner.hpp index 7899d970ab..92725dde53 100644 --- a/cpp/src/neighbors/ivf_flat/detail/jit_lto_kernels/interleaved_scan_planner.hpp +++ b/cpp/src/neighbors/ivf_flat/detail/jit_lto_kernels/interleaved_scan_planner.hpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/src/neighbors/ivf_pq/detail/jit_lto_kernels/compute_similarity_planner.hpp b/cpp/src/neighbors/ivf_pq/detail/jit_lto_kernels/compute_similarity_planner.hpp index 7152aaeebd..52b136018e 100644 --- a/cpp/src/neighbors/ivf_pq/detail/jit_lto_kernels/compute_similarity_planner.hpp +++ b/cpp/src/neighbors/ivf_pq/detail/jit_lto_kernels/compute_similarity_planner.hpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/src/neighbors/ivf_sq/detail/jit_lto_kernels/scan_planner.hpp b/cpp/src/neighbors/ivf_sq/detail/jit_lto_kernels/scan_planner.hpp index 5dc47dc612..21cf22f7f4 100644 --- a/cpp/src/neighbors/ivf_sq/detail/jit_lto_kernels/scan_planner.hpp +++ b/cpp/src/neighbors/ivf_sq/detail/jit_lto_kernels/scan_planner.hpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/tests/neighbors/distance_nn.cu b/cpp/tests/neighbors/distance_nn.cu index 6b17fc646b..1c4b92ec7e 100644 --- a/cpp/tests/neighbors/distance_nn.cu +++ b/cpp/tests/neighbors/distance_nn.cu @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/tests/neighbors/distance_nn_helper.cuh b/cpp/tests/neighbors/distance_nn_helper.cuh index dfa71f71a4..51028876ff 100644 --- a/cpp/tests/neighbors/distance_nn_helper.cuh +++ b/cpp/tests/neighbors/distance_nn_helper.cuh @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ From e28ecd9db77a90b71ab194563a5604992205be7b Mon Sep 17 00:00:00 2001 From: Divye Gala Date: Thu, 6 Aug 2026 19:55:47 -0400 Subject: [PATCH 30/82] Delete cpp/src/cluster/detail/.nfs000000001cb8b01800001a98 --- .../detail/.nfs000000001cb8b01800001a98 | 414 ------------------ 1 file changed, 414 deletions(-) delete mode 100644 cpp/src/cluster/detail/.nfs000000001cb8b01800001a98 diff --git a/cpp/src/cluster/detail/.nfs000000001cb8b01800001a98 b/cpp/src/cluster/detail/.nfs000000001cb8b01800001a98 deleted file mode 100644 index 63f5d20ba6..0000000000 --- a/cpp/src/cluster/detail/.nfs000000001cb8b01800001a98 +++ /dev/null @@ -1,414 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: Apache-2.0 - */ - -#include "../../distance/fused_distance_nn.cuh" -#include "../../distance/unfused_distance_nn.cuh" -#include "kmeans_common.cuh" - -#include - -namespace cuvs::cluster::kmeans::detail { - -namespace { - -template -__global__ void unpack_kvp_to_soa(IndexT* nearest_idx, - DataT* nearest_dist, - const raft::KeyValuePair* kvp, - IndexT n) -{ - IndexT i = blockIdx.x * blockDim.x + threadIdx.x; - if (i < n) { - if (nearest_idx != nullptr) { nearest_idx[i] = kvp[i].key; } - if (nearest_dist != nullptr) { nearest_dist[i] = kvp[i].value; } - } -} - -template -void unpack_kvp(raft::resources const& handle, - raft::device_vector_view nearest_idx, - raft::device_vector_view nearest_dist, - raft::device_vector_view, IndexT> kvp) -{ - auto stream = raft::resource::get_cuda_stream(handle); - auto n = static_cast(kvp.extent(0)); - int blks = static_cast((n + 255) / 256); - unpack_kvp_to_soa<<>>( - nearest_idx.data_handle(), nearest_dist.data_handle(), kvp.data_handle(), n); - RAFT_CUDA_TRY(cudaGetLastError()); -} - -} // namespace - -template -void minClusterAndDistanceCompute(raft::resources const& handle, - raft::device_matrix_view X, - raft::device_matrix_view centroids, - raft::device_vector_view nearest_idx, - raft::device_vector_view nearest_dist, - raft::device_vector_view L2NormX, - rmm::device_uvector& L2NormBuf_OR_DistBuf, - cuvs::distance::DistanceType metric, - int batch_samples, - int batch_centroids, - rmm::device_uvector& workspace) -{ - cudaStream_t stream = raft::resource::get_cuda_stream(handle); - auto n_samples = X.extent(0); - auto n_features = X.extent(1); - auto n_clusters = centroids.extent(0); - const bool is_l2_cos = metric == cuvs::distance::DistanceType::L2Expanded || - metric == cuvs::distance::DistanceType::L2SqrtExpanded || - metric == cuvs::distance::DistanceType::CosineExpanded; - const FusedDistancePath fused_path = - use_fused(handle, n_samples, n_clusters, n_features, metric); - - if (uses_fused_distance_nn(fused_path)) { - L2NormBuf_OR_DistBuf.resize(n_clusters, stream); - auto centroidsNorm = - raft::make_device_vector_view(L2NormBuf_OR_DistBuf.data(), n_clusters); - - if (is_l2_cos) { - if (metric == cuvs::distance::DistanceType::CosineExpanded) { - raft::linalg::norm( - handle, centroids, centroidsNorm, raft::sqrt_op{}); - } else { - raft::linalg::norm( - handle, centroids, centroidsNorm); - } - } - - auto centroidsNormConst = - raft::make_device_vector_view(L2NormBuf_OR_DistBuf.data(), n_clusters); - - raft::KeyValuePair* cutlass_kvp_scratch = nullptr; - rmm::device_uvector> temp_kvp(0, stream); - if (needs_cutlass_kvp_scratch(fused_path)) { - temp_kvp.resize(n_samples, stream); - cutlass_kvp_scratch = temp_kvp.data(); - workspace.resize(sizeof(int) * n_samples, stream); - } else if constexpr (std::is_same_v) { - // The cuTile kernel uses i32 internally and widens labels after the launch. - workspace.resize(sizeof(int) * static_cast(n_samples), stream); - } - - cuvs::distance::fusedDistanceNNMinReduce( - nearest_idx.data_handle(), - nearest_dist.data_handle(), - X.data_handle(), - centroids.data_handle(), - L2NormX.data_handle(), - centroidsNormConst.data_handle(), - n_samples, - n_clusters, - n_features, - needs_fused_mutex_workspace(fused_path) || std::is_same_v - ? (void*)workspace.data() - : nullptr, - metric != cuvs::distance::DistanceType::L2Expanded, - true, - true, - metric, - 0.0f, - cutlass_kvp_scratch, - stream); - } else if (is_l2_cos) { - L2NormBuf_OR_DistBuf.resize(n_clusters, stream); - auto centroidsNorm = - raft::make_device_vector_view(L2NormBuf_OR_DistBuf.data(), n_clusters); - - if (metric == cuvs::distance::DistanceType::CosineExpanded) { - raft::linalg::norm( - handle, centroids, centroidsNorm, raft::sqrt_op{}); - } else { - raft::linalg::norm( - handle, centroids, centroidsNorm); - } - - auto centroidsNormConst = - raft::make_device_vector_view(L2NormBuf_OR_DistBuf.data(), n_clusters); - - auto dataBatchSize = getDataBatchSize(batch_samples, n_samples); - auto centroidsBatchSize = getCentroidsBatchSize(batch_centroids, n_clusters); - - // The unfused reduction indexes its distance matrix with IndexT. - dataBatchSize = - std::min(dataBatchSize, std::numeric_limits::max() / centroidsBatchSize); - workspace.resize(sizeof(DataT) * dataBatchSize * centroidsBatchSize, stream); - - using KeyValueT = raft::KeyValuePair; - auto temp_kvp = raft::make_device_vector(handle, n_samples); - KeyValueT initial_value(0, std::numeric_limits::max()); - raft::matrix::fill(handle, temp_kvp.view(), initial_value); - - const bool tileCentroids = centroidsBatchSize < n_clusters; - rmm::device_uvector batchMinClusterAndDistance(tileCentroids ? dataBatchSize : 0, - stream); - - for (IndexT dIdx = 0; dIdx < n_samples;) { - auto ns = std::min(dataBatchSize, n_samples - dIdx); - auto minClusterAndDistanceView = raft::make_device_vector_view( - temp_kvp.data_handle() + dIdx, ns); - - for (IndexT cIdx = 0; cIdx < n_clusters;) { - auto nc = std::min(centroidsBatchSize, n_clusters - cIdx); - auto batchMin = tileCentroids ? batchMinClusterAndDistance.data() - : minClusterAndDistanceView.data_handle(); - - cuvs::distance::unfusedDistanceNNMinReduce( - handle, - batchMin, - X.data_handle() + dIdx * n_features, - centroids.data_handle() + cIdx * n_features, - L2NormX.data_handle() + dIdx, - centroidsNormConst.data_handle() + cIdx, - ns, - nc, - n_features, - (void*)workspace.data(), - metric != cuvs::distance::DistanceType::L2Expanded, - tileCentroids, - true, - metric, - 0.0f, - stream); - - if (tileCentroids) { - // Convert tile-local centroid indices and merge the tile minima. - auto batchMinView = raft::make_device_vector_view(batchMin, ns); - raft::linalg::map( - handle, - minClusterAndDistanceView, - [cIdx] __device__(KeyValueT current, KeyValueT batch) { - batch.key += cIdx; - return batch.value < current.value ? batch : current; - }, - raft::make_const_mdspan(minClusterAndDistanceView), - batchMinView); - } - cIdx += nc; - } - dIdx += ns; - } - - unpack_kvp(handle, nearest_idx, nearest_dist, raft::make_const_mdspan(temp_kvp.view())); - } else { - auto dataBatchSize = getDataBatchSize(batch_samples, n_samples); - auto centroidsBatchSize = getCentroidsBatchSize(batch_centroids, n_clusters); - - L2NormBuf_OR_DistBuf.resize(dataBatchSize * centroidsBatchSize, stream); - - auto pairwiseDistance = raft::make_device_matrix_view( - L2NormBuf_OR_DistBuf.data(), dataBatchSize, centroidsBatchSize); - - auto temp_kvp = - raft::make_device_vector, IndexT>(handle, n_samples); - raft::KeyValuePair initial_value(0, std::numeric_limits::max()); - raft::matrix::fill(handle, temp_kvp.view(), initial_value); - - for (IndexT dIdx = 0; dIdx < n_samples; dIdx += dataBatchSize) { - auto ns = std::min((IndexT)dataBatchSize, n_samples - dIdx); - - auto datasetView = raft::make_device_matrix_view( - X.data_handle() + (dIdx * n_features), ns, n_features); - - auto temp_kvp_view = raft::make_device_vector_view, IndexT>( - temp_kvp.data_handle() + dIdx, ns); - - for (IndexT cIdx = 0; cIdx < n_clusters; cIdx += centroidsBatchSize) { - auto nc = std::min((IndexT)centroidsBatchSize, n_clusters - cIdx); - - auto centroidsView = raft::make_device_matrix_view( - centroids.data_handle() + (cIdx * n_features), nc, n_features); - - auto pairwiseDistanceView = - raft::make_device_matrix_view(pairwiseDistance.data_handle(), ns, nc); - - pairwise_distance_kmeans( - handle, datasetView, centroidsView, pairwiseDistanceView, metric); - - raft::linalg::coalescedReduction( - temp_kvp_view.data_handle(), - pairwiseDistanceView.data_handle(), - pairwiseDistanceView.extent(1), - pairwiseDistanceView.extent(0), - initial_value, - stream, - true, - [=] __device__(const DataT val, const IndexT i) { - raft::KeyValuePair pair; - pair.key = cIdx + i; - pair.value = val; - return pair; - }, - raft::argmin_op{}, - raft::identity_op{}); - } - } - - unpack_kvp(handle, nearest_idx, nearest_dist, raft::make_const_mdspan(temp_kvp.view())); - } -} - -#define INSTANTIATE_MIN_CLUSTER_AND_DISTANCE(DataT, IndexT) \ - template void minClusterAndDistanceCompute( \ - raft::resources const& handle, \ - raft::device_matrix_view X, \ - raft::device_matrix_view centroids, \ - raft::device_vector_view nearest_idx, \ - raft::device_vector_view nearest_dist, \ - raft::device_vector_view L2NormX, \ - rmm::device_uvector& L2NormBuf_OR_DistBuf, \ - cuvs::distance::DistanceType metric, \ - int batch_samples, \ - int batch_centroids, \ - rmm::device_uvector& workspace); - -INSTANTIATE_MIN_CLUSTER_AND_DISTANCE(float, int64_t) -INSTANTIATE_MIN_CLUSTER_AND_DISTANCE(double, int64_t) -INSTANTIATE_MIN_CLUSTER_AND_DISTANCE(float, int) -INSTANTIATE_MIN_CLUSTER_AND_DISTANCE(double, int) - -#undef INSTANTIATE_MIN_CLUSTER_AND_DISTANCE - -template -void minClusterDistanceCompute(raft::resources const& handle, - raft::device_matrix_view X, - raft::device_matrix_view centroids, - raft::device_vector_view minClusterDistance, - raft::device_vector_view L2NormX, - rmm::device_uvector& L2NormBuf_OR_DistBuf, - cuvs::distance::DistanceType metric, - int batch_samples, - int batch_centroids, - rmm::device_uvector& workspace) -{ - cudaStream_t stream = raft::resource::get_cuda_stream(handle); - auto n_samples = X.extent(0); - auto n_features = X.extent(1); - auto n_clusters = centroids.extent(0); - - const bool is_l2_cos = metric == cuvs::distance::DistanceType::L2Expanded || - metric == cuvs::distance::DistanceType::L2SqrtExpanded || - metric == cuvs::distance::DistanceType::CosineExpanded; - - raft::matrix::fill(handle, minClusterDistance, std::numeric_limits::max()); - - const FusedDistancePath fused_path = - is_l2_cos ? use_fused(handle, n_samples, n_clusters, n_features, metric) - : FusedDistancePath::Unfused; - - if (uses_fused_distance_nn(fused_path)) { - L2NormBuf_OR_DistBuf.resize(n_clusters, stream); - auto centroidsNorm = - raft::make_device_vector_view(L2NormBuf_OR_DistBuf.data(), n_clusters); - - if (metric == cuvs::distance::DistanceType::CosineExpanded) { - raft::linalg::norm( - handle, - raft::make_device_matrix_view( - centroids.data_handle(), centroids.extent(0), centroids.extent(1)), - centroidsNorm, - raft::sqrt_op{}); - } else { - raft::linalg::norm( - handle, - raft::make_device_matrix_view( - centroids.data_handle(), centroids.extent(0), centroids.extent(1)), - centroidsNorm); - } - - raft::KeyValuePair* cutlass_kvp_scratch = nullptr; - rmm::device_uvector> temp_kvp(0, stream); - if (needs_cutlass_kvp_scratch(fused_path)) { - temp_kvp.resize(n_samples, stream); - cutlass_kvp_scratch = temp_kvp.data(); - workspace.resize(sizeof(int) * n_samples, stream); - } - - cuvs::distance::fusedDistanceNNMinReduce( - nullptr, - minClusterDistance.data_handle(), - X.data_handle(), - centroids.data_handle(), - L2NormX.data_handle(), - centroidsNorm.data_handle(), - n_samples, - n_clusters, - n_features, - needs_fused_mutex_workspace(fused_path) ? (void*)workspace.data() : nullptr, - metric != cuvs::distance::DistanceType::L2Expanded, - true, - true, - metric, - 0.0f, - cutlass_kvp_scratch, - stream); - } else { - auto dataBatchSize = getDataBatchSize(batch_samples, n_samples); - auto centroidsBatchSize = getCentroidsBatchSize(batch_centroids, n_clusters); - - L2NormBuf_OR_DistBuf.resize(dataBatchSize * centroidsBatchSize, stream); - - auto pairwiseDistance = raft::make_device_matrix_view( - L2NormBuf_OR_DistBuf.data(), dataBatchSize, centroidsBatchSize); - - for (IndexT dIdx = 0; dIdx < n_samples; dIdx += dataBatchSize) { - auto ns = std::min((IndexT)dataBatchSize, n_samples - dIdx); - - auto datasetView = raft::make_device_matrix_view( - X.data_handle() + dIdx * n_features, ns, n_features); - - auto minClusterDistanceView = - raft::make_device_vector_view(minClusterDistance.data_handle() + dIdx, ns); - - for (IndexT cIdx = 0; cIdx < n_clusters; cIdx += centroidsBatchSize) { - auto nc = std::min((IndexT)centroidsBatchSize, n_clusters - cIdx); - - auto centroidsView = raft::make_device_matrix_view( - centroids.data_handle() + cIdx * n_features, nc, n_features); - - auto pairwiseDistanceView = - raft::make_device_matrix_view(pairwiseDistance.data_handle(), ns, nc); - - pairwise_distance_kmeans( - handle, datasetView, centroidsView, pairwiseDistanceView, metric); - - raft::linalg::coalescedReduction(minClusterDistanceView.data_handle(), - pairwiseDistanceView.data_handle(), - pairwiseDistanceView.extent(1), - pairwiseDistanceView.extent(0), - std::numeric_limits::max(), - stream, - true, - raft::identity_op{}, - raft::min_op{}, - raft::identity_op{}); - } - } - } -} - -#define INSTANTIATE_MIN_CLUSTER_DISTANCE(DataT, IndexT) \ - template void minClusterDistanceCompute( \ - raft::resources const& handle, \ - raft::device_matrix_view X, \ - raft::device_matrix_view centroids, \ - raft::device_vector_view minClusterDistance, \ - raft::device_vector_view L2NormX, \ - rmm::device_uvector& L2NormBuf_OR_DistBuf, \ - cuvs::distance::DistanceType metric, \ - int batch_samples, \ - int batch_centroids, \ - rmm::device_uvector& workspace); - -INSTANTIATE_MIN_CLUSTER_DISTANCE(float, int64_t) -INSTANTIATE_MIN_CLUSTER_DISTANCE(double, int64_t) -INSTANTIATE_MIN_CLUSTER_DISTANCE(float, int) -INSTANTIATE_MIN_CLUSTER_DISTANCE(double, int) - -#undef INSTANTIATE_MIN_CLUSTER_DISTANCE - -} // namespace cuvs::cluster::kmeans::detail From 22b08454c87037a43051f8145631485406e230f2 Mon Sep 17 00:00:00 2001 From: divyegala Date: Fri, 7 Aug 2026 00:11:56 +0000 Subject: [PATCH 31/82] allow rattler to install from pypi/pip --- conda/recipes/libcuvs/recipe.yaml | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/conda/recipes/libcuvs/recipe.yaml b/conda/recipes/libcuvs/recipe.yaml index 27bfd5c14b..95bd29b208 100644 --- a/conda/recipes/libcuvs/recipe.yaml +++ b/conda/recipes/libcuvs/recipe.yaml @@ -33,7 +33,13 @@ cache: # cuTile and TileIRAS are PyPI-only build tools. Their CUDA components # are installed under the Python package namespace, alongside the conda toolchain. if [[ "${{ cuda_major }}" == "13" ]]; then - python -m pip install --no-cache-dir \ + # rattler-build disables package indexes and dependency resolution by default. + env -u PIP_NO_INDEX \ + -u PIP_NO_DEPENDENCIES \ + -u PIP_NO_DEPS \ + -u PIP_IGNORE_INSTALLED \ + python -m pip install --no-cache-dir \ + --index-url https://pypi.org/simple \ "cuda-tile" \ "cuda-toolkit[tileiras]==${{ cuda_version }}.*" fi From f1bf39dec378ccdc1396f311d334af434d845f49 Mon Sep 17 00:00:00 2001 From: divyegala Date: Mon, 10 Aug 2026 21:50:06 +0000 Subject: [PATCH 32/82] fix upstream merge --- ...mpute_inner_products_with_bitwise_block_sort_planner.hpp | 6 +++--- .../compute_inner_products_with_bitwise_planner.hpp | 6 +++--- ...ute_inner_products_with_lut16_opt_block_sort_planner.hpp | 6 +++--- .../compute_inner_products_with_lut16_opt_planner.hpp | 6 +++--- .../compute_inner_products_with_lut_block_sort_planner.hpp | 6 +++--- .../compute_inner_products_with_lut_planner.hpp | 6 +++--- 6 files changed, 18 insertions(+), 18 deletions(-) diff --git a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_bitwise_block_sort_planner.hpp b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_bitwise_block_sort_planner.hpp index dbe6ec5fd4..b7a18ed29c 100644 --- a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_bitwise_block_sort_planner.hpp +++ b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_bitwise_block_sort_planner.hpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ @@ -11,11 +11,11 @@ namespace cuvs::neighbors::ivf_rabitq::detail { -struct ComputeInnerProductsWithBitwiseBlockSortPlanner : AlgorithmPlanner { +struct ComputeInnerProductsWithBitwiseBlockSortPlanner : LTOAlgorithmPlanner { inline static LauncherJitCache launcher_jit_cache{}; ComputeInnerProductsWithBitwiseBlockSortPlanner() - : AlgorithmPlanner("compute_inner_products_with_bitwise_block_sort", launcher_jit_cache) + : LTOAlgorithmPlanner("compute_inner_products_with_bitwise_block_sort", launcher_jit_cache) { } diff --git a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_bitwise_planner.hpp b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_bitwise_planner.hpp index 3ec4809395..4e7d3fb4eb 100644 --- a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_bitwise_planner.hpp +++ b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_bitwise_planner.hpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ @@ -11,11 +11,11 @@ namespace cuvs::neighbors::ivf_rabitq::detail { -struct ComputeInnerProductsWithBitwisePlanner : AlgorithmPlanner { +struct ComputeInnerProductsWithBitwisePlanner : LTOAlgorithmPlanner { inline static LauncherJitCache launcher_jit_cache{}; ComputeInnerProductsWithBitwisePlanner() - : AlgorithmPlanner("compute_inner_products_with_bitwise", launcher_jit_cache) + : LTOAlgorithmPlanner("compute_inner_products_with_bitwise", launcher_jit_cache) { } diff --git a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut16_opt_block_sort_planner.hpp b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut16_opt_block_sort_planner.hpp index acf4e27e3a..8d8cc079b1 100644 --- a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut16_opt_block_sort_planner.hpp +++ b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut16_opt_block_sort_planner.hpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ @@ -11,11 +11,11 @@ namespace cuvs::neighbors::ivf_rabitq::detail { -struct ComputeInnerProductsWithLut16OptBlockSortPlanner : AlgorithmPlanner { +struct ComputeInnerProductsWithLut16OptBlockSortPlanner : LTOAlgorithmPlanner { inline static LauncherJitCache launcher_jit_cache{}; ComputeInnerProductsWithLut16OptBlockSortPlanner() - : AlgorithmPlanner("compute_inner_products_with_lut16_opt_block_sort", launcher_jit_cache) + : LTOAlgorithmPlanner("compute_inner_products_with_lut16_opt_block_sort", launcher_jit_cache) { } diff --git a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut16_opt_planner.hpp b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut16_opt_planner.hpp index 95e0694c9c..a8a8fd15ff 100644 --- a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut16_opt_planner.hpp +++ b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut16_opt_planner.hpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ @@ -11,11 +11,11 @@ namespace cuvs::neighbors::ivf_rabitq::detail { -struct ComputeInnerProductsWithLut16OptPlanner : AlgorithmPlanner { +struct ComputeInnerProductsWithLut16OptPlanner : LTOAlgorithmPlanner { inline static LauncherJitCache launcher_jit_cache{}; ComputeInnerProductsWithLut16OptPlanner() - : AlgorithmPlanner("compute_inner_products_with_lut16_opt", launcher_jit_cache) + : LTOAlgorithmPlanner("compute_inner_products_with_lut16_opt", launcher_jit_cache) { } diff --git a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut_block_sort_planner.hpp b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut_block_sort_planner.hpp index 9db3839880..42073ac8eb 100644 --- a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut_block_sort_planner.hpp +++ b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut_block_sort_planner.hpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ @@ -11,11 +11,11 @@ namespace cuvs::neighbors::ivf_rabitq::detail { -struct ComputeInnerProductsWithLutBlockSortPlanner : AlgorithmPlanner { +struct ComputeInnerProductsWithLutBlockSortPlanner : LTOAlgorithmPlanner { inline static LauncherJitCache launcher_jit_cache{}; ComputeInnerProductsWithLutBlockSortPlanner() - : AlgorithmPlanner("compute_inner_products_with_lut_block_sort", launcher_jit_cache) + : LTOAlgorithmPlanner("compute_inner_products_with_lut_block_sort", launcher_jit_cache) { } diff --git a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut_planner.hpp b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut_planner.hpp index 3559a9bee1..d7a37a8be1 100644 --- a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut_planner.hpp +++ b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut_planner.hpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ @@ -11,11 +11,11 @@ namespace cuvs::neighbors::ivf_rabitq::detail { -struct ComputeInnerProductsWithLutPlanner : AlgorithmPlanner { +struct ComputeInnerProductsWithLutPlanner : LTOAlgorithmPlanner { inline static LauncherJitCache launcher_jit_cache{}; ComputeInnerProductsWithLutPlanner() - : AlgorithmPlanner("compute_inner_products_with_lut", launcher_jit_cache) + : LTOAlgorithmPlanner("compute_inner_products_with_lut", launcher_jit_cache) { } From 25fe50b6c880764262a188422afdb452104a55b3 Mon Sep 17 00:00:00 2001 From: divyegala Date: Mon, 10 Aug 2026 23:07:32 +0000 Subject: [PATCH 33/82] fix test --- cpp/include/cuvs/detail/jit_lto/cutile_arch_tags.hpp | 6 ++++++ .../cutile/fused_1nn_cutile_matrix.json | 10 ++++++++++ .../fused_distance_nn/cutile/fused_1nn_planner.hpp | 6 ++++++ cpp/tests/cluster/kmeans_predict_batching.cu | 9 +++++++-- 4 files changed, 29 insertions(+), 2 deletions(-) diff --git a/cpp/include/cuvs/detail/jit_lto/cutile_arch_tags.hpp b/cpp/include/cuvs/detail/jit_lto/cutile_arch_tags.hpp index 95f7c81ecd..f88ee39c98 100644 --- a/cpp/include/cuvs/detail/jit_lto/cutile_arch_tags.hpp +++ b/cpp/include/cuvs/detail/jit_lto/cutile_arch_tags.hpp @@ -24,6 +24,11 @@ struct cutile_arch_8_6 { static constexpr int cc_minor = 6; }; +struct cutile_arch_8_9 { + static constexpr int cc_major = 8; + static constexpr int cc_minor = 9; +}; + struct cutile_arch_9_0 { static constexpr int cc_major = 9; static constexpr int cc_minor = 0; @@ -43,6 +48,7 @@ inline bool is_embedded_cubin_arch(int cc_major, int cc_minor) { if (cc_major == 8 && cc_minor == 0) { return true; } if (cc_major == 8 && cc_minor == 6) { return true; } + if (cc_major == 8 && cc_minor == 9) { return true; } if (cc_major == 9 && cc_minor == 0) { return true; } if (cc_major == 10 && cc_minor == 0) { return true; } if (cc_major == 12 && cc_minor == 0) { return true; } diff --git a/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_cutile_matrix.json b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_cutile_matrix.json index c94aec352d..9adf7ddfce 100644 --- a/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_cutile_matrix.json +++ b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_cutile_matrix.json @@ -71,6 +71,16 @@ "cc_minor": 6, "arch_tag": "cutile_arch_8_6" }, + { + "output_format": "cubin", + "artifact_ext": "cubin", + "artifact_basename": "@data_type@_@index_abbrev@_@abi_abbrev@_@gpu_code@", + "register": "cubin", + "gpu_code": "sm_89", + "cc_major": 8, + "cc_minor": 9, + "arch_tag": "cutile_arch_8_9" + }, { "output_format": "cubin", "artifact_ext": "cubin", diff --git a/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_planner.hpp b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_planner.hpp index e9e66a7a26..24389e6edc 100644 --- a/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_planner.hpp +++ b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_planner.hpp @@ -54,6 +54,7 @@ struct Fused1nnTilePlanner : TileAlgorithmPlanner { using cuvs::detail::jit_lto::cutile_arch_12_0; using cuvs::detail::jit_lto::cutile_arch_8_0; using cuvs::detail::jit_lto::cutile_arch_8_6; + using cuvs::detail::jit_lto::cutile_arch_8_9; using cuvs::detail::jit_lto::cutile_arch_9_0; constexpr bool is_relaxed = std::is_same_v; @@ -63,6 +64,9 @@ struct Fused1nnTilePlanner : TileAlgorithmPlanner { using Tile86 = std::conditional_t; + using Tile89 = std::conditional_t; using Tile90 = std::conditional_t; @@ -77,6 +81,8 @@ struct Fused1nnTilePlanner : TileAlgorithmPlanner { fragment_tag_fused_1nn_cubin>(); this->add_static_fragment< fragment_tag_fused_1nn_cubin>(); + this->add_static_fragment< + fragment_tag_fused_1nn_cubin>(); this->add_static_fragment< fragment_tag_fused_1nn_cubin>(); this->add_static_fragment< diff --git a/cpp/tests/cluster/kmeans_predict_batching.cu b/cpp/tests/cluster/kmeans_predict_batching.cu index e22cbc9dcd..b740c3f998 100644 --- a/cpp/tests/cluster/kmeans_predict_batching.cu +++ b/cpp/tests/cluster/kmeans_predict_batching.cu @@ -136,8 +136,13 @@ TEST(KMeansPredict, BatchParametersPreserveResultsAndReduceUnfusedAllocations) // predict selects fused or unfused 1-NN according to the architecture heuristic. The batching // parameters only affect the unfused path, so every GPU checks the results while allocation // reductions are required only when this problem shape dispatches to unfused 1-NN. - const bool uses_unfused_path = - !detail::use_fused(handle, test_n_samples, test_n_clusters, test_n_features); + const auto fused_path = + detail::use_fused(handle, + test_n_samples, + test_n_clusters, + test_n_features, + cuvs::distance::DistanceType::L2Expanded); + const bool uses_unfused_path = !detail::uses_fused_distance_nn(fused_path); for (std::size_t i = 1; i < batch_configs.size(); ++i) { auto config = batch_configs[i]; From f225d056140381f224b2c94e985adacd49b7589c Mon Sep 17 00:00:00 2001 From: divyegala Date: Tue, 11 Aug 2026 01:02:54 +0000 Subject: [PATCH 34/82] inertia calculation with unexpanded l2 metric --- cpp/src/cluster/detail/kmeans_common.cuh | 9 ++++++++- cpp/src/cluster/detail/kmeans_mg.cuh | 3 ++- cpp/src/cluster/kmeans.cuh | 14 ++++++++++---- 3 files changed, 20 insertions(+), 6 deletions(-) diff --git a/cpp/src/cluster/detail/kmeans_common.cuh b/cpp/src/cluster/detail/kmeans_common.cuh index 82aedceacd..578236b29b 100644 --- a/cpp/src/cluster/detail/kmeans_common.cuh +++ b/cpp/src/cluster/detail/kmeans_common.cuh @@ -381,8 +381,15 @@ void pairwise_distance_kmeans(raft::resources const& handle, DataT, raft::layout_c_contiguous, IndexT>(handle, X, centroids, pairwiseDistance); + } else if (metric == cuvs::distance::DistanceType::L2Unexpanded) { + cuvs::distance::distance(handle, X, centroids, pairwiseDistance); } else { - RAFT_FAIL("kmeans requires L2Expanded or L2SqrtExpanded distance, have %i", + RAFT_FAIL("kmeans requires L2Expanded, L2SqrtExpanded, or L2Unexpanded distance, have %i", static_cast(metric)); } } diff --git a/cpp/src/cluster/detail/kmeans_mg.cuh b/cpp/src/cluster/detail/kmeans_mg.cuh index 94555b50d8..c1f51026eb 100644 --- a/cpp/src/cluster/detail/kmeans_mg.cuh +++ b/cpp/src/cluster/detail/kmeans_mg.cuh @@ -557,7 +557,8 @@ void mnmg_fit( batch_data_view, rank_centroids_const, raft::make_device_scalar_view(batch_clustering_cost.data_handle()), - batch_sw); + batch_sw, + cuvs::distance::DistanceType::L2Unexpanded); raft::linalg::add(dev_res, raft::make_const_mdspan(clustering_cost.view()), diff --git a/cpp/src/cluster/kmeans.cuh b/cpp/src/cluster/kmeans.cuh index 6c1a8c0d2a..a58d4b439c 100644 --- a/cpp/src/cluster/kmeans.cuh +++ b/cpp/src/cluster/kmeans.cuh @@ -331,6 +331,7 @@ void min_cluster_distance(raft::resources const& handle, * @param[in] centroids Cluster centroids [n_clusters x n_features] * @param[out] cost Sum of squared distances to nearest centroid (device) * @param[in] sample_weight Optional per-sample weights [n_samples] + * @param[in] metric Squared-L2 implementation used to evaluate inertia */ template void cluster_cost( @@ -338,23 +339,28 @@ void cluster_cost( raft::device_matrix_view X, raft::device_matrix_view centroids, raft::device_scalar_view cost, - std::optional> sample_weight = std::nullopt) + std::optional> sample_weight = std::nullopt, + cuvs::distance::DistanceType metric = cuvs::distance::DistanceType::L2Unexpanded) { auto stream = raft::resource::get_cuda_stream(handle); auto n_clusters = centroids.extent(0); auto n_samples = X.extent(0); auto n_features = X.extent(1); + RAFT_EXPECTS(metric == cuvs::distance::DistanceType::L2Expanded || + metric == cuvs::distance::DistanceType::L2Unexpanded, + "cluster_cost requires a squared-L2 distance metric"); + rmm::device_uvector workspace(n_samples * sizeof(IndexT), stream); auto x_norms = raft::make_device_vector(handle, n_samples); - raft::linalg::norm(handle, X, x_norms.view()); + if (metric == cuvs::distance::DistanceType::L2Expanded) { + raft::linalg::norm(handle, X, x_norms.view()); + } auto min_cluster_distance = raft::make_device_vector(handle, n_samples); rmm::device_uvector l2_norm_or_distance_buffer(0, stream); - auto metric = cuvs::distance::DistanceType::L2Expanded; - cuvs::cluster::kmeans::min_cluster_distance( handle, X, From 52e6be66c6d5990e8e7bcbf43a18bc2059d6cf9a Mon Sep 17 00:00:00 2001 From: divyegala Date: Tue, 11 Aug 2026 01:12:39 +0000 Subject: [PATCH 35/82] correct from host to build deps --- conda/recipes/libcuvs/recipe.yaml | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/conda/recipes/libcuvs/recipe.yaml b/conda/recipes/libcuvs/recipe.yaml index 95bd29b208..adc0c88beb 100644 --- a/conda/recipes/libcuvs/recipe.yaml +++ b/conda/recipes/libcuvs/recipe.yaml @@ -84,6 +84,10 @@ cache: - cuda-version =${{ cuda_version }} - cmake ${{ cmake_version }} - ninja + - if: cuda_major == "13" + then: + - pip + - python - ${{ stdlib("c") }} host: - libnvjitlink-dev @@ -98,10 +102,6 @@ cache: - libcurand-dev - libcusolver-dev - libcusparse-dev - - if: cuda_major == "13" - then: - - pip - - python # These are used for bench-ann - openblas - if: linux64 From e6b36d13c8d591c61a482e801c91a7b00472fe0e Mon Sep 17 00:00:00 2001 From: divyegala Date: Tue, 11 Aug 2026 02:18:58 +0000 Subject: [PATCH 36/82] batch --- cpp/src/cluster/detail/kmeans_common.cuh | 16 +++++--- cpp/src/cluster/kmeans.cuh | 37 +++++++++++++++++++ .../detail/pairwise_matrix/dispatch-ext.cuh | 2 +- cpp/src/distance/distance-ext.cuh | 2 +- cpp/src/distance/distance.cu | 2 +- 5 files changed, 50 insertions(+), 9 deletions(-) diff --git a/cpp/src/cluster/detail/kmeans_common.cuh b/cpp/src/cluster/detail/kmeans_common.cuh index 578236b29b..3ece6b0e19 100644 --- a/cpp/src/cluster/detail/kmeans_common.cuh +++ b/cpp/src/cluster/detail/kmeans_common.cuh @@ -382,12 +382,16 @@ void pairwise_distance_kmeans(raft::resources const& handle, raft::layout_c_contiguous, IndexT>(handle, X, centroids, pairwiseDistance); } else if (metric == cuvs::distance::DistanceType::L2Unexpanded) { - cuvs::distance::distance(handle, X, centroids, pairwiseDistance); + if constexpr (std::is_same_v) { + cuvs::distance::distance(handle, X, centroids, pairwiseDistance); + } else { + RAFT_FAIL("L2Unexpanded KMeans distance requires int32-indexed batches"); + } } else { RAFT_FAIL("kmeans requires L2Expanded, L2SqrtExpanded, or L2Unexpanded distance, have %i", static_cast(metric)); diff --git a/cpp/src/cluster/kmeans.cuh b/cpp/src/cluster/kmeans.cuh index a58d4b439c..dc9e56c746 100644 --- a/cpp/src/cluster/kmeans.cuh +++ b/cpp/src/cluster/kmeans.cuh @@ -14,9 +14,13 @@ #include #include #include +#include #include +#include +#include #include +#include namespace cuvs::cluster::kmeans { @@ -351,6 +355,39 @@ void cluster_cost( metric == cuvs::distance::DistanceType::L2Unexpanded, "cluster_cost requires a squared-L2 distance metric"); + if constexpr (std::is_same_v) { + if (metric == cuvs::distance::DistanceType::L2Unexpanded) { + constexpr IndexT max_i32 = std::numeric_limits::max(); + RAFT_EXPECTS(n_clusters > 0 && n_clusters <= max_i32 && n_features <= max_i32, + "stable cluster_cost requires n_clusters and n_features to fit in int32"); + + raft::matrix::fill(handle, cost, DataT{0}); + auto batch_cost = raft::make_device_scalar(handle, DataT{0}); + auto centroids_i32 = raft::make_device_matrix_view( + centroids.data_handle(), static_cast(n_clusters), static_cast(n_features)); + const IndexT max_batch_rows = max_i32 / n_clusters; + + for (IndexT offset = 0; offset < n_samples; offset += max_batch_rows) { + const int batch_rows = static_cast(std::min(max_batch_rows, n_samples - offset)); + auto X_i32 = raft::make_device_matrix_view( + X.data_handle() + offset * n_features, batch_rows, static_cast(n_features)); + + std::optional> batch_weights = std::nullopt; + if (sample_weight.has_value()) { + batch_weights = raft::make_device_vector_view( + sample_weight->data_handle() + offset, batch_rows); + } + + raft::matrix::fill(handle, batch_cost.view(), DataT{0}); + cluster_cost( + handle, X_i32, centroids_i32, batch_cost.view(), batch_weights, metric); + raft::linalg::add( + cost.data_handle(), cost.data_handle(), batch_cost.data_handle(), 1, stream); + } + return; + } + } + rmm::device_uvector workspace(n_samples * sizeof(IndexT), stream); auto x_norms = raft::make_device_vector(handle, n_samples); diff --git a/cpp/src/distance/detail/pairwise_matrix/dispatch-ext.cuh b/cpp/src/distance/detail/pairwise_matrix/dispatch-ext.cuh index c93a2f3f2b..f7b47e132c 100644 --- a/cpp/src/distance/detail/pairwise_matrix/dispatch-ext.cuh +++ b/cpp/src/distance/detail/pairwise_matrix/dispatch-ext.cuh @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2023, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2023, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ #pragma once diff --git a/cpp/src/distance/distance-ext.cuh b/cpp/src/distance/distance-ext.cuh index e3841d2caa..1b9637420d 100644 --- a/cpp/src/distance/distance-ext.cuh +++ b/cpp/src/distance/distance-ext.cuh @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2018-2024, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2018-2024, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ #pragma once diff --git a/cpp/src/distance/distance.cu b/cpp/src/distance/distance.cu index 964f569ede..565b655142 100644 --- a/cpp/src/distance/distance.cu +++ b/cpp/src/distance/distance.cu @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2018-2024, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2018-2024, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ From d5c56ae5513f89282770bbb5bac25fb56b82d175 Mon Sep 17 00:00:00 2001 From: divyegala Date: Tue, 11 Aug 2026 02:37:41 +0000 Subject: [PATCH 37/82] use cutile from conda --- .../all_cuda-133_arch-aarch64.yaml | 6 ++---- .../environments/all_cuda-133_arch-x86_64.yaml | 6 ++---- .../bench_ann_cuda-133_arch-aarch64.yaml | 6 ++---- .../bench_ann_cuda-133_arch-x86_64.yaml | 6 ++---- conda/recipes/libcuvs/recipe.yaml | 18 ++---------------- .../modules/generate_cutile_kernels.cmake | 5 ++--- dependencies.yaml | 18 ++++++------------ 7 files changed, 18 insertions(+), 47 deletions(-) diff --git a/conda/environments/all_cuda-133_arch-aarch64.yaml b/conda/environments/all_cuda-133_arch-aarch64.yaml index bc20593702..c50c4c9352 100644 --- a/conda/environments/all_cuda-133_arch-aarch64.yaml +++ b/conda/environments/all_cuda-133_arch-aarch64.yaml @@ -15,8 +15,10 @@ dependencies: - cuda-nvrtc-dev - cuda-nvtx-dev - cuda-profiler-api +- cuda-tileiras=13.3.* - cuda-version=13.3 - cupy>=14.0.1,!=14.1.0 +- cutile-python - cxx-compiler - cython>=3.2.2 - dlpack>=0.8,<1.0 @@ -37,7 +39,6 @@ dependencies: - numpy>=2.0,<3.0 - openblas - openjdk=22.* -- pip - pre-commit - pylibraft==26.10.*,>=0.0.0a0 - pytest @@ -47,7 +48,4 @@ dependencies: - scikit-build-core>=0.11.0 - scikit-learn>=1.5 - sysroot_linux-aarch64==2.28 -- pip: - - cuda-tile - - cuda-toolkit[tileiras]==13.3.* name: all_cuda-133_arch-aarch64 diff --git a/conda/environments/all_cuda-133_arch-x86_64.yaml b/conda/environments/all_cuda-133_arch-x86_64.yaml index c1347e6c67..03844beb7e 100644 --- a/conda/environments/all_cuda-133_arch-x86_64.yaml +++ b/conda/environments/all_cuda-133_arch-x86_64.yaml @@ -15,8 +15,10 @@ dependencies: - cuda-nvrtc-dev - cuda-nvtx-dev - cuda-profiler-api +- cuda-tileiras=13.3.* - cuda-version=13.3 - cupy>=14.0.1,!=14.1.0 +- cutile-python - cxx-compiler - cython>=3.2.2 - dlpack>=0.8,<1.0 @@ -36,7 +38,6 @@ dependencies: - numpy>=2.0,<3.0 - openblas - openjdk=22.* -- pip - pre-commit - pylibraft==26.10.*,>=0.0.0a0 - pytest @@ -46,7 +47,4 @@ dependencies: - scikit-build-core>=0.11.0 - scikit-learn>=1.5 - sysroot_linux-64==2.28 -- pip: - - cuda-tile - - cuda-toolkit[tileiras]==13.3.* name: all_cuda-133_arch-x86_64 diff --git a/conda/environments/bench_ann_cuda-133_arch-aarch64.yaml b/conda/environments/bench_ann_cuda-133_arch-aarch64.yaml index 87e75a3e09..a3f1b00d1f 100644 --- a/conda/environments/bench_ann_cuda-133_arch-aarch64.yaml +++ b/conda/environments/bench_ann_cuda-133_arch-aarch64.yaml @@ -15,8 +15,10 @@ dependencies: - cuda-nvrtc-dev - cuda-nvtx-dev - cuda-profiler-api +- cuda-tileiras=13.3.* - cuda-version=13.3 - cupy>=14.0.1,!=14.1.0 +- cutile-python - cuvs==26.10.*,>=0.0.0a0 - cxx-compiler - cython>=3.2.2 @@ -38,7 +40,6 @@ dependencies: - openblas - opensearch-py>=2.4.0 - pandas -- pip - pylibraft==26.10.*,>=0.0.0a0 - pyyaml - rapids-build-backend>=0.4.0,<0.5.0 @@ -47,7 +48,4 @@ dependencies: - setuptools>=77.0.0 - sysroot_linux-aarch64==2.28 - wheel -- pip: - - cuda-tile - - cuda-toolkit[tileiras]==13.3.* name: bench_ann_cuda-133_arch-aarch64 diff --git a/conda/environments/bench_ann_cuda-133_arch-x86_64.yaml b/conda/environments/bench_ann_cuda-133_arch-x86_64.yaml index 36cba2b0f5..4bd1ce7df9 100644 --- a/conda/environments/bench_ann_cuda-133_arch-x86_64.yaml +++ b/conda/environments/bench_ann_cuda-133_arch-x86_64.yaml @@ -15,8 +15,10 @@ dependencies: - cuda-nvrtc-dev - cuda-nvtx-dev - cuda-profiler-api +- cuda-tileiras=13.3.* - cuda-version=13.3 - cupy>=14.0.1,!=14.1.0 +- cutile-python - cuvs==26.10.*,>=0.0.0a0 - cxx-compiler - cython>=3.2.2 @@ -41,7 +43,6 @@ dependencies: - openblas - opensearch-py>=2.4.0 - pandas -- pip - pylibraft==26.10.*,>=0.0.0a0 - pyyaml - rapids-build-backend>=0.4.0,<0.5.0 @@ -50,7 +51,4 @@ dependencies: - setuptools>=77.0.0 - sysroot_linux-64==2.28 - wheel -- pip: - - cuda-tile - - cuda-toolkit[tileiras]==13.3.* name: bench_ann_cuda-133_arch-x86_64 diff --git a/conda/recipes/libcuvs/recipe.yaml b/conda/recipes/libcuvs/recipe.yaml index adc0c88beb..220c48c2e2 100644 --- a/conda/recipes/libcuvs/recipe.yaml +++ b/conda/recipes/libcuvs/recipe.yaml @@ -30,20 +30,6 @@ cache: export CXXFLAGS=$(echo $CXXFLAGS | sed -E 's@\-fdebug\-prefix\-map[^ ]*@@g') set +x - # cuTile and TileIRAS are PyPI-only build tools. Their CUDA components - # are installed under the Python package namespace, alongside the conda toolchain. - if [[ "${{ cuda_major }}" == "13" ]]; then - # rattler-build disables package indexes and dependency resolution by default. - env -u PIP_NO_INDEX \ - -u PIP_NO_DEPENDENCIES \ - -u PIP_NO_DEPS \ - -u PIP_IGNORE_INSTALLED \ - python -m pip install --no-cache-dir \ - --index-url https://pypi.org/simple \ - "cuda-tile" \ - "cuda-toolkit[tileiras]==${{ cuda_version }}.*" - fi - ./build.sh libcuvs bench-ann tests --allgpuarch --mnmg-tests --build-metrics=compile_lib --incl-cache-stats --no-nvtx -n secrets: @@ -86,8 +72,8 @@ cache: - ninja - if: cuda_major == "13" then: - - pip - - python + - cutile-python + - cuda-tileiras =${{ cuda_version }}.* - ${{ stdlib("c") }} host: - libnvjitlink-dev diff --git a/cpp/cmake/modules/generate_cutile_kernels.cmake b/cpp/cmake/modules/generate_cutile_kernels.cmake index d37513d817..89389d5bea 100644 --- a/cpp/cmake/modules/generate_cutile_kernels.cmake +++ b/cpp/cmake/modules/generate_cutile_kernels.cmake @@ -64,9 +64,8 @@ function(_cutile_kernels_setup) ) if(NOT _cutile_import_result EQUAL 0) message( - FATAL_ERROR - "cuda.tile (cuTile Python) is required to build cuTile embedded kernels. " - "Install it in the active Python environment, e.g. pip install cuda-tile[tileiras]." + FATAL_ERROR "cuda.tile (cuTile Python) is required to build cuTile embedded kernels. " + "Install cutile-python and cuda-tileiras (conda), or cuda-tile[tileiras] (pip)." ) endif() diff --git a/dependencies.yaml b/dependencies.yaml index 258b7edfa8..6d5fa64d12 100644 --- a/dependencies.yaml +++ b/dependencies.yaml @@ -427,23 +427,17 @@ dependencies: - matrix: cuda: "13.3" packages: - - pip - - pip: - - cuda-tile - - cuda-toolkit[tileiras]==13.3.* + - cutile-python + - cuda-tileiras=13.3.* - matrix: cuda: "13.*" packages: - - pip - - pip: - - cuda-tile - - cuda-toolkit[tileiras]==13.* + - cutile-python + - cuda-tileiras - matrix: packages: - - pip - - pip: - - cuda-tile - - cuda-toolkit[tileiras]==13.* + - cutile-python + - cuda-tileiras - output_types: [requirements, pyproject] matrices: - matrix: From ccec474bc8542f3098ffc72a8402a43098165f5d Mon Sep 17 00:00:00 2001 From: divyegala Date: Tue, 11 Aug 2026 04:42:36 +0000 Subject: [PATCH 38/82] fix == --- conda/environments/all_cuda-133_arch-aarch64.yaml | 2 +- conda/environments/all_cuda-133_arch-x86_64.yaml | 2 +- conda/environments/bench_ann_cuda-133_arch-aarch64.yaml | 2 +- conda/environments/bench_ann_cuda-133_arch-x86_64.yaml | 2 +- conda/recipes/libcuvs/recipe.yaml | 2 +- dependencies.yaml | 2 +- 6 files changed, 6 insertions(+), 6 deletions(-) diff --git a/conda/environments/all_cuda-133_arch-aarch64.yaml b/conda/environments/all_cuda-133_arch-aarch64.yaml index c50c4c9352..481de23b33 100644 --- a/conda/environments/all_cuda-133_arch-aarch64.yaml +++ b/conda/environments/all_cuda-133_arch-aarch64.yaml @@ -15,7 +15,7 @@ dependencies: - cuda-nvrtc-dev - cuda-nvtx-dev - cuda-profiler-api -- cuda-tileiras=13.3.* +- cuda-tileiras==13.3.* - cuda-version=13.3 - cupy>=14.0.1,!=14.1.0 - cutile-python diff --git a/conda/environments/all_cuda-133_arch-x86_64.yaml b/conda/environments/all_cuda-133_arch-x86_64.yaml index 03844beb7e..6cbda9e374 100644 --- a/conda/environments/all_cuda-133_arch-x86_64.yaml +++ b/conda/environments/all_cuda-133_arch-x86_64.yaml @@ -15,7 +15,7 @@ dependencies: - cuda-nvrtc-dev - cuda-nvtx-dev - cuda-profiler-api -- cuda-tileiras=13.3.* +- cuda-tileiras==13.3.* - cuda-version=13.3 - cupy>=14.0.1,!=14.1.0 - cutile-python diff --git a/conda/environments/bench_ann_cuda-133_arch-aarch64.yaml b/conda/environments/bench_ann_cuda-133_arch-aarch64.yaml index a3f1b00d1f..d30f558c0c 100644 --- a/conda/environments/bench_ann_cuda-133_arch-aarch64.yaml +++ b/conda/environments/bench_ann_cuda-133_arch-aarch64.yaml @@ -15,7 +15,7 @@ dependencies: - cuda-nvrtc-dev - cuda-nvtx-dev - cuda-profiler-api -- cuda-tileiras=13.3.* +- cuda-tileiras==13.3.* - cuda-version=13.3 - cupy>=14.0.1,!=14.1.0 - cutile-python diff --git a/conda/environments/bench_ann_cuda-133_arch-x86_64.yaml b/conda/environments/bench_ann_cuda-133_arch-x86_64.yaml index 4bd1ce7df9..8c12e2dca5 100644 --- a/conda/environments/bench_ann_cuda-133_arch-x86_64.yaml +++ b/conda/environments/bench_ann_cuda-133_arch-x86_64.yaml @@ -15,7 +15,7 @@ dependencies: - cuda-nvrtc-dev - cuda-nvtx-dev - cuda-profiler-api -- cuda-tileiras=13.3.* +- cuda-tileiras==13.3.* - cuda-version=13.3 - cupy>=14.0.1,!=14.1.0 - cutile-python diff --git a/conda/recipes/libcuvs/recipe.yaml b/conda/recipes/libcuvs/recipe.yaml index 220c48c2e2..39eb23f4ae 100644 --- a/conda/recipes/libcuvs/recipe.yaml +++ b/conda/recipes/libcuvs/recipe.yaml @@ -73,7 +73,7 @@ cache: - if: cuda_major == "13" then: - cutile-python - - cuda-tileiras =${{ cuda_version }}.* + - cuda-tileiras ==${{ cuda_version }}.* - ${{ stdlib("c") }} host: - libnvjitlink-dev diff --git a/dependencies.yaml b/dependencies.yaml index 6d5fa64d12..a0f27865f3 100644 --- a/dependencies.yaml +++ b/dependencies.yaml @@ -428,7 +428,7 @@ dependencies: cuda: "13.3" packages: - cutile-python - - cuda-tileiras=13.3.* + - cuda-tileiras==13.3.* - matrix: cuda: "13.*" packages: From d63649e72fe16e8e680f177108f5bc7c9baf96ac Mon Sep 17 00:00:00 2001 From: divyegala Date: Tue, 11 Aug 2026 16:54:37 +0000 Subject: [PATCH 39/82] use sass within family --- .../cuvs/detail/jit_lto/cutile_arch_tags.hpp | 14 ++------- .../cuvs/detail/jit_lto/cutile_module.hpp | 29 +++++++++++++++---- .../cuvs/detail/jit_lto/tileir_compat.hpp | 6 ++-- .../detail/jit_lto/TileAlgorithmPlanner.cpp | 7 ++--- .../cutile/fused_1nn_cutile_matrix.json | 10 ------- .../cutile/fused_1nn_planner.hpp | 6 ---- 6 files changed, 32 insertions(+), 40 deletions(-) diff --git a/cpp/include/cuvs/detail/jit_lto/cutile_arch_tags.hpp b/cpp/include/cuvs/detail/jit_lto/cutile_arch_tags.hpp index f88ee39c98..2b378dac78 100644 --- a/cpp/include/cuvs/detail/jit_lto/cutile_arch_tags.hpp +++ b/cpp/include/cuvs/detail/jit_lto/cutile_arch_tags.hpp @@ -24,11 +24,6 @@ struct cutile_arch_8_6 { static constexpr int cc_minor = 6; }; -struct cutile_arch_8_9 { - static constexpr int cc_major = 8; - static constexpr int cc_minor = 9; -}; - struct cutile_arch_9_0 { static constexpr int cc_major = 9; static constexpr int cc_minor = 0; @@ -46,13 +41,8 @@ struct cutile_arch_12_0 { inline bool is_embedded_cubin_arch(int cc_major, int cc_minor) { - if (cc_major == 8 && cc_minor == 0) { return true; } - if (cc_major == 8 && cc_minor == 6) { return true; } - if (cc_major == 8 && cc_minor == 9) { return true; } - if (cc_major == 9 && cc_minor == 0) { return true; } - if (cc_major == 10 && cc_minor == 0) { return true; } - if (cc_major == 12 && cc_minor == 0) { return true; } - return false; + if (cc_minor < 0) { return false; } + return cc_major == 8 || cc_major == 9 || cc_major == 10 || cc_major == 12; } #else diff --git a/cpp/include/cuvs/detail/jit_lto/cutile_module.hpp b/cpp/include/cuvs/detail/jit_lto/cutile_module.hpp index 29f30ae15a..80029ba967 100644 --- a/cpp/include/cuvs/detail/jit_lto/cutile_module.hpp +++ b/cpp/include/cuvs/detail/jit_lto/cutile_module.hpp @@ -40,7 +40,28 @@ inline bool get_device_compute_capability(int& cc_major, int& cc_minor) return true; } -/** Selects a prebuilt cubin for the device CC, or embedded TileIR when the driver can JIT it. */ +/** + * Selects the newest compatible cubin in the device's compute-capability major family. + * + * CUDA cubins are forward compatible across minor revisions within a major family, so an SM 8.9 + * device can load SM 8.6 SASS and an SM 12.1 device can load SM 12.0 SASS. + */ +inline const CubinFragmentEntry* find_compatible_cubin_fragment( + int cc_major, + int cc_minor, + const std::vector>& cubin_fragments) +{ + const CubinFragmentEntry* best = nullptr; + for (const auto& fragment : cubin_fragments) { + if (fragment->get_cc_major() != cc_major || fragment->get_cc_minor() > cc_minor) { continue; } + if (best == nullptr || fragment->get_cc_minor() > best->get_cc_minor()) { + best = fragment.get(); + } + } + return best; +} + +/** Selects compatible prebuilt SASS for the device, or TileIR when the driver can JIT it. */ inline std::optional resolve_cutile_module_image( int cc_major, int cc_minor, @@ -48,10 +69,8 @@ inline std::optional resolve_cutile_module_image( const std::vector>& cubin_fragments, const TileIrBytecodeFragmentEntry* tileir_fragment) { - for (const auto& fragment : cubin_fragments) { - if (fragment->get_cc_major() == cc_major && fragment->get_cc_minor() == cc_minor) { - return CutileModuleImage{fragment->get_data(), fragment->get_length()}; - } + if (const auto* fragment = find_compatible_cubin_fragment(cc_major, cc_minor, cubin_fragments)) { + return CutileModuleImage{fragment->get_data(), fragment->get_length()}; } if (tileir_fragment != nullptr && tileir_fallback_available(driver_version)) { return CutileModuleImage{tileir_fragment->get_data(), tileir_fragment->get_length()}; diff --git a/cpp/include/cuvs/detail/jit_lto/tileir_compat.hpp b/cpp/include/cuvs/detail/jit_lto/tileir_compat.hpp index 029d53bfa0..bde7dab302 100644 --- a/cpp/include/cuvs/detail/jit_lto/tileir_compat.hpp +++ b/cpp/include/cuvs/detail/jit_lto/tileir_compat.hpp @@ -44,7 +44,7 @@ inline bool cutile_integration_enabled() return library_built_with_cutile() && runtime_cuda13_or_newer(); } -/** True when this build embeds a prebuilt cubin for the given compute capability. */ +/** True when this build embeds compatible SASS in the device's compute-capability major family. */ inline bool has_embedded_cubin_for_arch(int cc_major, int cc_minor) { return is_embedded_cubin_arch(cc_major, cc_minor); @@ -58,8 +58,8 @@ inline bool tileir_fallback_available(int driver_version) /** * True when a cuTile launch may be attempted for the given device: cuTile is enabled, the runtime - * is CUDA 13+, and either a matching embedded cubin exists (no driver JIT required) or the driver - * can JIT the embedded TileIR bytecode fallback. + * is CUDA 13+, and either compatible same-family SASS exists (no driver JIT required) or the + * driver can JIT the embedded TileIR bytecode fallback. */ #if CUVS_CUTILE_ENABLED inline bool cutile_launch_available_for_arch(int cc_major, int cc_minor, int driver_version) diff --git a/cpp/src/detail/jit_lto/TileAlgorithmPlanner.cpp b/cpp/src/detail/jit_lto/TileAlgorithmPlanner.cpp index 907b9fbc47..7195c532d7 100644 --- a/cpp/src/detail/jit_lto/TileAlgorithmPlanner.cpp +++ b/cpp/src/detail/jit_lto/TileAlgorithmPlanner.cpp @@ -49,10 +49,9 @@ CutileTileConfig TileAlgorithmPlanner::tile_config() const int cc_major = 0; int cc_minor = 0; if (cuvs::detail::jit_lto::get_device_compute_capability(cc_major, cc_minor)) { - for (const auto& fragment : cubin_fragments_) { - if (fragment->get_cc_major() == cc_major && fragment->get_cc_minor() == cc_minor) { - return tile_config_from_fragment(fragment.get(), entrypoint); - } + if (const auto* fragment = cuvs::detail::jit_lto::find_compatible_cubin_fragment( + cc_major, cc_minor, cubin_fragments_)) { + return tile_config_from_fragment(fragment, entrypoint); } } diff --git a/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_cutile_matrix.json b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_cutile_matrix.json index 9adf7ddfce..c94aec352d 100644 --- a/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_cutile_matrix.json +++ b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_cutile_matrix.json @@ -71,16 +71,6 @@ "cc_minor": 6, "arch_tag": "cutile_arch_8_6" }, - { - "output_format": "cubin", - "artifact_ext": "cubin", - "artifact_basename": "@data_type@_@index_abbrev@_@abi_abbrev@_@gpu_code@", - "register": "cubin", - "gpu_code": "sm_89", - "cc_major": 8, - "cc_minor": 9, - "arch_tag": "cutile_arch_8_9" - }, { "output_format": "cubin", "artifact_ext": "cubin", diff --git a/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_planner.hpp b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_planner.hpp index 24389e6edc..e9e66a7a26 100644 --- a/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_planner.hpp +++ b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_planner.hpp @@ -54,7 +54,6 @@ struct Fused1nnTilePlanner : TileAlgorithmPlanner { using cuvs::detail::jit_lto::cutile_arch_12_0; using cuvs::detail::jit_lto::cutile_arch_8_0; using cuvs::detail::jit_lto::cutile_arch_8_6; - using cuvs::detail::jit_lto::cutile_arch_8_9; using cuvs::detail::jit_lto::cutile_arch_9_0; constexpr bool is_relaxed = std::is_same_v; @@ -64,9 +63,6 @@ struct Fused1nnTilePlanner : TileAlgorithmPlanner { using Tile86 = std::conditional_t; - using Tile89 = std::conditional_t; using Tile90 = std::conditional_t; @@ -81,8 +77,6 @@ struct Fused1nnTilePlanner : TileAlgorithmPlanner { fragment_tag_fused_1nn_cubin>(); this->add_static_fragment< fragment_tag_fused_1nn_cubin>(); - this->add_static_fragment< - fragment_tag_fused_1nn_cubin>(); this->add_static_fragment< fragment_tag_fused_1nn_cubin>(); this->add_static_fragment< From b7ad9206d7bd0c8c3c10ded9d5e6006f5e2ae97e Mon Sep 17 00:00:00 2001 From: divyegala Date: Tue, 11 Aug 2026 18:11:26 +0000 Subject: [PATCH 40/82] address reviews --- .../all_cuda-133_arch-aarch64.yaml | 2 +- .../all_cuda-133_arch-x86_64.yaml | 2 +- .../bench_ann_cuda-133_arch-aarch64.yaml | 2 +- .../bench_ann_cuda-133_arch-x86_64.yaml | 2 +- conda/recipes/libcuvs/recipe.yaml | 2 +- cpp/src/cluster/detail/kmeans_balanced.cuh | 72 +++-- cpp/src/cluster/kmeans.cuh | 2 +- .../cutile/.nfs000000001a3bd0d600001bbe | 290 ++++++++++++++++++ .../cutile/fused_1nn_cutile_matrix.json | 69 ++++- .../cutile/fused_1nn_kernel.py | 5 +- .../cutile/fused_1nn_planner.hpp | 18 +- .../cutile/fused_1nn_tile.cu | 12 +- dependencies.yaml | 2 +- 13 files changed, 405 insertions(+), 75 deletions(-) create mode 100644 cpp/src/distance/detail/fused_distance_nn/cutile/.nfs000000001a3bd0d600001bbe diff --git a/conda/environments/all_cuda-133_arch-aarch64.yaml b/conda/environments/all_cuda-133_arch-aarch64.yaml index 481de23b33..aa50303efa 100644 --- a/conda/environments/all_cuda-133_arch-aarch64.yaml +++ b/conda/environments/all_cuda-133_arch-aarch64.yaml @@ -15,7 +15,7 @@ dependencies: - cuda-nvrtc-dev - cuda-nvtx-dev - cuda-profiler-api -- cuda-tileiras==13.3.* +- cuda-tileiras - cuda-version=13.3 - cupy>=14.0.1,!=14.1.0 - cutile-python diff --git a/conda/environments/all_cuda-133_arch-x86_64.yaml b/conda/environments/all_cuda-133_arch-x86_64.yaml index 6cbda9e374..47d0226234 100644 --- a/conda/environments/all_cuda-133_arch-x86_64.yaml +++ b/conda/environments/all_cuda-133_arch-x86_64.yaml @@ -15,7 +15,7 @@ dependencies: - cuda-nvrtc-dev - cuda-nvtx-dev - cuda-profiler-api -- cuda-tileiras==13.3.* +- cuda-tileiras - cuda-version=13.3 - cupy>=14.0.1,!=14.1.0 - cutile-python diff --git a/conda/environments/bench_ann_cuda-133_arch-aarch64.yaml b/conda/environments/bench_ann_cuda-133_arch-aarch64.yaml index d30f558c0c..1d5f229c8f 100644 --- a/conda/environments/bench_ann_cuda-133_arch-aarch64.yaml +++ b/conda/environments/bench_ann_cuda-133_arch-aarch64.yaml @@ -15,7 +15,7 @@ dependencies: - cuda-nvrtc-dev - cuda-nvtx-dev - cuda-profiler-api -- cuda-tileiras==13.3.* +- cuda-tileiras - cuda-version=13.3 - cupy>=14.0.1,!=14.1.0 - cutile-python diff --git a/conda/environments/bench_ann_cuda-133_arch-x86_64.yaml b/conda/environments/bench_ann_cuda-133_arch-x86_64.yaml index 8c12e2dca5..228d11c3d8 100644 --- a/conda/environments/bench_ann_cuda-133_arch-x86_64.yaml +++ b/conda/environments/bench_ann_cuda-133_arch-x86_64.yaml @@ -15,7 +15,7 @@ dependencies: - cuda-nvrtc-dev - cuda-nvtx-dev - cuda-profiler-api -- cuda-tileiras==13.3.* +- cuda-tileiras - cuda-version=13.3 - cupy>=14.0.1,!=14.1.0 - cutile-python diff --git a/conda/recipes/libcuvs/recipe.yaml b/conda/recipes/libcuvs/recipe.yaml index 39eb23f4ae..d9cb0ec5da 100644 --- a/conda/recipes/libcuvs/recipe.yaml +++ b/conda/recipes/libcuvs/recipe.yaml @@ -73,7 +73,7 @@ cache: - if: cuda_major == "13" then: - cutile-python - - cuda-tileiras ==${{ cuda_version }}.* + - cuda-tileiras - ${{ stdlib("c") }} host: - libnvjitlink-dev diff --git a/cpp/src/cluster/detail/kmeans_balanced.cuh b/cpp/src/cluster/detail/kmeans_balanced.cuh index 1311b9056a..8a84dbfd9f 100644 --- a/cpp/src/cluster/detail/kmeans_balanced.cuh +++ b/cpp/src/cluster/detail/kmeans_balanced.cuh @@ -5,6 +5,7 @@ #pragma once +#include "../kmeans.cuh" #include "kmeans_common.cuh" #include @@ -104,7 +105,7 @@ inline std::enable_if_t> predict_core( if constexpr (std::is_same_v) { auto labels_view = raft::make_device_vector_view(labels, n_rows); - cuvs::cluster::kmeans::detail::minClusterAndDistanceCompute( + cuvs::cluster::kmeans::min_cluster_and_distance( handle, X_view, centroids_view, @@ -119,18 +120,17 @@ inline std::enable_if_t> predict_core( } else { auto nearest_idx = raft::make_device_mdarray(handle, mr, raft::make_extents(n_rows)); - cuvs::cluster::kmeans::detail::minClusterAndDistanceCompute( - handle, - X_view, - centroids_view, - nearest_idx.view(), - nearest_dist.view(), - X_norm_view, - L2NormBuf_OR_DistBuf, - params.metric, - 0, - 0, - workspace); + cuvs::cluster::kmeans::min_cluster_and_distance(handle, + X_view, + centroids_view, + nearest_idx.view(), + nearest_dist.view(), + X_norm_view, + L2NormBuf_OR_DistBuf, + params.metric, + 0, + 0, + workspace); raft::copy( handle, raft::make_device_vector_view(labels, n_rows), nearest_idx.view()); } @@ -152,33 +152,31 @@ inline std::enable_if_t> predict_core( if constexpr (std::is_same_v) { auto labels_view = raft::make_device_vector_view(labels, n_rows); - cuvs::cluster::kmeans::detail::minClusterAndDistanceCompute( - handle, - X_view, - centroids_view, - labels_view, - nearest_dist.view(), - X_norm_view, - L2NormBuf_OR_DistBuf, - params.metric, - 0, - 0, - workspace); + cuvs::cluster::kmeans::min_cluster_and_distance(handle, + X_view, + centroids_view, + labels_view, + nearest_dist.view(), + X_norm_view, + L2NormBuf_OR_DistBuf, + params.metric, + 0, + 0, + workspace); } else { auto nearest_idx = raft::make_device_mdarray(handle, mr, raft::make_extents(n_rows)); - cuvs::cluster::kmeans::detail::minClusterAndDistanceCompute( - handle, - X_view, - centroids_view, - nearest_idx.view(), - nearest_dist.view(), - X_norm_view, - L2NormBuf_OR_DistBuf, - params.metric, - 0, - 0, - workspace); + cuvs::cluster::kmeans::min_cluster_and_distance(handle, + X_view, + centroids_view, + nearest_idx.view(), + nearest_dist.view(), + X_norm_view, + L2NormBuf_OR_DistBuf, + params.metric, + 0, + 0, + workspace); raft::copy(handle, raft::make_device_vector_view(labels, n_rows), nearest_idx.view()); diff --git a/cpp/src/cluster/kmeans.cuh b/cpp/src/cluster/kmeans.cuh index dc9e56c746..d635426229 100644 --- a/cpp/src/cluster/kmeans.cuh +++ b/cpp/src/cluster/kmeans.cuh @@ -483,7 +483,7 @@ void min_cluster_and_distance(raft::resources const& handle, raft::device_matrix_view centroids, raft::device_vector_view nearest_idx, raft::device_vector_view nearest_dist, - raft::device_vector_view L2NormX, + raft::device_vector_view L2NormX, rmm::device_uvector& L2NormBuf_OR_DistBuf, cuvs::distance::DistanceType metric, int batch_samples, diff --git a/cpp/src/distance/detail/fused_distance_nn/cutile/.nfs000000001a3bd0d600001bbe b/cpp/src/distance/detail/fused_distance_nn/cutile/.nfs000000001a3bd0d600001bbe new file mode 100644 index 0000000000..b858e6bdcf --- /dev/null +++ b/cpp/src/distance/detail/fused_distance_nn/cutile/.nfs000000001a3bd0d600001bbe @@ -0,0 +1,290 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include "fused_1nn_tile.hpp" + +#include "fused_1nn_planner.hpp" + +#include +#include +#include +#include +#include + +#include +#include +#include +#include + +namespace cuvs { +namespace distance { +namespace detail { + +namespace { + +template +bool launch_fused_1nn_tile(IdxT* nearest_idx, + DataT* nearest_dist, + const DataT* x, + const DataT* y, + const DataT* xn, + const DataT* yn, + IdxT m, + IdxT n, + IdxT k, + cuvs::distance::DistanceType metric, + bool is_sqrt, + cudaStream_t stream) +{ + if constexpr (!std::is_same_v && !std::is_same_v) { return false; } + + if (nearest_dist == nullptr) { return false; } + + Fused1nnTilePlanner planner; + planner.add_entrypoint(); + planner.add_tileir_fallback(); + const CutileTileConfig tile_cfg = planner.tile_config(); + auto launcher = planner.try_get_launcher(); + if (!launcher) { return false; } + + int metric_code; + bool apply_sqrt = false; + switch (metric) { + case cuvs::distance::DistanceType::InnerProduct: + metric_code = static_cast(cuvs::distance::DistanceType::InnerProduct); + break; + case cuvs::distance::DistanceType::L2Expanded: + case cuvs::distance::DistanceType::L2SqrtExpanded: + metric_code = static_cast(cuvs::distance::DistanceType::L2Expanded); + apply_sqrt = is_sqrt; + break; + case cuvs::distance::DistanceType::CosineExpanded: + metric_code = static_cast(cuvs::distance::DistanceType::CosineExpanded); + break; + default: return false; + } + + IdxT shape_x[2] = {m, k}; + IdxT stride_x[2] = {k, IdxT{1}}; + IdxT shape_y[2] = {n, k}; + IdxT stride_y[2] = {k, IdxT{1}}; + IdxT shape_xn = m; + IdxT stride_xn = IdxT{1}; + IdxT shape_yn = n; + IdxT stride_yn = IdxT{1}; + IdxT shape_idx = m; + IdxT stride_idx = IdxT{1}; + IdxT shape_dist = m; + IdxT stride_dist = IdxT{1}; + + IdxT M = m; + IdxT N = n; + IdxT K = k; + + void* x_ptr = const_cast(x); + void* y_ptr = const_cast(y); + void* xn_ptr = const_cast(xn); + void* yn_ptr = const_cast(yn); + // OutIdx must be a valid device pointer for the launch ABI; when store_idx is 0 the kernel + // does not write it (dist-only callers pass nearest_dist as a stand-in). + const IdxT store_idx = nearest_idx != nullptr ? IdxT{1} : IdxT{0}; + void* idx_ptr = + nearest_idx != nullptr ? static_cast(nearest_idx) : static_cast(nearest_dist); + void* dist_ptr = nearest_dist; + + const int tile_m = tile_cfg.tile_m; + dim3 grid((static_cast(m) + tile_m - 1) / tile_m, 1, 1); + dim3 block(1, 1, 1); + + using fused_1nn_cutile_kernel_t = void(void*, + IdxT, + IdxT, + IdxT, + IdxT, + void*, + IdxT, + IdxT, + IdxT, + IdxT, + void*, + IdxT, + IdxT, + void*, + IdxT, + IdxT, + void*, + IdxT, + IdxT, + void*, + IdxT, + IdxT, + IdxT, + IdxT, + IdxT, + IdxT, + IdxT, + int); + launcher->template dispatch(stream, + grid, + block, + 0, + x_ptr, + shape_x[0], + shape_x[1], + stride_x[0], + stride_x[1], + y_ptr, + shape_y[0], + shape_y[1], + stride_y[0], + stride_y[1], + xn_ptr, + shape_xn, + stride_xn, + yn_ptr, + shape_yn, + stride_yn, + idx_ptr, + shape_idx, + stride_idx, + dist_ptr, + shape_dist, + stride_dist, + M, + N, + K, + static_cast(apply_sqrt), + store_idx, + metric_code); + RAFT_CUDA_TRY(cudaGetLastError()); + return true; +} + +template +bool try_fused_1nn_tile_dispatch(IdxT* nearest_idx, + DataT* nearest_dist, + const DataT* x, + const DataT* y, + const DataT* xn, + const DataT* yn, + IdxT m, + IdxT n, + IdxT k, + cuvs::distance::DistanceType metric, + bool is_sqrt, + cudaStream_t stream) +{ + return launch_fused_1nn_tile( + nearest_idx, nearest_dist, x, y, xn, yn, m, n, k, metric, is_sqrt, stream); +} + +} // namespace + +template + requires is_fused_1nn_cutile_data_v +bool try_fused_1nn_tile(IdxT* nearest_idx, + DataT* nearest_dist, + const DataT* x, + const DataT* y, + const DataT* xn, + const DataT* yn, + IdxT m, + IdxT n, + IdxT k, + cuvs::distance::DistanceType metric, + bool is_sqrt, + void* index_workspace, + cudaStream_t stream) +{ + if (!cuvs::detail::jit_lto::cutile_launch_available_on_current_device()) { return false; } + static_assert(std::is_same_v || std::is_same_v); + + int cc_major = 0; + int cc_minor = 0; + if (!cuvs::detail::jit_lto::get_device_compute_capability(cc_major, cc_minor)) { return false; } + constexpr int tma_pitch_elements = 16 / sizeof(DataT); + const bool use_strict_abi = cc_major >= 9 && k % tma_pitch_elements == 0; + + if constexpr (std::is_same_v) { + if (use_strict_abi) { + return try_fused_1nn_tile_dispatch( + nearest_idx, nearest_dist, x, y, xn, yn, m, n, k, metric, is_sqrt, stream); + } + return try_fused_1nn_tile_dispatch( + nearest_idx, nearest_dist, x, y, xn, yn, m, n, k, metric, is_sqrt, stream); + } else { + constexpr int64_t max_i32 = std::numeric_limits::max(); + if (n > max_i32 || k > max_i32) { return false; } + if (nearest_idx != nullptr && index_workspace == nullptr) { return false; } + + auto* tmp_idx = static_cast(index_workspace); + for (int64_t offset = 0; offset < m; offset += max_i32) { + const int batch_m = static_cast(std::min(max_i32, m - offset)); + const auto* batch_x = x + static_cast(offset) * static_cast(k); + const auto* batch_xn = xn == nullptr ? nullptr : xn + offset; + auto* batch_dist = nearest_dist == nullptr ? nullptr : nearest_dist + offset; + + const bool launched = + use_strict_abi + ? try_fused_1nn_tile_dispatch(tmp_idx, + batch_dist, + batch_x, + y, + batch_xn, + yn, + batch_m, + static_cast(n), + static_cast(k), + metric, + is_sqrt, + stream) + : try_fused_1nn_tile_dispatch(tmp_idx, + batch_dist, + batch_x, + y, + batch_xn, + yn, + batch_m, + static_cast(n), + static_cast(k), + metric, + is_sqrt, + stream); + if (!launched) { return false; } + + if (nearest_idx != nullptr) { + raft::linalg::unaryOp( + nearest_idx + offset, tmp_idx, batch_m, raft::cast_op{}, stream); + } + } + return true; + } +} + +#define CUVS_INST_TRY_FUSED_1NN_TILE(DataT, IdxT) \ + template CUVS_EXPORT bool try_fused_1nn_tile(IdxT*, \ + DataT*, \ + const DataT*, \ + const DataT*, \ + const DataT*, \ + const DataT*, \ + IdxT, \ + IdxT, \ + IdxT, \ + cuvs::distance::DistanceType, \ + bool, \ + void*, \ + cudaStream_t) + +CUVS_INST_TRY_FUSED_1NN_TILE(float, int); +CUVS_INST_TRY_FUSED_1NN_TILE(float, int64_t); +CUVS_INST_TRY_FUSED_1NN_TILE(half, int); +CUVS_INST_TRY_FUSED_1NN_TILE(half, int64_t); + +#undef CUVS_INST_TRY_FUSED_1NN_TILE + +} // namespace detail +} // namespace distance +} // namespace cuvs diff --git a/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_cutile_matrix.json b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_cutile_matrix.json index c94aec352d..cce186a34d 100644 --- a/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_cutile_matrix.json +++ b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_cutile_matrix.json @@ -1,11 +1,6 @@ [ { "_abi": [ - { - "matrix_layout": "strict", - "abi_abbrev": "strict", - "abi_tag": "cutile_abi_strict" - }, { "matrix_layout": "relaxed", "abi_abbrev": "relaxed", @@ -41,16 +36,6 @@ } ], "_export": [ - { - "output_format": "cubin", - "artifact_ext": "cubin", - "artifact_basename": "@data_type@_@index_abbrev@_@abi_abbrev@_@gpu_code@", - "register": "cubin", - "gpu_code": "sm_100", - "cc_major": 10, - "cc_minor": 0, - "arch_tag": "cutile_arch_10_0" - }, { "output_format": "cubin", "artifact_ext": "cubin", @@ -70,6 +55,60 @@ "cc_major": 8, "cc_minor": 6, "arch_tag": "cutile_arch_8_6" + } + ] + }, + { + "_abi": [ + { + "matrix_layout": "strict", + "abi_abbrev": "strict", + "abi_tag": "cutile_abi_strict" + }, + { + "matrix_layout": "relaxed", + "abi_abbrev": "relaxed", + "abi_tag": "cutile_abi_relaxed" + } + ], + "_data": [ + { + "data_type": "half", + "data_abbrev": "h" + }, + { + "data_type": "float", + "data_abbrev": "f" + } + ], + "_metric": [ + { + "metric": "runtime" + } + ], + "_index": [ + { + "index_type": "int32", + "index_abbrev": "i32" + } + ], + "_tile": [ + { + "tile_m": 128, + "tile_n": 128, + "tile_k": 32 + } + ], + "_export": [ + { + "output_format": "cubin", + "artifact_ext": "cubin", + "artifact_basename": "@data_type@_@index_abbrev@_@abi_abbrev@_@gpu_code@", + "register": "cubin", + "gpu_code": "sm_100", + "cc_major": 10, + "cc_minor": 0, + "arch_tag": "cutile_arch_10_0" }, { "output_format": "cubin", diff --git a/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_kernel.py b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_kernel.py index 0937dc8116..ab79c72ccb 100644 --- a/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_kernel.py +++ b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_kernel.py @@ -55,7 +55,7 @@ def make_kernel( core_shape = (tile_m, tile_n) best_shape = (tile_m, 1) - @ct.kernel(occupancy=ct.ByTarget(sm_100=2, sm_120=2)) + @ct.kernel(occupancy=ct.ByTarget(sm_120=2)) def fused_1nn_kernel( A, B, @@ -120,13 +120,14 @@ def red_op(a_score, a_idx, b_score, b_idx): B_norm, index=(n,), shape=(tn,), padding_mode=zero_pad ) if metric_code == METRIC_L2_EXPANDED: + # L2 receives squared row norms; cosine receives L2 magnitudes. # The A norm is constant across centroids. Reduce # 0.5 * ||y||^2 - dot(x, y), then recover full L2 once. score = (0.5 * b_norm)[None, :] - accumulator else: # Defer the A-norm division until after selecting the # winning centroid. - score = -(accumulator / b_norm[None, :]) + score = accumulator / (-b_norm)[None, :] if n == num_tiles_n - 1: col = ct.arange(tn, dtype=ct.int16) diff --git a/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_planner.hpp b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_planner.hpp index e9e66a7a26..f3f1f19da4 100644 --- a/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_planner.hpp +++ b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_planner.hpp @@ -57,12 +57,6 @@ struct Fused1nnTilePlanner : TileAlgorithmPlanner { using cuvs::detail::jit_lto::cutile_arch_9_0; constexpr bool is_relaxed = std::is_same_v; - using Tile80 = std::conditional_t; - using Tile86 = std::conditional_t; using Tile90 = std::conditional_t; @@ -73,10 +67,14 @@ struct Fused1nnTilePlanner : TileAlgorithmPlanner { fused_1nn_matrix_tile_cutile_arch_12_0_relaxed, fused_1nn_matrix_tile_cutile_arch_12_0_strict>; - this->add_static_fragment< - fragment_tag_fused_1nn_cubin>(); - this->add_static_fragment< - fragment_tag_fused_1nn_cubin>(); + if constexpr (is_relaxed) { + using Tile80 = fused_1nn_matrix_tile_cutile_arch_8_0_relaxed; + using Tile86 = fused_1nn_matrix_tile_cutile_arch_8_6_relaxed; + this->add_static_fragment< + fragment_tag_fused_1nn_cubin>(); + this->add_static_fragment< + fragment_tag_fused_1nn_cubin>(); + } this->add_static_fragment< fragment_tag_fused_1nn_cubin>(); this->add_static_fragment< diff --git a/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_tile.cu b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_tile.cu index 622af314a7..a0df8f60c2 100644 --- a/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_tile.cu +++ b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_tile.cu @@ -201,9 +201,14 @@ bool try_fused_1nn_tile(IdxT* nearest_idx, if (!cuvs::detail::jit_lto::cutile_launch_available_on_current_device()) { return false; } static_assert(std::is_same_v || std::is_same_v); + int cc_major = 0; + int cc_minor = 0; + if (!cuvs::detail::jit_lto::get_device_compute_capability(cc_major, cc_minor)) { return false; } + constexpr int tma_pitch_elements = 16 / sizeof(DataT); + const bool use_strict_abi = cc_major >= 9 && k % tma_pitch_elements == 0; + if constexpr (std::is_same_v) { - constexpr int tma_pitch_elements = 16 / sizeof(DataT); - if (k % tma_pitch_elements == 0) { + if (use_strict_abi) { return try_fused_1nn_tile_dispatch( nearest_idx, nearest_dist, x, y, xn, yn, m, n, k, metric, is_sqrt, stream); } @@ -221,9 +226,8 @@ bool try_fused_1nn_tile(IdxT* nearest_idx, const auto* batch_xn = xn == nullptr ? nullptr : xn + offset; auto* batch_dist = nearest_dist == nullptr ? nullptr : nearest_dist + offset; - constexpr int tma_pitch_elements = 16 / sizeof(DataT); const bool launched = - k % tma_pitch_elements == 0 + use_strict_abi ? try_fused_1nn_tile_dispatch(tmp_idx, batch_dist, batch_x, diff --git a/dependencies.yaml b/dependencies.yaml index a0f27865f3..de4b796d2b 100644 --- a/dependencies.yaml +++ b/dependencies.yaml @@ -428,7 +428,7 @@ dependencies: cuda: "13.3" packages: - cutile-python - - cuda-tileiras==13.3.* + - cuda-tileiras - matrix: cuda: "13.*" packages: From 8f8e72188c0f611a5e05b97f0682171ac5c46a51 Mon Sep 17 00:00:00 2001 From: divyegala Date: Tue, 11 Aug 2026 18:15:06 +0000 Subject: [PATCH 41/82] delete resource --- .../cutile/.nfs000000001a3bd0d600001bbe | 290 ------------------ 1 file changed, 290 deletions(-) delete mode 100644 cpp/src/distance/detail/fused_distance_nn/cutile/.nfs000000001a3bd0d600001bbe diff --git a/cpp/src/distance/detail/fused_distance_nn/cutile/.nfs000000001a3bd0d600001bbe b/cpp/src/distance/detail/fused_distance_nn/cutile/.nfs000000001a3bd0d600001bbe deleted file mode 100644 index b858e6bdcf..0000000000 --- a/cpp/src/distance/detail/fused_distance_nn/cutile/.nfs000000001a3bd0d600001bbe +++ /dev/null @@ -1,290 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: Apache-2.0 - */ - -#include "fused_1nn_tile.hpp" - -#include "fused_1nn_planner.hpp" - -#include -#include -#include -#include -#include - -#include -#include -#include -#include - -namespace cuvs { -namespace distance { -namespace detail { - -namespace { - -template -bool launch_fused_1nn_tile(IdxT* nearest_idx, - DataT* nearest_dist, - const DataT* x, - const DataT* y, - const DataT* xn, - const DataT* yn, - IdxT m, - IdxT n, - IdxT k, - cuvs::distance::DistanceType metric, - bool is_sqrt, - cudaStream_t stream) -{ - if constexpr (!std::is_same_v && !std::is_same_v) { return false; } - - if (nearest_dist == nullptr) { return false; } - - Fused1nnTilePlanner planner; - planner.add_entrypoint(); - planner.add_tileir_fallback(); - const CutileTileConfig tile_cfg = planner.tile_config(); - auto launcher = planner.try_get_launcher(); - if (!launcher) { return false; } - - int metric_code; - bool apply_sqrt = false; - switch (metric) { - case cuvs::distance::DistanceType::InnerProduct: - metric_code = static_cast(cuvs::distance::DistanceType::InnerProduct); - break; - case cuvs::distance::DistanceType::L2Expanded: - case cuvs::distance::DistanceType::L2SqrtExpanded: - metric_code = static_cast(cuvs::distance::DistanceType::L2Expanded); - apply_sqrt = is_sqrt; - break; - case cuvs::distance::DistanceType::CosineExpanded: - metric_code = static_cast(cuvs::distance::DistanceType::CosineExpanded); - break; - default: return false; - } - - IdxT shape_x[2] = {m, k}; - IdxT stride_x[2] = {k, IdxT{1}}; - IdxT shape_y[2] = {n, k}; - IdxT stride_y[2] = {k, IdxT{1}}; - IdxT shape_xn = m; - IdxT stride_xn = IdxT{1}; - IdxT shape_yn = n; - IdxT stride_yn = IdxT{1}; - IdxT shape_idx = m; - IdxT stride_idx = IdxT{1}; - IdxT shape_dist = m; - IdxT stride_dist = IdxT{1}; - - IdxT M = m; - IdxT N = n; - IdxT K = k; - - void* x_ptr = const_cast(x); - void* y_ptr = const_cast(y); - void* xn_ptr = const_cast(xn); - void* yn_ptr = const_cast(yn); - // OutIdx must be a valid device pointer for the launch ABI; when store_idx is 0 the kernel - // does not write it (dist-only callers pass nearest_dist as a stand-in). - const IdxT store_idx = nearest_idx != nullptr ? IdxT{1} : IdxT{0}; - void* idx_ptr = - nearest_idx != nullptr ? static_cast(nearest_idx) : static_cast(nearest_dist); - void* dist_ptr = nearest_dist; - - const int tile_m = tile_cfg.tile_m; - dim3 grid((static_cast(m) + tile_m - 1) / tile_m, 1, 1); - dim3 block(1, 1, 1); - - using fused_1nn_cutile_kernel_t = void(void*, - IdxT, - IdxT, - IdxT, - IdxT, - void*, - IdxT, - IdxT, - IdxT, - IdxT, - void*, - IdxT, - IdxT, - void*, - IdxT, - IdxT, - void*, - IdxT, - IdxT, - void*, - IdxT, - IdxT, - IdxT, - IdxT, - IdxT, - IdxT, - IdxT, - int); - launcher->template dispatch(stream, - grid, - block, - 0, - x_ptr, - shape_x[0], - shape_x[1], - stride_x[0], - stride_x[1], - y_ptr, - shape_y[0], - shape_y[1], - stride_y[0], - stride_y[1], - xn_ptr, - shape_xn, - stride_xn, - yn_ptr, - shape_yn, - stride_yn, - idx_ptr, - shape_idx, - stride_idx, - dist_ptr, - shape_dist, - stride_dist, - M, - N, - K, - static_cast(apply_sqrt), - store_idx, - metric_code); - RAFT_CUDA_TRY(cudaGetLastError()); - return true; -} - -template -bool try_fused_1nn_tile_dispatch(IdxT* nearest_idx, - DataT* nearest_dist, - const DataT* x, - const DataT* y, - const DataT* xn, - const DataT* yn, - IdxT m, - IdxT n, - IdxT k, - cuvs::distance::DistanceType metric, - bool is_sqrt, - cudaStream_t stream) -{ - return launch_fused_1nn_tile( - nearest_idx, nearest_dist, x, y, xn, yn, m, n, k, metric, is_sqrt, stream); -} - -} // namespace - -template - requires is_fused_1nn_cutile_data_v -bool try_fused_1nn_tile(IdxT* nearest_idx, - DataT* nearest_dist, - const DataT* x, - const DataT* y, - const DataT* xn, - const DataT* yn, - IdxT m, - IdxT n, - IdxT k, - cuvs::distance::DistanceType metric, - bool is_sqrt, - void* index_workspace, - cudaStream_t stream) -{ - if (!cuvs::detail::jit_lto::cutile_launch_available_on_current_device()) { return false; } - static_assert(std::is_same_v || std::is_same_v); - - int cc_major = 0; - int cc_minor = 0; - if (!cuvs::detail::jit_lto::get_device_compute_capability(cc_major, cc_minor)) { return false; } - constexpr int tma_pitch_elements = 16 / sizeof(DataT); - const bool use_strict_abi = cc_major >= 9 && k % tma_pitch_elements == 0; - - if constexpr (std::is_same_v) { - if (use_strict_abi) { - return try_fused_1nn_tile_dispatch( - nearest_idx, nearest_dist, x, y, xn, yn, m, n, k, metric, is_sqrt, stream); - } - return try_fused_1nn_tile_dispatch( - nearest_idx, nearest_dist, x, y, xn, yn, m, n, k, metric, is_sqrt, stream); - } else { - constexpr int64_t max_i32 = std::numeric_limits::max(); - if (n > max_i32 || k > max_i32) { return false; } - if (nearest_idx != nullptr && index_workspace == nullptr) { return false; } - - auto* tmp_idx = static_cast(index_workspace); - for (int64_t offset = 0; offset < m; offset += max_i32) { - const int batch_m = static_cast(std::min(max_i32, m - offset)); - const auto* batch_x = x + static_cast(offset) * static_cast(k); - const auto* batch_xn = xn == nullptr ? nullptr : xn + offset; - auto* batch_dist = nearest_dist == nullptr ? nullptr : nearest_dist + offset; - - const bool launched = - use_strict_abi - ? try_fused_1nn_tile_dispatch(tmp_idx, - batch_dist, - batch_x, - y, - batch_xn, - yn, - batch_m, - static_cast(n), - static_cast(k), - metric, - is_sqrt, - stream) - : try_fused_1nn_tile_dispatch(tmp_idx, - batch_dist, - batch_x, - y, - batch_xn, - yn, - batch_m, - static_cast(n), - static_cast(k), - metric, - is_sqrt, - stream); - if (!launched) { return false; } - - if (nearest_idx != nullptr) { - raft::linalg::unaryOp( - nearest_idx + offset, tmp_idx, batch_m, raft::cast_op{}, stream); - } - } - return true; - } -} - -#define CUVS_INST_TRY_FUSED_1NN_TILE(DataT, IdxT) \ - template CUVS_EXPORT bool try_fused_1nn_tile(IdxT*, \ - DataT*, \ - const DataT*, \ - const DataT*, \ - const DataT*, \ - const DataT*, \ - IdxT, \ - IdxT, \ - IdxT, \ - cuvs::distance::DistanceType, \ - bool, \ - void*, \ - cudaStream_t) - -CUVS_INST_TRY_FUSED_1NN_TILE(float, int); -CUVS_INST_TRY_FUSED_1NN_TILE(float, int64_t); -CUVS_INST_TRY_FUSED_1NN_TILE(half, int); -CUVS_INST_TRY_FUSED_1NN_TILE(half, int64_t); - -#undef CUVS_INST_TRY_FUSED_1NN_TILE - -} // namespace detail -} // namespace distance -} // namespace cuvs From c92ed15f22c900f368b04c3d7138ea2e6cd3a0bf Mon Sep 17 00:00:00 2001 From: divyegala Date: Tue, 11 Aug 2026 19:06:59 +0000 Subject: [PATCH 42/82] python and cutile in build not host --- conda/recipes/libcuvs/recipe.yaml | 1 + .../modules/generate_cutile_kernels.cmake | 13 +- .../cutile/fused_1nn_cutile_matrix.json | 126 +++++++++++++++--- 3 files changed, 119 insertions(+), 21 deletions(-) diff --git a/conda/recipes/libcuvs/recipe.yaml b/conda/recipes/libcuvs/recipe.yaml index d9cb0ec5da..8386351ada 100644 --- a/conda/recipes/libcuvs/recipe.yaml +++ b/conda/recipes/libcuvs/recipe.yaml @@ -72,6 +72,7 @@ cache: - ninja - if: cuda_major == "13" then: + - python - cutile-python - cuda-tileiras - ${{ stdlib("c") }} diff --git a/cpp/cmake/modules/generate_cutile_kernels.cmake b/cpp/cmake/modules/generate_cutile_kernels.cmake index 89389d5bea..1eaf0c36b6 100644 --- a/cpp/cmake/modules/generate_cutile_kernels.cmake +++ b/cpp/cmake/modules/generate_cutile_kernels.cmake @@ -36,7 +36,17 @@ function(_cutile_kernels_setup) set(multi_value) cmake_parse_arguments(_CUTILE "${options}" "${one_value}" "${multi_value}" ${ARGN}) - find_package(Python3 REQUIRED COMPONENTS Interpreter) + if(DEFINED ENV{BUILD_PREFIX}) + find_program( + _cutile_build_python + NAMES python3 python + PATHS "$ENV{BUILD_PREFIX}/bin" + NO_DEFAULT_PATH REQUIRED NO_CACHE + ) + set(Python3_EXECUTABLE "${_cutile_build_python}") + else() + find_package(Python3 REQUIRED COMPONENTS Interpreter) + endif() find_package(CUDAToolkit REQUIRED) if(CUDAToolkit_VERSION VERSION_LESS 13.0) @@ -68,6 +78,7 @@ function(_cutile_kernels_setup) "Install cutile-python and cuda-tileiras (conda), or cuda-tile[tileiras] (pip)." ) endif() + message(STATUS "Using cuTile Python: ${Python3_EXECUTABLE}") set_property( DIRECTORY diff --git a/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_cutile_matrix.json b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_cutile_matrix.json index cce186a34d..b870f70597 100644 --- a/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_cutile_matrix.json +++ b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_cutile_matrix.json @@ -30,7 +30,7 @@ ], "_tile": [ { - "tile_m": 128, + "tile_m": 64, "tile_n": 128, "tile_k": 32 } @@ -63,12 +63,18 @@ { "matrix_layout": "strict", "abi_abbrev": "strict", - "abi_tag": "cutile_abi_strict" + "abi_tag": "cutile_abi_strict", + "tile_m": 128, + "tile_n": 128, + "tile_k": 32 }, { "matrix_layout": "relaxed", "abi_abbrev": "relaxed", - "abi_tag": "cutile_abi_relaxed" + "abi_tag": "cutile_abi_relaxed", + "tile_m": 128, + "tile_n": 256, + "tile_k": 16 } ], "_data": [ @@ -92,13 +98,6 @@ "index_abbrev": "i32" } ], - "_tile": [ - { - "tile_m": 128, - "tile_n": 128, - "tile_k": 32 - } - ], "_export": [ { "output_format": "cubin", @@ -109,17 +108,51 @@ "cc_major": 10, "cc_minor": 0, "arch_tag": "cutile_arch_10_0" + } + ] + }, + { + "_abi": [ + { + "matrix_layout": "strict", + "abi_abbrev": "strict", + "abi_tag": "cutile_abi_strict" }, { - "output_format": "cubin", - "artifact_ext": "cubin", - "artifact_basename": "@data_type@_@index_abbrev@_@abi_abbrev@_@gpu_code@", - "register": "cubin", - "gpu_code": "sm_90", - "cc_major": 9, - "cc_minor": 0, - "arch_tag": "cutile_arch_9_0" + "matrix_layout": "relaxed", + "abi_abbrev": "relaxed", + "abi_tag": "cutile_abi_relaxed" + } + ], + "_data": [ + { + "data_type": "half", + "data_abbrev": "h" }, + { + "data_type": "float", + "data_abbrev": "f" + } + ], + "_metric": [ + { + "metric": "runtime" + } + ], + "_index": [ + { + "index_type": "int32", + "index_abbrev": "i32" + } + ], + "_tile": [ + { + "tile_m": 128, + "tile_n": 128, + "tile_k": 32 + } + ], + "_export": [ { "output_format": "tileir_bytecode", "artifact_ext": "tilebc", @@ -136,8 +169,8 @@ "matrix_layout": "strict", "abi_abbrev": "strict", "abi_tag": "cutile_abi_strict", - "tile_m": 128, - "tile_n": 128, + "tile_m": 64, + "tile_n": 256, "tile_k": 32 }, { @@ -146,7 +179,60 @@ "abi_tag": "cutile_abi_relaxed", "tile_m": 128, "tile_n": 128, + "tile_k": 64 + } + ], + "_data": [ + { + "data_type": "half", + "data_abbrev": "h" + }, + { + "data_type": "float", + "data_abbrev": "f" + } + ], + "_metric": [ + { + "metric": "runtime" + } + ], + "_index": [ + { + "index_type": "int32", + "index_abbrev": "i32" + } + ], + "_export": [ + { + "output_format": "cubin", + "artifact_ext": "cubin", + "artifact_basename": "@data_type@_@index_abbrev@_@abi_abbrev@_@gpu_code@", + "register": "cubin", + "gpu_code": "sm_90", + "cc_major": 9, + "cc_minor": 0, + "arch_tag": "cutile_arch_9_0" + } + ] + }, + { + "_abi": [ + { + "matrix_layout": "strict", + "abi_abbrev": "strict", + "abi_tag": "cutile_abi_strict", + "tile_m": 64, + "tile_n": 128, "tile_k": 32 + }, + { + "matrix_layout": "relaxed", + "abi_abbrev": "relaxed", + "abi_tag": "cutile_abi_relaxed", + "tile_m": 64, + "tile_n": 128, + "tile_k": 64 } ], "_data": [ From 53d04f5af1e262b52081a40ee37096ed71590ae0 Mon Sep 17 00:00:00 2001 From: divyegala Date: Tue, 11 Aug 2026 20:37:28 +0000 Subject: [PATCH 43/82] correctly express python --- conda/recipes/libcuvs/recipe.yaml | 2 +- .../modules/compute_matrix_product.cmake | 33 +++++++++++++++---- .../modules/generate_cutile_kernels.cmake | 17 +++------- .../cutile/fused_1nn_tile.cu | 1 + 4 files changed, 34 insertions(+), 19 deletions(-) diff --git a/conda/recipes/libcuvs/recipe.yaml b/conda/recipes/libcuvs/recipe.yaml index 8386351ada..c65723421e 100644 --- a/conda/recipes/libcuvs/recipe.yaml +++ b/conda/recipes/libcuvs/recipe.yaml @@ -70,9 +70,9 @@ cache: - cuda-version =${{ cuda_version }} - cmake ${{ cmake_version }} - ninja + - python - if: cuda_major == "13" then: - - python - cutile-python - cuda-tileiras - ${{ stdlib("c") }} diff --git a/cpp/cmake/modules/compute_matrix_product.cmake b/cpp/cmake/modules/compute_matrix_product.cmake index 82a34f9242..fbbcd13b69 100644 --- a/cpp/cmake/modules/compute_matrix_product.cmake +++ b/cpp/cmake/modules/compute_matrix_product.cmake @@ -1,12 +1,31 @@ # ============================================================================= # cmake-format: off -# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # cmake-format: on # ============================================================================= include_guard(GLOBAL) +function(cuvs_find_build_python output_var) + if(DEFINED ENV{BUILD_PREFIX}) + find_program( + _cuvs_build_python + NAMES python3 python + PATHS "$ENV{BUILD_PREFIX}/bin" + NO_DEFAULT_PATH REQUIRED NO_CACHE + ) + set(_python_executable "${_cuvs_build_python}") + else() + find_package(Python3 REQUIRED COMPONENTS Interpreter) + set(_python_executable "${Python3_EXECUTABLE}") + endif() + set(${output_var} + "${_python_executable}" + PARENT_SCOPE + ) +endfunction() + function(compute_matrix_product output_var) set(options) set(one_value MATRIX_JSON_FILE MATRIX_JSON_STRING) @@ -14,19 +33,21 @@ function(compute_matrix_product output_var) cmake_parse_arguments(_JIT_LTO "${options}" "${one_value}" "${multi_value}" ${ARGN}) - find_package(Python3 REQUIRED COMPONENTS Interpreter) + cuvs_find_build_python(_matrix_python_executable) if(_JIT_LTO_MATRIX_JSON_FILE) execute_process( - COMMAND "${Python3_EXECUTABLE}" "${CMAKE_CURRENT_FUNCTION_LIST_DIR}/compute_matrix_product.py" - "${_JIT_LTO_MATRIX_JSON_FILE}" # + COMMAND + "${_matrix_python_executable}" + "${CMAKE_CURRENT_FUNCTION_LIST_DIR}/compute_matrix_product.py" + "${_JIT_LTO_MATRIX_JSON_FILE}" OUTPUT_VARIABLE output COMMAND_ERROR_IS_FATAL ANY ) else() execute_process( COMMAND "${CMAKE_COMMAND}" -E echo "${_JIT_LTO_MATRIX_JSON_STRING}" - COMMAND "${Python3_EXECUTABLE}" "${CMAKE_CURRENT_FUNCTION_LIST_DIR}/compute_matrix_product.py" - - + COMMAND "${_matrix_python_executable}" + "${CMAKE_CURRENT_FUNCTION_LIST_DIR}/compute_matrix_product.py" - OUTPUT_VARIABLE output COMMAND_ERROR_IS_FATAL ANY ) endif() diff --git a/cpp/cmake/modules/generate_cutile_kernels.cmake b/cpp/cmake/modules/generate_cutile_kernels.cmake index 1eaf0c36b6..0fbb69e567 100644 --- a/cpp/cmake/modules/generate_cutile_kernels.cmake +++ b/cpp/cmake/modules/generate_cutile_kernels.cmake @@ -36,17 +36,6 @@ function(_cutile_kernels_setup) set(multi_value) cmake_parse_arguments(_CUTILE "${options}" "${one_value}" "${multi_value}" ${ARGN}) - if(DEFINED ENV{BUILD_PREFIX}) - find_program( - _cutile_build_python - NAMES python3 python - PATHS "$ENV{BUILD_PREFIX}/bin" - NO_DEFAULT_PATH REQUIRED NO_CACHE - ) - set(Python3_EXECUTABLE "${_cutile_build_python}") - else() - find_package(Python3 REQUIRED COMPONENTS Interpreter) - endif() find_package(CUDAToolkit REQUIRED) if(CUDAToolkit_VERSION VERSION_LESS 13.0) @@ -61,6 +50,8 @@ function(_cutile_kernels_setup) return() endif() + cuvs_find_build_python(Python3_EXECUTABLE) + find_program( CUTILE_BIN2C NAMES bin2c @@ -203,7 +194,9 @@ function(process_cutile_matrix_entry source_list_var) set(multi_value FRAGMENT_TAG_HEADER_FILES) cmake_parse_arguments(_CUTILE "${options}" "${one_value}" "${multi_value}" ${ARGN}) - find_package(Python3 REQUIRED COMPONENTS Interpreter) + if(NOT Python3_EXECUTABLE) + cuvs_find_build_python(Python3_EXECUTABLE) + endif() populate_matrix_variables("${_CUTILE_MATRIX_JSON_ENTRY}") diff --git a/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_tile.cu b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_tile.cu index a0df8f60c2..caf5bbc446 100644 --- a/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_tile.cu +++ b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_tile.cu @@ -8,6 +8,7 @@ #include "fused_1nn_planner.hpp" #include +#include #include #include #include From 7f4c205f8118713ad2bc92065df15dd5ee2ea4d9 Mon Sep 17 00:00:00 2001 From: divyegala Date: Tue, 11 Aug 2026 23:51:35 +0000 Subject: [PATCH 44/82] fix c --- ci/build_standalone_c.sh | 3 +++ cpp/cmake/modules/compute_matrix_product.cmake | 3 +-- 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/ci/build_standalone_c.sh b/ci/build_standalone_c.sh index 94cffe1d64..72f4ad694a 100755 --- a/ci/build_standalone_c.sh +++ b/ci/build_standalone_c.sh @@ -41,6 +41,9 @@ source rapids-configure-sccache source rapids-datetime-string rapids-pip-retry install cmake +if [[ "${RAPIDS_CUDA_VERSION%%.*}" == "13" ]]; then + rapids-pip-retry install cuda-tile 'cuda-toolkit[tileiras]==13.*' +fi pyenv rehash rapids-print-env diff --git a/cpp/cmake/modules/compute_matrix_product.cmake b/cpp/cmake/modules/compute_matrix_product.cmake index fbbcd13b69..563b5e0428 100644 --- a/cpp/cmake/modules/compute_matrix_product.cmake +++ b/cpp/cmake/modules/compute_matrix_product.cmake @@ -12,8 +12,7 @@ function(cuvs_find_build_python output_var) find_program( _cuvs_build_python NAMES python3 python - PATHS "$ENV{BUILD_PREFIX}/bin" - NO_DEFAULT_PATH REQUIRED NO_CACHE + HINTS "$ENV{BUILD_PREFIX}/bin" REQUIRED NO_CACHE ) set(_python_executable "${_cuvs_build_python}") else() From 9715ac3b540798aed9ac211a2be489e331e89956 Mon Sep 17 00:00:00 2001 From: divyegala Date: Wed, 12 Aug 2026 19:33:06 +0000 Subject: [PATCH 45/82] use find package --- .../modules/compute_matrix_product.cmake | 31 +++---------------- .../modules/generate_cutile_kernels.cmake | 14 ++++++--- 2 files changed, 14 insertions(+), 31 deletions(-) diff --git a/cpp/cmake/modules/compute_matrix_product.cmake b/cpp/cmake/modules/compute_matrix_product.cmake index 563b5e0428..c91564f8e5 100644 --- a/cpp/cmake/modules/compute_matrix_product.cmake +++ b/cpp/cmake/modules/compute_matrix_product.cmake @@ -7,24 +7,6 @@ include_guard(GLOBAL) -function(cuvs_find_build_python output_var) - if(DEFINED ENV{BUILD_PREFIX}) - find_program( - _cuvs_build_python - NAMES python3 python - HINTS "$ENV{BUILD_PREFIX}/bin" REQUIRED NO_CACHE - ) - set(_python_executable "${_cuvs_build_python}") - else() - find_package(Python3 REQUIRED COMPONENTS Interpreter) - set(_python_executable "${Python3_EXECUTABLE}") - endif() - set(${output_var} - "${_python_executable}" - PARENT_SCOPE - ) -endfunction() - function(compute_matrix_product output_var) set(options) set(one_value MATRIX_JSON_FILE MATRIX_JSON_STRING) @@ -32,21 +14,18 @@ function(compute_matrix_product output_var) cmake_parse_arguments(_JIT_LTO "${options}" "${one_value}" "${multi_value}" ${ARGN}) - cuvs_find_build_python(_matrix_python_executable) + find_package(Python3 REQUIRED COMPONENTS Interpreter) if(_JIT_LTO_MATRIX_JSON_FILE) execute_process( - COMMAND - "${_matrix_python_executable}" - "${CMAKE_CURRENT_FUNCTION_LIST_DIR}/compute_matrix_product.py" - "${_JIT_LTO_MATRIX_JSON_FILE}" - OUTPUT_VARIABLE output COMMAND_ERROR_IS_FATAL ANY + COMMAND "${Python3_EXECUTABLE}" "${CMAKE_CURRENT_FUNCTION_LIST_DIR}/compute_matrix_product.py" + "${_JIT_LTO_MATRIX_JSON_FILE}" OUTPUT_VARIABLE output COMMAND_ERROR_IS_FATAL ANY ) else() execute_process( COMMAND "${CMAKE_COMMAND}" -E echo "${_JIT_LTO_MATRIX_JSON_STRING}" - COMMAND "${_matrix_python_executable}" - "${CMAKE_CURRENT_FUNCTION_LIST_DIR}/compute_matrix_product.py" - + COMMAND "${Python3_EXECUTABLE}" "${CMAKE_CURRENT_FUNCTION_LIST_DIR}/compute_matrix_product.py" + - OUTPUT_VARIABLE output COMMAND_ERROR_IS_FATAL ANY ) endif() diff --git a/cpp/cmake/modules/generate_cutile_kernels.cmake b/cpp/cmake/modules/generate_cutile_kernels.cmake index 0fbb69e567..185ac17e3d 100644 --- a/cpp/cmake/modules/generate_cutile_kernels.cmake +++ b/cpp/cmake/modules/generate_cutile_kernels.cmake @@ -50,7 +50,7 @@ function(_cutile_kernels_setup) return() endif() - cuvs_find_build_python(Python3_EXECUTABLE) + find_package(Python3 REQUIRED COMPONENTS Interpreter) find_program( CUTILE_BIN2C @@ -61,12 +61,16 @@ function(_cutile_kernels_setup) execute_process( COMMAND "${Python3_EXECUTABLE}" -c "import cuda.tile" RESULT_VARIABLE _cutile_import_result - OUTPUT_QUIET ERROR_QUIET + ERROR_VARIABLE _cutile_import_error + OUTPUT_QUIET ERROR_STRIP_TRAILING_WHITESPACE ) if(NOT _cutile_import_result EQUAL 0) message( - FATAL_ERROR "cuda.tile (cuTile Python) is required to build cuTile embedded kernels. " - "Install cutile-python and cuda-tileiras (conda), or cuda-tile[tileiras] (pip)." + FATAL_ERROR + "cuda.tile (cuTile Python) is required to build cuTile embedded kernels. " + "Install cutile-python and cuda-tileiras (conda), or cuda-tile[tileiras] (pip).\n" + "Interpreter: ${Python3_EXECUTABLE}\n" + "Import error: ${_cutile_import_error}" ) endif() message(STATUS "Using cuTile Python: ${Python3_EXECUTABLE}") @@ -195,7 +199,7 @@ function(process_cutile_matrix_entry source_list_var) cmake_parse_arguments(_CUTILE "${options}" "${one_value}" "${multi_value}" ${ARGN}) if(NOT Python3_EXECUTABLE) - cuvs_find_build_python(Python3_EXECUTABLE) + find_package(Python3 REQUIRED COMPONENTS Interpreter) endif() populate_matrix_variables("${_CUTILE_MATRIX_JSON_ENTRY}") From 9cdc42712664d856bc80272abd7ef8dca919dbb5 Mon Sep 17 00:00:00 2001 From: divyegala Date: Thu, 13 Aug 2026 01:02:28 +0000 Subject: [PATCH 46/82] use build prefix python --- cpp/cmake/modules/compute_matrix_product.cmake | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/cpp/cmake/modules/compute_matrix_product.cmake b/cpp/cmake/modules/compute_matrix_product.cmake index c91564f8e5..7221f1aa27 100644 --- a/cpp/cmake/modules/compute_matrix_product.cmake +++ b/cpp/cmake/modules/compute_matrix_product.cmake @@ -7,6 +7,12 @@ include_guard(GLOBAL) +if(DEFINED ENV{BUILD_PREFIX}) + set(Python3_ROOT_DIR "$ENV{BUILD_PREFIX}") + set(Python3_FIND_STRATEGY LOCATION) + set(Python3_FIND_VIRTUALENV STANDARD) +endif() + function(compute_matrix_product output_var) set(options) set(one_value MATRIX_JSON_FILE MATRIX_JSON_STRING) From 766eefcd6f1709738f3761115164f1df769a7032 Mon Sep 17 00:00:00 2001 From: divyegala Date: Thu, 13 Aug 2026 01:50:23 +0000 Subject: [PATCH 47/82] hint python early --- cpp/CMakeLists.txt | 7 +++++++ cpp/cmake/modules/compute_matrix_product.cmake | 6 ------ 2 files changed, 7 insertions(+), 6 deletions(-) diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt index 141fa2fdd1..2d7f5e3ade 100644 --- a/cpp/CMakeLists.txt +++ b/cpp/CMakeLists.txt @@ -4,6 +4,13 @@ # SPDX-License-Identifier: Apache-2.0 # cmake-format: on cmake_minimum_required(VERSION 4.0 FATAL_ERROR) + +if(DEFINED ENV{BUILD_PREFIX}) + set(Python3_ROOT_DIR "$ENV{BUILD_PREFIX}") + set(Python3_FIND_STRATEGY LOCATION) + set(Python3_FIND_VIRTUALENV STANDARD) +endif() + include(../cmake/rapids_config.cmake) include(rapids-cmake) include(rapids-cpm) diff --git a/cpp/cmake/modules/compute_matrix_product.cmake b/cpp/cmake/modules/compute_matrix_product.cmake index 7221f1aa27..c91564f8e5 100644 --- a/cpp/cmake/modules/compute_matrix_product.cmake +++ b/cpp/cmake/modules/compute_matrix_product.cmake @@ -7,12 +7,6 @@ include_guard(GLOBAL) -if(DEFINED ENV{BUILD_PREFIX}) - set(Python3_ROOT_DIR "$ENV{BUILD_PREFIX}") - set(Python3_FIND_STRATEGY LOCATION) - set(Python3_FIND_VIRTUALENV STANDARD) -endif() - function(compute_matrix_product output_var) set(options) set(one_value MATRIX_JSON_FILE MATRIX_JSON_STRING) From 810a582685750f5bff3bdd260dfcf921687edec2 Mon Sep 17 00:00:00 2001 From: divyegala Date: Thu, 13 Aug 2026 02:09:10 +0000 Subject: [PATCH 48/82] explicilty find python --- cpp/CMakeLists.txt | 20 +++++++++++++++++--- 1 file changed, 17 insertions(+), 3 deletions(-) diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt index 2d7f5e3ade..4cc8497ff0 100644 --- a/cpp/CMakeLists.txt +++ b/cpp/CMakeLists.txt @@ -6,9 +6,23 @@ cmake_minimum_required(VERSION 4.0 FATAL_ERROR) if(DEFINED ENV{BUILD_PREFIX}) - set(Python3_ROOT_DIR "$ENV{BUILD_PREFIX}") - set(Python3_FIND_STRATEGY LOCATION) - set(Python3_FIND_VIRTUALENV STANDARD) + file( + GLOB _cuvs_build_python_candidates + LIST_DIRECTORIES FALSE + "$ENV{BUILD_PREFIX}/bin/python3.[0-9]*" + ) + list(FILTER _cuvs_build_python_candidates EXCLUDE REGEX "-config$") + list(LENGTH _cuvs_build_python_candidates _cuvs_build_python_count) + if(NOT _cuvs_build_python_count EQUAL 1) + message(FATAL_ERROR "Expected one Python interpreter in $ENV{BUILD_PREFIX}/bin, found: " + "${_cuvs_build_python_candidates}" + ) + endif() + list(GET _cuvs_build_python_candidates 0 Python3_EXECUTABLE) + set(Python3_EXECUTABLE + "${Python3_EXECUTABLE}" + CACHE FILEPATH "Python interpreter from the conda build prefix" FORCE + ) endif() include(../cmake/rapids_config.cmake) From 85eb36724c36519480e8a65ab2fbce219f81ca10 Mon Sep 17 00:00:00 2001 From: divyegala Date: Tue, 18 Aug 2026 16:19:44 +0000 Subject: [PATCH 49/82] python fix --- cpp/CMakeLists.txt | 18 ++++++++++++++---- 1 file changed, 14 insertions(+), 4 deletions(-) diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt index 4cc8497ff0..c346ffce22 100644 --- a/cpp/CMakeLists.txt +++ b/cpp/CMakeLists.txt @@ -12,13 +12,23 @@ if(DEFINED ENV{BUILD_PREFIX}) "$ENV{BUILD_PREFIX}/bin/python3.[0-9]*" ) list(FILTER _cuvs_build_python_candidates EXCLUDE REGEX "-config$") - list(LENGTH _cuvs_build_python_candidates _cuvs_build_python_count) + if(NOT _cuvs_build_python_candidates) + message(FATAL_ERROR "No Python interpreter found in $ENV{BUILD_PREFIX}/bin") + endif() + set(_cuvs_build_python_executables) + foreach(_candidate IN LISTS _cuvs_build_python_candidates) + get_filename_component(_executable "${_candidate}" REALPATH) + list(APPEND _cuvs_build_python_executables "${_executable}") + endforeach() + list(REMOVE_DUPLICATES _cuvs_build_python_executables) + list(LENGTH _cuvs_build_python_executables _cuvs_build_python_count) if(NOT _cuvs_build_python_count EQUAL 1) - message(FATAL_ERROR "Expected one Python interpreter in $ENV{BUILD_PREFIX}/bin, found: " - "${_cuvs_build_python_candidates}" + message( + FATAL_ERROR "Python aliases in $ENV{BUILD_PREFIX}/bin resolve to different interpreters: " + "${_cuvs_build_python_executables}" ) endif() - list(GET _cuvs_build_python_candidates 0 Python3_EXECUTABLE) + list(GET _cuvs_build_python_executables 0 Python3_EXECUTABLE) set(Python3_EXECUTABLE "${Python3_EXECUTABLE}" CACHE FILEPATH "Python interpreter from the conda build prefix" FORCE From ac08848a6ce8fe49ebf39ae6ed971cbbb0c951ad Mon Sep 17 00:00:00 2001 From: divyegala Date: Tue, 18 Aug 2026 17:53:06 +0000 Subject: [PATCH 50/82] python only for cutile --- conda/recipes/libcuvs/recipe.yaml | 5 ++- cpp/CMakeLists.txt | 30 ---------------- .../modules/generate_cutile_kernels.cmake | 35 ++++++++++++++++++- 3 files changed, 38 insertions(+), 32 deletions(-) diff --git a/conda/recipes/libcuvs/recipe.yaml b/conda/recipes/libcuvs/recipe.yaml index c65723421e..c455ba6c15 100644 --- a/conda/recipes/libcuvs/recipe.yaml +++ b/conda/recipes/libcuvs/recipe.yaml @@ -70,7 +70,6 @@ cache: - cuda-version =${{ cuda_version }} - cmake ${{ cmake_version }} - ninja - - python - if: cuda_major == "13" then: - cutile-python @@ -395,6 +394,10 @@ outputs: - cuda-version =${{ cuda_version }} - cmake ${{ cmake_version }} - ninja + - if: cuda_major == "13" + then: + - cutile-python + - cuda-tileiras - ${{ stdlib("c") }} host: - ${{ pin_subpackage("libcuvs-headers", exact=True) }} diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt index 56c84a048b..cf057e9b5c 100644 --- a/cpp/CMakeLists.txt +++ b/cpp/CMakeLists.txt @@ -5,36 +5,6 @@ # cmake-format: on cmake_minimum_required(VERSION 4.0 FATAL_ERROR) -if(DEFINED ENV{BUILD_PREFIX}) - file( - GLOB _cuvs_build_python_candidates - LIST_DIRECTORIES FALSE - "$ENV{BUILD_PREFIX}/bin/python3.[0-9]*" - ) - list(FILTER _cuvs_build_python_candidates EXCLUDE REGEX "-config$") - if(NOT _cuvs_build_python_candidates) - message(FATAL_ERROR "No Python interpreter found in $ENV{BUILD_PREFIX}/bin") - endif() - set(_cuvs_build_python_executables) - foreach(_candidate IN LISTS _cuvs_build_python_candidates) - get_filename_component(_executable "${_candidate}" REALPATH) - list(APPEND _cuvs_build_python_executables "${_executable}") - endforeach() - list(REMOVE_DUPLICATES _cuvs_build_python_executables) - list(LENGTH _cuvs_build_python_executables _cuvs_build_python_count) - if(NOT _cuvs_build_python_count EQUAL 1) - message( - FATAL_ERROR "Python aliases in $ENV{BUILD_PREFIX}/bin resolve to different interpreters: " - "${_cuvs_build_python_executables}" - ) - endif() - list(GET _cuvs_build_python_executables 0 Python3_EXECUTABLE) - set(Python3_EXECUTABLE - "${Python3_EXECUTABLE}" - CACHE FILEPATH "Python interpreter from the conda build prefix" FORCE - ) -endif() - include(../cmake/rapids_config.cmake) include(rapids-cmake) include(rapids-cpm) diff --git a/cpp/cmake/modules/generate_cutile_kernels.cmake b/cpp/cmake/modules/generate_cutile_kernels.cmake index 185ac17e3d..d343985663 100644 --- a/cpp/cmake/modules/generate_cutile_kernels.cmake +++ b/cpp/cmake/modules/generate_cutile_kernels.cmake @@ -30,6 +30,39 @@ function(_cutile_fragment_tag_header_files output_var) ) endfunction() +function(_cutile_find_python output_var) + if(DEFINED ENV{BUILD_PREFIX}) + file( + GLOB _python_candidates + LIST_DIRECTORIES FALSE + "$ENV{BUILD_PREFIX}/bin/python3.[0-9]*" + ) + list(FILTER _python_candidates EXCLUDE REGEX "-config$") + foreach(_candidate IN LISTS _python_candidates) + get_filename_component(_executable "${_candidate}" REALPATH) + list(APPEND _python_executables "${_executable}") + endforeach() + list(REMOVE_DUPLICATES _python_executables) + list(LENGTH _python_executables _python_count) + if(NOT _python_count EQUAL 1) + message(FATAL_ERROR "Expected one Python installation in $ENV{BUILD_PREFIX}, found: " + "${_python_executables}" + ) + endif() + list(GET _python_executables 0 Python3_EXECUTABLE) + set(${output_var} + "${Python3_EXECUTABLE}" + PARENT_SCOPE + ) + return() + endif() + find_package(Python3 REQUIRED COMPONENTS Interpreter) + set(${output_var} + "${Python3_EXECUTABLE}" + PARENT_SCOPE + ) +endfunction() + function(_cutile_kernels_setup) set(options) set(one_value MATRIX_JSON_FILE OUTPUT_DIRECTORY) @@ -50,7 +83,7 @@ function(_cutile_kernels_setup) return() endif() - find_package(Python3 REQUIRED COMPONENTS Interpreter) + _cutile_find_python(Python3_EXECUTABLE) find_program( CUTILE_BIN2C From 513a1d5020c483bbd73c5168284aeac608c5a1f6 Mon Sep 17 00:00:00 2001 From: divyegala Date: Tue, 18 Aug 2026 18:54:10 +0000 Subject: [PATCH 51/82] try python again --- .../modules/compute_matrix_product.cmake | 26 ++++++++++--- .../modules/generate_cutile_kernels.cmake | 37 +------------------ 2 files changed, 23 insertions(+), 40 deletions(-) diff --git a/cpp/cmake/modules/compute_matrix_product.cmake b/cpp/cmake/modules/compute_matrix_product.cmake index c91564f8e5..c7ba84ad76 100644 --- a/cpp/cmake/modules/compute_matrix_product.cmake +++ b/cpp/cmake/modules/compute_matrix_product.cmake @@ -7,6 +7,19 @@ include_guard(GLOBAL) +function(cuvs_find_build_python output_var) + if(DEFINED ENV{BUILD_PREFIX}) + set(Python_ROOT_DIR "$ENV{BUILD_PREFIX}") + set(Python_FIND_STRATEGY LOCATION) + set(Python_FIND_VIRTUALENV STANDARD) + endif() + find_package(Python REQUIRED COMPONENTS Interpreter) + set(${output_var} + "${Python_EXECUTABLE}" + PARENT_SCOPE + ) +endfunction() + function(compute_matrix_product output_var) set(options) set(one_value MATRIX_JSON_FILE MATRIX_JSON_STRING) @@ -14,18 +27,21 @@ function(compute_matrix_product output_var) cmake_parse_arguments(_JIT_LTO "${options}" "${one_value}" "${multi_value}" ${ARGN}) - find_package(Python3 REQUIRED COMPONENTS Interpreter) + cuvs_find_build_python(_matrix_python_executable) if(_JIT_LTO_MATRIX_JSON_FILE) execute_process( - COMMAND "${Python3_EXECUTABLE}" "${CMAKE_CURRENT_FUNCTION_LIST_DIR}/compute_matrix_product.py" - "${_JIT_LTO_MATRIX_JSON_FILE}" OUTPUT_VARIABLE output COMMAND_ERROR_IS_FATAL ANY + COMMAND + "${_matrix_python_executable}" + "${CMAKE_CURRENT_FUNCTION_LIST_DIR}/compute_matrix_product.py" + "${_JIT_LTO_MATRIX_JSON_FILE}" + OUTPUT_VARIABLE output COMMAND_ERROR_IS_FATAL ANY ) else() execute_process( COMMAND "${CMAKE_COMMAND}" -E echo "${_JIT_LTO_MATRIX_JSON_STRING}" - COMMAND "${Python3_EXECUTABLE}" "${CMAKE_CURRENT_FUNCTION_LIST_DIR}/compute_matrix_product.py" - - + COMMAND "${_matrix_python_executable}" + "${CMAKE_CURRENT_FUNCTION_LIST_DIR}/compute_matrix_product.py" - OUTPUT_VARIABLE output COMMAND_ERROR_IS_FATAL ANY ) endif() diff --git a/cpp/cmake/modules/generate_cutile_kernels.cmake b/cpp/cmake/modules/generate_cutile_kernels.cmake index d343985663..ce6815531a 100644 --- a/cpp/cmake/modules/generate_cutile_kernels.cmake +++ b/cpp/cmake/modules/generate_cutile_kernels.cmake @@ -30,39 +30,6 @@ function(_cutile_fragment_tag_header_files output_var) ) endfunction() -function(_cutile_find_python output_var) - if(DEFINED ENV{BUILD_PREFIX}) - file( - GLOB _python_candidates - LIST_DIRECTORIES FALSE - "$ENV{BUILD_PREFIX}/bin/python3.[0-9]*" - ) - list(FILTER _python_candidates EXCLUDE REGEX "-config$") - foreach(_candidate IN LISTS _python_candidates) - get_filename_component(_executable "${_candidate}" REALPATH) - list(APPEND _python_executables "${_executable}") - endforeach() - list(REMOVE_DUPLICATES _python_executables) - list(LENGTH _python_executables _python_count) - if(NOT _python_count EQUAL 1) - message(FATAL_ERROR "Expected one Python installation in $ENV{BUILD_PREFIX}, found: " - "${_python_executables}" - ) - endif() - list(GET _python_executables 0 Python3_EXECUTABLE) - set(${output_var} - "${Python3_EXECUTABLE}" - PARENT_SCOPE - ) - return() - endif() - find_package(Python3 REQUIRED COMPONENTS Interpreter) - set(${output_var} - "${Python3_EXECUTABLE}" - PARENT_SCOPE - ) -endfunction() - function(_cutile_kernels_setup) set(options) set(one_value MATRIX_JSON_FILE OUTPUT_DIRECTORY) @@ -83,7 +50,7 @@ function(_cutile_kernels_setup) return() endif() - _cutile_find_python(Python3_EXECUTABLE) + cuvs_find_build_python(Python3_EXECUTABLE) find_program( CUTILE_BIN2C @@ -232,7 +199,7 @@ function(process_cutile_matrix_entry source_list_var) cmake_parse_arguments(_CUTILE "${options}" "${one_value}" "${multi_value}" ${ARGN}) if(NOT Python3_EXECUTABLE) - find_package(Python3 REQUIRED COMPONENTS Interpreter) + cuvs_find_build_python(Python3_EXECUTABLE) endif() populate_matrix_variables("${_CUTILE_MATRIX_JSON_ENTRY}") From e61047355ca065ea54392f235b843f16510c9ae6 Mon Sep 17 00:00:00 2001 From: divyegala Date: Tue, 18 Aug 2026 20:24:26 +0000 Subject: [PATCH 52/82] debug --- cpp/cmake/modules/compute_matrix_product.cmake | 2 ++ 1 file changed, 2 insertions(+) diff --git a/cpp/cmake/modules/compute_matrix_product.cmake b/cpp/cmake/modules/compute_matrix_product.cmake index c7ba84ad76..47cce1c6c0 100644 --- a/cpp/cmake/modules/compute_matrix_product.cmake +++ b/cpp/cmake/modules/compute_matrix_product.cmake @@ -13,7 +13,9 @@ function(cuvs_find_build_python output_var) set(Python_FIND_STRATEGY LOCATION) set(Python_FIND_VIRTUALENV STANDARD) endif() + set(CMAKE_FIND_DEBUG_MODE TRUE) find_package(Python REQUIRED COMPONENTS Interpreter) + set(CMAKE_FIND_DEBUG_MODE FALSE) set(${output_var} "${Python_EXECUTABLE}" PARENT_SCOPE From f3507bbe38cce3881f5178ccc78c678eb7fd6f9a Mon Sep 17 00:00:00 2001 From: divyegala Date: Tue, 18 Aug 2026 22:25:01 +0000 Subject: [PATCH 53/82] add explicit python back --- conda/recipes/libcuvs/recipe.yaml | 2 ++ 1 file changed, 2 insertions(+) diff --git a/conda/recipes/libcuvs/recipe.yaml b/conda/recipes/libcuvs/recipe.yaml index c455ba6c15..8f148573b3 100644 --- a/conda/recipes/libcuvs/recipe.yaml +++ b/conda/recipes/libcuvs/recipe.yaml @@ -70,6 +70,7 @@ cache: - cuda-version =${{ cuda_version }} - cmake ${{ cmake_version }} - ninja + - python - if: cuda_major == "13" then: - cutile-python @@ -394,6 +395,7 @@ outputs: - cuda-version =${{ cuda_version }} - cmake ${{ cmake_version }} - ninja + - python - if: cuda_major == "13" then: - cutile-python From a06747b8f9a23933bf1171193672cb75e4844869 Mon Sep 17 00:00:00 2001 From: divyegala Date: Tue, 18 Aug 2026 23:34:47 +0000 Subject: [PATCH 54/82] disable cmake paths --- cpp/cmake/modules/compute_matrix_product.cmake | 3 +++ 1 file changed, 3 insertions(+) diff --git a/cpp/cmake/modules/compute_matrix_product.cmake b/cpp/cmake/modules/compute_matrix_product.cmake index 47cce1c6c0..339866af11 100644 --- a/cpp/cmake/modules/compute_matrix_product.cmake +++ b/cpp/cmake/modules/compute_matrix_product.cmake @@ -12,6 +12,9 @@ function(cuvs_find_build_python output_var) set(Python_ROOT_DIR "$ENV{BUILD_PREFIX}") set(Python_FIND_STRATEGY LOCATION) set(Python_FIND_VIRTUALENV STANDARD) + set(CMAKE_FIND_USE_INSTALL_PREFIX FALSE) + set(CMAKE_FIND_USE_CMAKE_PATH FALSE) + set(CMAKE_FIND_USE_CMAKE_ENVIRONMENT_PATH FALSE) endif() set(CMAKE_FIND_DEBUG_MODE TRUE) find_package(Python REQUIRED COMPONENTS Interpreter) From 18750f39bbed0d16d82389a2705ac78a357dd7b4 Mon Sep 17 00:00:00 2001 From: divyegala Date: Wed, 19 Aug 2026 22:51:58 +0000 Subject: [PATCH 55/82] use Python_ROOT --- cpp/cmake/modules/compute_matrix_product.cmake | 7 +------ 1 file changed, 1 insertion(+), 6 deletions(-) diff --git a/cpp/cmake/modules/compute_matrix_product.cmake b/cpp/cmake/modules/compute_matrix_product.cmake index 339866af11..b5f3d06f86 100644 --- a/cpp/cmake/modules/compute_matrix_product.cmake +++ b/cpp/cmake/modules/compute_matrix_product.cmake @@ -9,12 +9,7 @@ include_guard(GLOBAL) function(cuvs_find_build_python output_var) if(DEFINED ENV{BUILD_PREFIX}) - set(Python_ROOT_DIR "$ENV{BUILD_PREFIX}") - set(Python_FIND_STRATEGY LOCATION) - set(Python_FIND_VIRTUALENV STANDARD) - set(CMAKE_FIND_USE_INSTALL_PREFIX FALSE) - set(CMAKE_FIND_USE_CMAKE_PATH FALSE) - set(CMAKE_FIND_USE_CMAKE_ENVIRONMENT_PATH FALSE) + set(Python_ROOT "$ENV{BUILD_PREFIX}") endif() set(CMAKE_FIND_DEBUG_MODE TRUE) find_package(Python REQUIRED COMPONENTS Interpreter) From ff5e3f958b714c398e3ef55b9992c359ae14246c Mon Sep 17 00:00:00 2001 From: divyegala Date: Thu, 20 Aug 2026 17:08:31 +0000 Subject: [PATCH 56/82] tf32 norms --- .../detail/minClusterDistanceCompute.cu | 160 ++++++++++++++---- .../cutile/fused_1nn_kernel.py | 46 +---- 2 files changed, 137 insertions(+), 69 deletions(-) diff --git a/cpp/src/cluster/detail/minClusterDistanceCompute.cu b/cpp/src/cluster/detail/minClusterDistanceCompute.cu index a02baca068..b230b8af16 100644 --- a/cpp/src/cluster/detail/minClusterDistanceCompute.cu +++ b/cpp/src/cluster/detail/minClusterDistanceCompute.cu @@ -7,12 +7,66 @@ #include "../../distance/unfused_distance_nn.cuh" #include "kmeans_common.cuh" +#include #include +#include +#include + namespace cuvs::cluster::kmeans::detail { namespace { +__device__ __forceinline__ float round_to_tf32(float value) +{ +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800 + return nvcuda::wmma::__float_to_tf32(value); +#else + return value; +#endif +} + +struct tf32_square_op { + template + __device__ float operator()(float value, IndexT) const + { + const float rounded = round_to_tf32(value); + return rounded * rounded; + } +}; + +template +void compute_tf32_row_norms(raft::resources const& handle, + const float* matrix, + float* norms, + IndexT n_rows, + IndexT n_cols, + bool take_sqrt) +{ + if (n_rows == 0) { return; } + auto matrix_view = raft::make_device_matrix_view(matrix, n_rows, n_cols); + auto norms_view = raft::make_device_vector_view(norms, n_rows); + if (take_sqrt) { + raft::linalg::coalesced_reduction(handle, + matrix_view, + norms_view, + 0.0f, + false, + tf32_square_op{}, + raft::add_op{}, + raft::sqrt_op{}); + } else { + raft::linalg::coalesced_reduction(handle, + matrix_view, + norms_view, + 0.0f, + false, + tf32_square_op{}, + raft::add_op{}, + raft::identity_op{}); + } +} + template __global__ void unpack_kvp_to_soa(IndexT* nearest_idx, DataT* nearest_dist, @@ -66,23 +120,45 @@ void minClusterAndDistanceCompute(raft::resources const& handle, use_fused(handle, n_samples, n_clusters, n_features, metric); if (uses_fused_distance_nn(fused_path)) { - L2NormBuf_OR_DistBuf.resize(n_clusters, stream); - auto centroidsNorm = - raft::make_device_vector_view(L2NormBuf_OR_DistBuf.data(), n_clusters); + const DataT* x_norm_ptr = L2NormX.data_handle(); + const DataT* centroids_norm_ptr; + if constexpr (std::is_same_v) { + if (fused_path == FusedDistancePath::FusedCutile && is_l2_cos) { + constexpr size_t norm_alignment = 16 / sizeof(float); + const size_t x_norm_storage = raft::alignTo(static_cast(n_samples), norm_alignment); + L2NormBuf_OR_DistBuf.resize(x_norm_storage + static_cast(n_clusters), stream); + auto* tf32_x_norms = L2NormBuf_OR_DistBuf.data(); + auto* tf32_centroid_norms = tf32_x_norms + x_norm_storage; + const bool take_sqrt = metric == cuvs::distance::DistanceType::CosineExpanded; + compute_tf32_row_norms( + handle, X.data_handle(), tf32_x_norms, n_samples, n_features, take_sqrt); + compute_tf32_row_norms( + handle, centroids.data_handle(), tf32_centroid_norms, n_clusters, n_features, take_sqrt); + x_norm_ptr = tf32_x_norms; + centroids_norm_ptr = tf32_centroid_norms; + } else { + L2NormBuf_OR_DistBuf.resize(n_clusters, stream); + centroids_norm_ptr = L2NormBuf_OR_DistBuf.data(); + } + } else { + L2NormBuf_OR_DistBuf.resize(n_clusters, stream); + centroids_norm_ptr = L2NormBuf_OR_DistBuf.data(); + } - if (is_l2_cos) { + if (!(fused_path == FusedDistancePath::FusedCutile && is_l2_cos && + std::is_same_v) && + is_l2_cos) { + auto centroids_norm = + raft::make_device_vector_view(L2NormBuf_OR_DistBuf.data(), n_clusters); if (metric == cuvs::distance::DistanceType::CosineExpanded) { raft::linalg::norm( - handle, centroids, centroidsNorm, raft::sqrt_op{}); + handle, centroids, centroids_norm, raft::sqrt_op{}); } else { raft::linalg::norm( - handle, centroids, centroidsNorm); + handle, centroids, centroids_norm); } } - auto centroidsNormConst = - raft::make_device_vector_view(L2NormBuf_OR_DistBuf.data(), n_clusters); - raft::KeyValuePair* cutlass_kvp_scratch = nullptr; rmm::device_uvector> temp_kvp(0, stream); if (needs_cutlass_kvp_scratch(fused_path)) { @@ -99,8 +175,8 @@ void minClusterAndDistanceCompute(raft::resources const& handle, nearest_dist.data_handle(), X.data_handle(), centroids.data_handle(), - L2NormX.data_handle(), - centroidsNormConst.data_handle(), + x_norm_ptr, + centroids_norm_ptr, n_samples, n_clusters, n_features, @@ -301,23 +377,49 @@ void minClusterDistanceCompute(raft::resources const& handle, : FusedDistancePath::Unfused; if (uses_fused_distance_nn(fused_path)) { - L2NormBuf_OR_DistBuf.resize(n_clusters, stream); - auto centroidsNorm = - raft::make_device_vector_view(L2NormBuf_OR_DistBuf.data(), n_clusters); - - if (metric == cuvs::distance::DistanceType::CosineExpanded) { - raft::linalg::norm( - handle, - raft::make_device_matrix_view( - centroids.data_handle(), centroids.extent(0), centroids.extent(1)), - centroidsNorm, - raft::sqrt_op{}); + const DataT* x_norm_ptr = L2NormX.data_handle(); + const DataT* centroids_norm_ptr; + if constexpr (std::is_same_v) { + if (fused_path == FusedDistancePath::FusedCutile && is_l2_cos) { + constexpr size_t norm_alignment = 16 / sizeof(float); + const size_t x_norm_storage = raft::alignTo(static_cast(n_samples), norm_alignment); + L2NormBuf_OR_DistBuf.resize(x_norm_storage + static_cast(n_clusters), stream); + auto* tf32_x_norms = L2NormBuf_OR_DistBuf.data(); + auto* tf32_centroid_norms = tf32_x_norms + x_norm_storage; + const bool take_sqrt = metric == cuvs::distance::DistanceType::CosineExpanded; + compute_tf32_row_norms( + handle, X.data_handle(), tf32_x_norms, n_samples, n_features, take_sqrt); + compute_tf32_row_norms( + handle, centroids.data_handle(), tf32_centroid_norms, n_clusters, n_features, take_sqrt); + x_norm_ptr = tf32_x_norms; + centroids_norm_ptr = tf32_centroid_norms; + } else { + L2NormBuf_OR_DistBuf.resize(n_clusters, stream); + centroids_norm_ptr = L2NormBuf_OR_DistBuf.data(); + } } else { - raft::linalg::norm( - handle, - raft::make_device_matrix_view( - centroids.data_handle(), centroids.extent(0), centroids.extent(1)), - centroidsNorm); + L2NormBuf_OR_DistBuf.resize(n_clusters, stream); + centroids_norm_ptr = L2NormBuf_OR_DistBuf.data(); + } + + if (!(fused_path == FusedDistancePath::FusedCutile && is_l2_cos && + std::is_same_v)) { + auto centroids_norm = + raft::make_device_vector_view(L2NormBuf_OR_DistBuf.data(), n_clusters); + if (metric == cuvs::distance::DistanceType::CosineExpanded) { + raft::linalg::norm( + handle, + raft::make_device_matrix_view( + centroids.data_handle(), centroids.extent(0), centroids.extent(1)), + centroids_norm, + raft::sqrt_op{}); + } else { + raft::linalg::norm( + handle, + raft::make_device_matrix_view( + centroids.data_handle(), centroids.extent(0), centroids.extent(1)), + centroids_norm); + } } raft::KeyValuePair* cutlass_kvp_scratch = nullptr; @@ -333,8 +435,8 @@ void minClusterDistanceCompute(raft::resources const& handle, minClusterDistance.data_handle(), X.data_handle(), centroids.data_handle(), - L2NormX.data_handle(), - centroidsNorm.data_handle(), + x_norm_ptr, + centroids_norm_ptr, n_samples, n_clusters, n_features, diff --git a/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_kernel.py b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_kernel.py index feeb01012d..2e14cc0d97 100644 --- a/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_kernel.py +++ b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_kernel.py @@ -98,18 +98,9 @@ def red_op(a_score, a_idx, b_score, b_idx): keepdims=True, ) - def sum_values(a, b): - return a + b - - use_tf32 = A.dtype == ct.float32 - if use_tf32: - rounded_a_norm = ct.full((tm, 1), 0, dtype=acc_dtype) - local_indices = ct.arange(tn, dtype=ct.int16)[None, :] for n in range(num_tiles_n): accumulator = ct.full((tm, tn), 0, dtype=acc_dtype) - if use_tf32: - rounded_b_norm = ct.full((1, tn), 0, dtype=acc_dtype) for k in range(num_tiles_k): dtype = ct.tfloat32 if A.dtype == ct.float32 else A.dtype a = ct.load( @@ -123,31 +114,13 @@ def sum_values(a, b): order=(1, 0), ).astype(dtype) accumulator = ct.mma(a, b_T, accumulator) - if use_tf32: - a_fp32 = a.astype(acc_dtype) - b_fp32 = b_T.astype(acc_dtype) - if n == 0: - rounded_a_norm += ct.reduce( - a_fp32 * a_fp32, 1, sum_values, 0.0, keepdims=True - ) - rounded_b_norm += ct.reduce( - b_fp32 * b_fp32, 0, sum_values, 0.0, keepdims=True - ) if metric_code == METRIC_INNER_PRODUCT: score = -accumulator else: - if use_tf32: - b_norm = rounded_b_norm.reshape((tn,)) - b_norm = ct.where( - metric_code == METRIC_COSINE_EXPANDED, - ct.sqrt(b_norm), - b_norm, - ) - else: - b_norm = ct.load( - B_norm, index=(n,), shape=(tn,), padding_mode=zero_pad - ) + b_norm = ct.load( + B_norm, index=(n,), shape=(tn,), padding_mode=zero_pad + ) if metric_code == METRIC_L2_EXPANDED: # L2 receives squared row norms; cosine receives L2 magnitudes. # The A norm is constant across centroids. Reduce @@ -172,16 +145,9 @@ def sum_values(a, b): if metric_code == METRIC_INNER_PRODUCT: out_dist = -best_dist else: - if use_tf32: - a_norm = ct.where( - metric_code == METRIC_COSINE_EXPANDED, - ct.sqrt(rounded_a_norm), - rounded_a_norm, - ) - else: - a_norm = ct.load( - A_norm, index=(bidm,), shape=(tm,), padding_mode=zero_pad - )[:, None] + a_norm = ct.load( + A_norm, index=(bidm,), shape=(tm,), padding_mode=zero_pad + )[:, None] if metric_code == METRIC_L2_EXPANDED: out_dist = a_norm + 2.0 * best_dist out_dist = ct.where( From 8d88dfd0a2e92f6e529b719866e1d6be22a250bf Mon Sep 17 00:00:00 2001 From: divyegala Date: Tue, 25 Aug 2026 04:14:27 +0000 Subject: [PATCH 57/82] new tile shapes --- .../modules/generate_cutile_kernels.cmake | 121 +++--- .../cuvs/detail/jit_lto/cutile_arch_tags.hpp | 5 + .../cutile/export_fused_1nn.py | 4 + .../cutile/fused_1nn_cutile_matrix.json | 365 +++++++++++++++--- .../cutile/fused_1nn_kernel.py | 6 +- .../cutile/fused_1nn_planner.hpp | 84 +++- 6 files changed, 464 insertions(+), 121 deletions(-) diff --git a/cpp/cmake/modules/generate_cutile_kernels.cmake b/cpp/cmake/modules/generate_cutile_kernels.cmake index 1d81963be8..f38f0a974b 100644 --- a/cpp/cmake/modules/generate_cutile_kernels.cmake +++ b/cpp/cmake/modules/generate_cutile_kernels.cmake @@ -97,7 +97,10 @@ function(_cutile_kernels_setup) ) endfunction() -macro(_cutile_append_matrix_tile_aliases entry abi_abbrev tile_m tile_n tile_k) +macro(_cutile_append_matrix_tile_aliases entry data_abbrev abi_abbrev tile_geometry) + list(GET tile_geometry 0 tile_m) + list(GET tile_geometry 1 tile_n) + list(GET tile_geometry 2 tile_k) string(JSON _cutile_export_len LENGTH "${entry}" "_export") set(_cutile_export_idx 0) while(_cutile_export_idx LESS _cutile_export_len) @@ -105,9 +108,9 @@ macro(_cutile_append_matrix_tile_aliases entry abi_abbrev tile_m tile_n tile_k) string(JSON _cutile_register GET "${_cutile_export_entry}" "register") if(_cutile_register STREQUAL "cubin") string(JSON _cutile_arch_tag GET "${_cutile_export_entry}" "arch_tag") - set(_cutile_alias_suffix "${_cutile_arch_tag}_${abi_abbrev}") + set(_cutile_alias_suffix "${data_abbrev}_${_cutile_arch_tag}_${abi_abbrev}") elseif(_cutile_register STREQUAL "tileir") - set(_cutile_alias_suffix "tileir_${abi_abbrev}") + set(_cutile_alias_suffix "${data_abbrev}_tileir_${abi_abbrev}") else() message(FATAL_ERROR "Unknown cuTile register kind '${_cutile_register}'") endif() @@ -145,30 +148,39 @@ function(_cutile_generate_matrix_tiles_header header_path matrix_json_file) string(JSON _default_tile_k GET "${_entry_tile}" "tile_k") endif() - string(JSON _abi_len LENGTH "${_entry}" "_abi") - set(_abi_idx 0) - while(_abi_idx LESS _abi_len) - string(JSON _abi_entry GET "${_entry}" "_abi" "${_abi_idx}") - string(JSON _abi_abbrev GET "${_abi_entry}" "abi_abbrev") - string(JSON _tile_m ERROR_VARIABLE _tile_m_error GET "${_abi_entry}" "tile_m") - string(JSON _tile_n ERROR_VARIABLE _tile_n_error GET "${_abi_entry}" "tile_n") - string(JSON _tile_k ERROR_VARIABLE _tile_k_error GET "${_abi_entry}" "tile_k") - if(_tile_m_error - OR _tile_n_error - OR _tile_k_error - ) - if(_entry_tile_error) - message(FATAL_ERROR "Missing cuTile geometry for ABI ${_abi_abbrev}") + string(JSON _data_len LENGTH "${_entry}" "_data") + set(_data_idx 0) + while(_data_idx LESS _data_len) + string(JSON _data_entry GET "${_entry}" "_data" "${_data_idx}") + string(JSON _data_abbrev GET "${_data_entry}" "data_abbrev") + + string(JSON _abi_len LENGTH "${_entry}" "_abi") + set(_abi_idx 0) + while(_abi_idx LESS _abi_len) + string(JSON _abi_entry GET "${_entry}" "_abi" "${_abi_idx}") + string(JSON _abi_abbrev GET "${_abi_entry}" "abi_abbrev") + string(JSON _tile_m ERROR_VARIABLE _tile_m_error GET "${_abi_entry}" "tile_m") + string(JSON _tile_n ERROR_VARIABLE _tile_n_error GET "${_abi_entry}" "tile_n") + string(JSON _tile_k ERROR_VARIABLE _tile_k_error GET "${_abi_entry}" "tile_k") + if(_tile_m_error + OR _tile_n_error + OR _tile_k_error + ) + if(_entry_tile_error) + message(FATAL_ERROR "Missing cuTile geometry for ${_data_abbrev}/${_abi_abbrev}") + endif() + set(_tile_m "${_default_tile_m}") + set(_tile_n "${_default_tile_n}") + set(_tile_k "${_default_tile_k}") endif() - set(_tile_m "${_default_tile_m}") - set(_tile_n "${_default_tile_n}") - set(_tile_k "${_default_tile_k}") - endif() - _cutile_append_matrix_tile_aliases( - "${_entry}" "${_abi_abbrev}" "${_tile_m}" "${_tile_n}" "${_tile_k}" - ) - math(EXPR _abi_idx "${_abi_idx} + 1") + set(_tile_geometry "${_tile_m};${_tile_n};${_tile_k}") + _cutile_append_matrix_tile_aliases( + "${_entry}" "${_data_abbrev}" "${_abi_abbrev}" "${_tile_geometry}" + ) + math(EXPR _abi_idx "${_abi_idx} + 1") + endwhile() + math(EXPR _data_idx "${_data_idx} + 1") endwhile() math(EXPR _entry_idx "${_entry_idx} + 1") endwhile() @@ -190,6 +202,40 @@ ${_tile_aliases} ) endfunction() +function(_cutile_make_python_args output_var) + set(_python_args + --format + "${output_format}" + --data-type + "${data_type}" + --metric + "${metric}" + --index-type + "${index_type}" + --tile-m + "${tile_m}" + --tile-n + "${tile_n}" + --tile-k + "${tile_k}" + --gpu-code + "${gpu_code}" + ) + if(DEFINED bytecode_version AND NOT "${bytecode_version}" STREQUAL "") + list(APPEND _python_args --bytecode-version "${bytecode_version}") + endif() + if(DEFINED matrix_layout AND NOT "${matrix_layout}" STREQUAL "") + list(APPEND _python_args --matrix-layout "${matrix_layout}") + endif() + if(DEFINED occupancy AND NOT "${occupancy}" STREQUAL "") + list(APPEND _python_args --occupancy "${occupancy}") + endif() + set(${output_var} + "${_python_args}" + PARENT_SCOPE + ) +endfunction() + function(process_cutile_matrix_entry source_list_var) set(options) set(one_value KERNEL_DIR KERNEL_BASENAME KERNEL_PYTHON EXPORT_SCRIPT OUTPUT_DIRECTORY @@ -227,30 +273,7 @@ function(process_cutile_matrix_entry source_list_var) set(_fragment_cpp "${_CUTILE_OUTPUT_DIRECTORY}/${_artifact_stem}_${register}.cpp") set(embedded_header_file "${_artifact_stem}_${register}.h") - set(_python_args - --format - "${output_format}" - --data-type - "${data_type}" - --metric - "${metric}" - --index-type - "${index_type}" - --tile-m - "${tile_m}" - --tile-n - "${tile_n}" - --tile-k - "${tile_k}" - --gpu-code - "${gpu_code}" - ) - if(DEFINED bytecode_version AND NOT "${bytecode_version}" STREQUAL "") - list(APPEND _python_args --bytecode-version "${bytecode_version}") - endif() - if(DEFINED matrix_layout AND NOT "${matrix_layout}" STREQUAL "") - list(APPEND _python_args --matrix-layout "${matrix_layout}") - endif() + _cutile_make_python_args(_python_args) set(_export_python_executable "${Python3_EXECUTABLE}") if(DEFINED python_executable AND NOT "${python_executable}" STREQUAL "") diff --git a/cpp/include/cuvs/detail/jit_lto/cutile_arch_tags.hpp b/cpp/include/cuvs/detail/jit_lto/cutile_arch_tags.hpp index 2b378dac78..ff7898fd31 100644 --- a/cpp/include/cuvs/detail/jit_lto/cutile_arch_tags.hpp +++ b/cpp/include/cuvs/detail/jit_lto/cutile_arch_tags.hpp @@ -24,6 +24,11 @@ struct cutile_arch_8_6 { static constexpr int cc_minor = 6; }; +struct cutile_arch_8_9 { + static constexpr int cc_major = 8; + static constexpr int cc_minor = 9; +}; + struct cutile_arch_9_0 { static constexpr int cc_major = 9; static constexpr int cc_minor = 0; diff --git a/cpp/src/distance/detail/fused_distance_nn/cutile/export_fused_1nn.py b/cpp/src/distance/detail/fused_distance_nn/cutile/export_fused_1nn.py index 2ddc717b60..4c2f2c8993 100644 --- a/cpp/src/distance/detail/fused_distance_nn/cutile/export_fused_1nn.py +++ b/cpp/src/distance/detail/fused_distance_nn/cutile/export_fused_1nn.py @@ -181,6 +181,7 @@ def export_binary( tile_k: int, gpu_code: str, matrix_layout: str = "strict", + occupancy: int | None = None, bytecode_version: str | None = None, ) -> str: kernel = make_kernel( @@ -192,6 +193,7 @@ def export_binary( index_type=index_type, gpu_code=gpu_code, matrix_layout=matrix_layout, + occupancy=occupancy, ) signature = _kernel_signature( data_type, @@ -244,6 +246,7 @@ def main() -> int: choices=("strict", "relaxed"), default="strict", ) + parser.add_argument("--occupancy", type=int) parser.add_argument( "--bytecode-version", default=DEFAULT_TILEIR_BYTECODE_VERSION ) @@ -261,6 +264,7 @@ def main() -> int: tile_k=args.tile_k, gpu_code=args.gpu_code, matrix_layout=args.matrix_layout, + occupancy=args.occupancy, bytecode_version=args.bytecode_version, ) ) diff --git a/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_cutile_matrix.json b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_cutile_matrix.json index b870f70597..8ccbcef0f1 100644 --- a/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_cutile_matrix.json +++ b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_cutile_matrix.json @@ -4,14 +4,22 @@ { "matrix_layout": "relaxed", "abi_abbrev": "relaxed", - "abi_tag": "cutile_abi_relaxed" + "abi_tag": "cutile_abi_relaxed", + "tile_m": 64, + "tile_n": 128, + "tile_k": 32 + }, + { + "matrix_layout": "strict", + "abi_abbrev": "strict", + "abi_tag": "cutile_abi_strict", + "tile_m": 64, + "tile_n": 128, + "tile_k": 32, + "occupancy": 2 } ], "_data": [ - { - "data_type": "half", - "data_abbrev": "h" - }, { "data_type": "float", "data_abbrev": "f" @@ -28,11 +36,74 @@ "index_abbrev": "i32" } ], - "_tile": [ + "_export": [ { - "tile_m": 64, + "output_format": "cubin", + "artifact_ext": "cubin", + "artifact_basename": "@data_type@_@index_abbrev@_@abi_abbrev@_@gpu_code@", + "register": "cubin", + "gpu_code": "sm_80", + "cc_major": 8, + "cc_minor": 0, + "arch_tag": "cutile_arch_8_0" + }, + { + "output_format": "cubin", + "artifact_ext": "cubin", + "artifact_basename": "@data_type@_@index_abbrev@_@abi_abbrev@_@gpu_code@", + "register": "cubin", + "gpu_code": "sm_86", + "cc_major": 8, + "cc_minor": 6, + "arch_tag": "cutile_arch_8_6" + }, + { + "output_format": "cubin", + "artifact_ext": "cubin", + "artifact_basename": "@data_type@_@index_abbrev@_@abi_abbrev@_@gpu_code@", + "register": "cubin", + "gpu_code": "sm_89", + "cc_major": 8, + "cc_minor": 9, + "arch_tag": "cutile_arch_8_9" + } + ] + }, + { + "_abi": [ + { + "matrix_layout": "relaxed", + "abi_abbrev": "relaxed", + "abi_tag": "cutile_abi_relaxed", + "tile_m": 128, "tile_n": 128, "tile_k": 32 + }, + { + "matrix_layout": "strict", + "abi_abbrev": "strict", + "abi_tag": "cutile_abi_strict", + "tile_m": 128, + "tile_n": 128, + "tile_k": 128, + "occupancy": 2 + } + ], + "_data": [ + { + "data_type": "half", + "data_abbrev": "h" + } + ], + "_metric": [ + { + "metric": "runtime" + } + ], + "_index": [ + { + "index_type": "int32", + "index_abbrev": "i32" } ], "_export": [ @@ -60,28 +131,75 @@ }, { "_abi": [ + { + "matrix_layout": "relaxed", + "abi_abbrev": "relaxed", + "abi_tag": "cutile_abi_relaxed", + "tile_m": 128, + "tile_n": 128, + "tile_k": 32, + "occupancy": 2 + }, { "matrix_layout": "strict", "abi_abbrev": "strict", "abi_tag": "cutile_abi_strict", "tile_m": 128, "tile_n": 128, - "tile_k": 32 - }, + "tile_k": 32, + "occupancy": 2 + } + ], + "_data": [ + { + "data_type": "half", + "data_abbrev": "h" + } + ], + "_metric": [ + { + "metric": "runtime" + } + ], + "_index": [ + { + "index_type": "int32", + "index_abbrev": "i32" + } + ], + "_export": [ + { + "output_format": "cubin", + "artifact_ext": "cubin", + "artifact_basename": "@data_type@_@index_abbrev@_@abi_abbrev@_@gpu_code@", + "register": "cubin", + "gpu_code": "sm_89", + "cc_major": 8, + "cc_minor": 9, + "arch_tag": "cutile_arch_8_9" + } + ] + }, + { + "_abi": [ { "matrix_layout": "relaxed", "abi_abbrev": "relaxed", "abi_tag": "cutile_abi_relaxed", "tile_m": 128, + "tile_n": 128, + "tile_k": 64 + }, + { + "matrix_layout": "strict", + "abi_abbrev": "strict", + "abi_tag": "cutile_abi_strict", + "tile_m": 64, "tile_n": 256, - "tile_k": 16 + "tile_k": 32 } ], "_data": [ - { - "data_type": "half", - "data_abbrev": "h" - }, { "data_type": "float", "data_abbrev": "f" @@ -104,34 +222,30 @@ "artifact_ext": "cubin", "artifact_basename": "@data_type@_@index_abbrev@_@abi_abbrev@_@gpu_code@", "register": "cubin", - "gpu_code": "sm_100", - "cc_major": 10, + "gpu_code": "sm_90", + "cc_major": 9, "cc_minor": 0, - "arch_tag": "cutile_arch_10_0" + "arch_tag": "cutile_arch_9_0" } ] }, { "_abi": [ - { - "matrix_layout": "strict", - "abi_abbrev": "strict", - "abi_tag": "cutile_abi_strict" - }, { "matrix_layout": "relaxed", "abi_abbrev": "relaxed", "abi_tag": "cutile_abi_relaxed" + }, + { + "matrix_layout": "strict", + "abi_abbrev": "strict", + "abi_tag": "cutile_abi_strict" } ], "_data": [ { "data_type": "half", "data_abbrev": "h" - }, - { - "data_type": "float", - "data_abbrev": "f" } ], "_metric": [ @@ -149,47 +263,95 @@ { "tile_m": 128, "tile_n": 128, - "tile_k": 32 + "tile_k": 128 } ], "_export": [ { - "output_format": "tileir_bytecode", - "artifact_ext": "tilebc", - "artifact_basename": "@data_type@_@index_abbrev@_@abi_abbrev@", - "register": "tileir", - "gpu_code": "sm_80", - "bytecode_version": "13.1" + "output_format": "cubin", + "artifact_ext": "cubin", + "artifact_basename": "@data_type@_@index_abbrev@_@abi_abbrev@_@gpu_code@", + "register": "cubin", + "gpu_code": "sm_90", + "cc_major": 9, + "cc_minor": 0, + "arch_tag": "cutile_arch_9_0" } ] }, { "_abi": [ + { + "matrix_layout": "relaxed", + "abi_abbrev": "relaxed", + "abi_tag": "cutile_abi_relaxed", + "tile_m": 128, + "tile_n": 256, + "tile_k": 16 + }, { "matrix_layout": "strict", "abi_abbrev": "strict", "abi_tag": "cutile_abi_strict", - "tile_m": 64, - "tile_n": 256, + "tile_m": 128, + "tile_n": 128, "tile_k": 32 - }, + } + ], + "_data": [ + { + "data_type": "float", + "data_abbrev": "f" + } + ], + "_metric": [ + { + "metric": "runtime" + } + ], + "_index": [ + { + "index_type": "int32", + "index_abbrev": "i32" + } + ], + "_export": [ + { + "output_format": "cubin", + "artifact_ext": "cubin", + "artifact_basename": "@data_type@_@index_abbrev@_@abi_abbrev@_@gpu_code@", + "register": "cubin", + "gpu_code": "sm_100", + "cc_major": 10, + "cc_minor": 0, + "arch_tag": "cutile_arch_10_0" + } + ] + }, + { + "_abi": [ { "matrix_layout": "relaxed", "abi_abbrev": "relaxed", "abi_tag": "cutile_abi_relaxed", "tile_m": 128, + "tile_n": 256, + "tile_k": 16, + "occupancy": 2 + }, + { + "matrix_layout": "strict", + "abi_abbrev": "strict", + "abi_tag": "cutile_abi_strict", + "tile_m": 128, "tile_n": 128, - "tile_k": 64 + "tile_k": 128 } ], "_data": [ { "data_type": "half", "data_abbrev": "h" - }, - { - "data_type": "float", - "data_abbrev": "f" } ], "_metric": [ @@ -209,40 +371,89 @@ "artifact_ext": "cubin", "artifact_basename": "@data_type@_@index_abbrev@_@abi_abbrev@_@gpu_code@", "register": "cubin", - "gpu_code": "sm_90", - "cc_major": 9, + "gpu_code": "sm_100", + "cc_major": 10, "cc_minor": 0, - "arch_tag": "cutile_arch_9_0" + "arch_tag": "cutile_arch_10_0" } ] }, { "_abi": [ + { + "matrix_layout": "relaxed", + "abi_abbrev": "relaxed", + "abi_tag": "cutile_abi_relaxed", + "tile_m": 64, + "tile_n": 128, + "tile_k": 64, + "occupancy": 2 + }, { "matrix_layout": "strict", "abi_abbrev": "strict", "abi_tag": "cutile_abi_strict", "tile_m": 64, "tile_n": 128, - "tile_k": 32 - }, + "tile_k": 32, + "occupancy": 2 + } + ], + "_data": [ + { + "data_type": "float", + "data_abbrev": "f" + } + ], + "_metric": [ + { + "metric": "runtime" + } + ], + "_index": [ + { + "index_type": "int32", + "index_abbrev": "i32" + } + ], + "_export": [ + { + "output_format": "cubin", + "artifact_ext": "cubin", + "artifact_basename": "@data_type@_@index_abbrev@_@abi_abbrev@_@gpu_code@", + "register": "cubin", + "gpu_code": "sm_120", + "cc_major": 12, + "cc_minor": 0, + "arch_tag": "cutile_arch_12_0" + } + ] + }, + { + "_abi": [ { "matrix_layout": "relaxed", "abi_abbrev": "relaxed", "abi_tag": "cutile_abi_relaxed", "tile_m": 64, "tile_n": 128, - "tile_k": 64 + "tile_k": 128, + "occupancy": 2 + }, + { + "matrix_layout": "strict", + "abi_abbrev": "strict", + "abi_tag": "cutile_abi_strict", + "tile_m": 64, + "tile_n": 256, + "tile_k": 64, + "occupancy": 2 } ], "_data": [ { "data_type": "half", "data_abbrev": "h" - }, - { - "data_type": "float", - "data_abbrev": "f" } ], "_metric": [ @@ -268,5 +479,57 @@ "arch_tag": "cutile_arch_12_0" } ] + }, + { + "_abi": [ + { + "matrix_layout": "relaxed", + "abi_abbrev": "relaxed", + "abi_tag": "cutile_abi_relaxed" + }, + { + "matrix_layout": "strict", + "abi_abbrev": "strict", + "abi_tag": "cutile_abi_strict" + } + ], + "_data": [ + { + "data_type": "half", + "data_abbrev": "h" + }, + { + "data_type": "float", + "data_abbrev": "f" + } + ], + "_metric": [ + { + "metric": "runtime" + } + ], + "_index": [ + { + "index_type": "int32", + "index_abbrev": "i32" + } + ], + "_tile": [ + { + "tile_m": 128, + "tile_n": 128, + "tile_k": 32 + } + ], + "_export": [ + { + "output_format": "tileir_bytecode", + "artifact_ext": "tilebc", + "artifact_basename": "@data_type@_@index_abbrev@_@abi_abbrev@", + "register": "tileir", + "gpu_code": "sm_80", + "bytecode_version": "13.1" + } + ] } ] diff --git a/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_kernel.py b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_kernel.py index 2e14cc0d97..dddd95954a 100644 --- a/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_kernel.py +++ b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_kernel.py @@ -38,6 +38,7 @@ def make_kernel( index_type: str = "int32", gpu_code: str = "sm_80", matrix_layout: str = "strict", + occupancy: int | None = None, ): """Build the flat-reduction runtime-metric cuTile kernel.""" if data_type not in ("half", "float"): @@ -54,8 +55,11 @@ def make_kernel( out_dist_dtype = ct.float16 if data_type == "half" else ct.float32 core_shape = (tile_m, tile_n) best_shape = (tile_m, 1) + kernel_options = {} + if occupancy is not None: + kernel_options["occupancy"] = ct.ByTarget(**{gpu_code: occupancy}) - @ct.kernel(occupancy=ct.ByTarget(sm_120=2)) + @ct.kernel(**kernel_options) def fused_1nn_kernel( A, B, diff --git a/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_planner.hpp b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_planner.hpp index 029fb87746..cceb9c0e67 100644 --- a/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_planner.hpp +++ b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_planner.hpp @@ -54,27 +54,66 @@ struct Fused1nnTilePlanner : cuvs::detail::jit_lto::TileAlgorithmPlanner { using cuvs::detail::jit_lto::cutile_arch_12_0; using cuvs::detail::jit_lto::cutile_arch_8_0; using cuvs::detail::jit_lto::cutile_arch_8_6; + using cuvs::detail::jit_lto::cutile_arch_8_9; using cuvs::detail::jit_lto::cutile_arch_9_0; constexpr bool is_relaxed = std::is_same_v; - using Tile90 = std::conditional_t; - using Tile100 = std::conditional_t; - using Tile120 = std::conditional_t; + constexpr bool is_float = std::is_same_v; + using Tile80 = + std::conditional_t, + std::conditional_t>; + using Tile86 = + std::conditional_t, + std::conditional_t>; + using Tile89 = + std::conditional_t, + std::conditional_t>; + using Tile90 = + std::conditional_t, + std::conditional_t>; + using Tile100 = + std::conditional_t, + std::conditional_t>; + using Tile120 = + std::conditional_t, + std::conditional_t>; - if constexpr (is_relaxed) { - using Tile80 = fused_1nn_matrix_tile_cutile_arch_8_0_relaxed; - using Tile86 = fused_1nn_matrix_tile_cutile_arch_8_6_relaxed; - this->add_static_fragment< - fragment_tag_fused_1nn_cubin>(); - this->add_static_fragment< - fragment_tag_fused_1nn_cubin>(); - } + this->add_static_fragment< + fragment_tag_fused_1nn_cubin>(); + this->add_static_fragment< + fragment_tag_fused_1nn_cubin>(); + this->add_static_fragment< + fragment_tag_fused_1nn_cubin>(); this->add_static_fragment< fragment_tag_fused_1nn_cubin>(); this->add_static_fragment< @@ -86,9 +125,14 @@ struct Fused1nnTilePlanner : cuvs::detail::jit_lto::TileAlgorithmPlanner { void add_tileir_fallback() { constexpr bool is_relaxed = std::is_same_v; - using TileIr = std::conditional_t; + constexpr bool is_float = std::is_same_v; + using TileIr = std::conditional_t, + std::conditional_t>; this->add_static_tileir_fragment< fragment_tag_fused_1nn_tileir>(); } From 4d7d01cdc86f709150412b98165d6b57d3b091b4 Mon Sep 17 00:00:00 2001 From: divyegala Date: Wed, 26 Aug 2026 21:45:52 +0000 Subject: [PATCH 58/82] sm86 for sm89 --- .../modules/generate_cutile_kernels.cmake | 7 ++--- .../cuvs/detail/jit_lto/cutile_arch_tags.hpp | 5 ---- .../cutile/fused_1nn_cutile_matrix.json | 26 +++---------------- .../cutile/fused_1nn_planner.hpp | 11 -------- 4 files changed, 7 insertions(+), 42 deletions(-) diff --git a/cpp/cmake/modules/generate_cutile_kernels.cmake b/cpp/cmake/modules/generate_cutile_kernels.cmake index f38f0a974b..12444c98fc 100644 --- a/cpp/cmake/modules/generate_cutile_kernels.cmake +++ b/cpp/cmake/modules/generate_cutile_kernels.cmake @@ -98,9 +98,10 @@ function(_cutile_kernels_setup) endfunction() macro(_cutile_append_matrix_tile_aliases entry data_abbrev abi_abbrev tile_geometry) - list(GET tile_geometry 0 tile_m) - list(GET tile_geometry 1 tile_n) - list(GET tile_geometry 2 tile_k) + set(_cutile_tile_geometry "${tile_geometry}") + list(GET _cutile_tile_geometry 0 tile_m) + list(GET _cutile_tile_geometry 1 tile_n) + list(GET _cutile_tile_geometry 2 tile_k) string(JSON _cutile_export_len LENGTH "${entry}" "_export") set(_cutile_export_idx 0) while(_cutile_export_idx LESS _cutile_export_len) diff --git a/cpp/include/cuvs/detail/jit_lto/cutile_arch_tags.hpp b/cpp/include/cuvs/detail/jit_lto/cutile_arch_tags.hpp index ff7898fd31..2b378dac78 100644 --- a/cpp/include/cuvs/detail/jit_lto/cutile_arch_tags.hpp +++ b/cpp/include/cuvs/detail/jit_lto/cutile_arch_tags.hpp @@ -24,11 +24,6 @@ struct cutile_arch_8_6 { static constexpr int cc_minor = 6; }; -struct cutile_arch_8_9 { - static constexpr int cc_major = 8; - static constexpr int cc_minor = 9; -}; - struct cutile_arch_9_0 { static constexpr int cc_major = 9; static constexpr int cc_minor = 0; diff --git a/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_cutile_matrix.json b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_cutile_matrix.json index 8ccbcef0f1..1c3eacc6bf 100644 --- a/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_cutile_matrix.json +++ b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_cutile_matrix.json @@ -56,16 +56,6 @@ "cc_major": 8, "cc_minor": 6, "arch_tag": "cutile_arch_8_6" - }, - { - "output_format": "cubin", - "artifact_ext": "cubin", - "artifact_basename": "@data_type@_@index_abbrev@_@abi_abbrev@_@gpu_code@", - "register": "cubin", - "gpu_code": "sm_89", - "cc_major": 8, - "cc_minor": 9, - "arch_tag": "cutile_arch_8_9" } ] }, @@ -116,16 +106,6 @@ "cc_major": 8, "cc_minor": 0, "arch_tag": "cutile_arch_8_0" - }, - { - "output_format": "cubin", - "artifact_ext": "cubin", - "artifact_basename": "@data_type@_@index_abbrev@_@abi_abbrev@_@gpu_code@", - "register": "cubin", - "gpu_code": "sm_86", - "cc_major": 8, - "cc_minor": 6, - "arch_tag": "cutile_arch_8_6" } ] }, @@ -173,10 +153,10 @@ "artifact_ext": "cubin", "artifact_basename": "@data_type@_@index_abbrev@_@abi_abbrev@_@gpu_code@", "register": "cubin", - "gpu_code": "sm_89", + "gpu_code": "sm_86", "cc_major": 8, - "cc_minor": 9, - "arch_tag": "cutile_arch_8_9" + "cc_minor": 6, + "arch_tag": "cutile_arch_8_6" } ] }, diff --git a/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_planner.hpp b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_planner.hpp index cceb9c0e67..0755788121 100644 --- a/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_planner.hpp +++ b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_planner.hpp @@ -54,7 +54,6 @@ struct Fused1nnTilePlanner : cuvs::detail::jit_lto::TileAlgorithmPlanner { using cuvs::detail::jit_lto::cutile_arch_12_0; using cuvs::detail::jit_lto::cutile_arch_8_0; using cuvs::detail::jit_lto::cutile_arch_8_6; - using cuvs::detail::jit_lto::cutile_arch_8_9; using cuvs::detail::jit_lto::cutile_arch_9_0; constexpr bool is_relaxed = std::is_same_v; @@ -75,14 +74,6 @@ struct Fused1nnTilePlanner : cuvs::detail::jit_lto::TileAlgorithmPlanner { std::conditional_t>; - using Tile89 = - std::conditional_t, - std::conditional_t>; using Tile90 = std::conditional_t>(); this->add_static_fragment< fragment_tag_fused_1nn_cubin>(); - this->add_static_fragment< - fragment_tag_fused_1nn_cubin>(); this->add_static_fragment< fragment_tag_fused_1nn_cubin>(); this->add_static_fragment< From 09d2838561fca13fdce924f55fd432e867944833 Mon Sep 17 00:00:00 2001 From: divyegala Date: Wed, 26 Aug 2026 22:18:24 +0000 Subject: [PATCH 59/82] sm80/86 shape divisible by --- .../cutile/export_fused_1nn.py | 32 +++++++++++++++---- .../cutile/fused_1nn_tile.cu | 7 ++-- 2 files changed, 28 insertions(+), 11 deletions(-) diff --git a/cpp/src/distance/detail/fused_distance_nn/cutile/export_fused_1nn.py b/cpp/src/distance/detail/fused_distance_nn/cutile/export_fused_1nn.py index 4c2f2c8993..011ff27636 100644 --- a/cpp/src/distance/detail/fused_distance_nn/cutile/export_fused_1nn.py +++ b/cpp/src/distance/detail/fused_distance_nn/cutile/export_fused_1nn.py @@ -58,20 +58,31 @@ def _elem_stride_divisible_for_tma(elem_dtype) -> tuple[int, int]: return (16 // bytes_per_elem, 1) +def _elem_shape_divisible_for_ldgsts(elem_dtype) -> tuple[int, int]: + """Matrix extent aligned to the same 16-byte row pitch enforced on strides.""" + bytes_per_elem = 2 if elem_dtype == ct.float16 else 4 + return (1, 16 // bytes_per_elem) + + def _cuvs_matrix_constraint( elem_dtype, *, index_dtype=ct.int32, require_tma_friendly_pitch: bool = True, + require_ldgsts_friendly_shape: bool = False, ): """Row-major device matrices for cuVS KMeans benchmarks. - Assumes raft/cupy-style contiguous layout: stride[-1]==1, stride[0]==D, - 16-byte base alignment, and row pitch 16-byte aligned (float32 D%4==0, - float16 D%8==0). Applies to both points and centroids matrices. + Assumes raft/cupy-style contiguous layout: stride[-1]==1, stride[0]==D, + 16-byte base alignment, and row pitch 16-byte aligned (float32 D%4==0, + float16 D%8==0). Applies to both points and centroids matrices. + + SM80/SM86 strict exports also express the row-pitch guarantee as + shape_divisible_by=(1, 4) for float32 or (1, 8) for float16. This + duplicates the stride constraint intentionally so the compiler selects + LDGSTS instead of LDG. Tail tiles remain masked in the kernel. - shape_divisible_by is (1, 1); tail tiles are masked in the kernel. - Odd D or general layouts need a separate relaxed export profile. + Odd D or general layouts need a separate relaxed export profile. """ return ArrayConstraint( elem_dtype, @@ -86,7 +97,11 @@ def _cuvs_matrix_constraint( if require_tma_friendly_pitch else (1, 1) ), - shape_divisible_by=(1, 1), + shape_divisible_by=( + _elem_shape_divisible_for_ldgsts(elem_dtype) + if require_ldgsts_friendly_shape + else (1, 1) + ), base_addr_divisible_by=16, ) @@ -127,6 +142,7 @@ def _kernel_signature( tile_m: int, tile_n: int, tile_k: int, + gpu_code: str, matrix_layout: str, ) -> KernelSignature: elem = _dtype_for(data_type) @@ -135,6 +151,9 @@ def _kernel_signature( elem, index_dtype=idx_dtype, require_tma_friendly_pitch=matrix_layout == "strict", + require_ldgsts_friendly_shape=( + matrix_layout == "strict" and gpu_code in ("sm_80", "sm_86") + ), ) norm_array = _cuvs_vector_constraint(elem, index_dtype=idx_dtype) idx_array = _cuvs_vector_constraint(idx_dtype, index_dtype=idx_dtype) @@ -202,6 +221,7 @@ def export_binary( tile_m, tile_n, tile_k, + gpu_code, matrix_layout, ) diff --git a/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_tile.cu b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_tile.cu index 3e1a46cdf1..fa3e4baafa 100644 --- a/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_tile.cu +++ b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_tile.cu @@ -199,11 +199,8 @@ bool try_fused_1nn_tile(IdxT* nearest_idx, if (!cuvs::detail::jit_lto::cutile_launch_available_on_current_device()) { return false; } static_assert(std::is_same_v || std::is_same_v); - int cc_major = 0; - int cc_minor = 0; - if (!cuvs::detail::jit_lto::get_device_compute_capability(cc_major, cc_minor)) { return false; } - constexpr int tma_pitch_elements = 16 / sizeof(DataT); - const bool use_strict_abi = cc_major >= 9 && k % tma_pitch_elements == 0; + constexpr int strict_pitch_elements = 16 / sizeof(DataT); + const bool use_strict_abi = k % strict_pitch_elements == 0; if constexpr (std::is_same_v) { if (use_strict_abi) { From 200aad6562d2ba92db73cc0e5dd1703395cbb506 Mon Sep 17 00:00:00 2001 From: divyegala Date: Thu, 27 Aug 2026 01:09:54 +0000 Subject: [PATCH 60/82] int64 fix --- c/src/cluster/kmeans.cpp | 12 ++++++++---- cpp/src/cluster/kmeans.cuh | 4 +++- 2 files changed, 11 insertions(+), 5 deletions(-) diff --git a/c/src/cluster/kmeans.cpp b/c/src/cluster/kmeans.cpp index 06e21fcda7..06cd8a1759 100644 --- a/c/src/cluster/kmeans.cpp +++ b/c/src/cluster/kmeans.cpp @@ -44,18 +44,22 @@ cuvs::cluster::kmeans::balanced_params convert_balanced_params(const cuvsKMeansP constexpr int64_t kKMeansInt32IndexMax = std::numeric_limits::max(); -bool dlpack_shape_exceeds_int32_index(const DLTensor& tensor) +bool dlpack_tensor_size_exceeds_int32_index(const DLTensor& tensor) { + int64_t size = 1; for (int i = 0; i < tensor.ndim; ++i) { - if (tensor.shape[i] > kKMeansInt32IndexMax) { return true; } + const int64_t extent = tensor.shape[i]; + if (extent <= 0) { return false; } + if (extent > kKMeansInt32IndexMax / size) { return true; } + size *= extent; } return false; } bool kmeans_tensor_shapes_use_int64_index(DLManagedTensor* X, DLManagedTensor* centroids) { - if (dlpack_shape_exceeds_int32_index(X->dl_tensor)) { return true; } - if (dlpack_shape_exceeds_int32_index(centroids->dl_tensor)) { return true; } + if (dlpack_tensor_size_exceeds_int32_index(X->dl_tensor)) { return true; } + if (dlpack_tensor_size_exceeds_int32_index(centroids->dl_tensor)) { return true; } return false; } diff --git a/cpp/src/cluster/kmeans.cuh b/cpp/src/cluster/kmeans.cuh index d635426229..5f7e5d7981 100644 --- a/cpp/src/cluster/kmeans.cuh +++ b/cpp/src/cluster/kmeans.cuh @@ -365,7 +365,9 @@ void cluster_cost( auto batch_cost = raft::make_device_scalar(handle, DataT{0}); auto centroids_i32 = raft::make_device_matrix_view( centroids.data_handle(), static_cast(n_clusters), static_cast(n_features)); - const IndexT max_batch_rows = max_i32 / n_clusters; + // The i32 path indexes both X[batch_rows, n_features] and its distance workspace + // [batch_rows, n_clusters]. + const IndexT max_batch_rows = max_i32 / std::max(n_clusters, n_features); for (IndexT offset = 0; offset < n_samples; offset += max_batch_rows) { const int batch_rows = static_cast(std::min(max_batch_rows, n_samples - offset)); From 146fccef7f1c9c1f9ebb3c2bc781aadcdb884067 Mon Sep 17 00:00:00 2001 From: divyegala Date: Thu, 27 Aug 2026 16:39:06 +0000 Subject: [PATCH 61/82] fix tests --- cpp/tests/cluster/kmeans.cu | 3 ++- cpp/tests/neighbors/distance_nn.cu | 2 ++ 2 files changed, 4 insertions(+), 1 deletion(-) diff --git a/cpp/tests/cluster/kmeans.cu b/cpp/tests/cluster/kmeans.cu index 59051484f4..aa7de8619f 100644 --- a/cpp/tests/cluster/kmeans.cu +++ b/cpp/tests/cluster/kmeans.cu @@ -679,7 +679,8 @@ TEST_P(KmeansFitBatchedTestF, Result) { prepareBlobInputs(); fitBatchedTest(); - ASSERT_TRUE(centroids_match); + // Batched FP32 centroid accumulation uses atomics, so its reduction order is not deterministic. + // Compare the resulting clustering and inertia rather than individual centroid coordinates. ASSERT_TRUE(score >= 0.99); ASSERT_TRUE(inertia_match); runInitSizeCompare(); diff --git a/cpp/tests/neighbors/distance_nn.cu b/cpp/tests/neighbors/distance_nn.cu index 376a60a36b..64c07b4058 100644 --- a/cpp/tests/neighbors/distance_nn.cu +++ b/cpp/tests/neighbors/distance_nn.cu @@ -211,9 +211,11 @@ const std::vector> input_fp32 = { template const std::vector> input_fp32_fused = [] { auto inputs = input_fp32; +#if CUVS_CUTILE_ENABLED inputs.insert( inputs.begin() + 6, NNInputs{512, 1024, 64, DistanceType::InnerProduct, false, uint64_t(31415926), 0.1}); +#endif inputs.push_back( NNInputs{1000, 8, 32, DistanceType::L2Expanded, false, uint64_t(31415926), 0.1}); inputs.push_back( From 87080b80a693cb7bba0fd98677cfbaa9cbb69705 Mon Sep 17 00:00:00 2001 From: divyegala Date: Fri, 28 Aug 2026 02:33:23 +0000 Subject: [PATCH 62/82] address review --- c/src/cluster/kmeans.cpp | 53 +++++---- c/tests/cluster/kmeans_c.cu | 11 ++ cpp/include/cuvs/cluster/kmeans.hpp | 35 +++++- cpp/src/cluster/detail/kmeans.cuh | 44 +++++--- cpp/src/cluster/detail/kmeans_balanced.cuh | 47 +++++--- cpp/src/cluster/detail/kmeans_common.cuh | 12 +-- .../detail/minClusterDistanceCompute.cu | 69 ++++++++++-- cpp/src/cluster/kmeans.cuh | 23 ++-- cpp/src/cluster/kmeans_impl.cuh | 6 +- cpp/src/cluster/kmeans_predict_double.cu | 44 ++++++-- cpp/src/cluster/kmeans_predict_float.cu | 44 ++++++-- cpp/src/distance/detail/fused_distance_nn.cuh | 7 +- .../cutile/fused_1nn_kernel.py | 23 ++-- .../cutile/fused_1nn_tile.cu | 101 ++++++++++++++++-- .../cutile/fused_1nn_tile.hpp | 35 ++++++ cpp/src/distance/fused_distance_nn-inl.cuh | 3 +- cpp/tests/neighbors/distance_nn.cu | 48 +++++++++ 17 files changed, 484 insertions(+), 121 deletions(-) diff --git a/c/src/cluster/kmeans.cpp b/c/src/cluster/kmeans.cpp index 06cd8a1759..cb5f045df0 100644 --- a/c/src/cluster/kmeans.cpp +++ b/c/src/cluster/kmeans.cpp @@ -71,7 +71,7 @@ bool kmeans_fit_uses_int64_index(DLManagedTensor* X, return kmeans_tensor_shapes_use_int64_index(X, centroids); } -bool kmeans_labels_use_int64_index(const DLTensor& labels) +bool kmeans_labels_are_int64(const DLTensor& labels) { return labels.dtype.code == kDLInt && labels.dtype.bits == 64; } @@ -86,16 +86,6 @@ void validate_kmeans_labels_dtype(const DLTensor& labels) labels.dtype.bits); } -bool kmeans_predict_uses_int64_index(DLManagedTensor* X, - DLManagedTensor* centroids, - DLManagedTensor* labels, - int n_clusters) -{ - validate_kmeans_labels_dtype(labels->dl_tensor); - if (kmeans_labels_use_int64_index(labels->dl_tensor)) { return true; } - return kmeans_fit_uses_int64_index(X, centroids, n_clusters); -} - template void _fit(cuvsResources_t res, const cuvsKMeansParams& params, @@ -372,23 +362,44 @@ extern "C" cuvsError_t cuvsKMeansPredict(cuvsResources_t res, { return cuvs::core::translate_exceptions([=] { auto dataset = X->dl_tensor; - const bool use_int64_index = - kmeans_predict_uses_int64_index(X, centroids, labels, params->n_clusters); + validate_kmeans_labels_dtype(labels->dl_tensor); + const bool use_int64_index = kmeans_fit_uses_int64_index(X, centroids, params->n_clusters); + const bool use_int64_labels = kmeans_labels_are_int64(labels->dl_tensor); if (dataset.dtype.code == kDLFloat && dataset.dtype.bits == 32) { if (use_int64_index) { - _predict( - res, *params, X, sample_weight, centroids, labels, normalize_weight, inertia); + if (use_int64_labels) { + _predict( + res, *params, X, sample_weight, centroids, labels, normalize_weight, inertia); + } else { + _predict( + res, *params, X, sample_weight, centroids, labels, normalize_weight, inertia); + } } else { - _predict( - res, *params, X, sample_weight, centroids, labels, normalize_weight, inertia); + if (use_int64_labels) { + _predict( + res, *params, X, sample_weight, centroids, labels, normalize_weight, inertia); + } else { + _predict( + res, *params, X, sample_weight, centroids, labels, normalize_weight, inertia); + } } } else if (dataset.dtype.code == kDLFloat && dataset.dtype.bits == 64) { if (use_int64_index) { - _predict( - res, *params, X, sample_weight, centroids, labels, normalize_weight, inertia); + if (use_int64_labels) { + _predict( + res, *params, X, sample_weight, centroids, labels, normalize_weight, inertia); + } else { + _predict( + res, *params, X, sample_weight, centroids, labels, normalize_weight, inertia); + } } else { - _predict( - res, *params, X, sample_weight, centroids, labels, normalize_weight, inertia); + if (use_int64_labels) { + _predict( + res, *params, X, sample_weight, centroids, labels, normalize_weight, inertia); + } else { + _predict( + res, *params, X, sample_weight, centroids, labels, normalize_weight, inertia); + } } } else { RAFT_FAIL("Unsupported dataset DLtensor dtype: %d and bits: %d", diff --git a/c/tests/cluster/kmeans_c.cu b/c/tests/cluster/kmeans_c.cu index 5caac1ba98..a1d3b1f64d 100644 --- a/c/tests/cluster/kmeans_c.cu +++ b/c/tests/cluster/kmeans_c.cu @@ -44,6 +44,7 @@ float kInitCentroids[kNClusters][kNFeatures] = { float kExpectedCentroids[kNClusters * kNFeatures] = {1.5f, 1.5f, 10.5f, 10.5f}; int32_t kExpectedLabels[kNSamples] = {0, 0, 0, 0, 1, 1, 1, 1}; +int64_t kExpectedLabels64[kNSamples] = {0, 0, 0, 0, 1, 1, 1, 1}; // 8 points, each at squared distance 0.5 from its cluster mean -> 4.0. constexpr double kExpectedInertia = 4.0; @@ -56,6 +57,7 @@ void test_fit_predict() rmm::device_uvector dataset_d(kNSamples * kNFeatures, stream); rmm::device_uvector centroids_d(kNClusters * kNFeatures, stream); rmm::device_uvector labels_d(kNSamples, stream); + rmm::device_uvector labels64_d(kNSamples, stream); raft::copy(dataset_d.data(), reinterpret_cast(kDataset), @@ -90,6 +92,9 @@ void test_fit_predict() DLManagedTensor labels_t{}; cuvs::core::to_dlpack( raft::make_device_vector_view(labels_d.data(), kNSamples), &labels_t); + DLManagedTensor labels64_t{}; + cuvs::core::to_dlpack( + raft::make_device_vector_view(labels64_d.data(), kNSamples), &labels64_t); double inertia = -1.0; int n_iter = -1; @@ -101,6 +106,9 @@ void test_fit_predict() ASSERT_EQ(cuvsKMeansPredict( res, params, &dataset_t, NULL, ¢roids_t, &labels_t, false, &predict_inertia), CUVS_SUCCESS); + ASSERT_EQ(cuvsKMeansPredict( + res, params, &dataset_t, NULL, ¢roids_t, &labels64_t, false, &predict_inertia), + CUVS_SUCCESS); ASSERT_EQ(cuvsKMeansClusterCost(res, &dataset_t, ¢roids_t, &cluster_cost), CUVS_SUCCESS); ASSERT_TRUE(cuvs::devArrMatchHost(kExpectedCentroids, @@ -109,12 +117,15 @@ void test_fit_predict() cuvs::CompareApprox(1e-4f))); ASSERT_TRUE(cuvs::devArrMatchHost( kExpectedLabels, labels_d.data(), kNSamples, cuvs::Compare())); + ASSERT_TRUE(cuvs::devArrMatchHost( + kExpectedLabels64, labels64_d.data(), kNSamples, cuvs::Compare())); EXPECT_GT(n_iter, 0); EXPECT_NEAR(inertia, kExpectedInertia, 1e-4); EXPECT_NEAR(predict_inertia, kExpectedInertia, 1e-4); EXPECT_NEAR(cluster_cost, kExpectedInertia, 1e-4); + labels64_t.deleter(&labels64_t); labels_t.deleter(&labels_t); centroids_t.deleter(¢roids_t); dataset_t.deleter(&dataset_t); diff --git a/cpp/include/cuvs/cluster/kmeans.hpp b/cpp/include/cuvs/cluster/kmeans.hpp index ceb42d1e71..251ff40f69 100644 --- a/cpp/include/cuvs/cluster/kmeans.hpp +++ b/cpp/include/cuvs/cluster/kmeans.hpp @@ -787,6 +787,15 @@ void predict(raft::resources const& handle, bool normalize_weight, raft::host_scalar_view inertia); +void predict(raft::resources const& handle, + const kmeans::params& params, + raft::device_matrix_view X, + std::optional> sample_weight, + raft::device_matrix_view centroids, + raft::device_vector_view labels, + bool normalize_weight, + raft::host_scalar_view inertia); + /** * @brief Predict the closest cluster each sample in X belongs to. * @@ -838,10 +847,10 @@ void predict(raft::resources const& handle, */ void predict(raft::resources const& handle, const kmeans::params& params, - raft::device_matrix_view X, - std::optional> sample_weight, - raft::device_matrix_view centroids, - raft::device_vector_view labels, + raft::device_matrix_view X, + std::optional> sample_weight, + raft::device_matrix_view centroids, + raft::device_vector_view labels, bool normalize_weight, raft::host_scalar_view inertia); @@ -952,6 +961,24 @@ void predict(raft::resources const& handle, * @param[out] inertia Sum of squared distances of samples to * their closest cluster center. */ +void predict(raft::resources const& handle, + const kmeans::params& params, + raft::device_matrix_view X, + std::optional> sample_weight, + raft::device_matrix_view centroids, + raft::device_vector_view labels, + bool normalize_weight, + raft::host_scalar_view inertia); + +void predict(raft::resources const& handle, + const kmeans::params& params, + raft::device_matrix_view X, + std::optional> sample_weight, + raft::device_matrix_view centroids, + raft::device_vector_view labels, + bool normalize_weight, + raft::host_scalar_view inertia); + void predict(raft::resources const& handle, const kmeans::params& params, raft::device_matrix_view X, diff --git a/cpp/src/cluster/detail/kmeans.cuh b/cpp/src/cluster/detail/kmeans.cuh index abb6b1a648..819545e8d1 100644 --- a/cpp/src/cluster/detail/kmeans.cuh +++ b/cpp/src/cluster/detail/kmeans.cuh @@ -51,6 +51,7 @@ #include #include #include +#include namespace cuvs::cluster::kmeans::detail { @@ -1022,13 +1023,13 @@ void fit(raft::resources const& handle, kmeans_fit(handle, params, X, sample_weight, centroids, inertia, n_iter); } -template +template void kmeans_predict(raft::resources const& handle, const cuvs::cluster::kmeans::params& pams, raft::device_matrix_view X, std::optional> sample_weight, raft::device_matrix_view centroids, - raft::device_vector_view labels, + raft::device_vector_view labels, bool normalize_weight, raft::host_scalar_view inertia) { @@ -1088,17 +1089,34 @@ void kmeans_predict(raft::resources const& handle, auto l2normx_view = raft::make_device_vector_view(L2NormX.data_handle(), n_samples); - cuvs::cluster::kmeans::detail::minClusterAndDistanceCompute(handle, - X, - centroids, - labels, - nearest_dist.view(), - l2normx_view, - L2NormBuf_OR_DistBuf, - pams.metric, - pams.batch_samples, - pams.batch_centroids, - workspace); + if constexpr (std::is_same_v) { + cuvs::cluster::kmeans::detail::minClusterAndDistanceCompute(handle, + X, + centroids, + labels, + nearest_dist.view(), + l2normx_view, + L2NormBuf_OR_DistBuf, + pams.metric, + pams.batch_samples, + pams.batch_centroids, + workspace); + } else { + auto index_labels = raft::make_device_vector(handle, n_samples); + cuvs::cluster::kmeans::detail::minClusterAndDistanceCompute(handle, + X, + centroids, + index_labels.view(), + nearest_dist.view(), + l2normx_view, + L2NormBuf_OR_DistBuf, + pams.metric, + pams.batch_samples, + pams.batch_centroids, + workspace); + raft::linalg::map( + handle, labels, raft::cast_op{}, raft::make_const_mdspan(index_labels.view())); + } rmm::device_scalar clusterCostD(stream); raft::linalg::map(handle, diff --git a/cpp/src/cluster/detail/kmeans_balanced.cuh b/cpp/src/cluster/detail/kmeans_balanced.cuh index e11302ce92..ebda7397a2 100644 --- a/cpp/src/cluster/detail/kmeans_balanced.cuh +++ b/cpp/src/cluster/detail/kmeans_balanced.cuh @@ -225,21 +225,22 @@ inline std::enable_if_t> predict_core( * * @tparam MathT type of the centroids and mapped data * @tparam IdxT index type + * @tparam LabelT label type * * @param[in] n_clusters number of clusters in kmeans clustering * @param[in] n_rows Number of samples in the dataset * @param[in] dim Number of features in the dataset * @param[in] metric Distance metric - * @param[in] needs_conversion Whether the data needs to be converted to MathT + * @param[in] data_is_math_type Whether the input data already uses MathT * @return A suggested minibatch size and the expected memory cost per-row (in bytes) */ -template +template auto calc_minibatch_size(const raft::resources& handle, IdxT n_clusters, IdxT n_rows, IdxT dim, cuvs::distance::DistanceType metric, - bool needs_conversion) -> std::tuple + bool data_is_math_type) -> std::tuple { n_clusters = std::max(1, n_clusters); @@ -248,12 +249,31 @@ auto calc_minibatch_size(const raft::resources& handle, switch (metric) { case distance::DistanceType::L2Expanded: case distance::DistanceType::L2SqrtExpanded: + case distance::DistanceType::CosineExpanded: case distance::DistanceType::InnerProduct: { - switch (use_fused(handle, n_rows, n_clusters, dim, metric)) { + const auto fused_path = use_fused(handle, n_rows, n_clusters, dim, metric); + + // min_cluster_and_distance always materializes the nearest distance for fused/L2 paths. + if (metric != distance::DistanceType::InnerProduct || + fused_path != FusedDistancePath::Unfused) { + mem_per_row += sizeof(MathT); + if constexpr (!std::is_same_v) { mem_per_row += sizeof(IdxT); } + } + if (metric != distance::DistanceType::InnerProduct) { + // predict may need a minibatch-sized input-norm buffer before entering predict_core. + mem_per_row += sizeof(MathT); + } + + switch (fused_path) { case FusedDistancePath::FusedCutile: - if constexpr (std::is_same_v) { - // cuTile computes labels with i32 and widens them after each launch. - mem_per_row += sizeof(int); + // Conservatively budget the fallback in case the eventual pointer-aware probe fails. + mem_per_row += sizeof(int); + mem_per_row += sizeof(raft::KeyValuePair); + if constexpr (std::is_same_v) { + if (metric != distance::DistanceType::InnerProduct) { + // TF32-compatible row norms are materialized for cuTile L2/cosine. + mem_per_row += sizeof(MathT); + } } break; case FusedDistancePath::FusedCutlass: @@ -264,6 +284,9 @@ auto calc_minibatch_size(const raft::resources& handle, case FusedDistancePath::Unfused: // unfused / GEMM+argmin path needs a full distance matrix row. mem_per_row += sizeof(MathT) * n_clusters; + if (metric != distance::DistanceType::InnerProduct) { + mem_per_row += sizeof(raft::KeyValuePair); + } break; } } break; @@ -274,7 +297,7 @@ auto calc_minibatch_size(const raft::resources& handle, } // If we need to convert to MathT, space required for the converted batch. - if (!needs_conversion) { mem_per_row += sizeof(MathT) * dim; } + if (!data_is_math_type) { mem_per_row += sizeof(MathT) * dim; } // Heuristic: calculate the minibatch size in order to use at most 80% or 512MB workspace memory. // We go below 1GB here as the allocation is mostly done in a single chunk which @@ -283,10 +306,6 @@ auto calc_minibatch_size(const raft::resources& handle, const auto available_ws_size = std::min((free_ws_size * size_t{8}) / size_t{10}, size_t{1} << 29); - // A fused implementation may require no per-row temporary workspace. In that case, - // process the complete input rather than dividing the available workspace by zero. - if (mem_per_row == 0) { return std::make_tuple(n_rows, mem_per_row); } - IdxT minibatch_size = std::max(IdxT{1}, static_cast(available_ws_size / mem_per_row)); minibatch_size = raft::round_down_safe(minibatch_size, IdxT{64}); @@ -467,7 +486,7 @@ void predict(const raft::resources& handle, raft::common::nvtx::range fun_scope( "predict(%zu, %u)", static_cast(n_rows), n_clusters); auto mem_res = mr.value_or(raft::resource::get_workspace_resource_ref(handle)); - auto [max_minibatch_size, _mem_per_row] = calc_minibatch_size( + auto [max_minibatch_size, _mem_per_row] = calc_minibatch_size( handle, n_clusters, n_rows, dim, params.metric, std::is_same_v); rmm::device_uvector cur_dataset( std::is_same_v ? 0 : max_minibatch_size * dim, stream, mem_res); @@ -1221,7 +1240,7 @@ void build_hierarchical(const raft::resources& handle, // TODO: Remove the explicit managed memory- we shouldn't be creating this on the user's behalf. rmm::mr::managed_memory_resource managed_memory; rmm::device_async_resource_ref device_memory = raft::resource::get_workspace_resource_ref(handle); - auto [max_minibatch_size, mem_per_row] = calc_minibatch_size( + auto [max_minibatch_size, mem_per_row] = calc_minibatch_size( handle, n_clusters, n_rows, dim, params.metric, std::is_same_v); // Precompute the L2 norm of the dataset if relevant and not yet computed. diff --git a/cpp/src/cluster/detail/kmeans_common.cuh b/cpp/src/cluster/detail/kmeans_common.cuh index 3ece6b0e19..095ab28603 100644 --- a/cpp/src/cluster/detail/kmeans_common.cuh +++ b/cpp/src/cluster/detail/kmeans_common.cuh @@ -66,7 +66,7 @@ inline constexpr bool is_cutile_fused_data_type_v = enum class FusedDistancePath : std::uint8_t { /** unfusedDistanceNNMinReduce or batched pairwise distance. */ Unfused = 0, - /** fusedDistanceNNMinReduce via cuTile; no CUTLASS mutex / KVP scratch. */ + /** fusedDistanceNNMinReduce via cuTile; scratch depends on the launchability probe. */ FusedCutile, /** fusedDistanceNNMinReduce via legacy CUTLASS; needs mutex workspace + KVP scratch. */ FusedCutlass, @@ -77,16 +77,6 @@ inline constexpr bool uses_fused_distance_nn(FusedDistancePath path) return path != FusedDistancePath::Unfused; } -inline constexpr bool needs_cutlass_kvp_scratch(FusedDistancePath path) -{ - return path == FusedDistancePath::FusedCutlass; -} - -inline constexpr bool needs_fused_mutex_workspace(FusedDistancePath path) -{ - return path == FusedDistancePath::FusedCutlass; -} - /** * @brief Selects the fused-distance assignment path for KMeans. * diff --git a/cpp/src/cluster/detail/minClusterDistanceCompute.cu b/cpp/src/cluster/detail/minClusterDistanceCompute.cu index b230b8af16..7570e15adb 100644 --- a/cpp/src/cluster/detail/minClusterDistanceCompute.cu +++ b/cpp/src/cluster/detail/minClusterDistanceCompute.cu @@ -116,8 +116,24 @@ void minClusterAndDistanceCompute(raft::resources const& handle, const bool is_l2_cos = metric == cuvs::distance::DistanceType::L2Expanded || metric == cuvs::distance::DistanceType::L2SqrtExpanded || metric == cuvs::distance::DistanceType::CosineExpanded; - const FusedDistancePath fused_path = + FusedDistancePath fused_path = use_fused(handle, n_samples, n_clusters, n_features, metric); + if constexpr (is_cutile_fused_data_type_v) { + if (fused_path == FusedDistancePath::FusedCutile && + metric == cuvs::distance::DistanceType::InnerProduct && + !cuvs::distance::detail::can_launch_fused_1nn_tile(nearest_idx.data_handle(), + nearest_dist.data_handle(), + X.data_handle(), + centroids.data_handle(), + static_cast(nullptr), + static_cast(nullptr), + n_samples, + n_clusters, + n_features, + metric)) { + fused_path = FusedDistancePath::Unfused; + } + } if (uses_fused_distance_nn(fused_path)) { const DataT* x_norm_ptr = L2NormX.data_handle(); @@ -159,14 +175,30 @@ void minClusterAndDistanceCompute(raft::resources const& handle, } } + bool cutile_ready = false; + if constexpr (is_cutile_fused_data_type_v) { + if (fused_path == FusedDistancePath::FusedCutile) { + cutile_ready = cuvs::distance::detail::can_launch_fused_1nn_tile(nearest_idx.data_handle(), + nearest_dist.data_handle(), + X.data_handle(), + centroids.data_handle(), + x_norm_ptr, + centroids_norm_ptr, + n_samples, + n_clusters, + n_features, + metric); + } + } + raft::KeyValuePair* cutlass_kvp_scratch = nullptr; rmm::device_uvector> temp_kvp(0, stream); - if (needs_cutlass_kvp_scratch(fused_path)) { + const bool needs_index_workspace = cutile_ready && std::is_same_v; + if (!cutile_ready) { temp_kvp.resize(n_samples, stream); cutlass_kvp_scratch = temp_kvp.data(); - workspace.resize(sizeof(int) * n_samples, stream); - } else if constexpr (std::is_same_v) { - // The cuTile kernel uses i32 internally and widens labels after the launch. + } + if (!cutile_ready || needs_index_workspace) { workspace.resize(sizeof(int) * static_cast(n_samples), stream); } @@ -180,9 +212,7 @@ void minClusterAndDistanceCompute(raft::resources const& handle, n_samples, n_clusters, n_features, - needs_fused_mutex_workspace(fused_path) || std::is_same_v - ? (void*)workspace.data() - : nullptr, + !cutile_ready || needs_index_workspace ? (void*)workspace.data() : nullptr, metric != cuvs::distance::DistanceType::L2Expanded, true, true, @@ -422,12 +452,29 @@ void minClusterDistanceCompute(raft::resources const& handle, } } + bool cutile_ready = false; + if constexpr (is_cutile_fused_data_type_v) { + if (fused_path == FusedDistancePath::FusedCutile) { + cutile_ready = + cuvs::distance::detail::can_launch_fused_1nn_tile(static_cast(nullptr), + minClusterDistance.data_handle(), + X.data_handle(), + centroids.data_handle(), + x_norm_ptr, + centroids_norm_ptr, + n_samples, + n_clusters, + n_features, + metric); + } + } + raft::KeyValuePair* cutlass_kvp_scratch = nullptr; rmm::device_uvector> temp_kvp(0, stream); - if (needs_cutlass_kvp_scratch(fused_path)) { + if (!cutile_ready) { temp_kvp.resize(n_samples, stream); cutlass_kvp_scratch = temp_kvp.data(); - workspace.resize(sizeof(int) * n_samples, stream); + workspace.resize(sizeof(int) * static_cast(n_samples), stream); } cuvs::distance::fusedDistanceNNMinReduce( @@ -440,7 +487,7 @@ void minClusterDistanceCompute(raft::resources const& handle, n_samples, n_clusters, n_features, - needs_fused_mutex_workspace(fused_path) ? (void*)workspace.data() : nullptr, + cutile_ready ? nullptr : (void*)workspace.data(), metric != cuvs::distance::DistanceType::L2Expanded, true, true, diff --git a/cpp/src/cluster/kmeans.cuh b/cpp/src/cluster/kmeans.cuh index 5f7e5d7981..572104f8eb 100644 --- a/cpp/src/cluster/kmeans.cuh +++ b/cpp/src/cluster/kmeans.cuh @@ -65,6 +65,7 @@ using KeyValueIndexOp = cuvs::cluster::kmeans::detail::KeyValueIndexOp +template void predict(raft::resources const& handle, const kmeans::params& params, raft::device_matrix_view X, std::optional> sample_weight, raft::device_matrix_view centroids, - raft::device_vector_view labels, + raft::device_vector_view labels, bool normalize_weight, raft::host_scalar_view inertia); -#define EXTERN_TEMPLATE_PREDICT(DataT, IndexT) \ - extern template void predict( \ +#define EXTERN_TEMPLATE_PREDICT(DataT, IndexT, LabelT) \ + extern template void predict( \ raft::resources const& handle, \ const kmeans::params& params, \ raft::device_matrix_view X, \ std::optional> sample_weight, \ raft::device_matrix_view centroids, \ - raft::device_vector_view labels, \ + raft::device_vector_view labels, \ bool normalize_weight, \ raft::host_scalar_view inertia); -EXTERN_TEMPLATE_PREDICT(double, int) -EXTERN_TEMPLATE_PREDICT(double, int64_t) -EXTERN_TEMPLATE_PREDICT(float, int) -EXTERN_TEMPLATE_PREDICT(float, int64_t) +EXTERN_TEMPLATE_PREDICT(double, int, int) +EXTERN_TEMPLATE_PREDICT(double, int, int64_t) +EXTERN_TEMPLATE_PREDICT(double, int64_t, int) +EXTERN_TEMPLATE_PREDICT(double, int64_t, int64_t) +EXTERN_TEMPLATE_PREDICT(float, int, int) +EXTERN_TEMPLATE_PREDICT(float, int, int64_t) +EXTERN_TEMPLATE_PREDICT(float, int64_t, int) +EXTERN_TEMPLATE_PREDICT(float, int64_t, int64_t) #undef EXTERN_TEMPLATE_PREDICT diff --git a/cpp/src/cluster/kmeans_impl.cuh b/cpp/src/cluster/kmeans_impl.cuh index 49eafac5f2..df3b5eedab 100644 --- a/cpp/src/cluster/kmeans_impl.cuh +++ b/cpp/src/cluster/kmeans_impl.cuh @@ -24,17 +24,17 @@ void fit(raft::resources const& handle, handle, params, X, sample_weight, centroids, inertia, n_iter); } -template +template void predict(raft::resources const& handle, const kmeans::params& params, raft::device_matrix_view X, std::optional> sample_weight, raft::device_matrix_view centroids, - raft::device_vector_view labels, + raft::device_vector_view labels, bool normalize_weight, raft::host_scalar_view inertia) { - cuvs::cluster::kmeans::detail::kmeans_predict( + cuvs::cluster::kmeans::detail::kmeans_predict( handle, params, X, sample_weight, centroids, labels, normalize_weight, inertia); } diff --git a/cpp/src/cluster/kmeans_predict_double.cu b/cpp/src/cluster/kmeans_predict_double.cu index 52d120a232..31531ceee5 100644 --- a/cpp/src/cluster/kmeans_predict_double.cu +++ b/cpp/src/cluster/kmeans_predict_double.cu @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2024-2025, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ @@ -9,19 +9,21 @@ namespace cuvs::cluster::kmeans { -#define INSTANTIATE_PREDICT(DataT, IndexT) \ - template void predict( \ +#define INSTANTIATE_PREDICT(DataT, IndexT, LabelT) \ + template void predict( \ raft::resources const& handle, \ const kmeans::params& params, \ raft::device_matrix_view X, \ std::optional> sample_weight, \ raft::device_matrix_view centroids, \ - raft::device_vector_view labels, \ + raft::device_vector_view labels, \ bool normalize_weight, \ raft::host_scalar_view inertia); -INSTANTIATE_PREDICT(double, int) -INSTANTIATE_PREDICT(double, int64_t) +INSTANTIATE_PREDICT(double, int, int) +INSTANTIATE_PREDICT(double, int, int64_t) +INSTANTIATE_PREDICT(double, int64_t, int) +INSTANTIATE_PREDICT(double, int64_t, int64_t) #undef INSTANTIATE_PREDICT @@ -35,7 +37,33 @@ void predict(raft::resources const& handle, raft::host_scalar_view inertia) { - cuvs::cluster::kmeans::predict( + cuvs::cluster::kmeans::predict( + handle, params, X, sample_weight, centroids, labels, normalize_weight, inertia); +} + +void predict(raft::resources const& handle, + const kmeans::params& params, + raft::device_matrix_view X, + std::optional> sample_weight, + raft::device_matrix_view centroids, + raft::device_vector_view labels, + bool normalize_weight, + raft::host_scalar_view inertia) +{ + cuvs::cluster::kmeans::predict( + handle, params, X, sample_weight, centroids, labels, normalize_weight, inertia); +} + +void predict(raft::resources const& handle, + const kmeans::params& params, + raft::device_matrix_view X, + std::optional> sample_weight, + raft::device_matrix_view centroids, + raft::device_vector_view labels, + bool normalize_weight, + raft::host_scalar_view inertia) +{ + cuvs::cluster::kmeans::predict( handle, params, X, sample_weight, centroids, labels, normalize_weight, inertia); } @@ -49,7 +77,7 @@ void predict(raft::resources const& handle, raft::host_scalar_view inertia) { - cuvs::cluster::kmeans::predict( + cuvs::cluster::kmeans::predict( handle, params, X, sample_weight, centroids, labels, normalize_weight, inertia); } } // namespace cuvs::cluster::kmeans diff --git a/cpp/src/cluster/kmeans_predict_float.cu b/cpp/src/cluster/kmeans_predict_float.cu index 30812aa141..27d33d2e0e 100644 --- a/cpp/src/cluster/kmeans_predict_float.cu +++ b/cpp/src/cluster/kmeans_predict_float.cu @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2022-2025, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ @@ -9,19 +9,21 @@ namespace cuvs::cluster::kmeans { -#define INSTANTIATE_PREDICT(DataT, IndexT) \ - template void predict( \ +#define INSTANTIATE_PREDICT(DataT, IndexT, LabelT) \ + template void predict( \ raft::resources const& handle, \ const kmeans::params& params, \ raft::device_matrix_view X, \ std::optional> sample_weight, \ raft::device_matrix_view centroids, \ - raft::device_vector_view labels, \ + raft::device_vector_view labels, \ bool normalize_weight, \ raft::host_scalar_view inertia); -INSTANTIATE_PREDICT(float, int) -INSTANTIATE_PREDICT(float, int64_t) +INSTANTIATE_PREDICT(float, int, int) +INSTANTIATE_PREDICT(float, int, int64_t) +INSTANTIATE_PREDICT(float, int64_t, int) +INSTANTIATE_PREDICT(float, int64_t, int64_t) #undef INSTANTIATE_PREDICT @@ -35,7 +37,33 @@ void predict(raft::resources const& handle, raft::host_scalar_view inertia) { - cuvs::cluster::kmeans::predict( + cuvs::cluster::kmeans::predict( + handle, params, X, sample_weight, centroids, labels, normalize_weight, inertia); +} + +void predict(raft::resources const& handle, + const kmeans::params& params, + raft::device_matrix_view X, + std::optional> sample_weight, + raft::device_matrix_view centroids, + raft::device_vector_view labels, + bool normalize_weight, + raft::host_scalar_view inertia) +{ + cuvs::cluster::kmeans::predict( + handle, params, X, sample_weight, centroids, labels, normalize_weight, inertia); +} + +void predict(raft::resources const& handle, + const kmeans::params& params, + raft::device_matrix_view X, + std::optional> sample_weight, + raft::device_matrix_view centroids, + raft::device_vector_view labels, + bool normalize_weight, + raft::host_scalar_view inertia) +{ + cuvs::cluster::kmeans::predict( handle, params, X, sample_weight, centroids, labels, normalize_weight, inertia); } @@ -49,7 +77,7 @@ void predict(raft::resources const& handle, raft::host_scalar_view inertia) { - cuvs::cluster::kmeans::predict( + cuvs::cluster::kmeans::predict( handle, params, X, sample_weight, centroids, labels, normalize_weight, inertia); } } // namespace cuvs::cluster::kmeans diff --git a/cpp/src/distance/detail/fused_distance_nn.cuh b/cpp/src/distance/detail/fused_distance_nn.cuh index dbc87f468d..07ccc11cf4 100644 --- a/cpp/src/distance/detail/fused_distance_nn.cuh +++ b/cpp/src/distance/detail/fused_distance_nn.cuh @@ -63,6 +63,12 @@ void fusedDistanceNNImpl(IdxT* nearest_idx, } } + // InnerProduct is a cuTile-only specialization of this fused primitive. Callers that cannot + // launch cuTile must select their unfused InnerProduct path instead of returning an untouched + // sentinel from the legacy L2/cosine implementation below. + RAFT_EXPECTS(metric != cuvs::distance::DistanceType::InnerProduct, + "Fused InnerProduct 1-NN requires a compatible cuTile launcher"); + RAFT_EXPECTS(cutlass_kvp_scratch != nullptr, "CUTLASS fused 1-NN requires a scratch KVP buffer"); if (initOutBuffer) { @@ -113,7 +119,6 @@ void fusedDistanceNNImpl(IdxT* nearest_idx, cutlass_kvp_scratch, stream); break; - case cuvs::distance::DistanceType::InnerProduct: break; default: assert("only cosine/l2 metric is supported with fusedDistanceNN\n"); break; } diff --git a/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_kernel.py b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_kernel.py index dddd95954a..dd3013c80b 100644 --- a/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_kernel.py +++ b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_kernel.py @@ -53,6 +53,7 @@ def make_kernel( acc_dtype = ct.float32 idx_dtype = _idx_dtype(index_type) out_dist_dtype = ct.float16 if data_type == "half" else ct.float32 + l2_clamp_precision = 1e-3 if data_type == "half" else 1e-6 core_shape = (tile_m, tile_n) best_shape = (tile_m, 1) kernel_options = {} @@ -83,6 +84,11 @@ def fused_1nn_kernel( num_tiles_k = ct.num_tiles(A, axis=1, shape=(tm, tk)) num_tiles_n = ct.num_tiles(B, axis=0, shape=(tn, tk)) zero_pad = ct.PaddingMode.ZERO + a_norm = ct.zeros(best_shape, acc_dtype) + if metric_code != METRIC_INNER_PRODUCT: + a_norm = ct.load( + A_norm, index=(bidm,), shape=(tm,), padding_mode=zero_pad + )[:, None] def reduce_scores(dists, indices): def red_op(a_score, a_idx, b_score, b_idx): @@ -126,10 +132,14 @@ def red_op(a_score, a_idx, b_score, b_idx): B_norm, index=(n,), shape=(tn,), padding_mode=zero_pad ) if metric_code == METRIC_L2_EXPANDED: - # L2 receives squared row norms; cosine receives L2 magnitudes. - # The A norm is constant across centroids. Reduce - # 0.5 * ||y||^2 - dot(x, y), then recover full L2 once. - score = (0.5 * b_norm)[None, :] - accumulator + # Match the existing expanded-L2 epilogue: clamp negative + # distances and tiny self-neighbor roundoff before argmin. + score = a_norm + b_norm[None, :] - 2.0 * accumulator + self_roundoff = (score * score < l2_clamp_precision) & ( + a_norm == b_norm[None, :] + ) + score = ct.where(score > 0.0, score, 0.0) + score = ct.where(self_roundoff, 0.0, score) else: # Defer the A-norm division until after selecting the # winning centroid. @@ -149,11 +159,8 @@ def red_op(a_score, a_idx, b_score, b_idx): if metric_code == METRIC_INNER_PRODUCT: out_dist = -best_dist else: - a_norm = ct.load( - A_norm, index=(bidm,), shape=(tm,), padding_mode=zero_pad - )[:, None] if metric_code == METRIC_L2_EXPANDED: - out_dist = a_norm + 2.0 * best_dist + out_dist = best_dist out_dist = ct.where( apply_sqrt != 0, ct.sqrt(out_dist), out_dist ) diff --git a/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_tile.cu b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_tile.cu index fa3e4baafa..eb576a7318 100644 --- a/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_tile.cu +++ b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_tile.cu @@ -25,6 +25,20 @@ namespace detail { namespace { +bool is_16_byte_aligned(const void* ptr) +{ + return ptr == nullptr || reinterpret_cast(ptr) % 16 == 0; +} + +template +bool has_fused_1nn_tile_launcher() +{ + Fused1nnTilePlanner planner; + planner.add_entrypoint(); + planner.add_tileir_fallback(); + return planner.try_get_launcher() != nullptr; +} + template bool launch_fused_1nn_tile(IdxT* nearest_idx, DataT* nearest_dist, @@ -46,9 +60,9 @@ bool launch_fused_1nn_tile(IdxT* nearest_idx, Fused1nnTilePlanner planner; planner.add_entrypoint(); planner.add_tileir_fallback(); - const cuvs::detail::jit_lto::CutileTileConfig tile_cfg = planner.tile_config(); - auto launcher = planner.try_get_launcher(); + auto launcher = planner.try_get_launcher(); if (!launcher) { return false; } + const cuvs::detail::jit_lto::CutileTileConfig tile_cfg = planner.tile_config(); int metric_code; bool apply_sqrt = false; @@ -180,6 +194,52 @@ bool try_fused_1nn_tile_dispatch(IdxT* nearest_idx, } // namespace +template + requires is_fused_1nn_cutile_data_v +bool can_launch_fused_1nn_tile(IdxT* nearest_idx, + DataT* nearest_dist, + const DataT* x, + const DataT* y, + const DataT* xn, + const DataT* yn, + IdxT m, + IdxT n, + IdxT k, + cuvs::distance::DistanceType metric) +{ + if (!cuvs::detail::jit_lto::cutile_launch_available_on_current_device()) { return false; } + static_assert(std::is_same_v || std::is_same_v); + + if (nearest_dist == nullptr || x == nullptr || y == nullptr || m <= 0 || n <= 0 || k <= 0) { + return false; + } + if (metric != cuvs::distance::DistanceType::InnerProduct && + metric != cuvs::distance::DistanceType::L2Expanded && + metric != cuvs::distance::DistanceType::L2SqrtExpanded && + metric != cuvs::distance::DistanceType::CosineExpanded) { + return false; + } + if (metric != cuvs::distance::DistanceType::InnerProduct && (xn == nullptr || yn == nullptr)) { + return false; + } + + // Both exported ABIs promise 16-byte base alignment for every array parameter. + if (!is_16_byte_aligned(nearest_dist) || !is_16_byte_aligned(x) || !is_16_byte_aligned(y) || + !is_16_byte_aligned(xn) || !is_16_byte_aligned(yn)) { + return false; + } + if constexpr (std::is_same_v) { + if (!is_16_byte_aligned(nearest_idx)) { return false; } + } else { + constexpr int64_t max_i32 = std::numeric_limits::max(); + if (n > max_i32 || k > max_i32) { return false; } + } + + constexpr int strict_pitch_elements = 16 / sizeof(DataT); + return k % strict_pitch_elements == 0 ? has_fused_1nn_tile_launcher() + : has_fused_1nn_tile_launcher(); +} + template requires is_fused_1nn_cutile_data_v bool try_fused_1nn_tile(IdxT* nearest_idx, @@ -196,8 +256,9 @@ bool try_fused_1nn_tile(IdxT* nearest_idx, void* index_workspace, cudaStream_t stream) { - if (!cuvs::detail::jit_lto::cutile_launch_available_on_current_device()) { return false; } - static_assert(std::is_same_v || std::is_same_v); + if (!can_launch_fused_1nn_tile(nearest_idx, nearest_dist, x, y, xn, yn, m, n, k, metric)) { + return false; + } constexpr int strict_pitch_elements = 16 / sizeof(DataT); const bool use_strict_abi = k % strict_pitch_elements == 0; @@ -210,13 +271,16 @@ bool try_fused_1nn_tile(IdxT* nearest_idx, return try_fused_1nn_tile_dispatch( nearest_idx, nearest_dist, x, y, xn, yn, m, n, k, metric, is_sqrt, stream); } else { - constexpr int64_t max_i32 = std::numeric_limits::max(); - if (n > max_i32 || k > max_i32) { return false; } if (nearest_idx != nullptr && index_workspace == nullptr) { return false; } + if (!is_16_byte_aligned(index_workspace)) { return false; } - auto* tmp_idx = static_cast(index_workspace); - for (int64_t offset = 0; offset < m; offset += max_i32) { - const int batch_m = static_cast(std::min(max_i32, m - offset)); + // Keep every chunk offset 16-byte aligned for x, xn, and nearest_dist. + constexpr int64_t max_i32 = std::numeric_limits::max(); + constexpr int64_t batch_alignment = 16 / sizeof(DataT); + constexpr int64_t max_batch_m = max_i32 - max_i32 % batch_alignment; + auto* tmp_idx = static_cast(index_workspace); + for (int64_t offset = 0; offset < m; offset += max_batch_m) { + const int batch_m = static_cast(std::min(max_batch_m, m - offset)); const auto* batch_x = x + static_cast(offset) * static_cast(k); const auto* batch_xn = xn == nullptr ? nullptr : xn + offset; auto* batch_dist = nearest_dist == nullptr ? nullptr : nearest_dist + offset; @@ -258,6 +322,25 @@ bool try_fused_1nn_tile(IdxT* nearest_idx, } } +#define CUVS_INST_CAN_LAUNCH_FUSED_1NN_TILE(DataT, IdxT) \ + template CUVS_EXPORT bool can_launch_fused_1nn_tile(IdxT*, \ + DataT*, \ + const DataT*, \ + const DataT*, \ + const DataT*, \ + const DataT*, \ + IdxT, \ + IdxT, \ + IdxT, \ + cuvs::distance::DistanceType) + +CUVS_INST_CAN_LAUNCH_FUSED_1NN_TILE(float, int); +CUVS_INST_CAN_LAUNCH_FUSED_1NN_TILE(float, int64_t); +CUVS_INST_CAN_LAUNCH_FUSED_1NN_TILE(half, int); +CUVS_INST_CAN_LAUNCH_FUSED_1NN_TILE(half, int64_t); + +#undef CUVS_INST_CAN_LAUNCH_FUSED_1NN_TILE + #define CUVS_INST_TRY_FUSED_1NN_TILE(DataT, IdxT) \ template CUVS_EXPORT bool try_fused_1nn_tile(IdxT*, \ DataT*, \ diff --git a/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_tile.hpp b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_tile.hpp index 7cdbabd411..3adf58c3d7 100644 --- a/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_tile.hpp +++ b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_tile.hpp @@ -25,6 +25,26 @@ inline constexpr bool is_fused_1nn_cutile_data_v = std::is_same_v || std::is_same_v; #if CUVS_CUTILE_ENABLED +/** + * Return whether the supplied problem can use cuTile without fallback scratch. + * + * The result includes runtime/device support, exported ABI constraints, and launcher construction. + * A successful probe populates the shared launcher cache used by try_fused_1nn_tile. + * An int64 output index still requires an int32 workspace with `m` elements. + */ +template + requires is_fused_1nn_cutile_data_v +bool can_launch_fused_1nn_tile(IdxT* nearest_idx, + DataT* nearest_dist, + const DataT* x, + const DataT* y, + const DataT* xn, + const DataT* yn, + IdxT m, + IdxT n, + IdxT k, + cuvs::distance::DistanceType metric); + template requires is_fused_1nn_cutile_data_v bool try_fused_1nn_tile(IdxT* nearest_idx, @@ -41,6 +61,21 @@ bool try_fused_1nn_tile(IdxT* nearest_idx, void* index_workspace, cudaStream_t stream); #else +template +bool can_launch_fused_1nn_tile(IdxT*, + DataT*, + const DataT*, + const DataT*, + const DataT*, + const DataT*, + IdxT, + IdxT, + IdxT, + cuvs::distance::DistanceType) +{ + return false; +} + template bool try_fused_1nn_tile(IdxT*, DataT*, diff --git a/cpp/src/distance/fused_distance_nn-inl.cuh b/cpp/src/distance/fused_distance_nn-inl.cuh index ccaef64319..26a931754e 100644 --- a/cpp/src/distance/fused_distance_nn-inl.cuh +++ b/cpp/src/distance/fused_distance_nn-inl.cuh @@ -213,7 +213,8 @@ void fusedDistanceNN(IdxT* nearest_idx, * @param[out] nearest_idx Nearest neighbor index per row, length `m` (required). * @param[out] nearest_dist Minimum distance per row, length `m` (optional, may be null). * @param[in] cutlass_kvp_scratch Temp KVP buffer, length `m`; required when CUTLASS/SIMT runs. - * Unused when cuTile handles the launch. + * It may be omitted only after + * detail::can_launch_fused_1nn_tile succeeds. */ template void fusedDistanceNNMinReduce(IdxT* nearest_idx, diff --git a/cpp/tests/neighbors/distance_nn.cu b/cpp/tests/neighbors/distance_nn.cu index 64c07b4058..1bd800a6cf 100644 --- a/cpp/tests/neighbors/distance_nn.cu +++ b/cpp/tests/neighbors/distance_nn.cu @@ -264,6 +264,54 @@ TEST_P(NNTest_fp16_unfused, test) INSTANTIATE_TEST_CASE_P(NNTest, NNTest_fp16_unfused, ::testing::ValuesIn(input_fp16)); +TEST(Fused1nn, ExpandedL2ClampsNegativeRoundoff) +{ + raft::resources handle; + auto stream = raft::resource::get_cuda_stream(handle); + constexpr int k = 64; + + auto x = raft::make_device_matrix(handle, 1, k); + auto y = raft::make_device_matrix(handle, 1, k); + auto x_norm = raft::make_device_vector(handle, 1); + auto y_norm = raft::make_device_vector(handle, 1); + auto out_idx = raft::make_device_vector(handle, 1); + auto out_dist = raft::make_device_vector(handle, 1); + auto out_kvp = raft::make_device_vector, int>(handle, 1); + auto workspace = raft::make_device_vector(handle, 1); + + raft::matrix::fill(handle, x.view(), 1.0006f); + raft::copy(y.data_handle(), x.data_handle(), k, stream); + raft::linalg::rowNorm( + x_norm.data_handle(), x.data_handle(), k, 1, stream); + raft::copy(y_norm.data_handle(), x_norm.data_handle(), 1, stream); + + cuvs::distance::fusedDistanceNNMinReduce(out_idx.data_handle(), + out_dist.data_handle(), + x.data_handle(), + y.data_handle(), + x_norm.data_handle(), + y_norm.data_handle(), + 1, + 1, + k, + workspace.data_handle(), + true, + true, + true, + DistanceType::L2SqrtExpanded, + 0.0f, + out_kvp.data_handle(), + stream); + + int actual_idx; + float actual_dist; + raft::update_host(&actual_idx, out_idx.data_handle(), 1, stream); + raft::update_host(&actual_dist, out_dist.data_handle(), 1, stream); + raft::resource::sync_stream(handle); + EXPECT_EQ(actual_idx, 0); + EXPECT_EQ(actual_dist, 0.0f); +} + template const std::vector> input_int8 = { {4096, 4096, 64, DistanceType::L2Expanded, false, uint64_t(31415926), 0.1}, From ab012a68d0e855a67a78d8917d5cde1b8edf4d02 Mon Sep 17 00:00:00 2001 From: divyegala Date: Fri, 28 Aug 2026 03:35:33 +0000 Subject: [PATCH 63/82] fix fp16 norms --- cpp/src/distance/detail/fused_distance_nn.cuh | 4 +- .../cutile/export_fused_1nn.py | 3 +- .../cutile/fused_1nn_kernel.py | 2 +- .../cutile/fused_1nn_tile.cu | 69 ++++++++++--------- .../cutile/fused_1nn_tile.hpp | 20 +++--- 5 files changed, 53 insertions(+), 45 deletions(-) diff --git a/cpp/src/distance/detail/fused_distance_nn.cuh b/cpp/src/distance/detail/fused_distance_nn.cuh index 07ccc11cf4..669bb05f1f 100644 --- a/cpp/src/distance/detail/fused_distance_nn.cuh +++ b/cpp/src/distance/detail/fused_distance_nn.cuh @@ -23,6 +23,7 @@ #include // size_t #include // std::numeric_limits +#include namespace cuvs { namespace distance { @@ -54,7 +55,8 @@ void fusedDistanceNNImpl(IdxT* nearest_idx, typedef raft::KeyValuePair KVP; constexpr auto maxVal = std::numeric_limits::max(); - if constexpr (is_fused_1nn_cutile_data_v) { + if constexpr (is_fused_1nn_cutile_data_v && + std::is_same_v, DataT>) { if constexpr (cuvs::detail::jit_lto::library_built_with_cutile()) { if (try_fused_1nn_tile( nearest_idx, nearest_dist, x, y, xn, yn, m, n, k, metric, sqrt, workspace, stream)) { diff --git a/cpp/src/distance/detail/fused_distance_nn/cutile/export_fused_1nn.py b/cpp/src/distance/detail/fused_distance_nn/cutile/export_fused_1nn.py index 011ff27636..81411e72f7 100644 --- a/cpp/src/distance/detail/fused_distance_nn/cutile/export_fused_1nn.py +++ b/cpp/src/distance/detail/fused_distance_nn/cutile/export_fused_1nn.py @@ -155,7 +155,8 @@ def _kernel_signature( matrix_layout == "strict" and gpu_code in ("sm_80", "sm_86") ), ) - norm_array = _cuvs_vector_constraint(elem, index_dtype=idx_dtype) + norm_elem = ct.float32 if data_type == "half" else elem + norm_array = _cuvs_vector_constraint(norm_elem, index_dtype=idx_dtype) idx_array = _cuvs_vector_constraint(idx_dtype, index_dtype=idx_dtype) dist_array = _cuvs_vector_constraint(elem, index_dtype=idx_dtype) diff --git a/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_kernel.py b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_kernel.py index dd3013c80b..87713b5ef1 100644 --- a/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_kernel.py +++ b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_kernel.py @@ -88,7 +88,7 @@ def fused_1nn_kernel( if metric_code != METRIC_INNER_PRODUCT: a_norm = ct.load( A_norm, index=(bidm,), shape=(tm,), padding_mode=zero_pad - )[:, None] + )[:, None].astype(acc_dtype) def reduce_scores(dists, indices): def red_op(a_score, a_idx, b_score, b_idx): diff --git a/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_tile.cu b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_tile.cu index eb576a7318..db1ee7bdb4 100644 --- a/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_tile.cu +++ b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_tile.cu @@ -44,8 +44,8 @@ bool launch_fused_1nn_tile(IdxT* nearest_idx, DataT* nearest_dist, const DataT* x, const DataT* y, - const DataT* xn, - const DataT* yn, + const fused_1nn_cutile_norm_t* xn, + const fused_1nn_cutile_norm_t* yn, IdxT m, IdxT n, IdxT k, @@ -100,8 +100,8 @@ bool launch_fused_1nn_tile(IdxT* nearest_idx, void* x_ptr = const_cast(x); void* y_ptr = const_cast(y); - void* xn_ptr = const_cast(xn); - void* yn_ptr = const_cast(yn); + void* xn_ptr = const_cast*>(xn); + void* yn_ptr = const_cast*>(yn); const IdxT store_idx = nearest_idx != nullptr ? IdxT{1} : IdxT{0}; void* idx_ptr = nearest_idx; void* dist_ptr = nearest_dist; @@ -179,8 +179,8 @@ bool try_fused_1nn_tile_dispatch(IdxT* nearest_idx, DataT* nearest_dist, const DataT* x, const DataT* y, - const DataT* xn, - const DataT* yn, + const fused_1nn_cutile_norm_t* xn, + const fused_1nn_cutile_norm_t* yn, IdxT m, IdxT n, IdxT k, @@ -200,8 +200,8 @@ bool can_launch_fused_1nn_tile(IdxT* nearest_idx, DataT* nearest_dist, const DataT* x, const DataT* y, - const DataT* xn, - const DataT* yn, + const fused_1nn_cutile_norm_t* xn, + const fused_1nn_cutile_norm_t* yn, IdxT m, IdxT n, IdxT k, @@ -246,8 +246,8 @@ bool try_fused_1nn_tile(IdxT* nearest_idx, DataT* nearest_dist, const DataT* x, const DataT* y, - const DataT* xn, - const DataT* yn, + const fused_1nn_cutile_norm_t* xn, + const fused_1nn_cutile_norm_t* yn, IdxT m, IdxT n, IdxT k, @@ -322,17 +322,18 @@ bool try_fused_1nn_tile(IdxT* nearest_idx, } } -#define CUVS_INST_CAN_LAUNCH_FUSED_1NN_TILE(DataT, IdxT) \ - template CUVS_EXPORT bool can_launch_fused_1nn_tile(IdxT*, \ - DataT*, \ - const DataT*, \ - const DataT*, \ - const DataT*, \ - const DataT*, \ - IdxT, \ - IdxT, \ - IdxT, \ - cuvs::distance::DistanceType) +#define CUVS_INST_CAN_LAUNCH_FUSED_1NN_TILE(DataT, IdxT) \ + template CUVS_EXPORT bool can_launch_fused_1nn_tile( \ + IdxT*, \ + DataT*, \ + const DataT*, \ + const DataT*, \ + const fused_1nn_cutile_norm_t*, \ + const fused_1nn_cutile_norm_t*, \ + IdxT, \ + IdxT, \ + IdxT, \ + cuvs::distance::DistanceType) CUVS_INST_CAN_LAUNCH_FUSED_1NN_TILE(float, int); CUVS_INST_CAN_LAUNCH_FUSED_1NN_TILE(float, int64_t); @@ -341,19 +342,19 @@ CUVS_INST_CAN_LAUNCH_FUSED_1NN_TILE(half, int64_t); #undef CUVS_INST_CAN_LAUNCH_FUSED_1NN_TILE -#define CUVS_INST_TRY_FUSED_1NN_TILE(DataT, IdxT) \ - template CUVS_EXPORT bool try_fused_1nn_tile(IdxT*, \ - DataT*, \ - const DataT*, \ - const DataT*, \ - const DataT*, \ - const DataT*, \ - IdxT, \ - IdxT, \ - IdxT, \ - cuvs::distance::DistanceType, \ - bool, \ - void*, \ +#define CUVS_INST_TRY_FUSED_1NN_TILE(DataT, IdxT) \ + template CUVS_EXPORT bool try_fused_1nn_tile(IdxT*, \ + DataT*, \ + const DataT*, \ + const DataT*, \ + const fused_1nn_cutile_norm_t*, \ + const fused_1nn_cutile_norm_t*, \ + IdxT, \ + IdxT, \ + IdxT, \ + cuvs::distance::DistanceType, \ + bool, \ + void*, \ cudaStream_t) CUVS_INST_TRY_FUSED_1NN_TILE(float, int); diff --git a/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_tile.hpp b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_tile.hpp index 3adf58c3d7..6aeecf8aae 100644 --- a/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_tile.hpp +++ b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_tile.hpp @@ -24,6 +24,10 @@ template inline constexpr bool is_fused_1nn_cutile_data_v = std::is_same_v || std::is_same_v; +// Tensor-core products accumulate in FP32; FP16 norms must remain FP32 through the epilogue. +template +using fused_1nn_cutile_norm_t = std::conditional_t, float, DataT>; + #if CUVS_CUTILE_ENABLED /** * Return whether the supplied problem can use cuTile without fallback scratch. @@ -38,8 +42,8 @@ bool can_launch_fused_1nn_tile(IdxT* nearest_idx, DataT* nearest_dist, const DataT* x, const DataT* y, - const DataT* xn, - const DataT* yn, + const fused_1nn_cutile_norm_t* xn, + const fused_1nn_cutile_norm_t* yn, IdxT m, IdxT n, IdxT k, @@ -51,8 +55,8 @@ bool try_fused_1nn_tile(IdxT* nearest_idx, DataT* nearest_dist, const DataT* x, const DataT* y, - const DataT* xn, - const DataT* yn, + const fused_1nn_cutile_norm_t* xn, + const fused_1nn_cutile_norm_t* yn, IdxT m, IdxT n, IdxT k, @@ -66,8 +70,8 @@ bool can_launch_fused_1nn_tile(IdxT*, DataT*, const DataT*, const DataT*, - const DataT*, - const DataT*, + const fused_1nn_cutile_norm_t*, + const fused_1nn_cutile_norm_t*, IdxT, IdxT, IdxT, @@ -81,8 +85,8 @@ bool try_fused_1nn_tile(IdxT*, DataT*, const DataT*, const DataT*, - const DataT*, - const DataT*, + const fused_1nn_cutile_norm_t*, + const fused_1nn_cutile_norm_t*, IdxT, IdxT, IdxT, From 7aa42b8508473d63163ca33441309e0866e65fb3 Mon Sep 17 00:00:00 2001 From: divyegala Date: Tue, 1 Sep 2026 01:28:08 +0000 Subject: [PATCH 64/82] clamp only the winner --- .../cutile/fused_1nn_kernel.py | 25 ++++++++----------- 1 file changed, 10 insertions(+), 15 deletions(-) diff --git a/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_kernel.py b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_kernel.py index 87713b5ef1..ad1ba8fbea 100644 --- a/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_kernel.py +++ b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_kernel.py @@ -53,7 +53,6 @@ def make_kernel( acc_dtype = ct.float32 idx_dtype = _idx_dtype(index_type) out_dist_dtype = ct.float16 if data_type == "half" else ct.float32 - l2_clamp_precision = 1e-3 if data_type == "half" else 1e-6 core_shape = (tile_m, tile_n) best_shape = (tile_m, 1) kernel_options = {} @@ -84,11 +83,6 @@ def fused_1nn_kernel( num_tiles_k = ct.num_tiles(A, axis=1, shape=(tm, tk)) num_tiles_n = ct.num_tiles(B, axis=0, shape=(tn, tk)) zero_pad = ct.PaddingMode.ZERO - a_norm = ct.zeros(best_shape, acc_dtype) - if metric_code != METRIC_INNER_PRODUCT: - a_norm = ct.load( - A_norm, index=(bidm,), shape=(tm,), padding_mode=zero_pad - )[:, None].astype(acc_dtype) def reduce_scores(dists, indices): def red_op(a_score, a_idx, b_score, b_idx): @@ -132,14 +126,9 @@ def red_op(a_score, a_idx, b_score, b_idx): B_norm, index=(n,), shape=(tn,), padding_mode=zero_pad ) if metric_code == METRIC_L2_EXPANDED: - # Match the existing expanded-L2 epilogue: clamp negative - # distances and tiny self-neighbor roundoff before argmin. - score = a_norm + b_norm[None, :] - 2.0 * accumulator - self_roundoff = (score * score < l2_clamp_precision) & ( - a_norm == b_norm[None, :] - ) - score = ct.where(score > 0.0, score, 0.0) - score = ct.where(self_roundoff, 0.0, score) + # The A norm is constant across centroids. Excluding it + # avoids cancellation in the score used by argmin. + score = (0.5 * b_norm)[None, :] - accumulator else: # Defer the A-norm division until after selecting the # winning centroid. @@ -159,8 +148,14 @@ def red_op(a_score, a_idx, b_score, b_idx): if metric_code == METRIC_INNER_PRODUCT: out_dist = -best_dist else: + a_norm = ct.load( + A_norm, index=(bidm,), shape=(tm,), padding_mode=zero_pad + )[:, None] if metric_code == METRIC_L2_EXPANDED: - out_dist = best_dist + out_dist = a_norm + 2.0 * best_dist + # Separately reduced norms and the MMA can reconstruct a + # slightly negative distance; clamp before an optional sqrt. + out_dist = ct.where(out_dist > 0.0, out_dist, 0.0) out_dist = ct.where( apply_sqrt != 0, ct.sqrt(out_dist), out_dist ) From 45c66a83696c95e945f649b734a71d294866e73b Mon Sep 17 00:00:00 2001 From: divyegala Date: Tue, 1 Sep 2026 20:53:34 +0000 Subject: [PATCH 65/82] regression and test fixes --- .../detail/minClusterDistanceCompute.cu | 151 ++++++++---------- .../cutile/fused_1nn_tile.cu | 42 ++++- .../cutile/fused_1nn_tile.hpp | 24 +++ .../predicated_tile_iterator_reduced_vec.h | 4 +- .../neighbors/ivf_flat/ivf_flat_search.cuh | 24 ++- 5 files changed, 154 insertions(+), 91 deletions(-) diff --git a/cpp/src/cluster/detail/minClusterDistanceCompute.cu b/cpp/src/cluster/detail/minClusterDistanceCompute.cu index 7570e15adb..e09cba9e70 100644 --- a/cpp/src/cluster/detail/minClusterDistanceCompute.cu +++ b/cpp/src/cluster/detail/minClusterDistanceCompute.cu @@ -118,76 +118,68 @@ void minClusterAndDistanceCompute(raft::resources const& handle, metric == cuvs::distance::DistanceType::CosineExpanded; FusedDistancePath fused_path = use_fused(handle, n_samples, n_clusters, n_features, metric); + bool cutile_ready = false; if constexpr (is_cutile_fused_data_type_v) { - if (fused_path == FusedDistancePath::FusedCutile && - metric == cuvs::distance::DistanceType::InnerProduct && - !cuvs::distance::detail::can_launch_fused_1nn_tile(nearest_idx.data_handle(), - nearest_dist.data_handle(), - X.data_handle(), - centroids.data_handle(), - static_cast(nullptr), - static_cast(nullptr), - n_samples, - n_clusters, - n_features, - metric)) { - fused_path = FusedDistancePath::Unfused; + if (fused_path == FusedDistancePath::FusedCutile) { + cutile_ready = cuvs::distance::detail::can_launch_fused_1nn_tile(nearest_idx.data_handle(), + nearest_dist.data_handle(), + X.data_handle(), + centroids.data_handle(), + n_samples, + n_clusters, + n_features, + metric); + if (!cutile_ready) { + fused_path = metric == cuvs::distance::DistanceType::InnerProduct + ? FusedDistancePath::Unfused + : FusedDistancePath::FusedCutlass; + } } } if (uses_fused_distance_nn(fused_path)) { - const DataT* x_norm_ptr = L2NormX.data_handle(); - const DataT* centroids_norm_ptr; - if constexpr (std::is_same_v) { - if (fused_path == FusedDistancePath::FusedCutile && is_l2_cos) { - constexpr size_t norm_alignment = 16 / sizeof(float); - const size_t x_norm_storage = raft::alignTo(static_cast(n_samples), norm_alignment); - L2NormBuf_OR_DistBuf.resize(x_norm_storage + static_cast(n_clusters), stream); - auto* tf32_x_norms = L2NormBuf_OR_DistBuf.data(); - auto* tf32_centroid_norms = tf32_x_norms + x_norm_storage; - const bool take_sqrt = metric == cuvs::distance::DistanceType::CosineExpanded; - compute_tf32_row_norms( - handle, X.data_handle(), tf32_x_norms, n_samples, n_features, take_sqrt); - compute_tf32_row_norms( - handle, centroids.data_handle(), tf32_centroid_norms, n_clusters, n_features, take_sqrt); - x_norm_ptr = tf32_x_norms; - centroids_norm_ptr = tf32_centroid_norms; + const DataT* x_norm_ptr = nullptr; + const DataT* centroids_norm_ptr = nullptr; + if (is_l2_cos) { + x_norm_ptr = L2NormX.data_handle(); + if constexpr (std::is_same_v) { + if (cutile_ready) { + constexpr size_t norm_alignment = 16 / sizeof(float); + const size_t x_norm_storage = + raft::alignTo(static_cast(n_samples), norm_alignment); + L2NormBuf_OR_DistBuf.resize(x_norm_storage + static_cast(n_clusters), stream); + auto* tf32_x_norms = L2NormBuf_OR_DistBuf.data(); + auto* tf32_centroid_norms = tf32_x_norms + x_norm_storage; + const bool take_sqrt = metric == cuvs::distance::DistanceType::CosineExpanded; + compute_tf32_row_norms( + handle, X.data_handle(), tf32_x_norms, n_samples, n_features, take_sqrt); + compute_tf32_row_norms(handle, + centroids.data_handle(), + tf32_centroid_norms, + n_clusters, + n_features, + take_sqrt); + x_norm_ptr = tf32_x_norms; + centroids_norm_ptr = tf32_centroid_norms; + } else { + L2NormBuf_OR_DistBuf.resize(n_clusters, stream); + centroids_norm_ptr = L2NormBuf_OR_DistBuf.data(); + } } else { L2NormBuf_OR_DistBuf.resize(n_clusters, stream); centroids_norm_ptr = L2NormBuf_OR_DistBuf.data(); } - } else { - L2NormBuf_OR_DistBuf.resize(n_clusters, stream); - centroids_norm_ptr = L2NormBuf_OR_DistBuf.data(); - } - - if (!(fused_path == FusedDistancePath::FusedCutile && is_l2_cos && - std::is_same_v) && - is_l2_cos) { - auto centroids_norm = - raft::make_device_vector_view(L2NormBuf_OR_DistBuf.data(), n_clusters); - if (metric == cuvs::distance::DistanceType::CosineExpanded) { - raft::linalg::norm( - handle, centroids, centroids_norm, raft::sqrt_op{}); - } else { - raft::linalg::norm( - handle, centroids, centroids_norm); - } - } - bool cutile_ready = false; - if constexpr (is_cutile_fused_data_type_v) { - if (fused_path == FusedDistancePath::FusedCutile) { - cutile_ready = cuvs::distance::detail::can_launch_fused_1nn_tile(nearest_idx.data_handle(), - nearest_dist.data_handle(), - X.data_handle(), - centroids.data_handle(), - x_norm_ptr, - centroids_norm_ptr, - n_samples, - n_clusters, - n_features, - metric); + if (!cutile_ready) { + auto centroids_norm = + raft::make_device_vector_view(L2NormBuf_OR_DistBuf.data(), n_clusters); + if (metric == cuvs::distance::DistanceType::CosineExpanded) { + raft::linalg::norm( + handle, centroids, centroids_norm, raft::sqrt_op{}); + } else { + raft::linalg::norm( + handle, centroids, centroids_norm); + } } } @@ -402,15 +394,30 @@ void minClusterDistanceCompute(raft::resources const& handle, raft::matrix::fill(handle, minClusterDistance, std::numeric_limits::max()); - const FusedDistancePath fused_path = + FusedDistancePath fused_path = is_l2_cos ? use_fused(handle, n_samples, n_clusters, n_features, metric) : FusedDistancePath::Unfused; + bool cutile_ready = false; + if constexpr (is_cutile_fused_data_type_v) { + if (fused_path == FusedDistancePath::FusedCutile) { + cutile_ready = + cuvs::distance::detail::can_launch_fused_1nn_tile(static_cast(nullptr), + minClusterDistance.data_handle(), + X.data_handle(), + centroids.data_handle(), + n_samples, + n_clusters, + n_features, + metric); + if (!cutile_ready) { fused_path = FusedDistancePath::FusedCutlass; } + } + } if (uses_fused_distance_nn(fused_path)) { const DataT* x_norm_ptr = L2NormX.data_handle(); const DataT* centroids_norm_ptr; if constexpr (std::is_same_v) { - if (fused_path == FusedDistancePath::FusedCutile && is_l2_cos) { + if (cutile_ready) { constexpr size_t norm_alignment = 16 / sizeof(float); const size_t x_norm_storage = raft::alignTo(static_cast(n_samples), norm_alignment); L2NormBuf_OR_DistBuf.resize(x_norm_storage + static_cast(n_clusters), stream); @@ -432,8 +439,7 @@ void minClusterDistanceCompute(raft::resources const& handle, centroids_norm_ptr = L2NormBuf_OR_DistBuf.data(); } - if (!(fused_path == FusedDistancePath::FusedCutile && is_l2_cos && - std::is_same_v)) { + if (!cutile_ready) { auto centroids_norm = raft::make_device_vector_view(L2NormBuf_OR_DistBuf.data(), n_clusters); if (metric == cuvs::distance::DistanceType::CosineExpanded) { @@ -452,23 +458,6 @@ void minClusterDistanceCompute(raft::resources const& handle, } } - bool cutile_ready = false; - if constexpr (is_cutile_fused_data_type_v) { - if (fused_path == FusedDistancePath::FusedCutile) { - cutile_ready = - cuvs::distance::detail::can_launch_fused_1nn_tile(static_cast(nullptr), - minClusterDistance.data_handle(), - X.data_handle(), - centroids.data_handle(), - x_norm_ptr, - centroids_norm_ptr, - n_samples, - n_clusters, - n_features, - metric); - } - } - raft::KeyValuePair* cutlass_kvp_scratch = nullptr; rmm::device_uvector> temp_kvp(0, stream); if (!cutile_ready) { diff --git a/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_tile.cu b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_tile.cu index db1ee7bdb4..5593949454 100644 --- a/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_tile.cu +++ b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_tile.cu @@ -200,8 +200,6 @@ bool can_launch_fused_1nn_tile(IdxT* nearest_idx, DataT* nearest_dist, const DataT* x, const DataT* y, - const fused_1nn_cutile_norm_t* xn, - const fused_1nn_cutile_norm_t* yn, IdxT m, IdxT n, IdxT k, @@ -219,13 +217,8 @@ bool can_launch_fused_1nn_tile(IdxT* nearest_idx, metric != cuvs::distance::DistanceType::CosineExpanded) { return false; } - if (metric != cuvs::distance::DistanceType::InnerProduct && (xn == nullptr || yn == nullptr)) { - return false; - } - // Both exported ABIs promise 16-byte base alignment for every array parameter. - if (!is_16_byte_aligned(nearest_dist) || !is_16_byte_aligned(x) || !is_16_byte_aligned(y) || - !is_16_byte_aligned(xn) || !is_16_byte_aligned(yn)) { + if (!is_16_byte_aligned(nearest_dist) || !is_16_byte_aligned(x) || !is_16_byte_aligned(y)) { return false; } if constexpr (std::is_same_v) { @@ -240,6 +233,28 @@ bool can_launch_fused_1nn_tile(IdxT* nearest_idx, : has_fused_1nn_tile_launcher(); } +template + requires is_fused_1nn_cutile_data_v +bool can_launch_fused_1nn_tile(IdxT* nearest_idx, + DataT* nearest_dist, + const DataT* x, + const DataT* y, + const fused_1nn_cutile_norm_t* xn, + const fused_1nn_cutile_norm_t* yn, + IdxT m, + IdxT n, + IdxT k, + cuvs::distance::DistanceType metric) +{ + if (!can_launch_fused_1nn_tile(nearest_idx, nearest_dist, x, y, m, n, k, metric)) { + return false; + } + if (metric != cuvs::distance::DistanceType::InnerProduct && (xn == nullptr || yn == nullptr)) { + return false; + } + return is_16_byte_aligned(xn) && is_16_byte_aligned(yn); +} + template requires is_fused_1nn_cutile_data_v bool try_fused_1nn_tile(IdxT* nearest_idx, @@ -322,6 +337,17 @@ bool try_fused_1nn_tile(IdxT* nearest_idx, } } +#define CUVS_INST_CAN_LAUNCH_FUSED_1NN_TILE_PREFLIGHT(DataT, IdxT) \ + template CUVS_EXPORT bool can_launch_fused_1nn_tile( \ + IdxT*, DataT*, const DataT*, const DataT*, IdxT, IdxT, IdxT, cuvs::distance::DistanceType) + +CUVS_INST_CAN_LAUNCH_FUSED_1NN_TILE_PREFLIGHT(float, int); +CUVS_INST_CAN_LAUNCH_FUSED_1NN_TILE_PREFLIGHT(float, int64_t); +CUVS_INST_CAN_LAUNCH_FUSED_1NN_TILE_PREFLIGHT(half, int); +CUVS_INST_CAN_LAUNCH_FUSED_1NN_TILE_PREFLIGHT(half, int64_t); + +#undef CUVS_INST_CAN_LAUNCH_FUSED_1NN_TILE_PREFLIGHT + #define CUVS_INST_CAN_LAUNCH_FUSED_1NN_TILE(DataT, IdxT) \ template CUVS_EXPORT bool can_launch_fused_1nn_tile( \ IdxT*, \ diff --git a/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_tile.hpp b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_tile.hpp index 6aeecf8aae..245a2c6afb 100644 --- a/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_tile.hpp +++ b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_tile.hpp @@ -38,6 +38,23 @@ using fused_1nn_cutile_norm_t = std::conditional_t, */ template requires is_fused_1nn_cutile_data_v +bool can_launch_fused_1nn_tile(IdxT* nearest_idx, + DataT* nearest_dist, + const DataT* x, + const DataT* y, + IdxT m, + IdxT n, + IdxT k, + cuvs::distance::DistanceType metric); + +/** + * Return whether the supplied problem and existing norm buffers can use cuTile. + * + * The overload without norm pointers is a preflight probe for callers that allocate aligned norm + * buffers only after the remaining launch requirements have been validated. + */ +template + requires is_fused_1nn_cutile_data_v bool can_launch_fused_1nn_tile(IdxT* nearest_idx, DataT* nearest_dist, const DataT* x, @@ -65,6 +82,13 @@ bool try_fused_1nn_tile(IdxT* nearest_idx, void* index_workspace, cudaStream_t stream); #else +template +bool can_launch_fused_1nn_tile( + IdxT*, DataT*, const DataT*, const DataT*, IdxT, IdxT, IdxT, cuvs::distance::DistanceType) +{ + return false; +} + template bool can_launch_fused_1nn_tile(IdxT*, DataT*, diff --git a/cpp/src/distance/detail/fused_distance_nn/predicated_tile_iterator_reduced_vec.h b/cpp/src/distance/detail/fused_distance_nn/predicated_tile_iterator_reduced_vec.h index 8d16c72c04..6b089493e5 100644 --- a/cpp/src/distance/detail/fused_distance_nn/predicated_tile_iterator_reduced_vec.h +++ b/cpp/src/distance/detail/fused_distance_nn/predicated_tile_iterator_reduced_vec.h @@ -437,8 +437,10 @@ class PredicatedTileIteratorReducedVec { __syncthreads(); if (row < total_rows) { + volatile Element* gmem_ptr = reinterpret_cast(first_tile_byte_pointer_); + if ((block_start_row_first_tile_ + row) < extent_row_) { - user_params.red_op_.merge(block_start_row_first_tile_ + row, row_local_min); + user_params.red_op_(block_start_row_first_tile_ + row, (gmem_ptr + row), row_local_min); } } diff --git a/cpp/src/neighbors/ivf_flat/ivf_flat_search.cuh b/cpp/src/neighbors/ivf_flat/ivf_flat_search.cuh index 960d48c818..e6f54dd457 100644 --- a/cpp/src/neighbors/ivf_flat/ivf_flat_search.cuh +++ b/cpp/src/neighbors/ivf_flat/ivf_flat_search.cuh @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2022-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ @@ -292,6 +292,28 @@ void search_impl(raft::resources const& handle, cuvs::selection::SelectAlgo::kAuto, num_samples_vector); } + if (manage_local_topk && grid_dim_x > 1) { + // The merge select initializes unused output slots with (kDummy, idx=0), even when the + // corresponding local inputs carried an out-of-range sentinel. Restore the sentinel before + // neighbor postprocessing so an underfilled result cannot turn idx=0 into a duplicate DB ID. + AccT dummy_out = effective_metric == cuvs::distance::DistanceType::CosineExpanded + ? raft::lower_bound() + : (select_min ? raft::upper_bound() : raft::lower_bound()); + if (effective_metric == cuvs::distance::DistanceType::L2SqrtExpanded || + effective_metric == cuvs::distance::DistanceType::L2SqrtUnexpanded) { + dummy_out = raft::sqrt_op{}(dummy_out); + } else if (effective_metric == cuvs::distance::DistanceType::CosineExpanded) { + dummy_out = AccT{1} - dummy_out; + } + raft::linalg::map_offset(handle, + raft::make_device_vector_view( + neighbors_uint32, std::size_t(n_queries) * std::size_t(k)), + [neighbors_uint32, distances, dummy_out] __device__(std::size_t i) { + return distances[i] == dummy_out + ? std::numeric_limits::max() + : neighbors_uint32[i]; + }); + } if (!manage_local_topk) { // post process distances && neighbor IDs ivf::detail::postprocess_distances( From b79f4bff012b473519ac72e7b04fd7609c96dd85 Mon Sep 17 00:00:00 2001 From: divyegala Date: Wed, 2 Sep 2026 02:48:07 +0000 Subject: [PATCH 66/82] self review --- cpp/src/cluster/detail/kmeans_common.cuh | 40 +++++++++--- .../detail/minClusterDistanceCompute.cu | 18 +++-- .../cutile/fused_1nn_tile.cu | 18 ++--- .../cutile/fused_1nn_tile.hpp | 21 +++++- cpp/tests/cluster/kmeans_predict_batching.cu | 14 ++++ cpp/tests/neighbors/distance_nn.cu | 65 ++++++++++++++++++- 6 files changed, 145 insertions(+), 31 deletions(-) diff --git a/cpp/src/cluster/detail/kmeans_common.cuh b/cpp/src/cluster/detail/kmeans_common.cuh index 095ab28603..7a92a48fe7 100644 --- a/cpp/src/cluster/detail/kmeans_common.cuh +++ b/cpp/src/cluster/detail/kmeans_common.cuh @@ -77,6 +77,34 @@ inline constexpr bool uses_fused_distance_nn(FusedDistancePath path) return path != FusedDistancePath::Unfused; } +/** + * @brief Select the pre-cuTile assignment path. + * + * This is also the fallback after a pointer-aware cuTile launch probe fails. + */ +template +constexpr FusedDistancePath use_legacy_fused(int cc_major, + IdxT m, + IdxT n, + cuvs::distance::DistanceType metric) +{ + if (metric == cuvs::distance::DistanceType::InnerProduct) { return FusedDistancePath::Unfused; } + + if (cc_major <= 8) { return FusedDistancePath::FusedCutlass; } + if (cc_major == 9 && (m >= 4096 || n >= 4096)) { return FusedDistancePath::FusedCutlass; } + return FusedDistancePath::Unfused; +} + +template +FusedDistancePath use_legacy_fused(const raft::resources& handle, + IdxT m, + IdxT n, + cuvs::distance::DistanceType metric) +{ + const auto prop = raft::resource::get_device_properties(handle); + return use_legacy_fused(prop.major, m, n, metric); +} + /** * @brief Selects the fused-distance assignment path for KMeans. * @@ -90,8 +118,6 @@ FusedDistancePath use_fused( const raft::resources& handle, IdxT m, IdxT n, IdxT k, cuvs::distance::DistanceType metric) { (void)k; - cudaDeviceProp prop; - prop = raft::resource::get_device_properties(handle); if constexpr (is_cutile_fused_data_type_v) { if constexpr (cuvs::detail::jit_lto::library_built_with_cutile()) { @@ -102,16 +128,10 @@ FusedDistancePath use_fused( return FusedDistancePath::FusedCutile; } } - if (metric == cuvs::distance::DistanceType::InnerProduct) { return FusedDistancePath::Unfused; } - if (prop.major <= 8) { return FusedDistancePath::FusedCutlass; } - if (prop.major == 9 && (m >= 4096 || n >= 4096)) { return FusedDistancePath::FusedCutlass; } - return FusedDistancePath::Unfused; + return use_legacy_fused(handle, m, n, metric); } - if (prop.major >= 10) { return FusedDistancePath::Unfused; } - if (prop.major <= 8) { return FusedDistancePath::FusedCutlass; } - if (prop.major == 9 && (m >= 4096 || n >= 4096)) { return FusedDistancePath::FusedCutlass; } - return FusedDistancePath::Unfused; + return use_legacy_fused(handle, m, n, metric); } template diff --git a/cpp/src/cluster/detail/minClusterDistanceCompute.cu b/cpp/src/cluster/detail/minClusterDistanceCompute.cu index e09cba9e70..8fe1571f47 100644 --- a/cpp/src/cluster/detail/minClusterDistanceCompute.cu +++ b/cpp/src/cluster/detail/minClusterDistanceCompute.cu @@ -129,11 +129,7 @@ void minClusterAndDistanceCompute(raft::resources const& handle, n_clusters, n_features, metric); - if (!cutile_ready) { - fused_path = metric == cuvs::distance::DistanceType::InnerProduct - ? FusedDistancePath::Unfused - : FusedDistancePath::FusedCutlass; - } + if (!cutile_ready) { fused_path = use_legacy_fused(handle, n_samples, n_clusters, metric); } } } @@ -189,9 +185,11 @@ void minClusterAndDistanceCompute(raft::resources const& handle, if (!cutile_ready) { temp_kvp.resize(n_samples, stream); cutlass_kvp_scratch = temp_kvp.data(); - } - if (!cutile_ready || needs_index_workspace) { workspace.resize(sizeof(int) * static_cast(n_samples), stream); + } else if (needs_index_workspace) { + const auto workspace_rows = + cuvs::distance::detail::fused_1nn_cutile_index_workspace_rows(n_samples); + workspace.resize(sizeof(int) * workspace_rows, stream); } cuvs::distance::fusedDistanceNNMinReduce( @@ -206,7 +204,7 @@ void minClusterAndDistanceCompute(raft::resources const& handle, n_features, !cutile_ready || needs_index_workspace ? (void*)workspace.data() : nullptr, metric != cuvs::distance::DistanceType::L2Expanded, - true, + false, true, metric, 0.0f, @@ -409,7 +407,7 @@ void minClusterDistanceCompute(raft::resources const& handle, n_clusters, n_features, metric); - if (!cutile_ready) { fused_path = FusedDistancePath::FusedCutlass; } + if (!cutile_ready) { fused_path = use_legacy_fused(handle, n_samples, n_clusters, metric); } } } @@ -478,7 +476,7 @@ void minClusterDistanceCompute(raft::resources const& handle, n_features, cutile_ready ? nullptr : (void*)workspace.data(), metric != cuvs::distance::DistanceType::L2Expanded, - true, + false, true, metric, 0.0f, diff --git a/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_tile.cu b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_tile.cu index 5593949454..836fc891a4 100644 --- a/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_tile.cu +++ b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_tile.cu @@ -290,15 +290,14 @@ bool try_fused_1nn_tile(IdxT* nearest_idx, if (!is_16_byte_aligned(index_workspace)) { return false; } // Keep every chunk offset 16-byte aligned for x, xn, and nearest_dist. - constexpr int64_t max_i32 = std::numeric_limits::max(); - constexpr int64_t batch_alignment = 16 / sizeof(DataT); - constexpr int64_t max_batch_m = max_i32 - max_i32 % batch_alignment; - auto* tmp_idx = static_cast(index_workspace); - for (int64_t offset = 0; offset < m; offset += max_batch_m) { - const int batch_m = static_cast(std::min(max_batch_m, m - offset)); - const auto* batch_x = x + static_cast(offset) * static_cast(k); - const auto* batch_xn = xn == nullptr ? nullptr : xn + offset; - auto* batch_dist = nearest_dist == nullptr ? nullptr : nearest_dist + offset; + constexpr int64_t max_batch_m = fused_1nn_cutile_max_batch_m; + auto* tmp_idx = static_cast(index_workspace); + for (int64_t offset = 0; offset < m;) { + const int64_t batch_m64 = std::min(max_batch_m, m - offset); + const int batch_m = static_cast(batch_m64); + const auto* batch_x = x + static_cast(offset) * static_cast(k); + const auto* batch_xn = xn == nullptr ? nullptr : xn + offset; + auto* batch_dist = nearest_dist == nullptr ? nullptr : nearest_dist + offset; const bool launched = use_strict_abi @@ -332,6 +331,7 @@ bool try_fused_1nn_tile(IdxT* nearest_idx, raft::linalg::unaryOp( nearest_idx + offset, tmp_idx, batch_m, raft::cast_op{}, stream); } + offset += batch_m64; } return true; } diff --git a/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_tile.hpp b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_tile.hpp index 245a2c6afb..5a4b339425 100644 --- a/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_tile.hpp +++ b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_tile.hpp @@ -5,6 +5,9 @@ #pragma once +#include +#include +#include #include #include @@ -28,13 +31,29 @@ inline constexpr bool is_fused_1nn_cutile_data_v = template using fused_1nn_cutile_norm_t = std::conditional_t, float, DataT>; +template +inline constexpr int64_t fused_1nn_cutile_max_batch_m = [] { + constexpr int64_t max_i32 = std::numeric_limits::max(); + constexpr int64_t batch_alignment = 16 / sizeof(DataT); + return max_i32 - max_i32 % batch_alignment; +}(); + +template +constexpr size_t fused_1nn_cutile_index_workspace_rows(IdxT m) +{ + const auto rows = static_cast(m); + if (rows <= 0) { return 0; } + return static_cast( + rows < fused_1nn_cutile_max_batch_m ? rows : fused_1nn_cutile_max_batch_m); +} + #if CUVS_CUTILE_ENABLED /** * Return whether the supplied problem can use cuTile without fallback scratch. * * The result includes runtime/device support, exported ABI constraints, and launcher construction. * A successful probe populates the shared launcher cache used by try_fused_1nn_tile. - * An int64 output index still requires an int32 workspace with `m` elements. + * An int64 output index still requires an int32 workspace sized to the largest launch chunk. */ template requires is_fused_1nn_cutile_data_v diff --git a/cpp/tests/cluster/kmeans_predict_batching.cu b/cpp/tests/cluster/kmeans_predict_batching.cu index b740c3f998..d54595254f 100644 --- a/cpp/tests/cluster/kmeans_predict_batching.cu +++ b/cpp/tests/cluster/kmeans_predict_batching.cu @@ -157,4 +157,18 @@ TEST(KMeansPredict, BatchParametersPreserveResultsAndReduceUnfusedAllocations) } } +TEST(KMeansPredict, ProbeFailureUsesLegacyArchitectureFallback) +{ + using cuvs::distance::DistanceType; + using detail::FusedDistancePath; + + for (auto metric : {DistanceType::L2Expanded, DistanceType::CosineExpanded}) { + EXPECT_EQ(detail::use_legacy_fused(8, 1024, 1024, metric), FusedDistancePath::FusedCutlass); + EXPECT_EQ(detail::use_legacy_fused(9, 4096, 1024, metric), FusedDistancePath::FusedCutlass); + EXPECT_EQ(detail::use_legacy_fused(9, 1024, 1024, metric), FusedDistancePath::Unfused); + EXPECT_EQ(detail::use_legacy_fused(10, 16384, 16384, metric), FusedDistancePath::Unfused); + EXPECT_EQ(detail::use_legacy_fused(12, 16384, 16384, metric), FusedDistancePath::Unfused); + } +} + } // namespace cuvs::cluster::kmeans diff --git a/cpp/tests/neighbors/distance_nn.cu b/cpp/tests/neighbors/distance_nn.cu index 1bd800a6cf..ec8624c0c9 100644 --- a/cpp/tests/neighbors/distance_nn.cu +++ b/cpp/tests/neighbors/distance_nn.cu @@ -125,7 +125,7 @@ class NNTest : public ::testing::TestWithParam> { k, (void*)workspace.data_handle(), sqrt, - true, + false, true, metric, 0.0, @@ -312,6 +312,69 @@ TEST(Fused1nn, ExpandedL2ClampsNegativeRoundoff) EXPECT_EQ(actual_dist, 0.0f); } +TEST(Fused1nn, Int64IndexWorkspaceUsesLargestChunk) +{ + constexpr int64_t max_batch_m_float = cuvs::distance::detail::fused_1nn_cutile_max_batch_m; + constexpr int64_t max_batch_m_half = cuvs::distance::detail::fused_1nn_cutile_max_batch_m; + EXPECT_EQ(max_batch_m_float, 2147483644); + EXPECT_EQ(max_batch_m_half, 2147483640); + EXPECT_EQ(cuvs::distance::detail::fused_1nn_cutile_index_workspace_rows(1024), 1024); + EXPECT_EQ(cuvs::distance::detail::fused_1nn_cutile_index_workspace_rows( + std::numeric_limits::max()), + static_cast(max_batch_m_float)); + EXPECT_EQ(cuvs::distance::detail::fused_1nn_cutile_index_workspace_rows( + std::numeric_limits::max()), + static_cast(max_batch_m_half)); +} + +#if CUVS_CUTILE_ENABLED +TEST(Fused1nn, PointerAwareProbeRejectsMisalignedArrays) +{ + raft::resources handle; + constexpr int m = 32; + constexpr int n = 32; + constexpr int k = 64; + + auto x = raft::make_device_vector(handle, m * k + 1); + auto y = raft::make_device_vector(handle, n * k); + auto x_norm = raft::make_device_vector(handle, m + 1); + auto y_norm = raft::make_device_vector(handle, n); + auto out_idx = raft::make_device_vector(handle, m); + auto out_dist = raft::make_device_vector(handle, m); + + for (auto metric : {DistanceType::L2Expanded, DistanceType::CosineExpanded}) { + EXPECT_FALSE(cuvs::distance::detail::can_launch_fused_1nn_tile(out_idx.data_handle(), + out_dist.data_handle(), + x.data_handle() + 1, + y.data_handle(), + m, + n, + k, + metric)); + + if (cuvs::distance::detail::can_launch_fused_1nn_tile(out_idx.data_handle(), + out_dist.data_handle(), + x.data_handle(), + y.data_handle(), + m, + n, + k, + metric)) { + EXPECT_FALSE(cuvs::distance::detail::can_launch_fused_1nn_tile(out_idx.data_handle(), + out_dist.data_handle(), + x.data_handle(), + y.data_handle(), + x_norm.data_handle() + 1, + y_norm.data_handle(), + m, + n, + k, + metric)); + } + } +} +#endif + template const std::vector> input_int8 = { {4096, 4096, 64, DistanceType::L2Expanded, false, uint64_t(31415926), 0.1}, From ff9dae12a493e2ec0cc9424bf37cd33512372e79 Mon Sep 17 00:00:00 2001 From: divyegala Date: Wed, 2 Sep 2026 04:22:14 +0000 Subject: [PATCH 67/82] preserve fallback fidelity better --- cpp/src/cluster/detail/kmeans.cuh | 132 +++++---- cpp/src/cluster/detail/kmeans_common.cuh | 259 ++++++++++++------ cpp/src/cluster/detail/kmeans_mg.cuh | 11 +- .../detail/minClusterDistanceCompute.cu | 242 ++++++++++++---- cpp/src/distance/detail/fused_distance_nn.cuh | 108 ++++---- .../cutile/fused_1nn_tile.cu | 52 ++-- .../cutile/fused_1nn_tile.hpp | 17 ++ .../fused_distance_nn/fused_cosine_nn.cuh | 14 +- .../detail/fused_distance_nn/fused_l2_nn.cuh | 14 +- cpp/src/distance/fused_distance_nn-inl.cuh | 8 +- 10 files changed, 572 insertions(+), 285 deletions(-) diff --git a/cpp/src/cluster/detail/kmeans.cuh b/cpp/src/cluster/detail/kmeans.cuh index 819545e8d1..1a71003fd0 100644 --- a/cpp/src/cluster/detail/kmeans.cuh +++ b/cpp/src/cluster/detail/kmeans.cuh @@ -685,8 +685,7 @@ void kmeans_fit( DataT* cur_centroids_ptr = cur_centroids_buf.data(); DataT* new_centroids_ptr = new_centroids_buf.data(); - auto nearest_idx = raft::make_device_vector(handle, device_buffer_samples); - auto nearest_dist = raft::make_device_vector(handle, device_buffer_samples); + rmm::device_uvector assignment_storage(0, stream); auto L2NormBatch = raft::make_device_vector(handle, device_buffer_samples); auto batch_weights_buf = raft::make_device_vector(handle, device_buffer_samples); rmm::device_uvector L2NormBuf_OR_DistBuf(0, stream); @@ -856,11 +855,6 @@ void kmeans_fit( auto batch_weights_view = cur_batch_weights(static_cast(data_batch.offset()), wt_data, cur_batch_size); - auto nearest_idx_view = - raft::make_device_vector_view(nearest_idx.data_handle(), cur_batch_size); - auto nearest_dist_view = - raft::make_device_vector_view(nearest_dist.data_handle(), cur_batch_size); - if constexpr (!data_on_device) { if (need_compute_norms) { if (!norms_cached) { @@ -888,8 +882,7 @@ void kmeans_fit( metric, iter_params.batch_samples, iter_params.batch_centroids, - nearest_idx_view, - nearest_dist_view, + assignment_storage, l2_const_view, L2NormBuf_OR_DistBuf, ws, @@ -1077,7 +1070,6 @@ void kmeans_predict(raft::resources const& handle, raft::make_const_mdspan(weight.view())); } - auto nearest_dist = raft::make_device_vector(handle, n_samples); rmm::device_uvector L2NormBuf_OR_DistBuf(0, stream); // L2 norm of X: ||x||^2 @@ -1089,49 +1081,87 @@ void kmeans_predict(raft::resources const& handle, auto l2normx_view = raft::make_device_vector_view(L2NormX.data_handle(), n_samples); - if constexpr (std::is_same_v) { - cuvs::cluster::kmeans::detail::minClusterAndDistanceCompute(handle, - X, - centroids, - labels, - nearest_dist.view(), - l2normx_view, - L2NormBuf_OR_DistBuf, - pams.metric, - pams.batch_samples, - pams.batch_centroids, - workspace); - } else { - auto index_labels = raft::make_device_vector(handle, n_samples); - cuvs::cluster::kmeans::detail::minClusterAndDistanceCompute(handle, - X, - centroids, - index_labels.view(), - nearest_dist.view(), - l2normx_view, - L2NormBuf_OR_DistBuf, - pams.metric, - pams.batch_samples, - pams.batch_centroids, - workspace); - raft::linalg::map( - handle, labels, raft::cast_op{}, raft::make_const_mdspan(index_labels.view())); - } - rmm::device_scalar clusterCostD(stream); - raft::linalg::map(handle, - nearest_dist.view(), - raft::mul_op{}, - raft::make_const_mdspan(nearest_dist.view()), - raft::make_const_mdspan(weight.view())); + auto path = select_min_cluster_distance_path(handle, X, centroids, metric); + if constexpr (std::is_same_v && std::is_same_v) { + if (path == FusedDistancePath::FusedCutile && + reinterpret_cast(labels.data_handle()) % 16 != 0) { + path = use_legacy_fused(handle, n_samples, centroids.extent(0), metric); + } + } - cuvs::cluster::kmeans::detail::computeClusterCost( - handle, - nearest_dist.view(), - workspace, - raft::make_device_scalar_view(clusterCostD.data()), - raft::identity_op{}, - raft::add_op{}); + if (path == FusedDistancePath::FusedCutile) { + auto nearest_dist = raft::make_device_vector(handle, n_samples); + if constexpr (std::is_same_v) { + minClusterAndDistanceCompute(handle, + X, + centroids, + labels, + nearest_dist.view(), + l2normx_view, + L2NormBuf_OR_DistBuf, + metric, + pams.batch_samples, + pams.batch_centroids, + workspace); + } else { + auto index_labels = raft::make_device_vector(handle, n_samples); + minClusterAndDistanceCompute(handle, + X, + centroids, + index_labels.view(), + nearest_dist.view(), + l2normx_view, + L2NormBuf_OR_DistBuf, + metric, + pams.batch_samples, + pams.batch_centroids, + workspace); + raft::linalg::map( + handle, labels, raft::cast_op{}, raft::make_const_mdspan(index_labels.view())); + } + raft::linalg::map(handle, + nearest_dist.view(), + raft::mul_op{}, + raft::make_const_mdspan(nearest_dist.view()), + raft::make_const_mdspan(weight.view())); + computeClusterCost(handle, + nearest_dist.view(), + workspace, + raft::make_device_scalar_view(clusterCostD.data()), + raft::identity_op{}, + raft::add_op{}); + } else { + using KvpT = raft::KeyValuePair; + auto nearest = raft::make_device_vector(handle, n_samples); + minClusterAndDistanceComputeKvp(handle, + X, + centroids, + nearest.view(), + l2normx_view, + L2NormBuf_OR_DistBuf, + metric, + pams.batch_samples, + pams.batch_centroids, + workspace, + path); + auto* nearest_ptr = nearest.data_handle(); + raft::linalg::map_offset(handle, labels, [nearest_ptr] __device__(IndexT i) { + return static_cast(nearest_ptr[i].key); + }); + auto* weights = weight.data_handle(); + cuda::counting_iterator indices(IndexT{0}); + cuda::transform_iterator weighted_dist(indices, + [nearest_ptr, weights] __device__(IndexT i) -> DataT { + return nearest_ptr[i].value * weights[i]; + }); + computeClusterCostFromIterator(handle, + weighted_dist, + n_samples, + workspace, + raft::make_device_scalar_view(clusterCostD.data()), + raft::add_op{}); + } inertia[0] = clusterCostD.value(stream); } diff --git a/cpp/src/cluster/detail/kmeans_common.cuh b/cpp/src/cluster/detail/kmeans_common.cuh index 7a92a48fe7..b4dfe3511c 100644 --- a/cpp/src/cluster/detail/kmeans_common.cuh +++ b/cpp/src/cluster/detail/kmeans_common.cuh @@ -259,28 +259,22 @@ IndexT getCentroidsBatchSize(int batch_centroids, IndexT n_local_clusters) return (minVal == 0) ? n_local_clusters : minVal; } -template -void computeClusterCost(raft::resources const& handle, - raft::device_vector_view minClusterDistance, - rmm::device_uvector& workspace, - raft::device_scalar_view clusterCost, - MainOpT main_op, - ReductionOpT reduction_op) +template +void computeClusterCostFromIterator(raft::resources const& handle, + InputIteratorT input, + IndexT n, + rmm::device_uvector& workspace, + raft::device_scalar_view clusterCost, + ReductionOpT reduction_op) { cudaStream_t stream = raft::resource::get_cuda_stream(handle); - cuda::transform_iterator itr(minClusterDistance.data_handle(), main_op); - size_t temp_storage_bytes = 0; RAFT_CUDA_TRY(cub::DeviceReduce::Reduce(nullptr, temp_storage_bytes, - itr, + input, clusterCost.data_handle(), - minClusterDistance.size(), + n, reduction_op, OutputT(), stream)); @@ -289,14 +283,31 @@ void computeClusterCost(raft::resources const& handle, RAFT_CUDA_TRY(cub::DeviceReduce::Reduce(workspace.data(), temp_storage_bytes, - itr, + input, clusterCost.data_handle(), - minClusterDistance.size(), + n, reduction_op, OutputT(), stream)); } +template +void computeClusterCost(raft::resources const& handle, + raft::device_vector_view minClusterDistance, + rmm::device_uvector& workspace, + raft::device_scalar_view clusterCost, + MainOpT main_op, + ReductionOpT reduction_op) +{ + cuda::transform_iterator input(minClusterDistance.data_handle(), main_op); + computeClusterCostFromIterator( + handle, input, minClusterDistance.size(), workspace, clusterCost, reduction_op); +} + template void sampleCentroids(raft::resources const& handle, raft::device_matrix_view X, @@ -442,6 +453,13 @@ void shuffleAndGather(raft::resources const& handle, } // Calculates nearest centroid index and distance for every sample in input 'X'. +template +FusedDistancePath select_min_cluster_distance_path( + raft::resources const& handle, + raft::device_matrix_view X, + raft::device_matrix_view centroids, + cuvs::distance::DistanceType metric); + template void minClusterAndDistanceCompute(raft::resources const& handle, raft::device_matrix_view X, @@ -455,6 +473,20 @@ void minClusterAndDistanceCompute(raft::resources const& handle, int batch_centroids, rmm::device_uvector& workspace); +template +void minClusterAndDistanceComputeKvp( + raft::resources const& handle, + raft::device_matrix_view X, + raft::device_matrix_view centroids, + raft::device_vector_view, IndexT> nearest, + raft::device_vector_view L2NormX, + rmm::device_uvector& L2NormBuf_OR_DistBuf, + cuvs::distance::DistanceType metric, + int batch_samples, + int batch_centroids, + rmm::device_uvector& workspace, + FusedDistancePath path); + #define EXTERN_TEMPLATE_MIN_CLUSTER_AND_DISTANCE(DataT, IndexT) \ extern template void minClusterAndDistanceCompute( \ raft::resources const& handle, \ @@ -519,35 +551,54 @@ void countSamplesInCluster(raft::resources const& handle, { cudaStream_t stream = raft::resource::get_cuda_stream(handle); auto n_samples = X.extent(0); - auto n_features = X.extent(1); auto n_clusters = centroids.extent(0); - // stores (key, value) pair corresponding to each sample where - // - key is the index of nearest cluster - // - value is the distance to the nearest cluster - auto nearest_idx = raft::make_device_vector(handle, n_samples); - auto nearest_dist = raft::make_device_vector(handle, n_samples); rmm::device_uvector L2NormBuf_OR_DistBuf(0, stream); - - cuvs::cluster::kmeans::detail::minClusterAndDistanceCompute( - handle, - X, - (raft::device_matrix_view)centroids, - nearest_idx.view(), - nearest_dist.view(), - L2NormX, - L2NormBuf_OR_DistBuf, - params.metric, - params.batch_samples, - params.batch_centroids, - workspace); - - countLabels(handle, - nearest_idx.data_handle(), - sampleCountInCluster.data_handle(), - (IndexT)n_samples, - (IndexT)n_clusters, - workspace); + auto centroids_const = raft::make_const_mdspan(centroids); + auto path = select_min_cluster_distance_path(handle, X, centroids_const, params.metric); + + auto count_labels = [&](auto labels) { + countLabels(handle, + labels, + sampleCountInCluster.data_handle(), + static_cast(n_samples), + static_cast(n_clusters), + workspace); + }; + + if (path == FusedDistancePath::FusedCutile) { + auto nearest_idx = raft::make_device_vector(handle, n_samples); + auto nearest_dist = raft::make_device_vector(handle, n_samples); + minClusterAndDistanceCompute(handle, + X, + centroids_const, + nearest_idx.view(), + nearest_dist.view(), + L2NormX, + L2NormBuf_OR_DistBuf, + params.metric, + params.batch_samples, + params.batch_centroids, + workspace); + count_labels(nearest_idx.data_handle()); + } else { + using KvpT = raft::KeyValuePair; + auto nearest = raft::make_device_vector(handle, n_samples); + minClusterAndDistanceComputeKvp(handle, + X, + centroids_const, + nearest.view(), + L2NormX, + L2NormBuf_OR_DistBuf, + params.metric, + params.batch_samples, + params.batch_centroids, + workspace, + path); + auto labels = + thrust::make_transform_iterator(nearest.data_handle(), KeyValueIndexOp{}); + count_labels(labels); + } } /** @@ -728,8 +779,7 @@ __device__ void check_convergence(raft::device_scalar_view clusteri * @param[in] batch_samples_param Batch-samples param forwarded to minClusterAndDistanceCompute * @param[in] batch_centroids_param Batch-centroids param forwarded to * minClusterAndDistanceCompute - * @param[inout] nearest_idx Nearest cluster index per sample [batch_size] - * @param[inout] nearest_dist Nearest distance per sample [batch_size] + * @param[inout] assignment_storage Native SoA or KVP assignment storage for one batch * @param[in] L2NormBatch Precomputed data norms [batch_size] * @param[inout] L2NormBuf_OR_DistBuf Resizable scratch * @param[inout] workspace Resizable scratch @@ -745,8 +795,7 @@ void process_batch(raft::resources const& handle, cuvs::distance::DistanceType metric, int batch_samples_param, int batch_centroids_param, - raft::device_vector_view nearest_idx, - raft::device_vector_view nearest_dist, + rmm::device_uvector& assignment_storage, raft::device_vector_view L2NormBatch, rmm::device_uvector& L2NormBuf_OR_DistBuf, rmm::device_uvector& workspace, @@ -755,44 +804,86 @@ void process_batch(raft::resources const& handle, raft::device_scalar_view clustering_cost, rmm::device_uvector& batch_workspace) { - cudaStream_t stream = raft::resource::get_cuda_stream(handle); + cudaStream_t stream = raft::resource::get_cuda_stream(handle); + const auto n_samples = batch_data.extent(0); + const auto path = select_min_cluster_distance_path(handle, batch_data, centroids, metric); + auto batch_cost = raft::make_device_scalar(handle, DataT{0}); + + if (path == FusedDistancePath::FusedCutile) { + const auto dist_offset = + raft::alignTo(sizeof(IndexT) * static_cast(n_samples), size_t{16}); + assignment_storage.resize(dist_offset + sizeof(DataT) * static_cast(n_samples), stream); + auto* nearest_idx = reinterpret_cast(assignment_storage.data()); + auto* nearest_dist = reinterpret_cast(assignment_storage.data() + dist_offset); + auto nearest_idx_view = raft::make_device_vector_view(nearest_idx, n_samples); + auto nearest_dist_view = raft::make_device_vector_view(nearest_dist, n_samples); + minClusterAndDistanceCompute(handle, + batch_data, + centroids, + nearest_idx_view, + nearest_dist_view, + L2NormBatch, + L2NormBuf_OR_DistBuf, + metric, + batch_samples_param, + batch_centroids_param, + workspace); + compute_centroid_adjustments(handle, + batch_data, + batch_weights, + nearest_idx, + static_cast(centroid_sums.extent(0)), + centroid_sums, + weight_per_cluster, + batch_workspace, + /*reset_sums=*/false); + auto weighted_dist = raft::make_device_vector(handle, n_samples); + raft::linalg::map(handle, + weighted_dist.view(), + raft::mul_op{}, + raft::make_const_mdspan(nearest_dist_view), + raft::make_const_mdspan(batch_weights)); + computeClusterCost(handle, + weighted_dist.view(), + workspace, + batch_cost.view(), + raft::identity_op{}, + raft::add_op{}); + } else { + using KvpT = raft::KeyValuePair; + assignment_storage.resize(sizeof(KvpT) * static_cast(n_samples), stream); + auto* nearest = reinterpret_cast(assignment_storage.data()); + auto nearest_view = raft::make_device_vector_view(nearest, n_samples); + minClusterAndDistanceComputeKvp(handle, + batch_data, + centroids, + nearest_view, + L2NormBatch, + L2NormBuf_OR_DistBuf, + metric, + batch_samples_param, + batch_centroids_param, + workspace, + path); + auto labels = thrust::make_transform_iterator(nearest, KeyValueIndexOp{}); + compute_centroid_adjustments(handle, + batch_data, + batch_weights, + labels, + static_cast(centroid_sums.extent(0)), + centroid_sums, + weight_per_cluster, + batch_workspace, + /*reset_sums=*/false); + auto* weights = batch_weights.data_handle(); + cuda::counting_iterator indices(IndexT{0}); + cuda::transform_iterator weighted_dist( + indices, + [nearest, weights] __device__(IndexT i) -> DataT { return nearest[i].value * weights[i]; }); + computeClusterCostFromIterator( + handle, weighted_dist, n_samples, workspace, batch_cost.view(), raft::add_op{}); + } - minClusterAndDistanceCompute(handle, - batch_data, - centroids, - nearest_idx, - nearest_dist, - L2NormBatch, - L2NormBuf_OR_DistBuf, - metric, - batch_samples_param, - batch_centroids_param, - workspace); - - compute_centroid_adjustments(handle, - batch_data, - batch_weights, - nearest_idx.data_handle(), - static_cast(centroid_sums.extent(0)), - centroid_sums, - weight_per_cluster, - batch_workspace, - /*reset_sums=*/false); - - auto weighted_dist = raft::make_device_vector(handle, nearest_dist.extent(0)); - raft::linalg::map(handle, - weighted_dist.view(), - raft::mul_op{}, - raft::make_const_mdspan(nearest_dist), - raft::make_const_mdspan(batch_weights)); - - auto batch_cost = raft::make_device_scalar(handle, DataT{0}); - computeClusterCost(handle, - weighted_dist.view(), - workspace, - batch_cost.view(), - raft::identity_op{}, - raft::add_op{}); raft::linalg::add(clustering_cost.data_handle(), clustering_cost.data_handle(), batch_cost.data_handle(), diff --git a/cpp/src/cluster/detail/kmeans_mg.cuh b/cpp/src/cluster/detail/kmeans_mg.cuh index c1f51026eb..1de8a401a7 100644 --- a/cpp/src/cluster/detail/kmeans_mg.cuh +++ b/cpp/src/cluster/detail/kmeans_mg.cuh @@ -214,8 +214,7 @@ void mnmg_fit( auto sqrd_norm_error_dev = raft::make_device_scalar(dev_res, DataT{0}); IndexT alloc_batch_size = device_buffer_samples; auto batch_weights = raft::make_device_vector(dev_res, alloc_batch_size); - auto nearest_idx = raft::make_device_vector(dev_res, alloc_batch_size); - auto nearest_dist = raft::make_device_vector(dev_res, alloc_batch_size); + rmm::device_uvector assignment_storage(0, stream); auto L2NormBatch = raft::make_device_vector(dev_res, data_on_device ? IndexT{0} : alloc_batch_size); rmm::device_uvector L2NormBuf_OR_DistBuf(0, stream); @@ -448,11 +447,6 @@ void mnmg_fit( L2NormBatch_const = raft::make_const_mdspan(norm_slice); } - auto nearest_idx_view = raft::make_device_vector_view( - nearest_idx.data_handle(), current_batch_size); - auto nearest_dist_view = raft::make_device_vector_view( - nearest_dist.data_handle(), current_batch_size); - cuvs::cluster::kmeans::detail::process_batch( dev_res, batch_data_view, @@ -461,8 +455,7 @@ void mnmg_fit( metric, params.batch_samples, params.batch_centroids, - nearest_idx_view, - nearest_dist_view, + assignment_storage, L2NormBatch_const, L2NormBuf_OR_DistBuf, workspace, diff --git a/cpp/src/cluster/detail/minClusterDistanceCompute.cu b/cpp/src/cluster/detail/minClusterDistanceCompute.cu index 8fe1571f47..a62feec252 100644 --- a/cpp/src/cluster/detail/minClusterDistanceCompute.cu +++ b/cpp/src/cluster/detail/minClusterDistanceCompute.cu @@ -82,32 +82,56 @@ __global__ void unpack_kvp_to_soa(IndexT* nearest_idx, template void unpack_kvp(raft::resources const& handle, - raft::device_vector_view nearest_idx, - raft::device_vector_view nearest_dist, - raft::device_vector_view, IndexT> kvp) + IndexT* nearest_idx, + DataT* nearest_dist, + const raft::KeyValuePair* kvp, + IndexT n) { auto stream = raft::resource::get_cuda_stream(handle); - auto n = static_cast(kvp.extent(0)); int blks = static_cast((n + 255) / 256); - unpack_kvp_to_soa<<>>( - nearest_idx.data_handle(), nearest_dist.data_handle(), kvp.data_handle(), n); + unpack_kvp_to_soa<<>>(nearest_idx, nearest_dist, kvp, n); RAFT_CUDA_TRY(cudaGetLastError()); } } // namespace template -void minClusterAndDistanceCompute(raft::resources const& handle, - raft::device_matrix_view X, - raft::device_matrix_view centroids, - raft::device_vector_view nearest_idx, - raft::device_vector_view nearest_dist, - raft::device_vector_view L2NormX, - rmm::device_uvector& L2NormBuf_OR_DistBuf, - cuvs::distance::DistanceType metric, - int batch_samples, - int batch_centroids, - rmm::device_uvector& workspace) +FusedDistancePath select_min_cluster_distance_path( + raft::resources const& handle, + raft::device_matrix_view X, + raft::device_matrix_view centroids, + cuvs::distance::DistanceType metric) +{ + auto path = + use_fused(handle, X.extent(0), centroids.extent(0), X.extent(1), metric); + if constexpr (is_cutile_fused_data_type_v) { + if (path == FusedDistancePath::FusedCutile && + !cuvs::distance::detail::can_launch_fused_1nn_tile(X.data_handle(), + centroids.data_handle(), + X.extent(0), + centroids.extent(0), + X.extent(1), + metric)) { + path = use_legacy_fused(handle, X.extent(0), centroids.extent(0), metric); + } + } + return path; +} + +template +void min_cluster_and_distance_compute_impl(raft::resources const& handle, + raft::device_matrix_view X, + raft::device_matrix_view centroids, + IndexT* nearest_idx, + DataT* nearest_dist, + raft::KeyValuePair* native_kvp, + raft::device_vector_view L2NormX, + rmm::device_uvector& L2NormBuf_OR_DistBuf, + cuvs::distance::DistanceType metric, + int batch_samples, + int batch_centroids, + rmm::device_uvector& workspace, + FusedDistancePath fused_path) { cudaStream_t stream = raft::resource::get_cuda_stream(handle); auto n_samples = X.extent(0); @@ -116,22 +140,9 @@ void minClusterAndDistanceCompute(raft::resources const& handle, const bool is_l2_cos = metric == cuvs::distance::DistanceType::L2Expanded || metric == cuvs::distance::DistanceType::L2SqrtExpanded || metric == cuvs::distance::DistanceType::CosineExpanded; - FusedDistancePath fused_path = - use_fused(handle, n_samples, n_clusters, n_features, metric); - bool cutile_ready = false; - if constexpr (is_cutile_fused_data_type_v) { - if (fused_path == FusedDistancePath::FusedCutile) { - cutile_ready = cuvs::distance::detail::can_launch_fused_1nn_tile(nearest_idx.data_handle(), - nearest_dist.data_handle(), - X.data_handle(), - centroids.data_handle(), - n_samples, - n_clusters, - n_features, - metric); - if (!cutile_ready) { fused_path = use_legacy_fused(handle, n_samples, n_clusters, metric); } - } - } + const bool cutile_ready = fused_path == FusedDistancePath::FusedCutile; + RAFT_EXPECTS(!cutile_ready || native_kvp == nullptr, + "cuTile fused 1-NN requires separate index and distance outputs"); if (uses_fused_distance_nn(fused_path)) { const DataT* x_norm_ptr = nullptr; @@ -183,8 +194,12 @@ void minClusterAndDistanceCompute(raft::resources const& handle, rmm::device_uvector> temp_kvp(0, stream); const bool needs_index_workspace = cutile_ready && std::is_same_v; if (!cutile_ready) { - temp_kvp.resize(n_samples, stream); - cutlass_kvp_scratch = temp_kvp.data(); + if (native_kvp != nullptr) { + cutlass_kvp_scratch = native_kvp; + } else { + temp_kvp.resize(n_samples, stream); + cutlass_kvp_scratch = temp_kvp.data(); + } workspace.resize(sizeof(int) * static_cast(n_samples), stream); } else if (needs_index_workspace) { const auto workspace_rows = @@ -193,8 +208,8 @@ void minClusterAndDistanceCompute(raft::resources const& handle, } cuvs::distance::fusedDistanceNNMinReduce( - nearest_idx.data_handle(), - nearest_dist.data_handle(), + nearest_idx, + nearest_dist, X.data_handle(), centroids.data_handle(), x_norm_ptr, @@ -235,9 +250,11 @@ void minClusterAndDistanceCompute(raft::resources const& handle, workspace.resize(sizeof(DataT) * dataBatchSize * centroidsBatchSize, stream); using KeyValueT = raft::KeyValuePair; - auto temp_kvp = raft::make_device_vector(handle, n_samples); + rmm::device_uvector temp_kvp(native_kvp == nullptr ? n_samples : 0, stream); + auto* kvp_output = native_kvp == nullptr ? temp_kvp.data() : native_kvp; + auto kvp_output_view = raft::make_device_vector_view(kvp_output, n_samples); KeyValueT initial_value(0, std::numeric_limits::max()); - raft::matrix::fill(handle, temp_kvp.view(), initial_value); + raft::matrix::fill(handle, kvp_output_view, initial_value); const bool tileCentroids = centroidsBatchSize < n_clusters; rmm::device_uvector batchMinClusterAndDistance(tileCentroids ? dataBatchSize : 0, @@ -246,7 +263,7 @@ void minClusterAndDistanceCompute(raft::resources const& handle, for (IndexT dIdx = 0; dIdx < n_samples;) { auto ns = std::min(dataBatchSize, n_samples - dIdx); auto minClusterAndDistanceView = - raft::make_device_vector_view(temp_kvp.data_handle() + dIdx, ns); + raft::make_device_vector_view(kvp_output + dIdx, ns); for (IndexT cIdx = 0; cIdx < n_clusters;) { auto nc = std::min(centroidsBatchSize, n_clusters - cIdx); @@ -289,7 +306,9 @@ void minClusterAndDistanceCompute(raft::resources const& handle, dIdx += ns; } - unpack_kvp(handle, nearest_idx, nearest_dist, raft::make_const_mdspan(temp_kvp.view())); + if (native_kvp == nullptr) { + unpack_kvp(handle, nearest_idx, nearest_dist, kvp_output, n_samples); + } } else { auto dataBatchSize = getDataBatchSize(batch_samples, n_samples); auto centroidsBatchSize = getCentroidsBatchSize(batch_centroids, n_clusters); @@ -299,10 +318,12 @@ void minClusterAndDistanceCompute(raft::resources const& handle, auto pairwiseDistance = raft::make_device_matrix_view( L2NormBuf_OR_DistBuf.data(), dataBatchSize, centroidsBatchSize); - auto temp_kvp = - raft::make_device_vector, IndexT>(handle, n_samples); - raft::KeyValuePair initial_value(0, std::numeric_limits::max()); - raft::matrix::fill(handle, temp_kvp.view(), initial_value); + using KeyValueT = raft::KeyValuePair; + rmm::device_uvector temp_kvp(native_kvp == nullptr ? n_samples : 0, stream); + auto* kvp_output = native_kvp == nullptr ? temp_kvp.data() : native_kvp; + auto kvp_output_view = raft::make_device_vector_view(kvp_output, n_samples); + KeyValueT initial_value(0, std::numeric_limits::max()); + raft::matrix::fill(handle, kvp_output_view, initial_value); for (IndexT dIdx = 0; dIdx < n_samples; dIdx += dataBatchSize) { auto ns = std::min((IndexT)dataBatchSize, n_samples - dIdx); @@ -310,8 +331,7 @@ void minClusterAndDistanceCompute(raft::resources const& handle, auto datasetView = raft::make_device_matrix_view( X.data_handle() + (dIdx * n_features), ns, n_features); - auto temp_kvp_view = raft::make_device_vector_view, IndexT>( - temp_kvp.data_handle() + dIdx, ns); + auto temp_kvp_view = raft::make_device_vector_view(kvp_output + dIdx, ns); for (IndexT cIdx = 0; cIdx < n_clusters; cIdx += centroidsBatchSize) { auto nc = std::min((IndexT)centroidsBatchSize, n_clusters - cIdx); @@ -334,7 +354,7 @@ void minClusterAndDistanceCompute(raft::resources const& handle, stream, true, [=] __device__(const DataT val, const IndexT i) { - raft::KeyValuePair pair; + KeyValueT pair; pair.key = cIdx + i; pair.value = val; return pair; @@ -344,10 +364,86 @@ void minClusterAndDistanceCompute(raft::resources const& handle, } } - unpack_kvp(handle, nearest_idx, nearest_dist, raft::make_const_mdspan(temp_kvp.view())); + if (native_kvp == nullptr) { + unpack_kvp(handle, nearest_idx, nearest_dist, kvp_output, n_samples); + } } } +template +void minClusterAndDistanceCompute(raft::resources const& handle, + raft::device_matrix_view X, + raft::device_matrix_view centroids, + raft::device_vector_view nearest_idx, + raft::device_vector_view nearest_dist, + raft::device_vector_view L2NormX, + rmm::device_uvector& L2NormBuf_OR_DistBuf, + cuvs::distance::DistanceType metric, + int batch_samples, + int batch_centroids, + rmm::device_uvector& workspace) +{ + auto path = select_min_cluster_distance_path(handle, X, centroids, metric); + if constexpr (is_cutile_fused_data_type_v) { + if (path == FusedDistancePath::FusedCutile && + !cuvs::distance::detail::can_launch_fused_1nn_tile(nearest_idx.data_handle(), + nearest_dist.data_handle(), + X.data_handle(), + centroids.data_handle(), + X.extent(0), + centroids.extent(0), + X.extent(1), + metric)) { + path = use_legacy_fused(handle, X.extent(0), centroids.extent(0), metric); + } + } + + min_cluster_and_distance_compute_impl(handle, + X, + centroids, + nearest_idx.data_handle(), + nearest_dist.data_handle(), + nullptr, + L2NormX, + L2NormBuf_OR_DistBuf, + metric, + batch_samples, + batch_centroids, + workspace, + path); +} + +template +void minClusterAndDistanceComputeKvp( + raft::resources const& handle, + raft::device_matrix_view X, + raft::device_matrix_view centroids, + raft::device_vector_view, IndexT> nearest, + raft::device_vector_view L2NormX, + rmm::device_uvector& L2NormBuf_OR_DistBuf, + cuvs::distance::DistanceType metric, + int batch_samples, + int batch_centroids, + rmm::device_uvector& workspace, + FusedDistancePath path) +{ + RAFT_EXPECTS(path != FusedDistancePath::FusedCutile, + "cuTile fused 1-NN cannot write native KVP output"); + min_cluster_and_distance_compute_impl(handle, + X, + centroids, + nullptr, + nullptr, + nearest.data_handle(), + L2NormX, + L2NormBuf_OR_DistBuf, + metric, + batch_samples, + batch_centroids, + workspace, + path); +} + #define INSTANTIATE_MIN_CLUSTER_AND_DISTANCE(DataT, IndexT) \ template void minClusterAndDistanceCompute( \ raft::resources const& handle, \ @@ -369,6 +465,41 @@ INSTANTIATE_MIN_CLUSTER_AND_DISTANCE(double, int) #undef INSTANTIATE_MIN_CLUSTER_AND_DISTANCE +#define INSTANTIATE_SELECT_MIN_CLUSTER_PATH(DataT, IndexT) \ + template FusedDistancePath select_min_cluster_distance_path( \ + raft::resources const&, \ + raft::device_matrix_view, \ + raft::device_matrix_view, \ + cuvs::distance::DistanceType); + +INSTANTIATE_SELECT_MIN_CLUSTER_PATH(float, int64_t) +INSTANTIATE_SELECT_MIN_CLUSTER_PATH(double, int64_t) +INSTANTIATE_SELECT_MIN_CLUSTER_PATH(float, int) +INSTANTIATE_SELECT_MIN_CLUSTER_PATH(double, int) + +#undef INSTANTIATE_SELECT_MIN_CLUSTER_PATH + +#define INSTANTIATE_MIN_CLUSTER_AND_DISTANCE_KVP(DataT, IndexT) \ + template void minClusterAndDistanceComputeKvp( \ + raft::resources const&, \ + raft::device_matrix_view, \ + raft::device_matrix_view, \ + raft::device_vector_view, IndexT>, \ + raft::device_vector_view, \ + rmm::device_uvector&, \ + cuvs::distance::DistanceType, \ + int, \ + int, \ + rmm::device_uvector&, \ + FusedDistancePath); + +INSTANTIATE_MIN_CLUSTER_AND_DISTANCE_KVP(float, int64_t) +INSTANTIATE_MIN_CLUSTER_AND_DISTANCE_KVP(double, int64_t) +INSTANTIATE_MIN_CLUSTER_AND_DISTANCE_KVP(float, int) +INSTANTIATE_MIN_CLUSTER_AND_DISTANCE_KVP(double, int) + +#undef INSTANTIATE_MIN_CLUSTER_AND_DISTANCE_KVP + template void minClusterDistanceCompute(raft::resources const& handle, raft::device_matrix_view X, @@ -390,8 +521,6 @@ void minClusterDistanceCompute(raft::resources const& handle, metric == cuvs::distance::DistanceType::L2SqrtExpanded || metric == cuvs::distance::DistanceType::CosineExpanded; - raft::matrix::fill(handle, minClusterDistance, std::numeric_limits::max()); - FusedDistancePath fused_path = is_l2_cos ? use_fused(handle, n_samples, n_clusters, n_features, metric) : FusedDistancePath::Unfused; @@ -456,13 +585,7 @@ void minClusterDistanceCompute(raft::resources const& handle, } } - raft::KeyValuePair* cutlass_kvp_scratch = nullptr; - rmm::device_uvector> temp_kvp(0, stream); - if (!cutile_ready) { - temp_kvp.resize(n_samples, stream); - cutlass_kvp_scratch = temp_kvp.data(); - workspace.resize(sizeof(int) * static_cast(n_samples), stream); - } + if (!cutile_ready) { workspace.resize(sizeof(int) * static_cast(n_samples), stream); } cuvs::distance::fusedDistanceNNMinReduce( nullptr, @@ -480,9 +603,10 @@ void minClusterDistanceCompute(raft::resources const& handle, true, metric, 0.0f, - cutlass_kvp_scratch, + nullptr, stream); } else { + raft::matrix::fill(handle, minClusterDistance, std::numeric_limits::max()); auto dataBatchSize = getDataBatchSize(batch_samples, n_samples); auto centroidsBatchSize = getCentroidsBatchSize(batch_centroids, n_clusters); diff --git a/cpp/src/distance/detail/fused_distance_nn.cuh b/cpp/src/distance/detail/fused_distance_nn.cuh index 669bb05f1f..33bf6fc953 100644 --- a/cpp/src/distance/detail/fused_distance_nn.cuh +++ b/cpp/src/distance/detail/fused_distance_nn.cuh @@ -71,60 +71,66 @@ void fusedDistanceNNImpl(IdxT* nearest_idx, RAFT_EXPECTS(metric != cuvs::distance::DistanceType::InnerProduct, "Fused InnerProduct 1-NN requires a compatible cuTile launcher"); - RAFT_EXPECTS(cutlass_kvp_scratch != nullptr, "CUTLASS fused 1-NN requires a scratch KVP buffer"); - - if (initOutBuffer) { - initFused1nnOutput(nearest_idx, nearest_dist, m, std::numeric_limits::max(), stream); - } - - MinAndDistanceReduceOpImpl cutlass_redOp; - cutlass_redOp.out_kvp = cutlass_kvp_scratch; - initialize( - cutlass_kvp_scratch, m, maxVal, cutlass_redOp, stream); - RAFT_CUDA_TRY(cudaMemsetAsync(workspace, 0, sizeof(int) * m, stream)); - switch (metric) { - case cuvs::distance::DistanceType::CosineExpanded: - fusedCosineNN(nearest_idx, - nearest_dist, - x, - y, - xn, - yn, - m, - n, - k, - workspace, - cutlass_redOp, - pairRedOp, - sqrt, - cutlass_kvp_scratch, - stream); - break; - case cuvs::distance::DistanceType::L2SqrtExpanded: - case cuvs::distance::DistanceType::L2Expanded: - fusedL2NNImpl(nearest_idx, - nearest_dist, - x, - y, - xn, - yn, - m, - n, - k, - workspace, - cutlass_redOp, - pairRedOp, - sqrt, - false, - cutlass_kvp_scratch, - stream); - break; - default: assert("only cosine/l2 metric is supported with fusedDistanceNN\n"); break; - } + auto launch_legacy = [&](OutT* out, auto cutlass_red_op) { + switch (metric) { + case cuvs::distance::DistanceType::CosineExpanded: + fusedCosineNN(nearest_idx, + nearest_dist, + x, + y, + xn, + yn, + m, + n, + k, + workspace, + cutlass_red_op, + pairRedOp, + sqrt, + out, + stream); + break; + case cuvs::distance::DistanceType::L2SqrtExpanded: + case cuvs::distance::DistanceType::L2Expanded: + fusedL2NNImpl(nearest_idx, + nearest_dist, + x, + y, + xn, + yn, + m, + n, + k, + workspace, + cutlass_red_op, + pairRedOp, + sqrt, + false, + out, + stream); + break; + default: assert("only cosine/l2 metric is supported with fusedDistanceNN\n"); break; + } + }; - unpackFused1nnKvpToSoa(nearest_idx, nearest_dist, cutlass_kvp_scratch, m, stream); + MinAndDistanceReduceOpImpl cutlass_red_op; + if (cutlass_kvp_scratch != nullptr) { + if (initOutBuffer) { initFused1nnOutput(nearest_idx, nearest_dist, m, maxVal, stream); } + cutlass_red_op.out_kvp = cutlass_kvp_scratch; + initialize( + cutlass_kvp_scratch, m, maxVal, cutlass_red_op, stream); + launch_legacy(cutlass_kvp_scratch, cutlass_red_op); + unpackFused1nnKvpToSoa(nearest_idx, nearest_dist, cutlass_kvp_scratch, m, stream); + } else { + RAFT_EXPECTS(nearest_idx == nullptr && nearest_dist != nullptr, + "Direct CUTLASS output supports distance-only results"); + cutlass_red_op.out_dist = nearest_dist; + initialize( + nearest_dist, m, maxVal, cutlass_red_op, stream); + launch_legacy(nearest_dist, cutlass_red_op); + } } } // namespace detail diff --git a/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_tile.cu b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_tile.cu index 836fc891a4..b93931c977 100644 --- a/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_tile.cu +++ b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_tile.cu @@ -196,21 +196,13 @@ bool try_fused_1nn_tile_dispatch(IdxT* nearest_idx, template requires is_fused_1nn_cutile_data_v -bool can_launch_fused_1nn_tile(IdxT* nearest_idx, - DataT* nearest_dist, - const DataT* x, - const DataT* y, - IdxT m, - IdxT n, - IdxT k, - cuvs::distance::DistanceType metric) +bool can_launch_fused_1nn_tile( + const DataT* x, const DataT* y, IdxT m, IdxT n, IdxT k, cuvs::distance::DistanceType metric) { if (!cuvs::detail::jit_lto::cutile_launch_available_on_current_device()) { return false; } static_assert(std::is_same_v || std::is_same_v); - if (nearest_dist == nullptr || x == nullptr || y == nullptr || m <= 0 || n <= 0 || k <= 0) { - return false; - } + if (x == nullptr || y == nullptr || m <= 0 || n <= 0 || k <= 0) { return false; } if (metric != cuvs::distance::DistanceType::InnerProduct && metric != cuvs::distance::DistanceType::L2Expanded && metric != cuvs::distance::DistanceType::L2SqrtExpanded && @@ -218,12 +210,8 @@ bool can_launch_fused_1nn_tile(IdxT* nearest_idx, return false; } - if (!is_16_byte_aligned(nearest_dist) || !is_16_byte_aligned(x) || !is_16_byte_aligned(y)) { - return false; - } - if constexpr (std::is_same_v) { - if (!is_16_byte_aligned(nearest_idx)) { return false; } - } else { + if (!is_16_byte_aligned(x) || !is_16_byte_aligned(y)) { return false; } + if constexpr (std::is_same_v) { constexpr int64_t max_i32 = std::numeric_limits::max(); if (n > max_i32 || k > max_i32) { return false; } } @@ -233,6 +221,25 @@ bool can_launch_fused_1nn_tile(IdxT* nearest_idx, : has_fused_1nn_tile_launcher(); } +template + requires is_fused_1nn_cutile_data_v +bool can_launch_fused_1nn_tile(IdxT* nearest_idx, + DataT* nearest_dist, + const DataT* x, + const DataT* y, + IdxT m, + IdxT n, + IdxT k, + cuvs::distance::DistanceType metric) +{ + if (!can_launch_fused_1nn_tile(x, y, m, n, k, metric)) { return false; } + if (nearest_dist == nullptr || !is_16_byte_aligned(nearest_dist)) { return false; } + if constexpr (std::is_same_v) { + if (!is_16_byte_aligned(nearest_idx)) { return false; } + } + return true; +} + template requires is_fused_1nn_cutile_data_v bool can_launch_fused_1nn_tile(IdxT* nearest_idx, @@ -337,6 +344,17 @@ bool try_fused_1nn_tile(IdxT* nearest_idx, } } +#define CUVS_INST_CAN_LAUNCH_FUSED_1NN_TILE_INPUTS(DataT, IdxT) \ + template CUVS_EXPORT bool can_launch_fused_1nn_tile( \ + const DataT*, const DataT*, IdxT, IdxT, IdxT, cuvs::distance::DistanceType) + +CUVS_INST_CAN_LAUNCH_FUSED_1NN_TILE_INPUTS(float, int); +CUVS_INST_CAN_LAUNCH_FUSED_1NN_TILE_INPUTS(float, int64_t); +CUVS_INST_CAN_LAUNCH_FUSED_1NN_TILE_INPUTS(half, int); +CUVS_INST_CAN_LAUNCH_FUSED_1NN_TILE_INPUTS(half, int64_t); + +#undef CUVS_INST_CAN_LAUNCH_FUSED_1NN_TILE_INPUTS + #define CUVS_INST_CAN_LAUNCH_FUSED_1NN_TILE_PREFLIGHT(DataT, IdxT) \ template CUVS_EXPORT bool can_launch_fused_1nn_tile( \ IdxT*, DataT*, const DataT*, const DataT*, IdxT, IdxT, IdxT, cuvs::distance::DistanceType) diff --git a/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_tile.hpp b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_tile.hpp index 5a4b339425..add5f7a94e 100644 --- a/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_tile.hpp +++ b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_tile.hpp @@ -48,6 +48,16 @@ constexpr size_t fused_1nn_cutile_index_workspace_rows(IdxT m) } #if CUVS_CUTILE_ENABLED +/** + * Return whether the input problem has a compatible cuTile launcher. + * + * This output-independent probe lets callers select native result storage before allocating it. + */ +template + requires is_fused_1nn_cutile_data_v +bool can_launch_fused_1nn_tile( + const DataT* x, const DataT* y, IdxT m, IdxT n, IdxT k, cuvs::distance::DistanceType metric); + /** * Return whether the supplied problem can use cuTile without fallback scratch. * @@ -101,6 +111,13 @@ bool try_fused_1nn_tile(IdxT* nearest_idx, void* index_workspace, cudaStream_t stream); #else +template +bool can_launch_fused_1nn_tile( + const DataT*, const DataT*, IdxT, IdxT, IdxT, cuvs::distance::DistanceType) +{ + return false; +} + template bool can_launch_fused_1nn_tile( IdxT*, DataT*, const DataT*, const DataT*, IdxT, IdxT, IdxT, cuvs::distance::DistanceType) diff --git a/cpp/src/distance/detail/fused_distance_nn/fused_cosine_nn.cuh b/cpp/src/distance/detail/fused_distance_nn/fused_cosine_nn.cuh index 43059a681c..5dea6c2256 100644 --- a/cpp/src/distance/detail/fused_distance_nn/fused_cosine_nn.cuh +++ b/cpp/src/distance/detail/fused_distance_nn/fused_cosine_nn.cuh @@ -24,7 +24,12 @@ namespace distance { namespace detail { -template +template void fusedCosineNN(IdxT* nearest_idx, DataT* nearest_dist, const DataT* x, @@ -38,14 +43,13 @@ void fusedCosineNN(IdxT* nearest_idx, ReduceOpT redOp, KVPReduceOpT pairRedOp, bool sqrt, - raft::KeyValuePair* cutlass_out, + OutT* cutlass_out, cudaStream_t stream) { typedef Policy P; dim3 blk(P::Nthreads); constexpr auto maxVal = std::numeric_limits::max(); - typedef raft::KeyValuePair KVPair; if (cutlass_out == nullptr) { initFused1nnOutput(nearest_idx, nearest_dist, m, maxVal, stream); @@ -59,7 +63,7 @@ void fusedCosineNN(IdxT* nearest_idx, raft::identity_op fin_op{}; auto kernel = fusedDistanceNNkernel +template void fusedL2NNImpl(IdxT* nearest_idx, DataT* nearest_dist, const DataT* x, @@ -39,14 +44,13 @@ void fusedL2NNImpl(IdxT* nearest_idx, KVPReduceOpT pairRedOp, bool sqrt, bool initOutBuffer, - raft::KeyValuePair* cutlass_out, + OutT* cutlass_out, cudaStream_t stream) { typedef Policy P; dim3 blk(P::Nthreads); constexpr auto maxVal = std::numeric_limits::max(); - typedef raft::KeyValuePair KVPair; if (initOutBuffer && cutlass_out == nullptr) { initFused1nnOutput(nearest_idx, nearest_dist, m, maxVal, stream); @@ -60,7 +64,7 @@ void fusedL2NNImpl(IdxT* nearest_idx, raft::identity_op fin_op{}; auto kernel = fusedDistanceNNkernel void fusedDistanceNNMinReduce(IdxT* nearest_idx, From cafd051a4605686709daf8ad90f905110eb57a0a Mon Sep 17 00:00:00 2001 From: divyegala Date: Wed, 2 Sep 2026 20:49:37 +0000 Subject: [PATCH 68/82] working through path unification --- .../cuvs/detail/jit_lto/cutile_module.hpp | 37 +- .../cuvs/detail/jit_lto/tileir_compat.hpp | 3 + cpp/src/cluster/detail/kmeans_balanced.cuh | 393 ++++++++++++------ cpp/src/cluster/detail/kmeans_common.cuh | 36 +- .../detail/minClusterDistanceCompute.cu | 71 +++- .../detail/jit_lto/TileAlgorithmPlanner.cpp | 22 +- cpp/src/distance/detail/fused_distance_nn.cuh | 26 +- .../cutile/export_fused_1nn.py | 15 +- .../cutile/fused_1nn_tile.cu | 60 ++- cpp/src/distance/fused_distance_nn-inl.cuh | 18 +- cpp/tests/cluster/kmeans_balanced.cu | 50 +++ cpp/tests/neighbors/distance_nn.cu | 148 ++++++- 12 files changed, 697 insertions(+), 182 deletions(-) diff --git a/cpp/include/cuvs/detail/jit_lto/cutile_module.hpp b/cpp/include/cuvs/detail/jit_lto/cutile_module.hpp index ae46b523e1..4e39450ba1 100644 --- a/cpp/include/cuvs/detail/jit_lto/cutile_module.hpp +++ b/cpp/include/cuvs/detail/jit_lto/cutile_module.hpp @@ -78,15 +78,44 @@ inline std::optional resolve_cutile_module_image( return std::nullopt; } -inline std::shared_ptr load_cutile_launcher( +inline bool is_expected_cutile_unavailable(cudaError_t status) +{ + switch (status) { + case cudaErrorInvalidDeviceFunction: + case cudaErrorInvalidPtx: + case cudaErrorNoKernelImageForDevice: + case cudaErrorSymbolNotFound: + case cudaErrorUnsupportedPtxVersion: + case cudaErrorCallRequiresNewerDriver: + case cudaErrorSharedObjectSymbolNotFound: + case cudaErrorSharedObjectInitFailed: + case cudaErrorJitCompilerNotFound: return true; + default: return false; + } +} + +/** + * Loads a cuTile launcher, returning null for an expected module/JIT compatibility rejection. + * Unexpected CUDA failures retain the normal RAFT exception behavior. + */ +inline std::shared_ptr try_load_cutile_launcher( const CutileModuleImage& image, const std::string& kernel_symbol) { cudaLibrary_t library{}; - RAFT_CUDA_TRY( - cudaLibraryLoadData(&library, image.data, nullptr, nullptr, 0, nullptr, nullptr, 0)); + auto load_status = + cudaLibraryLoadData(&library, image.data, nullptr, nullptr, 0, nullptr, nullptr, 0); + if (load_status != cudaSuccess) { + if (is_expected_cutile_unavailable(load_status)) { return nullptr; } + RAFT_CUDA_TRY(load_status); + } cudaKernel_t kernel{}; - RAFT_CUDA_TRY(cudaLibraryGetKernel(&kernel, library, kernel_symbol.c_str())); + load_status = cudaLibraryGetKernel(&kernel, library, kernel_symbol.c_str()); + if (load_status != cudaSuccess) { + RAFT_CUDA_TRY(cudaLibraryUnload(library)); + if (is_expected_cutile_unavailable(load_status)) { return nullptr; } + RAFT_CUDA_TRY(load_status); + } return std::make_shared(kernel, library); } diff --git a/cpp/include/cuvs/detail/jit_lto/tileir_compat.hpp b/cpp/include/cuvs/detail/jit_lto/tileir_compat.hpp index bde7dab302..a03a7a78fc 100644 --- a/cpp/include/cuvs/detail/jit_lto/tileir_compat.hpp +++ b/cpp/include/cuvs/detail/jit_lto/tileir_compat.hpp @@ -65,6 +65,9 @@ inline bool tileir_fallback_available(int driver_version) inline bool cutile_launch_available_for_arch(int cc_major, int cc_minor, int driver_version) { if (!runtime_cuda13_or_newer()) { return false; } + // The exported fused-1NN kernels require Ampere-or-newer tensor-core semantics, and the current + // integration is validated only through the SM12 family. + if (cc_major < 8 || cc_major > 12) { return false; } if (has_embedded_cubin_for_arch(cc_major, cc_minor)) { return true; } return tileir_fallback_available(driver_version); } diff --git a/cpp/src/cluster/detail/kmeans_balanced.cuh b/cpp/src/cluster/detail/kmeans_balanced.cuh index ebda7397a2..65db51f93d 100644 --- a/cpp/src/cluster/detail/kmeans_balanced.cuh +++ b/cpp/src/cluster/detail/kmeans_balanced.cuh @@ -44,6 +44,7 @@ #include #include +#include #include #include #include @@ -53,6 +54,93 @@ namespace cuvs::cluster::kmeans::detail { +template +bool predict_core_min_cluster(const raft::resources& handle, + raft::device_matrix_view X, + raft::device_matrix_view centroids, + raft::device_vector_view X_norm, + cuvs::distance::DistanceType metric, + LabelT* labels, + rmm::device_uvector& L2NormBuf_OR_DistBuf, + rmm::device_uvector& workspace, + rmm::device_async_resource_ref mr, + const MathT* cutile_x_norm) +{ + auto n_rows = X.extent(0); + auto path = select_min_cluster_distance_path(handle, X, centroids, metric); + + if constexpr (std::is_same_v && std::is_same_v) { + if (path == FusedDistancePath::FusedCutile && + reinterpret_cast(labels) % 16 != 0) { + path = use_legacy_fused(handle, n_rows, centroids.extent(0), metric); + } + } + + if (path == FusedDistancePath::FusedCutile) { + auto nearest_dist = + raft::make_device_mdarray(handle, mr, raft::make_extents(n_rows)); + + if constexpr (std::is_same_v) { + auto labels_view = raft::make_device_vector_view(labels, n_rows); + minClusterAndDistanceCompute(handle, + X, + centroids, + labels_view, + nearest_dist.view(), + X_norm, + L2NormBuf_OR_DistBuf, + metric, + 0, + 0, + workspace, + cutile_x_norm); + } else { + auto nearest_idx = + raft::make_device_mdarray(handle, mr, raft::make_extents(n_rows)); + minClusterAndDistanceCompute(handle, + X, + centroids, + nearest_idx.view(), + nearest_dist.view(), + X_norm, + L2NormBuf_OR_DistBuf, + metric, + 0, + 0, + workspace, + cutile_x_norm); + raft::copy( + handle, raft::make_device_vector_view(labels, n_rows), nearest_idx.view()); + } + return true; + } + + // The generic pairwise-distance assignment path does not implement InnerProduct. Preserve the + // historical GEMM-plus-argmin fallback in predict_core after any final cuTile rejection. + if (metric == cuvs::distance::DistanceType::InnerProduct) { return false; } + + using KvpT = raft::KeyValuePair; + auto nearest = + raft::make_device_mdarray(handle, mr, raft::make_extents(n_rows)); + minClusterAndDistanceComputeKvp(handle, + X, + centroids, + nearest.view(), + X_norm, + L2NormBuf_OR_DistBuf, + metric, + 0, + 0, + workspace, + path); + auto* nearest_ptr = nearest.data_handle(); + raft::linalg::map_offset( + handle, + raft::make_device_vector_view(labels, n_rows), + [nearest_ptr] __device__(IdxT i) -> LabelT { return static_cast(nearest_ptr[i].key); }); + return true; +} + /** * @brief Predict labels for the dataset; floating-point types only. * @@ -83,6 +171,7 @@ inline std::enable_if_t> predict_core( IdxT dim, const MathT* dataset, const MathT* dataset_norm, + const MathT* dataset_cutile_norm, IdxT n_rows, LabelT* labels, rmm::device_async_resource_ref mr) @@ -100,89 +189,39 @@ inline std::enable_if_t> predict_core( raft::make_device_matrix_view(centers, n_clusters, dim); auto X_norm_view = raft::make_device_vector_view(dataset_norm, n_rows); - auto nearest_dist = - raft::make_device_mdarray(handle, mr, raft::make_extents(n_rows)); - - if constexpr (std::is_same_v) { - auto labels_view = raft::make_device_vector_view(labels, n_rows); - cuvs::cluster::kmeans::min_cluster_and_distance( - handle, - X_view, - centroids_view, - labels_view, - nearest_dist.view(), - X_norm_view, - L2NormBuf_OR_DistBuf, - params.metric, - 0, // batch_samples (unused for fused reduction) - 0, // batch_centroids (unused for fused reduction) - workspace); - } else { - auto nearest_idx = - raft::make_device_mdarray(handle, mr, raft::make_extents(n_rows)); - cuvs::cluster::kmeans::min_cluster_and_distance(handle, - X_view, - centroids_view, - nearest_idx.view(), - nearest_dist.view(), - X_norm_view, - L2NormBuf_OR_DistBuf, - params.metric, - 0, - 0, - workspace); - raft::copy( - handle, raft::make_device_vector_view(labels, n_rows), nearest_idx.view()); - } + predict_core_min_cluster(handle, + X_view, + centroids_view, + X_norm_view, + params.metric, + labels, + L2NormBuf_OR_DistBuf, + workspace, + mr, + dataset_cutile_norm); break; } case cuvs::distance::DistanceType::InnerProduct: { - if (uses_fused_distance_nn( - use_fused(handle, n_rows, n_clusters, dim, params.metric))) { - rmm::device_uvector L2NormBuf_OR_DistBuf(0, stream, mr); - rmm::device_uvector workspace(0, stream, mr); - - auto X_view = raft::make_device_matrix_view(dataset, n_rows, dim); - auto centroids_view = - raft::make_device_matrix_view(centers, n_clusters, dim); - auto X_norm_view = raft::make_device_vector_view(dataset_norm, n_rows); - - auto nearest_dist = - raft::make_device_mdarray(handle, mr, raft::make_extents(n_rows)); - - if constexpr (std::is_same_v) { - auto labels_view = raft::make_device_vector_view(labels, n_rows); - cuvs::cluster::kmeans::min_cluster_and_distance(handle, - X_view, - centroids_view, - labels_view, - nearest_dist.view(), - X_norm_view, - L2NormBuf_OR_DistBuf, - params.metric, - 0, - 0, - workspace); - } else { - auto nearest_idx = - raft::make_device_mdarray(handle, mr, raft::make_extents(n_rows)); - cuvs::cluster::kmeans::min_cluster_and_distance(handle, - X_view, - centroids_view, - nearest_idx.view(), - nearest_dist.view(), - X_norm_view, - L2NormBuf_OR_DistBuf, - params.metric, - 0, - 0, - workspace); - raft::copy(handle, - raft::make_device_vector_view(labels, n_rows), - nearest_idx.view()); - } - } else { - rmm::device_uvector distances(n_rows * n_clusters, stream, mr); + rmm::device_uvector L2NormBuf_OR_DistBuf(0, stream, mr); + rmm::device_uvector workspace(0, stream, mr); + + auto X_view = raft::make_device_matrix_view(dataset, n_rows, dim); + auto centroids_view = + raft::make_device_matrix_view(centers, n_clusters, dim); + auto X_norm_view = raft::make_device_vector_view(dataset_norm, n_rows); + + if (!predict_core_min_cluster(handle, + X_view, + centroids_view, + X_norm_view, + params.metric, + labels, + L2NormBuf_OR_DistBuf, + workspace, + mr, + dataset_cutile_norm)) { + rmm::device_uvector distances( + static_cast(n_rows) * static_cast(n_clusters), stream, mr); MathT alpha = -1.0; MathT beta = 0.0; @@ -245,7 +284,16 @@ auto calc_minibatch_size(const raft::resources& handle, n_clusters = std::max(1, n_clusters); // Estimate memory needs per row (i.e element of the batch). - size_t mem_per_row = 0; + auto saturating_add = [](size_t a, size_t b) { + return b > std::numeric_limits::max() - a ? std::numeric_limits::max() : a + b; + }; + auto saturating_multiply = [](size_t a, size_t b) { + return b != 0 && a > std::numeric_limits::max() / b ? std::numeric_limits::max() + : a * b; + }; + size_t common_mem_per_row = 0; + size_t path_mem_per_row = 0; + size_t fixed_bytes = 0; switch (metric) { case distance::DistanceType::L2Expanded: case distance::DistanceType::L2SqrtExpanded: @@ -253,62 +301,88 @@ auto calc_minibatch_size(const raft::resources& handle, case distance::DistanceType::InnerProduct: { const auto fused_path = use_fused(handle, n_rows, n_clusters, dim, metric); - // min_cluster_and_distance always materializes the nearest distance for fused/L2 paths. - if (metric != distance::DistanceType::InnerProduct || - fused_path != FusedDistancePath::Unfused) { - mem_per_row += sizeof(MathT); - if constexpr (!std::is_same_v) { mem_per_row += sizeof(IdxT); } - } if (metric != distance::DistanceType::InnerProduct) { // predict may need a minibatch-sized input-norm buffer before entering predict_core. - mem_per_row += sizeof(MathT); + common_mem_per_row += sizeof(MathT); + fixed_bytes = saturating_multiply(sizeof(MathT), static_cast(n_clusters)); } - switch (fused_path) { - case FusedDistancePath::FusedCutile: - // Conservatively budget the fallback in case the eventual pointer-aware probe fails. - mem_per_row += sizeof(int); - mem_per_row += sizeof(raft::KeyValuePair); - if constexpr (std::is_same_v) { + auto path_bytes_per_row = [&](FusedDistancePath path) { + size_t bytes = 0; + switch (path) { + case FusedDistancePath::FusedCutile: + // cuTile writes separate distance and index arrays. + bytes += sizeof(MathT); + if constexpr (!std::is_same_v) { bytes += sizeof(IdxT); } + if constexpr (std::is_same_v) { + // cuTile converts chunk-local int32 indices to int64 output indices. + bytes += sizeof(int); + } + if constexpr (std::is_same_v) { + if (metric != distance::DistanceType::InnerProduct) { + // TF32-compatible row norms are materialized for cuTile L2/cosine. + bytes += sizeof(MathT); + } + } + break; + case FusedDistancePath::FusedCutlass: + // CUTLASS writes native KVP assignments and uses one mutex per row. + bytes += sizeof(int); + bytes += sizeof(raft::KeyValuePair); + break; + case FusedDistancePath::Unfused: + // Unfused assignment needs a full distance matrix row. + bytes = saturating_add( + bytes, saturating_multiply(sizeof(MathT), static_cast(n_clusters))); if (metric != distance::DistanceType::InnerProduct) { - // TF32-compatible row norms are materialized for cuTile L2/cosine. - mem_per_row += sizeof(MathT); + bytes += sizeof(raft::KeyValuePair); } - } - break; - case FusedDistancePath::FusedCutlass: - // fusedDistanceNNMinReduce CUTLASS fallback: mutex workspace + scratch KVP per row. - mem_per_row += sizeof(int); - mem_per_row += sizeof(raft::KeyValuePair); - break; - case FusedDistancePath::Unfused: - // unfused / GEMM+argmin path needs a full distance matrix row. - mem_per_row += sizeof(MathT) * n_clusters; - if (metric != distance::DistanceType::InnerProduct) { - mem_per_row += sizeof(raft::KeyValuePair); - } - break; + break; + } + return bytes; + }; + + path_mem_per_row = path_bytes_per_row(fused_path); + if (fused_path == FusedDistancePath::FusedCutile) { + const auto prop = raft::resource::get_device_properties(handle); + const auto fallback = use_legacy_fused(prop.major, n_rows, n_clusters, metric); + path_mem_per_row = std::max(path_mem_per_row, path_bytes_per_row(fallback)); + + // Hopper changes from unfused to CUTLASS at 4096 rows. A final minibatch can fall below + // that threshold even when full batches use CUTLASS, so reserve for both real outcomes. + if (prop.major == 9 && n_clusters < IdxT{4096} && n_rows >= IdxT{4096}) { + path_mem_per_row = + std::max(path_mem_per_row, path_bytes_per_row(FusedDistancePath::Unfused)); + } } } break; // Other metrics require storing a distance matrix. default: { - mem_per_row += sizeof(MathT) * n_clusters; + path_mem_per_row = saturating_multiply(sizeof(MathT), static_cast(n_clusters)); } } // If we need to convert to MathT, space required for the converted batch. - if (!data_is_math_type) { mem_per_row += sizeof(MathT) * dim; } + if (!data_is_math_type) { + common_mem_per_row = saturating_add( + common_mem_per_row, saturating_multiply(sizeof(MathT), static_cast(dim))); + } + const size_t mem_per_row = saturating_add(common_mem_per_row, path_mem_per_row); // Heuristic: calculate the minibatch size in order to use at most 80% or 512MB workspace memory. // We go below 1GB here as the allocation is mostly done in a single chunk which // is problematic if e.g. a pool allocator manages its own chunks <= 1GB. const auto free_ws_size = raft::resource::get_workspace_free_bytes(handle); - const auto available_ws_size = + const auto workspace_limit = std::min((free_ws_size * size_t{8}) / size_t{10}, size_t{1} << 29); + const auto available_ws_size = + workspace_limit > fixed_bytes ? workspace_limit - fixed_bytes : size_t{1}; IdxT minibatch_size = std::max(IdxT{1}, static_cast(available_ws_size / mem_per_row)); - minibatch_size = raft::round_down_safe(minibatch_size, IdxT{64}); + if (minibatch_size >= IdxT{64}) { + minibatch_size = raft::round_down_safe(minibatch_size, IdxT{64}); + } minibatch_size = std::min(minibatch_size, n_rows); return std::make_tuple(minibatch_size, mem_per_row); } @@ -448,6 +522,34 @@ void compute_norm(const raft::resources& handle, norm_fin_op); } +/** Computes TF32-compatible norms for cuTile, converting to float when necessary. */ +template +void compute_cutile_norm(const raft::resources& handle, + float* dataset_norm, + const T* dataset, + IdxT dim, + IdxT n_rows, + MappingOpT mapping_op, + bool take_sqrt, + rmm::device_async_resource_ref mr) +{ + auto stream = raft::resource::get_cuda_stream(handle); + rmm::device_uvector mapped_dataset(0, stream, mr); + const float* dataset_ptr = nullptr; + if constexpr (std::is_same_v) { + dataset_ptr = dataset; + } else { + mapped_dataset.resize(static_cast(n_rows) * static_cast(dim), stream); + raft::linalg::map( + handle, + raft::make_device_vector_view(dataset, n_rows * dim), + raft::make_device_vector_view(mapped_dataset.data(), n_rows * dim), + mapping_op); + dataset_ptr = mapped_dataset.data(); + } + computeCutileRowNorms(handle, dataset_ptr, dataset_norm, n_rows, dim, take_sqrt); +} + /** * @brief Predict labels for the dataset. * @@ -480,7 +582,8 @@ void predict(const raft::resources& handle, LabelT* labels, MappingOpT mapping_op, std::optional mr = std::nullopt, - const MathT* dataset_norm = nullptr) + const MathT* dataset_norm = nullptr, + const MathT* dataset_cutile_norm = nullptr) { auto stream = raft::resource::get_cuda_stream(handle); raft::common::nvtx::range fun_scope( @@ -543,6 +646,7 @@ void predict(const raft::resources& handle, dim, cur_dataset_ptr, dataset_norm_ptr, + dataset_cutile_norm == nullptr ? nullptr : dataset_cutile_norm + offset, minibatch_size, labels + offset, mem_res); @@ -872,6 +976,7 @@ void balancing_em_iters(const raft::resources& handle, IdxT dim, const T* dataset, const MathT* dataset_norm, + const MathT* dataset_cutile_norm, IdxT n_rows, IdxT n_clusters, MathT* cluster_centers, @@ -940,7 +1045,8 @@ void balancing_em_iters(const raft::resources& handle, cluster_labels, mapping_op, device_memory, - dataset_norm); + dataset_norm, + dataset_cutile_norm); // M: Maximization step - calculate optimal cluster centers calc_centers_and_sizes(handle, cluster_centers, @@ -974,7 +1080,8 @@ void build_clusters(const raft::resources& handle, CounterT* cluster_sizes, MappingOpT mapping_op, rmm::device_async_resource_ref device_memory, - const MathT* dataset_norm = nullptr) + const MathT* dataset_norm = nullptr, + const MathT* dataset_cutile_norm = nullptr) { auto stream = raft::resource::get_cuda_stream(handle); // "randomly" initialize labels @@ -1004,6 +1111,7 @@ void build_clusters(const raft::resources& handle, dim, dataset, dataset_norm, + dataset_cutile_norm, n_rows, n_clusters, cluster_centers, @@ -1105,6 +1213,7 @@ auto build_fine_clusters(const raft::resources& handle, IdxT dim, const T* dataset_mptr, const MathT* dataset_norm_mptr, + const MathT* dataset_cutile_norm_mptr, const LabelT* labels_mptr, IdxT n_rows, const IdxT* fine_clusters_nums, @@ -1125,9 +1234,12 @@ auto build_fine_clusters(const raft::resources& handle, auto large_ws = raft::resource::get_large_workspace_resource_ref(handle); rmm::device_uvector mc_trainset_buf(mesocluster_size_max * dim, stream, large_ws); rmm::device_uvector mc_trainset_norm_buf(mesocluster_size_max, stream, device_memory); - auto mc_trainset_ids = mc_trainset_ids_buf.data(); - auto mc_trainset = mc_trainset_buf.data(); - auto mc_trainset_norm = mc_trainset_norm_buf.data(); + rmm::device_uvector mc_trainset_cutile_norm_buf( + dataset_cutile_norm_mptr == nullptr ? 0 : mesocluster_size_max, stream, device_memory); + auto mc_trainset_ids = mc_trainset_ids_buf.data(); + auto mc_trainset = mc_trainset_buf.data(); + auto mc_trainset_norm = mc_trainset_norm_buf.data(); + auto mc_trainset_cutile_norm = mc_trainset_cutile_norm_buf.data(); // label (cluster ID) of each vector rmm::device_uvector mc_trainset_labels(mesocluster_size_max, stream, device_memory); @@ -1171,6 +1283,13 @@ auto build_fine_clusters(const raft::resources& handle, mc_trainset_ids + k, dataset_norm_mptr, mc_trainset_norm); + if (dataset_cutile_norm_mptr != nullptr) { + thrust::gather(raft::resource::get_thrust_policy(handle), + mc_trainset_ids, + mc_trainset_ids + k, + dataset_cutile_norm_mptr, + mc_trainset_cutile_norm); + } } build_clusters(handle, @@ -1184,7 +1303,8 @@ auto build_fine_clusters(const raft::resources& handle, mc_trainset_csizes_tmp.data(), mapping_op, device_memory, - mc_trainset_norm); + mc_trainset_norm, + dataset_cutile_norm_mptr == nullptr ? nullptr : mc_trainset_cutile_norm); raft::copy(handle, raft::make_device_vector_view(cluster_centers + (dim * fine_clusters_csum[i]), @@ -1245,7 +1365,9 @@ void build_hierarchical(const raft::resources& handle, // Precompute the L2 norm of the dataset if relevant and not yet computed. rmm::device_uvector dataset_norm_buf(0, stream, device_memory); - const MathT* dataset_norm = nullptr; + rmm::device_uvector dataset_cutile_norm_buf(0, stream, device_memory); + const MathT* dataset_norm = nullptr; + const MathT* dataset_cutile_norm = nullptr; if ((params.metric == cuvs::distance::DistanceType::L2Expanded || params.metric == cuvs::distance::DistanceType::L2SqrtExpanded || params.metric == cuvs::distance::DistanceType::CosineExpanded)) { @@ -1272,6 +1394,28 @@ void build_hierarchical(const raft::resources& handle, device_memory); } dataset_norm = (const MathT*)dataset_norm_buf.data(); + + if constexpr (std::is_same_v) { + if (use_fused(handle, n_rows, n_clusters, dim, params.metric) == + FusedDistancePath::FusedCutile) { + dataset_cutile_norm_buf.resize(n_rows, stream); + const bool take_sqrt = params.metric == cuvs::distance::DistanceType::CosineExpanded; + raft::common::nvtx::range cached_input_norm_scope( + "cutile_cached_input_norm"); + for (IdxT offset = 0; offset < n_rows; offset += max_minibatch_size) { + const IdxT minibatch_size = std::min(max_minibatch_size, n_rows - offset); + compute_cutile_norm(handle, + dataset_cutile_norm_buf.data() + offset, + dataset + dim * offset, + dim, + minibatch_size, + mapping_op, + take_sqrt, + device_memory); + } + dataset_cutile_norm = dataset_cutile_norm_buf.data(); + } + } } /* Temporary workaround to cub::DeviceHistogram not supporting any type that isn't natively @@ -1295,7 +1439,8 @@ void build_hierarchical(const raft::resources& handle, mesocluster_sizes_buf.data(), mapping_op, device_memory, - dataset_norm); + dataset_norm, + dataset_cutile_norm); } auto mesocluster_sizes = mesocluster_sizes_buf.data(); @@ -1327,6 +1472,7 @@ void build_hierarchical(const raft::resources& handle, dim, dataset, dataset_norm, + dataset_cutile_norm, mesocluster_labels, n_rows, fine_clusters_nums.data(), @@ -1365,6 +1511,7 @@ void build_hierarchical(const raft::resources& handle, dim, dataset, dataset_norm, + dataset_cutile_norm, n_rows, n_clusters, cluster_centers, diff --git a/cpp/src/cluster/detail/kmeans_common.cuh b/cpp/src/cluster/detail/kmeans_common.cuh index b4dfe3511c..05bfcda36b 100644 --- a/cpp/src/cluster/detail/kmeans_common.cuh +++ b/cpp/src/cluster/detail/kmeans_common.cuh @@ -68,7 +68,7 @@ enum class FusedDistancePath : std::uint8_t { Unfused = 0, /** fusedDistanceNNMinReduce via cuTile; scratch depends on the launchability probe. */ FusedCutile, - /** fusedDistanceNNMinReduce via legacy CUTLASS; needs mutex workspace + KVP scratch. */ + /** Legacy CUTLASS fused 1-NN, with native KVP assignment output and mutex workspace. */ FusedCutlass, }; @@ -471,7 +471,8 @@ void minClusterAndDistanceCompute(raft::resources const& handle, cuvs::distance::DistanceType metric, int batch_samples, int batch_centroids, - rmm::device_uvector& workspace); + rmm::device_uvector& workspace, + const DataT* cutile_x_norm = nullptr); template void minClusterAndDistanceComputeKvp( @@ -499,7 +500,8 @@ void minClusterAndDistanceComputeKvp( cuvs::distance::DistanceType metric, \ int batch_samples, \ int batch_centroids, \ - rmm::device_uvector& workspace); + rmm::device_uvector& workspace, \ + const DataT* cutile_x_norm); EXTERN_TEMPLATE_MIN_CLUSTER_AND_DISTANCE(float, int64_t) EXTERN_TEMPLATE_MIN_CLUSTER_AND_DISTANCE(float, int) @@ -508,6 +510,14 @@ EXTERN_TEMPLATE_MIN_CLUSTER_AND_DISTANCE(double, int) #undef EXTERN_TEMPLATE_MIN_CLUSTER_AND_DISTANCE +template +void computeCutileRowNorms(raft::resources const& handle, + const float* matrix, + float* norms, + IndexT n_rows, + IndexT n_cols, + bool take_sqrt); + template void minClusterDistanceCompute(raft::resources const& handle, raft::device_matrix_view X, @@ -837,18 +847,14 @@ void process_batch(raft::resources const& handle, weight_per_cluster, batch_workspace, /*reset_sums=*/false); - auto weighted_dist = raft::make_device_vector(handle, n_samples); - raft::linalg::map(handle, - weighted_dist.view(), - raft::mul_op{}, - raft::make_const_mdspan(nearest_dist_view), - raft::make_const_mdspan(batch_weights)); - computeClusterCost(handle, - weighted_dist.view(), - workspace, - batch_cost.view(), - raft::identity_op{}, - raft::add_op{}); + auto* weights = batch_weights.data_handle(); + cuda::counting_iterator indices(IndexT{0}); + cuda::transform_iterator weighted_dist(indices, + [nearest_dist, weights] __device__(IndexT i) -> DataT { + return nearest_dist[i] * weights[i]; + }); + computeClusterCostFromIterator( + handle, weighted_dist, n_samples, workspace, batch_cost.view(), raft::add_op{}); } else { using KvpT = raft::KeyValuePair; assignment_storage.resize(sizeof(KvpT) * static_cast(n_samples), stream); diff --git a/cpp/src/cluster/detail/minClusterDistanceCompute.cu b/cpp/src/cluster/detail/minClusterDistanceCompute.cu index a62feec252..c7d0dea340 100644 --- a/cpp/src/cluster/detail/minClusterDistanceCompute.cu +++ b/cpp/src/cluster/detail/minClusterDistanceCompute.cu @@ -3,6 +3,7 @@ * SPDX-License-Identifier: Apache-2.0 */ +#include "../../core/nvtx.hpp" #include "../../distance/fused_distance_nn.cuh" #include "../../distance/unfused_distance_nn.cuh" #include "kmeans_common.cuh" @@ -95,6 +96,17 @@ void unpack_kvp(raft::resources const& handle, } // namespace +template +void computeCutileRowNorms(raft::resources const& handle, + const float* matrix, + float* norms, + IndexT n_rows, + IndexT n_cols, + bool take_sqrt) +{ + compute_tf32_row_norms(handle, matrix, norms, n_rows, n_cols, take_sqrt); +} + template FusedDistancePath select_min_cluster_distance_path( raft::resources const& handle, @@ -131,7 +143,8 @@ void min_cluster_and_distance_compute_impl(raft::resources const& handle, int batch_samples, int batch_centroids, rmm::device_uvector& workspace, - FusedDistancePath fused_path) + FusedDistancePath fused_path, + const DataT* cutile_x_norm) { cudaStream_t stream = raft::resource::get_cuda_stream(handle); auto n_samples = X.extent(0); @@ -152,20 +165,31 @@ void min_cluster_and_distance_compute_impl(raft::resources const& handle, if constexpr (std::is_same_v) { if (cutile_ready) { constexpr size_t norm_alignment = 16 / sizeof(float); - const size_t x_norm_storage = - raft::alignTo(static_cast(n_samples), norm_alignment); - L2NormBuf_OR_DistBuf.resize(x_norm_storage + static_cast(n_clusters), stream); - auto* tf32_x_norms = L2NormBuf_OR_DistBuf.data(); - auto* tf32_centroid_norms = tf32_x_norms + x_norm_storage; - const bool take_sqrt = metric == cuvs::distance::DistanceType::CosineExpanded; - compute_tf32_row_norms( - handle, X.data_handle(), tf32_x_norms, n_samples, n_features, take_sqrt); - compute_tf32_row_norms(handle, - centroids.data_handle(), - tf32_centroid_norms, - n_clusters, - n_features, - take_sqrt); + const bool take_sqrt = metric == cuvs::distance::DistanceType::CosineExpanded; + size_t centroid_offset = 0; + if (cutile_x_norm == nullptr) { + centroid_offset = raft::alignTo(static_cast(n_samples), norm_alignment); + } + L2NormBuf_OR_DistBuf.resize(centroid_offset + static_cast(n_clusters), stream); + auto* tf32_x_norms = cutile_x_norm == nullptr ? L2NormBuf_OR_DistBuf.data() + : const_cast(cutile_x_norm); + auto* tf32_centroid_norms = L2NormBuf_OR_DistBuf.data() + centroid_offset; + if (cutile_x_norm == nullptr) { + raft::common::nvtx::range input_norm_scope( + "cutile_input_norm"); + compute_tf32_row_norms( + handle, X.data_handle(), tf32_x_norms, n_samples, n_features, take_sqrt); + } + { + raft::common::nvtx::range center_norm_scope( + "cutile_center_norm"); + compute_tf32_row_norms(handle, + centroids.data_handle(), + tf32_centroid_norms, + n_clusters, + n_features, + take_sqrt); + } x_norm_ptr = tf32_x_norms; centroids_norm_ptr = tf32_centroid_norms; } else { @@ -381,7 +405,8 @@ void minClusterAndDistanceCompute(raft::resources const& handle, cuvs::distance::DistanceType metric, int batch_samples, int batch_centroids, - rmm::device_uvector& workspace) + rmm::device_uvector& workspace, + const DataT* cutile_x_norm) { auto path = select_min_cluster_distance_path(handle, X, centroids, metric); if constexpr (is_cutile_fused_data_type_v) { @@ -410,7 +435,8 @@ void minClusterAndDistanceCompute(raft::resources const& handle, batch_samples, batch_centroids, workspace, - path); + path, + cutile_x_norm); } template @@ -441,7 +467,8 @@ void minClusterAndDistanceComputeKvp( batch_samples, batch_centroids, workspace, - path); + path, + nullptr); } #define INSTANTIATE_MIN_CLUSTER_AND_DISTANCE(DataT, IndexT) \ @@ -456,7 +483,8 @@ void minClusterAndDistanceComputeKvp( cuvs::distance::DistanceType metric, \ int batch_samples, \ int batch_centroids, \ - rmm::device_uvector& workspace); + rmm::device_uvector& workspace, \ + const DataT* cutile_x_norm); INSTANTIATE_MIN_CLUSTER_AND_DISTANCE(float, int64_t) INSTANTIATE_MIN_CLUSTER_AND_DISTANCE(double, int64_t) @@ -465,6 +493,11 @@ INSTANTIATE_MIN_CLUSTER_AND_DISTANCE(double, int) #undef INSTANTIATE_MIN_CLUSTER_AND_DISTANCE +template void computeCutileRowNorms( + raft::resources const&, const float*, float*, int, int, bool); +template void computeCutileRowNorms( + raft::resources const&, const float*, float*, int64_t, int64_t, bool); + #define INSTANTIATE_SELECT_MIN_CLUSTER_PATH(DataT, IndexT) \ template FusedDistancePath select_min_cluster_distance_path( \ raft::resources const&, \ diff --git a/cpp/src/detail/jit_lto/TileAlgorithmPlanner.cpp b/cpp/src/detail/jit_lto/TileAlgorithmPlanner.cpp index d3fee7e9dc..aa6fe1dd5d 100644 --- a/cpp/src/detail/jit_lto/TileAlgorithmPlanner.cpp +++ b/cpp/src/detail/jit_lto/TileAlgorithmPlanner.cpp @@ -83,6 +83,26 @@ std::string TileAlgorithmPlanner::get_planner_key() const key += fragment->get_key(); } if (tileir_fragment_) { key += tileir_fragment_->get_key(); } + + int device = -1; + int cc_major = -1; + int cc_minor = -1; + int driver_version = -1; + if (cudaGetDevice(&device) == cudaSuccess && + cuvs::detail::jit_lto::get_device_compute_capability(cc_major, cc_minor)) { + key += ":device=" + std::to_string(device); + key += ":cc=" + std::to_string(cc_major) + "." + std::to_string(cc_minor); + if (const auto* fragment = cuvs::detail::jit_lto::find_compatible_cubin_fragment( + cc_major, cc_minor, cubin_fragments_)) { + key += ":cubin=" + std::to_string(fragment->get_cc_major()) + "." + + std::to_string(fragment->get_cc_minor()); + } else { + key += ":tileir"; + } + if (cudaDriverGetVersion(&driver_version) == cudaSuccess) { + key += ":driver=" + std::to_string(driver_version); + } + } return key; } @@ -119,7 +139,7 @@ std::shared_ptr TileAlgorithmPlanner::build() cc_major, cc_minor, driver_version, cubin_fragments_, tileir_fragment_.get()); if (!image) { return nullptr; } - return cuvs::detail::jit_lto::load_cutile_launcher(*image, entrypoint_); + return cuvs::detail::jit_lto::try_load_cutile_launcher(*image, entrypoint_); } } // namespace cuvs::detail::jit_lto diff --git a/cpp/src/distance/detail/fused_distance_nn.cuh b/cpp/src/distance/detail/fused_distance_nn.cuh index 33bf6fc953..7bb9b7a4f9 100644 --- a/cpp/src/distance/detail/fused_distance_nn.cuh +++ b/cpp/src/distance/detail/fused_distance_nn.cuh @@ -30,13 +30,18 @@ namespace distance { namespace detail { -template +template void fusedDistanceNNImpl(IdxT* nearest_idx, DataT* nearest_dist, const DataT* x, const DataT* y, - const DataT* xn, - const DataT* yn, + const NormT* xn, + const NormT* yn, IdxT m, IdxT n, IdxT k, @@ -56,7 +61,7 @@ void fusedDistanceNNImpl(IdxT* nearest_idx, constexpr auto maxVal = std::numeric_limits::max(); if constexpr (is_fused_1nn_cutile_data_v && - std::is_same_v, DataT>) { + std::is_same_v, NormT>) { if constexpr (cuvs::detail::jit_lto::library_built_with_cutile()) { if (try_fused_1nn_tile( nearest_idx, nearest_dist, x, y, xn, yn, m, n, k, metric, sqrt, workspace, stream)) { @@ -70,6 +75,11 @@ void fusedDistanceNNImpl(IdxT* nearest_idx, // sentinel from the legacy L2/cosine implementation below. RAFT_EXPECTS(metric != cuvs::distance::DistanceType::InnerProduct, "Fused InnerProduct 1-NN requires a compatible cuTile launcher"); + RAFT_EXPECTS((std::is_same_v), + "Mixed-precision norm inputs require a compatible cuTile launcher"); + + const auto* legacy_xn = reinterpret_cast(xn); + const auto* legacy_yn = reinterpret_cast(yn); RAFT_CUDA_TRY(cudaMemsetAsync(workspace, 0, sizeof(int) * m, stream)); @@ -80,8 +90,8 @@ void fusedDistanceNNImpl(IdxT* nearest_idx, nearest_dist, x, y, - xn, - yn, + legacy_xn, + legacy_yn, m, n, k, @@ -98,8 +108,8 @@ void fusedDistanceNNImpl(IdxT* nearest_idx, nearest_dist, x, y, - xn, - yn, + legacy_xn, + legacy_yn, m, n, k, diff --git a/cpp/src/distance/detail/fused_distance_nn/cutile/export_fused_1nn.py b/cpp/src/distance/detail/fused_distance_nn/cutile/export_fused_1nn.py index 81411e72f7..2f126fde4f 100644 --- a/cpp/src/distance/detail/fused_distance_nn/cutile/export_fused_1nn.py +++ b/cpp/src/distance/detail/fused_distance_nn/cutile/export_fused_1nn.py @@ -89,7 +89,8 @@ def _cuvs_matrix_constraint( ndim=2, index_dtype=index_dtype, stride_lower_bound_incl=(0, None), - alias_groups=(), + # Dataset and centroid views are read-only and may legally share storage. + alias_groups=("read_only_inputs",), may_alias_internally=False, stride_constant=(None, 1), stride_divisible_by=( @@ -106,14 +107,16 @@ def _cuvs_matrix_constraint( ) -def _cuvs_vector_constraint(elem_dtype, *, index_dtype=ct.int32): +def _cuvs_vector_constraint( + elem_dtype, *, index_dtype=ct.int32, alias_groups=() +): """1-D device vectors: contiguous, 16-byte base. Length need not be divisible by 16.""" return ArrayConstraint( elem_dtype, ndim=1, index_dtype=index_dtype, stride_lower_bound_incl=(None,), - alias_groups=(), + alias_groups=alias_groups, may_alias_internally=False, stride_constant=(1,), stride_divisible_by=(1,), @@ -156,7 +159,11 @@ def _kernel_signature( ), ) norm_elem = ct.float32 if data_type == "half" else elem - norm_array = _cuvs_vector_constraint(norm_elem, index_dtype=idx_dtype) + norm_array = _cuvs_vector_constraint( + norm_elem, + index_dtype=idx_dtype, + alias_groups=("read_only_inputs",), + ) idx_array = _cuvs_vector_constraint(idx_dtype, index_dtype=idx_dtype) dist_array = _cuvs_vector_constraint(elem, index_dtype=idx_dtype) diff --git a/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_tile.cu b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_tile.cu index b93931c977..870c794789 100644 --- a/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_tile.cu +++ b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_tile.cu @@ -30,6 +30,30 @@ bool is_16_byte_aligned(const void* ptr) return ptr == nullptr || reinterpret_cast(ptr) % 16 == 0; } +bool byte_ranges_overlap(const void* lhs, size_t lhs_bytes, const void* rhs, size_t rhs_bytes) +{ + if (lhs == nullptr || rhs == nullptr || lhs_bytes == 0 || rhs_bytes == 0) { return false; } + const auto lhs_begin = reinterpret_cast(lhs); + const auto rhs_begin = reinterpret_cast(rhs); + if (lhs_bytes > std::numeric_limits::max() - lhs_begin || + rhs_bytes > std::numeric_limits::max() - rhs_begin) { + return true; + } + return lhs_begin < rhs_begin + rhs_bytes && rhs_begin < lhs_begin + lhs_bytes; +} + +template +size_t checked_tensor_bytes(IdxT rows, IdxT cols, size_t element_size) +{ + const auto rows_u = static_cast(rows); + const auto cols_u = static_cast(cols); + constexpr auto max_size = std::numeric_limits::max(); + if (cols_u != 0 && rows_u > max_size / cols_u) { return max_size; } + const auto elements = rows_u * cols_u; + if (element_size != 0 && elements > max_size / element_size) { return max_size; } + return static_cast(elements) * element_size; +} + template bool has_fused_1nn_tile_launcher() { @@ -237,6 +261,17 @@ bool can_launch_fused_1nn_tile(IdxT* nearest_idx, if constexpr (std::is_same_v) { if (!is_16_byte_aligned(nearest_idx)) { return false; } } + const auto x_bytes = checked_tensor_bytes(m, k, sizeof(DataT)); + const auto y_bytes = checked_tensor_bytes(n, k, sizeof(DataT)); + const auto dist_bytes = checked_tensor_bytes(m, IdxT{1}, sizeof(DataT)); + const auto idx_bytes = checked_tensor_bytes(m, IdxT{1}, sizeof(IdxT)); + if (byte_ranges_overlap(nearest_dist, dist_bytes, x, x_bytes) || + byte_ranges_overlap(nearest_dist, dist_bytes, y, y_bytes) || + byte_ranges_overlap(nearest_idx, idx_bytes, x, x_bytes) || + byte_ranges_overlap(nearest_idx, idx_bytes, y, y_bytes) || + byte_ranges_overlap(nearest_idx, idx_bytes, nearest_dist, dist_bytes)) { + return false; + } return true; } @@ -259,7 +294,15 @@ bool can_launch_fused_1nn_tile(IdxT* nearest_idx, if (metric != cuvs::distance::DistanceType::InnerProduct && (xn == nullptr || yn == nullptr)) { return false; } - return is_16_byte_aligned(xn) && is_16_byte_aligned(yn); + if (!is_16_byte_aligned(xn) || !is_16_byte_aligned(yn)) { return false; } + const auto xn_bytes = checked_tensor_bytes(m, IdxT{1}, sizeof(*xn)); + const auto yn_bytes = checked_tensor_bytes(n, IdxT{1}, sizeof(*yn)); + const auto dist_bytes = checked_tensor_bytes(m, IdxT{1}, sizeof(DataT)); + const auto idx_bytes = checked_tensor_bytes(m, IdxT{1}, sizeof(IdxT)); + return !byte_ranges_overlap(nearest_dist, dist_bytes, xn, xn_bytes) && + !byte_ranges_overlap(nearest_dist, dist_bytes, yn, yn_bytes) && + !byte_ranges_overlap(nearest_idx, idx_bytes, xn, xn_bytes) && + !byte_ranges_overlap(nearest_idx, idx_bytes, yn, yn_bytes); } template @@ -295,6 +338,21 @@ bool try_fused_1nn_tile(IdxT* nearest_idx, } else { if (nearest_idx != nullptr && index_workspace == nullptr) { return false; } if (!is_16_byte_aligned(index_workspace)) { return false; } + const auto workspace_bytes = checked_tensor_bytes(m, IdxT{1}, sizeof(int)); + const auto x_bytes = checked_tensor_bytes(m, k, sizeof(DataT)); + const auto y_bytes = checked_tensor_bytes(n, k, sizeof(DataT)); + const auto norm_x_bytes = checked_tensor_bytes(m, IdxT{1}, sizeof(*xn)); + const auto norm_y_bytes = checked_tensor_bytes(n, IdxT{1}, sizeof(*yn)); + const auto dist_bytes = checked_tensor_bytes(m, IdxT{1}, sizeof(DataT)); + const auto idx_bytes = checked_tensor_bytes(m, IdxT{1}, sizeof(IdxT)); + if (byte_ranges_overlap(index_workspace, workspace_bytes, x, x_bytes) || + byte_ranges_overlap(index_workspace, workspace_bytes, y, y_bytes) || + byte_ranges_overlap(index_workspace, workspace_bytes, xn, norm_x_bytes) || + byte_ranges_overlap(index_workspace, workspace_bytes, yn, norm_y_bytes) || + byte_ranges_overlap(index_workspace, workspace_bytes, nearest_dist, dist_bytes) || + byte_ranges_overlap(index_workspace, workspace_bytes, nearest_idx, idx_bytes)) { + return false; + } // Keep every chunk offset 16-byte aligned for x, xn, and nearest_dist. constexpr int64_t max_batch_m = fused_1nn_cutile_max_batch_m; diff --git a/cpp/src/distance/fused_distance_nn-inl.cuh b/cpp/src/distance/fused_distance_nn-inl.cuh index cbf90c8025..afb66a402a 100644 --- a/cpp/src/distance/fused_distance_nn-inl.cuh +++ b/cpp/src/distance/fused_distance_nn-inl.cuh @@ -29,13 +29,13 @@ namespace distance { * @{ */ -template +template void fusedDistanceNN(IdxT* nearest_idx, DataT* nearest_dist, const DataT* x, const DataT* y, - const DataT* xn, - const DataT* yn, + const NormT* xn, + const NormT* yn, IdxT m, IdxT n, IdxT k, @@ -60,6 +60,7 @@ void fusedDistanceNN(IdxT* nearest_idx, if (is_skinny) { detail::fusedDistanceNNImpl< DataT, + NormT, IdxT, typename raft::linalg::Policy4x4Skinny::Policy, ReduceOpT>(nearest_idx, @@ -84,6 +85,7 @@ void fusedDistanceNN(IdxT* nearest_idx, } else { detail::fusedDistanceNNImpl< DataT, + NormT, IdxT, typename raft::linalg::Policy4x4::Policy, ReduceOpT>(nearest_idx, @@ -110,6 +112,7 @@ void fusedDistanceNN(IdxT* nearest_idx, if (is_skinny) { detail::fusedDistanceNNImpl< DataT, + NormT, IdxT, typename raft::linalg::Policy4x4Skinny::Policy, ReduceOpT>(nearest_idx, @@ -134,6 +137,7 @@ void fusedDistanceNN(IdxT* nearest_idx, } else { detail::fusedDistanceNNImpl< DataT, + NormT, IdxT, typename raft::linalg::Policy4x4::Policy, ReduceOpT>(nearest_idx, @@ -159,6 +163,7 @@ void fusedDistanceNN(IdxT* nearest_idx, } else { if (is_skinny) { detail::fusedDistanceNNImpl::Policy, ReduceOpT>(nearest_idx, @@ -182,6 +187,7 @@ void fusedDistanceNN(IdxT* nearest_idx, stream); } else { detail::fusedDistanceNNImpl::Policy, ReduceOpT>(nearest_idx, @@ -216,13 +222,13 @@ void fusedDistanceNN(IdxT* nearest_idx, * output. Distance-only output may pass null and write directly * to `nearest_dist`. */ -template +template void fusedDistanceNNMinReduce(IdxT* nearest_idx, DataT* nearest_dist, const DataT* x, const DataT* y, - const DataT* xn, - const DataT* yn, + const NormT* xn, + const NormT* yn, IdxT m, IdxT n, IdxT k, diff --git a/cpp/tests/cluster/kmeans_balanced.cu b/cpp/tests/cluster/kmeans_balanced.cu index 10d4ce9cbb..dbe8798de2 100644 --- a/cpp/tests/cluster/kmeans_balanced.cu +++ b/cpp/tests/cluster/kmeans_balanced.cu @@ -274,4 +274,54 @@ KB_TEST((KmeansBalancedTest, KmeansBalancedTestFI8I32I32_SEP, inputsf_i32); +#if CUVS_CUTILE_ENABLED +TEST(KmeansBalancedPredict, CutilePointerRejectionPreservesMetricFallback) +{ + raft::resources handle; + auto stream = raft::resource::get_cuda_stream(handle); + constexpr int64_t n_rows = 4; + constexpr int64_t n_cols = 4; + constexpr int64_t n_clusters = 2; + + const std::vector h_x{ + 1.0f, 0.0f, 0.2f, 0.0f, 0.0f, 1.0f, 0.0f, 0.2f, 0.8f, 0.1f, 0.0f, 0.0f, 0.1f, 0.9f, 0.0f, 0.0f}; + const std::vector h_centroids{1.0f, 0.0f, 0.0f, 0.0f, 0.0f, 1.0f, 0.0f, 0.0f}; + + rmm::device_uvector x_aligned(h_x.size(), stream); + rmm::device_uvector x_misaligned(h_x.size() + 1, stream); + rmm::device_uvector centroids(h_centroids.size(), stream); + rmm::device_uvector labels_ref(n_rows, stream); + rmm::device_uvector labels_misaligned(n_rows + 1, stream); + raft::update_device(x_aligned.data(), h_x.data(), h_x.size(), stream); + raft::update_device(x_misaligned.data() + 1, h_x.data(), h_x.size(), stream); + raft::update_device(centroids.data(), h_centroids.data(), h_centroids.size(), stream); + + for (auto metric : {distance::DistanceType::L2Expanded, + distance::DistanceType::CosineExpanded, + distance::DistanceType::InnerProduct}) { + cluster::kmeans::balanced_params params; + params.metric = metric; + cluster::kmeans::predict( + handle, + params, + raft::make_device_matrix_view(x_aligned.data(), n_rows, n_cols), + raft::make_device_matrix_view(centroids.data(), n_clusters, n_cols), + raft::make_device_vector_view(labels_ref.data(), n_rows)); + cluster::kmeans::predict( + handle, + params, + raft::make_device_matrix_view(x_misaligned.data() + 1, n_rows, n_cols), + raft::make_device_matrix_view(centroids.data(), n_clusters, n_cols), + raft::make_device_vector_view(labels_misaligned.data() + 1, n_rows)); + + std::vector h_ref(n_rows); + std::vector h_actual(n_rows); + raft::update_host(h_ref.data(), labels_ref.data(), n_rows, stream); + raft::update_host(h_actual.data(), labels_misaligned.data() + 1, n_rows, stream); + raft::resource::sync_stream(handle, stream); + EXPECT_EQ(h_actual, h_ref); + } +} +#endif + } // namespace cuvs diff --git a/cpp/tests/neighbors/distance_nn.cu b/cpp/tests/neighbors/distance_nn.cu index ec8624c0c9..d2b28b70a1 100644 --- a/cpp/tests/neighbors/distance_nn.cu +++ b/cpp/tests/neighbors/distance_nn.cu @@ -9,6 +9,9 @@ #include "../../src/distance/fused_distance_nn.cuh" #include "../../src/distance/unfused_distance_nn.cuh" +#include +#include + #include #include #include @@ -312,6 +315,146 @@ TEST(Fused1nn, ExpandedL2ClampsNegativeRoundoff) EXPECT_EQ(actual_dist, 0.0f); } +TEST(Fused1nn, CutileAvailabilityRejectsUnsupportedArchitectures) +{ + EXPECT_FALSE(cuvs::detail::jit_lto::cutile_launch_available_for_arch(7, 5, 13000)); + EXPECT_FALSE(cuvs::detail::jit_lto::cutile_launch_available_for_arch(13, 0, 13000)); +} + +#if CUVS_CUTILE_ENABLED +TEST(Fused1nn, ExpectedModuleCompatibilityErrorsAreRecoverable) +{ + using cuvs::detail::jit_lto::is_expected_cutile_unavailable; + EXPECT_TRUE(is_expected_cutile_unavailable(cudaErrorInvalidDeviceFunction)); + EXPECT_TRUE(is_expected_cutile_unavailable(cudaErrorInvalidPtx)); + EXPECT_TRUE(is_expected_cutile_unavailable(cudaErrorNoKernelImageForDevice)); + EXPECT_TRUE(is_expected_cutile_unavailable(cudaErrorSymbolNotFound)); + EXPECT_TRUE(is_expected_cutile_unavailable(cudaErrorUnsupportedPtxVersion)); + EXPECT_TRUE(is_expected_cutile_unavailable(cudaErrorCallRequiresNewerDriver)); + EXPECT_FALSE(is_expected_cutile_unavailable(cudaErrorMemoryAllocation)); + EXPECT_FALSE(is_expected_cutile_unavailable(cudaErrorIllegalAddress)); +} + +template +void run_half_cutile_contract_case(int k) +{ + raft::resources handle; + auto stream = raft::resource::get_cuda_stream(handle); + constexpr IdxT m = 2; + constexpr IdxT n = 2; + + std::vector h_x(static_cast(m) * k, __float2half(0.0f)); + std::vector h_y(static_cast(n) * k, __float2half(0.0f)); + h_x[0] = __float2half(1.0f); + h_x[k + 1] = __float2half(1.0f); + h_y[0] = __float2half(1.0f); + h_y[k + 1] = __float2half(1.0f); + + rmm::device_uvector x(h_x.size(), stream); + rmm::device_uvector y(h_y.size(), stream); + rmm::device_uvector x_norm(m, stream); + rmm::device_uvector y_norm(n, stream); + rmm::device_uvector out_idx(m, stream); + rmm::device_uvector out_dist(m, stream); + rmm::device_uvector workspace(m, stream); + raft::update_device(x.data(), h_x.data(), h_x.size(), stream); + raft::update_device(y.data(), h_y.data(), h_y.size(), stream); + const std::vector h_norms(m, 1.0f); + raft::update_device(x_norm.data(), h_norms.data(), m, stream); + raft::update_device(y_norm.data(), h_norms.data(), n, stream); + + if constexpr (std::is_same_v) { + EXPECT_FALSE((cuvs::distance::detail::try_fused_1nn_tile(out_idx.data(), + out_dist.data(), + x.data(), + y.data(), + x_norm.data(), + y_norm.data(), + m, + n, + static_cast(k), + DistanceType::L2Expanded, + false, + x.data(), + stream))); + } + + for (auto metric : {DistanceType::L2Expanded, + DistanceType::L2SqrtExpanded, + DistanceType::CosineExpanded, + DistanceType::InnerProduct}) { + ASSERT_TRUE(( + cuvs::distance::detail::try_fused_1nn_tile(out_idx.data(), + out_dist.data(), + x.data(), + y.data(), + x_norm.data(), + y_norm.data(), + m, + n, + static_cast(k), + metric, + metric == DistanceType::L2SqrtExpanded, + workspace.data(), + stream))); + std::vector h_idx(m); + raft::update_host(h_idx.data(), out_idx.data(), m, stream); + raft::resource::sync_stream(handle, stream); + EXPECT_EQ(h_idx[0], IdxT{0}); + EXPECT_EQ(h_idx[1], IdxT{1}); + } + + ASSERT_TRUE((cuvs::distance::detail::try_fused_1nn_tile(out_idx.data(), + out_dist.data(), + x.data(), + x.data(), + x_norm.data(), + x_norm.data(), + m, + m, + static_cast(k), + DistanceType::L2Expanded, + false, + workspace.data(), + stream))); + std::vector h_alias_idx(m); + raft::update_host(h_alias_idx.data(), out_idx.data(), m, stream); + raft::resource::sync_stream(handle, stream); + EXPECT_EQ(h_alias_idx[0], IdxT{0}); + EXPECT_EQ(h_alias_idx[1], IdxT{1}); + + raft::copy(y.data() + k, y.data(), k, stream); + ASSERT_TRUE((cuvs::distance::detail::try_fused_1nn_tile(out_idx.data(), + out_dist.data(), + x.data(), + y.data(), + x_norm.data(), + y_norm.data(), + m, + n, + static_cast(k), + DistanceType::L2Expanded, + false, + workspace.data(), + stream))); + IdxT h_tie_idx; + half h_tie_dist; + raft::update_host(&h_tie_idx, out_idx.data(), 1, stream); + raft::update_host(&h_tie_dist, out_dist.data(), 1, stream); + raft::resource::sync_stream(handle, stream); + EXPECT_TRUE(h_tie_idx == IdxT{0} || h_tie_idx == IdxT{1}); + EXPECT_EQ(__half2float(h_tie_dist), 0.0f); +} + +TEST(Fused1nn, HalfUsesFloatNormsAcrossAbisAndIndexTypes) +{ + run_half_cutile_contract_case(8); + run_half_cutile_contract_case(7); + run_half_cutile_contract_case(8); + run_half_cutile_contract_case(7); +} +#endif + TEST(Fused1nn, Int64IndexWorkspaceUsesLargestChunk) { constexpr int64_t max_batch_m_float = cuvs::distance::detail::fused_1nn_cutile_max_batch_m; @@ -342,7 +485,8 @@ TEST(Fused1nn, PointerAwareProbeRejectsMisalignedArrays) auto out_idx = raft::make_device_vector(handle, m); auto out_dist = raft::make_device_vector(handle, m); - for (auto metric : {DistanceType::L2Expanded, DistanceType::CosineExpanded}) { + for (auto metric : + {DistanceType::L2Expanded, DistanceType::CosineExpanded, DistanceType::InnerProduct}) { EXPECT_FALSE(cuvs::distance::detail::can_launch_fused_1nn_tile(out_idx.data_handle(), out_dist.data_handle(), x.data_handle() + 1, @@ -370,6 +514,8 @@ TEST(Fused1nn, PointerAwareProbeRejectsMisalignedArrays) n, k, metric)); + EXPECT_FALSE(cuvs::distance::detail::can_launch_fused_1nn_tile( + out_idx.data_handle(), x.data_handle(), x.data_handle(), y.data_handle(), m, n, k, metric)); } } } From 638a2ecb4f6a8c18d14a26ead2afc9d11728a0e4 Mon Sep 17 00:00:00 2001 From: divyegala Date: Thu, 3 Sep 2026 00:36:05 +0000 Subject: [PATCH 69/82] refactor backend choice to primitive --- ci/build_standalone_c.sh | 3 - .../modules/compute_matrix_product.cmake | 2 - cpp/src/cluster/detail/kmeans.cuh | 19 +- cpp/src/cluster/detail/kmeans_balanced.cuh | 62 +++- cpp/src/cluster/detail/kmeans_common.cuh | 92 +++-- .../detail/minClusterDistanceCompute.cu | 168 ++++++--- .../detail/jit_lto/TileAlgorithmPlanner.cpp | 1 - cpp/src/distance/detail/fused_distance_nn.cuh | 43 ++- .../cutile/export_fused_1nn.py | 4 +- .../jit_lto_kernels/interleaved_scan_impl.cuh | 19 +- .../neighbors/ivf_flat/ivf_flat_search.cuh | 24 +- .../ivf_rabitq/gpu_index/quantizer_gpu.cu | 10 +- cpp/tests/cluster/kmeans.cu | 3 +- cpp/tests/cluster/kmeans_balanced.cu | 50 --- cpp/tests/cluster/kmeans_predict_batching.cu | 14 - cpp/tests/neighbors/distance_nn.cu | 325 ++---------------- cpp/tests/neighbors/distance_nn_helper.cuh | 95 ++--- 17 files changed, 336 insertions(+), 598 deletions(-) diff --git a/ci/build_standalone_c.sh b/ci/build_standalone_c.sh index f008a44f5f..2b8e0863f9 100755 --- a/ci/build_standalone_c.sh +++ b/ci/build_standalone_c.sh @@ -41,9 +41,6 @@ source rapids-configure-sccache source rapids-datetime-string rapids-pip-retry install cmake -if [[ "${RAPIDS_CUDA_VERSION%%.*}" == "13" ]]; then - rapids-pip-retry install cuda-tile 'cuda-toolkit[tileiras]==13.*' -fi pyenv rehash rapids-print-env diff --git a/cpp/cmake/modules/compute_matrix_product.cmake b/cpp/cmake/modules/compute_matrix_product.cmake index b5f3d06f86..60b96113f0 100644 --- a/cpp/cmake/modules/compute_matrix_product.cmake +++ b/cpp/cmake/modules/compute_matrix_product.cmake @@ -11,9 +11,7 @@ function(cuvs_find_build_python output_var) if(DEFINED ENV{BUILD_PREFIX}) set(Python_ROOT "$ENV{BUILD_PREFIX}") endif() - set(CMAKE_FIND_DEBUG_MODE TRUE) find_package(Python REQUIRED COMPONENTS Interpreter) - set(CMAKE_FIND_DEBUG_MODE FALSE) set(${output_var} "${Python_EXECUTABLE}" PARENT_SCOPE diff --git a/cpp/src/cluster/detail/kmeans.cuh b/cpp/src/cluster/detail/kmeans.cuh index 1a71003fd0..208784762d 100644 --- a/cpp/src/cluster/detail/kmeans.cuh +++ b/cpp/src/cluster/detail/kmeans.cuh @@ -1082,15 +1082,10 @@ void kmeans_predict(raft::resources const& handle, auto l2normx_view = raft::make_device_vector_view(L2NormX.data_handle(), n_samples); rmm::device_scalar clusterCostD(stream); - auto path = select_min_cluster_distance_path(handle, X, centroids, metric); - if constexpr (std::is_same_v && std::is_same_v) { - if (path == FusedDistancePath::FusedCutile && - reinterpret_cast(labels.data_handle()) % 16 != 0) { - path = use_legacy_fused(handle, n_samples, centroids.extent(0), metric); - } - } + const auto requirements = get_fused_1nn_requirements( + handle, X, centroids, metric, pams.batch_samples, pams.batch_centroids); - if (path == FusedDistancePath::FusedCutile) { + if (requirements.output_layout == Fused1nnOutputLayout::Soa) { auto nearest_dist = raft::make_device_vector(handle, n_samples); if constexpr (std::is_same_v) { minClusterAndDistanceCompute(handle, @@ -1103,7 +1098,8 @@ void kmeans_predict(raft::resources const& handle, metric, pams.batch_samples, pams.batch_centroids, - workspace); + workspace, + requirements); } else { auto index_labels = raft::make_device_vector(handle, n_samples); minClusterAndDistanceCompute(handle, @@ -1116,7 +1112,8 @@ void kmeans_predict(raft::resources const& handle, metric, pams.batch_samples, pams.batch_centroids, - workspace); + workspace, + requirements); raft::linalg::map( handle, labels, raft::cast_op{}, raft::make_const_mdspan(index_labels.view())); } @@ -1144,7 +1141,7 @@ void kmeans_predict(raft::resources const& handle, pams.batch_samples, pams.batch_centroids, workspace, - path); + requirements); auto* nearest_ptr = nearest.data_handle(); raft::linalg::map_offset(handle, labels, [nearest_ptr] __device__(IndexT i) { return static_cast(nearest_ptr[i].key); diff --git a/cpp/src/cluster/detail/kmeans_balanced.cuh b/cpp/src/cluster/detail/kmeans_balanced.cuh index 65db51f93d..1cab6a0ade 100644 --- a/cpp/src/cluster/detail/kmeans_balanced.cuh +++ b/cpp/src/cluster/detail/kmeans_balanced.cuh @@ -67,21 +67,57 @@ bool predict_core_min_cluster(const raft::resources& handle, const MathT* cutile_x_norm) { auto n_rows = X.extent(0); - auto path = select_min_cluster_distance_path(handle, X, centroids, metric); - - if constexpr (std::is_same_v && std::is_same_v) { - if (path == FusedDistancePath::FusedCutile && - reinterpret_cast(labels) % 16 != 0) { - path = use_legacy_fused(handle, n_rows, centroids.extent(0), metric); + const auto requirements = get_fused_1nn_requirements(handle, X, centroids, metric); + + if (requirements.output_layout == Fused1nnOutputLayout::Soa) { + // cuTile indices are int32. For 32-bit labels, request that native representation directly; + // cluster labels are never negative, so no conversion or copy is required. + if constexpr (std::is_same_v || std::is_same_v) { + if (n_rows <= static_cast(std::numeric_limits::max())) { + using CutileIdxT = int; + const auto cutile_rows = static_cast(n_rows); + const auto cutile_cols = static_cast(X.extent(1)); + const auto cutile_clusters = static_cast(centroids.extent(0)); + auto cutile_X = raft::make_device_matrix_view( + X.data_handle(), cutile_rows, cutile_cols); + auto cutile_centroids = raft::make_device_matrix_view( + centroids.data_handle(), cutile_clusters, cutile_cols); + const auto cutile_requirements = + get_fused_1nn_requirements(handle, cutile_X, cutile_centroids, metric); + RAFT_EXPECTS(cutile_requirements.output_layout == Fused1nnOutputLayout::Soa, + "resolved cuTile plan changed while adapting 32-bit labels"); + + auto nearest_dist = + raft::make_device_mdarray(handle, mr, raft::make_extents(cutile_rows)); + auto* cutile_labels = reinterpret_cast(labels); + auto labels_view = + raft::make_device_vector_view(cutile_labels, cutile_rows); + auto cutile_norm = raft::make_device_vector_view( + X_norm.data_handle(), cutile_rows); + minClusterAndDistanceCompute(handle, + cutile_X, + cutile_centroids, + labels_view, + nearest_dist.view(), + cutile_norm, + L2NormBuf_OR_DistBuf, + metric, + 0, + 0, + workspace, + cutile_requirements, + cutile_x_norm); + return true; + } } - } - if (path == FusedDistancePath::FusedCutile) { + constexpr bool label_storage_is_cutile_index = std::is_same_v; auto nearest_dist = raft::make_device_mdarray(handle, mr, raft::make_extents(n_rows)); - if constexpr (std::is_same_v) { - auto labels_view = raft::make_device_vector_view(labels, n_rows); + if constexpr (label_storage_is_cutile_index) { + auto* cutile_labels = reinterpret_cast(labels); + auto labels_view = raft::make_device_vector_view(cutile_labels, n_rows); minClusterAndDistanceCompute(handle, X, centroids, @@ -93,6 +129,7 @@ bool predict_core_min_cluster(const raft::resources& handle, 0, 0, workspace, + requirements, cutile_x_norm); } else { auto nearest_idx = @@ -108,6 +145,7 @@ bool predict_core_min_cluster(const raft::resources& handle, 0, 0, workspace, + requirements, cutile_x_norm); raft::copy( handle, raft::make_device_vector_view(labels, n_rows), nearest_idx.view()); @@ -132,7 +170,7 @@ bool predict_core_min_cluster(const raft::resources& handle, 0, 0, workspace, - path); + requirements); auto* nearest_ptr = nearest.data_handle(); raft::linalg::map_offset( handle, @@ -1400,8 +1438,6 @@ void build_hierarchical(const raft::resources& handle, FusedDistancePath::FusedCutile) { dataset_cutile_norm_buf.resize(n_rows, stream); const bool take_sqrt = params.metric == cuvs::distance::DistanceType::CosineExpanded; - raft::common::nvtx::range cached_input_norm_scope( - "cutile_cached_input_norm"); for (IdxT offset = 0; offset < n_rows; offset += max_minibatch_size) { const IdxT minibatch_size = std::min(max_minibatch_size, n_rows - offset); compute_cutile_norm(handle, diff --git a/cpp/src/cluster/detail/kmeans_common.cuh b/cpp/src/cluster/detail/kmeans_common.cuh index 05bfcda36b..6591e3bf77 100644 --- a/cpp/src/cluster/detail/kmeans_common.cuh +++ b/cpp/src/cluster/detail/kmeans_common.cuh @@ -5,6 +5,7 @@ #pragma once #include "../../distance/distance.cuh" +#include "../../distance/fused_distance_nn.cuh" #include #include #include @@ -62,14 +63,40 @@ template inline constexpr bool is_cutile_fused_data_type_v = std::is_same_v || std::is_same_v; -/** Which fused-distance implementation minCluster* will use (or Unfused). */ -enum class FusedDistancePath : std::uint8_t { - /** unfusedDistanceNNMinReduce or batched pairwise distance. */ - Unfused = 0, - /** fusedDistanceNNMinReduce via cuTile; scratch depends on the launchability probe. */ - FusedCutile, - /** Legacy CUTLASS fused 1-NN, with native KVP assignment output and mutex workspace. */ - FusedCutlass, +using FusedDistancePath = cuvs::distance::detail::Fused1nnBackend; + +/** Native result representation selected by fused 1-NN. */ +enum class Fused1nnOutputLayout : std::uint8_t { + /** Separate index and distance arrays (cuTile). */ + Soa, + /** Native key/value-pair array (CUTLASS/SIMT and the unfused reducer). */ + Kvp, +}; + +/** Norm representation required by the resolved fused-1NN implementation. */ +enum class Fused1nnNormPolicy : std::uint8_t { + Default, + Tf32, +}; + +/** + * Resolved fused-1NN storage and execution requirements. + * + * KMeans uses the result layout and byte counts to reuse its existing raw buffers. The fused + * implementation owns execution of the resolved path. + */ +template +struct Fused1nnRequirements { + FusedDistancePath path{}; + Fused1nnOutputLayout output_layout{}; + Fused1nnNormPolicy norm_policy{}; + size_t result_bytes{}; + size_t result_alignment{}; + size_t distance_offset{}; + size_t workspace_bytes{}; + size_t workspace_alignment{}; + IndexT sample_tile{}; + IndexT centroid_tile{}; }; inline constexpr bool uses_fused_distance_nn(FusedDistancePath path) @@ -88,11 +115,7 @@ constexpr FusedDistancePath use_legacy_fused(int cc_major, IdxT n, cuvs::distance::DistanceType metric) { - if (metric == cuvs::distance::DistanceType::InnerProduct) { return FusedDistancePath::Unfused; } - - if (cc_major <= 8) { return FusedDistancePath::FusedCutlass; } - if (cc_major == 9 && (m >= 4096 || n >= 4096)) { return FusedDistancePath::FusedCutlass; } - return FusedDistancePath::Unfused; + return cuvs::distance::detail::fused_1nn_legacy_backend(cc_major, m, n, metric); } template @@ -102,7 +125,7 @@ FusedDistancePath use_legacy_fused(const raft::resources& handle, cuvs::distance::DistanceType metric) { const auto prop = raft::resource::get_device_properties(handle); - return use_legacy_fused(prop.major, m, n, metric); + return cuvs::distance::detail::fused_1nn_legacy_backend(prop.major, m, n, metric); } /** @@ -454,11 +477,13 @@ void shuffleAndGather(raft::resources const& handle, // Calculates nearest centroid index and distance for every sample in input 'X'. template -FusedDistancePath select_min_cluster_distance_path( +Fused1nnRequirements get_fused_1nn_requirements( raft::resources const& handle, raft::device_matrix_view X, raft::device_matrix_view centroids, - cuvs::distance::DistanceType metric); + cuvs::distance::DistanceType metric, + int batch_samples = 0, + int batch_centroids = 0); template void minClusterAndDistanceCompute(raft::resources const& handle, @@ -472,6 +497,7 @@ void minClusterAndDistanceCompute(raft::resources const& handle, int batch_samples, int batch_centroids, rmm::device_uvector& workspace, + const Fused1nnRequirements& requirements, const DataT* cutile_x_norm = nullptr); template @@ -486,7 +512,7 @@ void minClusterAndDistanceComputeKvp( int batch_samples, int batch_centroids, rmm::device_uvector& workspace, - FusedDistancePath path); + const Fused1nnRequirements& requirements); #define EXTERN_TEMPLATE_MIN_CLUSTER_AND_DISTANCE(DataT, IndexT) \ extern template void minClusterAndDistanceCompute( \ @@ -501,6 +527,7 @@ void minClusterAndDistanceComputeKvp( int batch_samples, \ int batch_centroids, \ rmm::device_uvector& workspace, \ + const Fused1nnRequirements& requirements, \ const DataT* cutile_x_norm); EXTERN_TEMPLATE_MIN_CLUSTER_AND_DISTANCE(float, int64_t) @@ -565,7 +592,8 @@ void countSamplesInCluster(raft::resources const& handle, rmm::device_uvector L2NormBuf_OR_DistBuf(0, stream); auto centroids_const = raft::make_const_mdspan(centroids); - auto path = select_min_cluster_distance_path(handle, X, centroids_const, params.metric); + const auto requirements = get_fused_1nn_requirements( + handle, X, centroids_const, params.metric, params.batch_samples, params.batch_centroids); auto count_labels = [&](auto labels) { countLabels(handle, @@ -576,7 +604,7 @@ void countSamplesInCluster(raft::resources const& handle, workspace); }; - if (path == FusedDistancePath::FusedCutile) { + if (requirements.output_layout == Fused1nnOutputLayout::Soa) { auto nearest_idx = raft::make_device_vector(handle, n_samples); auto nearest_dist = raft::make_device_vector(handle, n_samples); minClusterAndDistanceCompute(handle, @@ -589,7 +617,8 @@ void countSamplesInCluster(raft::resources const& handle, params.metric, params.batch_samples, params.batch_centroids, - workspace); + workspace, + requirements); count_labels(nearest_idx.data_handle()); } else { using KvpT = raft::KeyValuePair; @@ -604,7 +633,7 @@ void countSamplesInCluster(raft::resources const& handle, params.batch_samples, params.batch_centroids, workspace, - path); + requirements); auto labels = thrust::make_transform_iterator(nearest.data_handle(), KeyValueIndexOp{}); count_labels(labels); @@ -816,13 +845,15 @@ void process_batch(raft::resources const& handle, { cudaStream_t stream = raft::resource::get_cuda_stream(handle); const auto n_samples = batch_data.extent(0); - const auto path = select_min_cluster_distance_path(handle, batch_data, centroids, metric); + const auto requirements = get_fused_1nn_requirements( + handle, batch_data, centroids, metric, batch_samples_param, batch_centroids_param); auto batch_cost = raft::make_device_scalar(handle, DataT{0}); - if (path == FusedDistancePath::FusedCutile) { - const auto dist_offset = - raft::alignTo(sizeof(IndexT) * static_cast(n_samples), size_t{16}); - assignment_storage.resize(dist_offset + sizeof(DataT) * static_cast(n_samples), stream); + if (requirements.output_layout == Fused1nnOutputLayout::Soa) { + const auto dist_offset = requirements.distance_offset; + if (assignment_storage.size() < requirements.result_bytes) { + assignment_storage.resize(requirements.result_bytes, stream); + } auto* nearest_idx = reinterpret_cast(assignment_storage.data()); auto* nearest_dist = reinterpret_cast(assignment_storage.data() + dist_offset); auto nearest_idx_view = raft::make_device_vector_view(nearest_idx, n_samples); @@ -837,7 +868,8 @@ void process_batch(raft::resources const& handle, metric, batch_samples_param, batch_centroids_param, - workspace); + workspace, + requirements); compute_centroid_adjustments(handle, batch_data, batch_weights, @@ -857,7 +889,9 @@ void process_batch(raft::resources const& handle, handle, weighted_dist, n_samples, workspace, batch_cost.view(), raft::add_op{}); } else { using KvpT = raft::KeyValuePair; - assignment_storage.resize(sizeof(KvpT) * static_cast(n_samples), stream); + if (assignment_storage.size() < requirements.result_bytes) { + assignment_storage.resize(requirements.result_bytes, stream); + } auto* nearest = reinterpret_cast(assignment_storage.data()); auto nearest_view = raft::make_device_vector_view(nearest, n_samples); minClusterAndDistanceComputeKvp(handle, @@ -870,7 +904,7 @@ void process_batch(raft::resources const& handle, batch_samples_param, batch_centroids_param, workspace, - path); + requirements); auto labels = thrust::make_transform_iterator(nearest, KeyValueIndexOp{}); compute_centroid_adjustments(handle, batch_data, diff --git a/cpp/src/cluster/detail/minClusterDistanceCompute.cu b/cpp/src/cluster/detail/minClusterDistanceCompute.cu index c7d0dea340..11be57e19a 100644 --- a/cpp/src/cluster/detail/minClusterDistanceCompute.cu +++ b/cpp/src/cluster/detail/minClusterDistanceCompute.cu @@ -3,7 +3,6 @@ * SPDX-License-Identifier: Apache-2.0 */ -#include "../../core/nvtx.hpp" #include "../../distance/fused_distance_nn.cuh" #include "../../distance/unfused_distance_nn.cuh" #include "kmeans_common.cuh" @@ -108,26 +107,70 @@ void computeCutileRowNorms(raft::resources const& handle, } template -FusedDistancePath select_min_cluster_distance_path( +Fused1nnRequirements get_fused_1nn_requirements( raft::resources const& handle, raft::device_matrix_view X, raft::device_matrix_view centroids, - cuvs::distance::DistanceType metric) + cuvs::distance::DistanceType metric, + int batch_samples, + int batch_centroids) { - auto path = - use_fused(handle, X.extent(0), centroids.extent(0), X.extent(1), metric); - if constexpr (is_cutile_fused_data_type_v) { - if (path == FusedDistancePath::FusedCutile && - !cuvs::distance::detail::can_launch_fused_1nn_tile(X.data_handle(), - centroids.data_handle(), - X.extent(0), - centroids.extent(0), - X.extent(1), - metric)) { - path = use_legacy_fused(handle, X.extent(0), centroids.extent(0), metric); + const auto path = cuvs::distance::detail::resolve_fused_1nn_backend( + handle, + X.data_handle(), + centroids.data_handle(), + X.extent(0), + centroids.extent(0), + X.extent(1), + metric); + + Fused1nnRequirements requirements{}; + requirements.path = path; + requirements.sample_tile = getDataBatchSize(batch_samples, X.extent(0)); + requirements.centroid_tile = getCentroidsBatchSize(batch_centroids, centroids.extent(0)); + requirements.workspace_alignment = alignof(int); + + if (path == FusedDistancePath::FusedCutile) { + requirements.output_layout = Fused1nnOutputLayout::Soa; + requirements.norm_policy = std::is_same_v + ? Fused1nnNormPolicy::Tf32 + : Fused1nnNormPolicy::Default; + requirements.result_alignment = 16; + requirements.distance_offset = + raft::alignTo(sizeof(IndexT) * static_cast(X.extent(0)), size_t{16}); + requirements.result_bytes = requirements.distance_offset + + sizeof(DataT) * static_cast(X.extent(0)); + if constexpr (std::is_same_v) { + requirements.workspace_bytes = + sizeof(int) * cuvs::distance::detail::fused_1nn_cutile_index_workspace_rows( + X.extent(0)); + } + } else { + requirements.output_layout = Fused1nnOutputLayout::Kvp; + requirements.norm_policy = Fused1nnNormPolicy::Default; + requirements.result_alignment = alignof(raft::KeyValuePair); + requirements.result_bytes = + sizeof(raft::KeyValuePair) * static_cast(X.extent(0)); + if (path == FusedDistancePath::FusedCutlass) { + requirements.workspace_bytes = sizeof(int) * static_cast(X.extent(0)); + } else if (path == FusedDistancePath::Unfused && + (metric == cuvs::distance::DistanceType::L2Expanded || + metric == cuvs::distance::DistanceType::L2SqrtExpanded || + metric == cuvs::distance::DistanceType::CosineExpanded)) { + auto sample_tile = requirements.sample_tile; + const auto centroid_tile = requirements.centroid_tile; + sample_tile = std::min(sample_tile, std::numeric_limits::max() / centroid_tile); + const size_t distance_bytes = sizeof(DataT) * static_cast(sample_tile) * + static_cast(centroid_tile); + requirements.workspace_alignment = std::max(alignof(DataT), alignof(raft::KeyValuePair)); + requirements.workspace_bytes = raft::alignTo(distance_bytes, alignof(raft::KeyValuePair)); + if (centroid_tile < centroids.extent(0)) { + requirements.workspace_bytes += + sizeof(raft::KeyValuePair) * static_cast(sample_tile); + } } } - return path; + return requirements; } template @@ -143,7 +186,7 @@ void min_cluster_and_distance_compute_impl(raft::resources const& handle, int batch_samples, int batch_centroids, rmm::device_uvector& workspace, - FusedDistancePath fused_path, + const Fused1nnRequirements& requirements, const DataT* cutile_x_norm) { cudaStream_t stream = raft::resource::get_cuda_stream(handle); @@ -153,7 +196,11 @@ void min_cluster_and_distance_compute_impl(raft::resources const& handle, const bool is_l2_cos = metric == cuvs::distance::DistanceType::L2Expanded || metric == cuvs::distance::DistanceType::L2SqrtExpanded || metric == cuvs::distance::DistanceType::CosineExpanded; + const auto fused_path = requirements.path; const bool cutile_ready = fused_path == FusedDistancePath::FusedCutile; + if (workspace.size() < requirements.workspace_bytes) { + workspace.resize(requirements.workspace_bytes, stream); + } RAFT_EXPECTS(!cutile_ready || native_kvp == nullptr, "cuTile fused 1-NN requires separate index and distance outputs"); @@ -175,14 +222,10 @@ void min_cluster_and_distance_compute_impl(raft::resources const& handle, : const_cast(cutile_x_norm); auto* tf32_centroid_norms = L2NormBuf_OR_DistBuf.data() + centroid_offset; if (cutile_x_norm == nullptr) { - raft::common::nvtx::range input_norm_scope( - "cutile_input_norm"); compute_tf32_row_norms( handle, X.data_handle(), tf32_x_norms, n_samples, n_features, take_sqrt); } { - raft::common::nvtx::range center_norm_scope( - "cutile_center_norm"); compute_tf32_row_norms(handle, centroids.data_handle(), tf32_centroid_norms, @@ -224,11 +267,15 @@ void min_cluster_and_distance_compute_impl(raft::resources const& handle, temp_kvp.resize(n_samples, stream); cutlass_kvp_scratch = temp_kvp.data(); } - workspace.resize(sizeof(int) * static_cast(n_samples), stream); + if (workspace.size() < sizeof(int) * static_cast(n_samples)) { + workspace.resize(sizeof(int) * static_cast(n_samples), stream); + } } else if (needs_index_workspace) { const auto workspace_rows = cuvs::distance::detail::fused_1nn_cutile_index_workspace_rows(n_samples); - workspace.resize(sizeof(int) * workspace_rows, stream); + if (workspace.size() < sizeof(int) * workspace_rows) { + workspace.resize(sizeof(int) * workspace_rows, stream); + } } cuvs::distance::fusedDistanceNNMinReduce( @@ -271,7 +318,6 @@ void min_cluster_and_distance_compute_impl(raft::resources const& handle, // The unfused reduction indexes its distance matrix with IndexT. dataBatchSize = std::min(dataBatchSize, std::numeric_limits::max() / centroidsBatchSize); - workspace.resize(sizeof(DataT) * dataBatchSize * centroidsBatchSize, stream); using KeyValueT = raft::KeyValuePair; rmm::device_uvector temp_kvp(native_kvp == nullptr ? n_samples : 0, stream); @@ -281,8 +327,17 @@ void min_cluster_and_distance_compute_impl(raft::resources const& handle, raft::matrix::fill(handle, kvp_output_view, initial_value); const bool tileCentroids = centroidsBatchSize < n_clusters; - rmm::device_uvector batchMinClusterAndDistance(tileCentroids ? dataBatchSize : 0, - stream); + const size_t distance_workspace_bytes = + sizeof(DataT) * static_cast(dataBatchSize) * static_cast(centroidsBatchSize); + const size_t batch_min_offset = raft::alignTo(distance_workspace_bytes, alignof(KeyValueT)); + const size_t required_workspace_bytes = + batch_min_offset + (tileCentroids ? sizeof(KeyValueT) * static_cast(dataBatchSize) : 0); + if (workspace.size() < required_workspace_bytes) { + workspace.resize(required_workspace_bytes, stream); + } + auto* batch_min_storage = tileCentroids + ? reinterpret_cast(workspace.data() + batch_min_offset) + : nullptr; for (IndexT dIdx = 0; dIdx < n_samples;) { auto ns = std::min(dataBatchSize, n_samples - dIdx); @@ -291,7 +346,7 @@ void min_cluster_and_distance_compute_impl(raft::resources const& handle, for (IndexT cIdx = 0; cIdx < n_clusters;) { auto nc = std::min(centroidsBatchSize, n_clusters - cIdx); - auto batchMin = tileCentroids ? batchMinClusterAndDistance.data() + auto batchMin = tileCentroids ? batch_min_storage : minClusterAndDistanceView.data_handle(); cuvs::distance::unfusedDistanceNNMinReduce( @@ -406,20 +461,22 @@ void minClusterAndDistanceCompute(raft::resources const& handle, int batch_samples, int batch_centroids, rmm::device_uvector& workspace, + const Fused1nnRequirements& requirements, const DataT* cutile_x_norm) { - auto path = select_min_cluster_distance_path(handle, X, centroids, metric); + RAFT_EXPECTS(requirements.output_layout == Fused1nnOutputLayout::Soa, + "resolved fused 1-NN plan requires KVP output"); if constexpr (is_cutile_fused_data_type_v) { - if (path == FusedDistancePath::FusedCutile && - !cuvs::distance::detail::can_launch_fused_1nn_tile(nearest_idx.data_handle(), - nearest_dist.data_handle(), - X.data_handle(), - centroids.data_handle(), - X.extent(0), - centroids.extent(0), - X.extent(1), - metric)) { - path = use_legacy_fused(handle, X.extent(0), centroids.extent(0), metric); + if (requirements.path == FusedDistancePath::FusedCutile) { + RAFT_EXPECTS(cuvs::distance::detail::can_launch_fused_1nn_tile(nearest_idx.data_handle(), + nearest_dist.data_handle(), + X.data_handle(), + centroids.data_handle(), + X.extent(0), + centroids.extent(0), + X.extent(1), + metric), + "resolved cuTile plan has incompatible output storage"); } } @@ -435,7 +492,7 @@ void minClusterAndDistanceCompute(raft::resources const& handle, batch_samples, batch_centroids, workspace, - path, + requirements, cutile_x_norm); } @@ -451,10 +508,10 @@ void minClusterAndDistanceComputeKvp( int batch_samples, int batch_centroids, rmm::device_uvector& workspace, - FusedDistancePath path) + const Fused1nnRequirements& requirements) { - RAFT_EXPECTS(path != FusedDistancePath::FusedCutile, - "cuTile fused 1-NN cannot write native KVP output"); + RAFT_EXPECTS(requirements.output_layout == Fused1nnOutputLayout::Kvp, + "resolved fused 1-NN plan requires separate output arrays"); min_cluster_and_distance_compute_impl(handle, X, centroids, @@ -467,7 +524,7 @@ void minClusterAndDistanceComputeKvp( batch_samples, batch_centroids, workspace, - path, + requirements, nullptr); } @@ -484,6 +541,7 @@ void minClusterAndDistanceComputeKvp( int batch_samples, \ int batch_centroids, \ rmm::device_uvector& workspace, \ + const Fused1nnRequirements& requirements, \ const DataT* cutile_x_norm); INSTANTIATE_MIN_CLUSTER_AND_DISTANCE(float, int64_t) @@ -498,19 +556,21 @@ template void computeCutileRowNorms( template void computeCutileRowNorms( raft::resources const&, const float*, float*, int64_t, int64_t, bool); -#define INSTANTIATE_SELECT_MIN_CLUSTER_PATH(DataT, IndexT) \ - template FusedDistancePath select_min_cluster_distance_path( \ - raft::resources const&, \ - raft::device_matrix_view, \ - raft::device_matrix_view, \ - cuvs::distance::DistanceType); +#define INSTANTIATE_FUSED_1NN_REQUIREMENTS(DataT, IndexT) \ + template Fused1nnRequirements get_fused_1nn_requirements( \ + raft::resources const&, \ + raft::device_matrix_view, \ + raft::device_matrix_view, \ + cuvs::distance::DistanceType, \ + int, \ + int); -INSTANTIATE_SELECT_MIN_CLUSTER_PATH(float, int64_t) -INSTANTIATE_SELECT_MIN_CLUSTER_PATH(double, int64_t) -INSTANTIATE_SELECT_MIN_CLUSTER_PATH(float, int) -INSTANTIATE_SELECT_MIN_CLUSTER_PATH(double, int) +INSTANTIATE_FUSED_1NN_REQUIREMENTS(float, int64_t) +INSTANTIATE_FUSED_1NN_REQUIREMENTS(double, int64_t) +INSTANTIATE_FUSED_1NN_REQUIREMENTS(float, int) +INSTANTIATE_FUSED_1NN_REQUIREMENTS(double, int) -#undef INSTANTIATE_SELECT_MIN_CLUSTER_PATH +#undef INSTANTIATE_FUSED_1NN_REQUIREMENTS #define INSTANTIATE_MIN_CLUSTER_AND_DISTANCE_KVP(DataT, IndexT) \ template void minClusterAndDistanceComputeKvp( \ @@ -524,7 +584,7 @@ INSTANTIATE_SELECT_MIN_CLUSTER_PATH(double, int) int, \ int, \ rmm::device_uvector&, \ - FusedDistancePath); + const Fused1nnRequirements&); INSTANTIATE_MIN_CLUSTER_AND_DISTANCE_KVP(float, int64_t) INSTANTIATE_MIN_CLUSTER_AND_DISTANCE_KVP(double, int64_t) diff --git a/cpp/src/detail/jit_lto/TileAlgorithmPlanner.cpp b/cpp/src/detail/jit_lto/TileAlgorithmPlanner.cpp index aa6fe1dd5d..d23d337817 100644 --- a/cpp/src/detail/jit_lto/TileAlgorithmPlanner.cpp +++ b/cpp/src/detail/jit_lto/TileAlgorithmPlanner.cpp @@ -57,7 +57,6 @@ std::shared_ptr TileAlgorithmPlanner::try_get_launcher return it->second; } - RAFT_LOG_DEBUG("Building launcher for kernel entrypoint: %s", entrypoint_.c_str()); auto launcher = this->build(); if (!launcher) { launcher_cache_.build_failed.insert(launch_key); diff --git a/cpp/src/distance/detail/fused_distance_nn.cuh b/cpp/src/distance/detail/fused_distance_nn.cuh index 7bb9b7a4f9..e1ec935a15 100644 --- a/cpp/src/distance/detail/fused_distance_nn.cuh +++ b/cpp/src/distance/detail/fused_distance_nn.cuh @@ -17,6 +17,8 @@ #include #include // raft::KeyValuePair #include // raft::identity_op +#include +#include #include // Policy #include // raft::util::arch::SM_* #include // raft::ceildiv, raft::shfl @@ -30,6 +32,43 @@ namespace distance { namespace detail { +/** Backend selected by the fused 1-NN primitive. */ +enum class Fused1nnBackend : std::uint8_t { + Unfused = 0, + FusedCutile, + FusedCutlass, +}; + +template +constexpr Fused1nnBackend fused_1nn_legacy_backend( + int cc_major, IdxT m, IdxT n, cuvs::distance::DistanceType metric) +{ + if (metric == cuvs::distance::DistanceType::InnerProduct) { return Fused1nnBackend::Unfused; } + if (cc_major <= 8 || (cc_major == 9 && (m >= 4096 || n >= 4096))) { + return Fused1nnBackend::FusedCutlass; + } + return Fused1nnBackend::Unfused; +} + +/** Resolve the executable fused 1-NN backend before callers allocate result storage. */ +template +Fused1nnBackend resolve_fused_1nn_backend(const raft::resources& handle, + const DataT* x, + const DataT* y, + IdxT m, + IdxT n, + IdxT k, + cuvs::distance::DistanceType metric) +{ + if constexpr (is_fused_1nn_cutile_data_v) { + if (can_launch_fused_1nn_tile(x, y, m, n, k, metric)) { + return Fused1nnBackend::FusedCutile; + } + } + const auto prop = raft::resource::get_device_properties(handle); + return fused_1nn_legacy_backend(prop.major, m, n, metric); +} + template ( cutlass_kvp_scratch, m, maxVal, cutlass_red_op, stream); launch_legacy(cutlass_kvp_scratch, cutlass_red_op); - unpackFused1nnKvpToSoa(nearest_idx, nearest_dist, cutlass_kvp_scratch, m, stream); + if (nearest_idx != nullptr || nearest_dist != nullptr) { + unpackFused1nnKvpToSoa(nearest_idx, nearest_dist, cutlass_kvp_scratch, m, stream); + } } else { RAFT_EXPECTS(nearest_idx == nullptr && nearest_dist != nullptr, "Direct CUTLASS output supports distance-only results"); diff --git a/cpp/src/distance/detail/fused_distance_nn/cutile/export_fused_1nn.py b/cpp/src/distance/detail/fused_distance_nn/cutile/export_fused_1nn.py index 2f126fde4f..1fdf313d18 100644 --- a/cpp/src/distance/detail/fused_distance_nn/cutile/export_fused_1nn.py +++ b/cpp/src/distance/detail/fused_distance_nn/cutile/export_fused_1nn.py @@ -280,8 +280,7 @@ def main() -> int: ) args = parser.parse_args() - print( - export_binary( + export_binary( args.output_file, output_format=args.format, data_type=args.data_type, @@ -294,7 +293,6 @@ def main() -> int: matrix_layout=args.matrix_layout, occupancy=args.occupancy, bytecode_version=args.bytecode_version, - ) ) return 0 diff --git a/cpp/src/neighbors/ivf_flat/detail/jit_lto_kernels/interleaved_scan_impl.cuh b/cpp/src/neighbors/ivf_flat/detail/jit_lto_kernels/interleaved_scan_impl.cuh index d28cbb6773..66a13a6cba 100644 --- a/cpp/src/neighbors/ivf_flat/detail/jit_lto_kernels/interleaved_scan_impl.cuh +++ b/cpp/src/neighbors/ivf_flat/detail/jit_lto_kernels/interleaved_scan_impl.cuh @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ @@ -179,12 +179,7 @@ __device__ __forceinline__ void interleaved_scan_impl(const uint32_t query_smem_ } if constexpr (kManageLocalTopK) { - // A filtered or padded record must not carry an in-range sample offset. If fewer than k - // valid records remain, the dummy can reach the output queue; using the end offset makes - // postprocess_neighbors translate it to kOutOfBoundsRecord instead of a duplicate valid - // database index. - const uint32_t sample_ix = valid ? sample_offset + vec_id : chunk_indices[n_probes - 1]; - queue.add(val, sample_ix); + queue.add(val, sample_offset + vec_id); } else { if (vec_id < list_length) distances[sample_offset + vec_id] = val; } @@ -207,16 +202,6 @@ __device__ __forceinline__ void interleaved_scan_impl(const uint32_t query_smem_ __syncthreads(); queue.done(interleaved_scan_kernel_smem); queue.store(distances, neighbors, [](auto val) { return post_process(val); }); - - // block_sort initializes slots that never received a candidate with (kDummy, idx=0). Scrub - // those too so a completely empty/filtered probe set cannot turn the internal index zero into - // a real database ID during neighbor postprocessing. - if (threadIdx.x < raft::WarpSize) { - const auto dummy_out = post_process(local_topk_t::queue_t::kDummy); - for (uint32_t i = threadIdx.x; i < k; i += raft::WarpSize) { - if (distances[i] == dummy_out) { neighbors[i] = chunk_indices[n_probes - 1]; } - } - } } } diff --git a/cpp/src/neighbors/ivf_flat/ivf_flat_search.cuh b/cpp/src/neighbors/ivf_flat/ivf_flat_search.cuh index e6f54dd457..960d48c818 100644 --- a/cpp/src/neighbors/ivf_flat/ivf_flat_search.cuh +++ b/cpp/src/neighbors/ivf_flat/ivf_flat_search.cuh @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2022-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ @@ -292,28 +292,6 @@ void search_impl(raft::resources const& handle, cuvs::selection::SelectAlgo::kAuto, num_samples_vector); } - if (manage_local_topk && grid_dim_x > 1) { - // The merge select initializes unused output slots with (kDummy, idx=0), even when the - // corresponding local inputs carried an out-of-range sentinel. Restore the sentinel before - // neighbor postprocessing so an underfilled result cannot turn idx=0 into a duplicate DB ID. - AccT dummy_out = effective_metric == cuvs::distance::DistanceType::CosineExpanded - ? raft::lower_bound() - : (select_min ? raft::upper_bound() : raft::lower_bound()); - if (effective_metric == cuvs::distance::DistanceType::L2SqrtExpanded || - effective_metric == cuvs::distance::DistanceType::L2SqrtUnexpanded) { - dummy_out = raft::sqrt_op{}(dummy_out); - } else if (effective_metric == cuvs::distance::DistanceType::CosineExpanded) { - dummy_out = AccT{1} - dummy_out; - } - raft::linalg::map_offset(handle, - raft::make_device_vector_view( - neighbors_uint32, std::size_t(n_queries) * std::size_t(k)), - [neighbors_uint32, distances, dummy_out] __device__(std::size_t i) { - return distances[i] == dummy_out - ? std::numeric_limits::max() - : neighbors_uint32[i]; - }); - } if (!manage_local_topk) { // post process distances && neighbor IDs ivf::detail::postprocess_distances( diff --git a/cpp/src/neighbors/ivf_rabitq/gpu_index/quantizer_gpu.cu b/cpp/src/neighbors/ivf_rabitq/gpu_index/quantizer_gpu.cu index c7701ae6ef..af0876acca 100644 --- a/cpp/src/neighbors/ivf_rabitq/gpu_index/quantizer_gpu.cu +++ b/cpp/src/neighbors/ivf_rabitq/gpu_index/quantizer_gpu.cu @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ @@ -519,8 +519,6 @@ void data_transformation_batch_opt(const float* d_data, // 5. Save the rotated centroid: copy CP into d_rotated_c. raft::copy(d_rotated_c, d_CP, D, stream); - if (num_points == 0) { return; } - // 6. Launch the single FUSED kernel for subtract, normalize, and binarize. const unsigned int FusedBlockSize = 256; // A good default, can be tuned. dim3 gridDim(num_points); @@ -637,8 +635,6 @@ void DataQuantizerGPU::quantize_batch_opt(const float* d_data, D, handle_); - if (num_points == 0) { return; } - rabitq_codes_and_factors_fused(d_rotated_c, d_bin_XP.data_handle(), d_XP.data_handle(), @@ -716,8 +712,6 @@ void data_transformation_batch_opt_contiguous(const float* d_contiguous_data, // 5. Save the rotated centroid: copy CP into d_rotated_c. raft::copy(d_rotated_c, d_CP, D, stream); - if (num_points == 0) { return; } - // 6. Launch the single FUSED kernel for subtract, normalize, and binarize. const unsigned int FusedBlockSize = 256; // A good default, can be tuned. dim3 gridDim(num_points); @@ -761,8 +755,6 @@ void DataQuantizerGPU::quantize_batch_opt_contiguous(const float* d_contiguous_d D, handle_); - if (num_points == 0) { return; } - rabitq_codes_and_factors_fused(d_rotated_c, d_bin_XP.data_handle(), d_XP.data_handle(), diff --git a/cpp/tests/cluster/kmeans.cu b/cpp/tests/cluster/kmeans.cu index aa7de8619f..59051484f4 100644 --- a/cpp/tests/cluster/kmeans.cu +++ b/cpp/tests/cluster/kmeans.cu @@ -679,8 +679,7 @@ TEST_P(KmeansFitBatchedTestF, Result) { prepareBlobInputs(); fitBatchedTest(); - // Batched FP32 centroid accumulation uses atomics, so its reduction order is not deterministic. - // Compare the resulting clustering and inertia rather than individual centroid coordinates. + ASSERT_TRUE(centroids_match); ASSERT_TRUE(score >= 0.99); ASSERT_TRUE(inertia_match); runInitSizeCompare(); diff --git a/cpp/tests/cluster/kmeans_balanced.cu b/cpp/tests/cluster/kmeans_balanced.cu index dbe8798de2..10d4ce9cbb 100644 --- a/cpp/tests/cluster/kmeans_balanced.cu +++ b/cpp/tests/cluster/kmeans_balanced.cu @@ -274,54 +274,4 @@ KB_TEST((KmeansBalancedTest, KmeansBalancedTestFI8I32I32_SEP, inputsf_i32); -#if CUVS_CUTILE_ENABLED -TEST(KmeansBalancedPredict, CutilePointerRejectionPreservesMetricFallback) -{ - raft::resources handle; - auto stream = raft::resource::get_cuda_stream(handle); - constexpr int64_t n_rows = 4; - constexpr int64_t n_cols = 4; - constexpr int64_t n_clusters = 2; - - const std::vector h_x{ - 1.0f, 0.0f, 0.2f, 0.0f, 0.0f, 1.0f, 0.0f, 0.2f, 0.8f, 0.1f, 0.0f, 0.0f, 0.1f, 0.9f, 0.0f, 0.0f}; - const std::vector h_centroids{1.0f, 0.0f, 0.0f, 0.0f, 0.0f, 1.0f, 0.0f, 0.0f}; - - rmm::device_uvector x_aligned(h_x.size(), stream); - rmm::device_uvector x_misaligned(h_x.size() + 1, stream); - rmm::device_uvector centroids(h_centroids.size(), stream); - rmm::device_uvector labels_ref(n_rows, stream); - rmm::device_uvector labels_misaligned(n_rows + 1, stream); - raft::update_device(x_aligned.data(), h_x.data(), h_x.size(), stream); - raft::update_device(x_misaligned.data() + 1, h_x.data(), h_x.size(), stream); - raft::update_device(centroids.data(), h_centroids.data(), h_centroids.size(), stream); - - for (auto metric : {distance::DistanceType::L2Expanded, - distance::DistanceType::CosineExpanded, - distance::DistanceType::InnerProduct}) { - cluster::kmeans::balanced_params params; - params.metric = metric; - cluster::kmeans::predict( - handle, - params, - raft::make_device_matrix_view(x_aligned.data(), n_rows, n_cols), - raft::make_device_matrix_view(centroids.data(), n_clusters, n_cols), - raft::make_device_vector_view(labels_ref.data(), n_rows)); - cluster::kmeans::predict( - handle, - params, - raft::make_device_matrix_view(x_misaligned.data() + 1, n_rows, n_cols), - raft::make_device_matrix_view(centroids.data(), n_clusters, n_cols), - raft::make_device_vector_view(labels_misaligned.data() + 1, n_rows)); - - std::vector h_ref(n_rows); - std::vector h_actual(n_rows); - raft::update_host(h_ref.data(), labels_ref.data(), n_rows, stream); - raft::update_host(h_actual.data(), labels_misaligned.data() + 1, n_rows, stream); - raft::resource::sync_stream(handle, stream); - EXPECT_EQ(h_actual, h_ref); - } -} -#endif - } // namespace cuvs diff --git a/cpp/tests/cluster/kmeans_predict_batching.cu b/cpp/tests/cluster/kmeans_predict_batching.cu index d54595254f..b740c3f998 100644 --- a/cpp/tests/cluster/kmeans_predict_batching.cu +++ b/cpp/tests/cluster/kmeans_predict_batching.cu @@ -157,18 +157,4 @@ TEST(KMeansPredict, BatchParametersPreserveResultsAndReduceUnfusedAllocations) } } -TEST(KMeansPredict, ProbeFailureUsesLegacyArchitectureFallback) -{ - using cuvs::distance::DistanceType; - using detail::FusedDistancePath; - - for (auto metric : {DistanceType::L2Expanded, DistanceType::CosineExpanded}) { - EXPECT_EQ(detail::use_legacy_fused(8, 1024, 1024, metric), FusedDistancePath::FusedCutlass); - EXPECT_EQ(detail::use_legacy_fused(9, 4096, 1024, metric), FusedDistancePath::FusedCutlass); - EXPECT_EQ(detail::use_legacy_fused(9, 1024, 1024, metric), FusedDistancePath::Unfused); - EXPECT_EQ(detail::use_legacy_fused(10, 16384, 16384, metric), FusedDistancePath::Unfused); - EXPECT_EQ(detail::use_legacy_fused(12, 16384, 16384, metric), FusedDistancePath::Unfused); - } -} - } // namespace cuvs::cluster::kmeans diff --git a/cpp/tests/neighbors/distance_nn.cu b/cpp/tests/neighbors/distance_nn.cu index d2b28b70a1..df13981279 100644 --- a/cpp/tests/neighbors/distance_nn.cu +++ b/cpp/tests/neighbors/distance_nn.cu @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ @@ -9,16 +9,11 @@ #include "../../src/distance/fused_distance_nn.cuh" #include "../../src/distance/unfused_distance_nn.cuh" -#include -#include - #include #include #include #include -#include - namespace cuvs::neighbors { enum class ImplType { fused, unfused }; @@ -47,7 +42,9 @@ __global__ void fill_int8(int8_t* buff, int len, int seed_offset) template class NNTest : public ::testing::TestWithParam> { public: - using RefOutT = raft::KeyValuePair; + using OutT = std::conditional_t>; NNTest() : params_{::testing::TestWithParam>::GetParam()}, m{params_.m}, @@ -60,10 +57,8 @@ class NNTest : public ::testing::TestWithParam> { y{raft::make_device_matrix(handle, n, k)}, x_norm{raft::make_device_vector(handle, m)}, y_norm{raft::make_device_vector(handle, n)}, - out_idx{raft::make_device_vector(handle, m)}, - out_dist{raft::make_device_vector(handle, m)}, - out_kvp{raft::make_device_vector(handle, m)}, - ref_out{raft::make_device_vector(handle, m)} + out{raft::make_device_vector(handle, m)}, + ref_out{raft::make_device_vector(handle, m)} { } @@ -99,11 +94,15 @@ class NNTest : public ::testing::TestWithParam> { workspace_size = m * n * sizeof(AccT); } - raft::matrix::fill(handle, raft::make_device_matrix_view(out_idx.data_handle(), m, 1), IdxT{0}); - raft::matrix::fill( - handle, raft::make_device_matrix_view(out_dist.data_handle(), m, 1), AccT{0}); - raft::matrix::fill( - handle, raft::make_device_matrix_view(ref_out.data_handle(), m, 1), RefOutT{0, 0}); + // Reset buffer + if constexpr (std::is_same_v>) { + // OutT is a RAFT KeyValuePair + raft::matrix::fill( + handle, raft::make_device_matrix_view(out.data_handle(), m, 1), OutT{0, 0}); + } else { + // OutT is a scalar type + raft::matrix::fill(handle, raft::make_device_matrix_view(out.data_handle(), m, 1), OutT{0}); + } raft::resource::sync_stream(handle, stream); } @@ -112,13 +111,13 @@ class NNTest : public ::testing::TestWithParam> { raft::device_vector workspace = raft::make_device_vector(handle, workspace_size); - ref_nn( + ref_nn( ref_out.data_handle(), x.data_handle(), y.data_handle(), m, n, k, sqrt, metric, stream); if constexpr (impl == ImplType::fused) { if constexpr (std::is_same_v) { - cuvs::distance::fusedDistanceNNMinReduce(out_idx.data_handle(), - out_dist.data_handle(), + cuvs::distance::fusedDistanceNNMinReduce(nullptr, + out.data_handle(), x.data_handle(), y.data_handle(), x_norm.data_handle(), @@ -128,20 +127,20 @@ class NNTest : public ::testing::TestWithParam> { k, (void*)workspace.data_handle(), sqrt, - false, + true, true, metric, 0.0, - out_kvp.data_handle(), + nullptr, stream); } else { static_assert(sizeof(DataT) == 0, "fusedDistanceNNMinReduce is not implemented for datatype other than float"); } } else if constexpr (impl == ImplType::unfused) { - cuvs::distance::unfusedDistanceNNMinReduce( + cuvs::distance::unfusedDistanceNNMinReduce( handle, - out_kvp.data_handle(), + out.data_handle(), x.data_handle(), y.data_handle(), x_norm.data_handle(), @@ -161,16 +160,7 @@ class NNTest : public ::testing::TestWithParam> { void compare() { - if constexpr (impl == ImplType::fused) { - vector_compare_soa( - handle, ref_out.data_handle(), out_idx.data_handle(), out_dist.data_handle(), m, summary); - // FP32 Tensor Core inputs are rounded to TF32, so near-tied candidates may select a - // different valid nearest neighbor than the full-FP32 reference. - const auto allowed_misses = std::max(1, (m + 499) / 500); - ASSERT_LE(summary.n_misses, allowed_misses) << summary; - } else { - vector_compare(handle, ref_out.data_handle(), out_kvp.data_handle(), m, summary); - } + vector_compare(handle, ref_out.data_handle(), out.data_handle(), m, summary); ASSERT_TRUE(summary.max_diff < params_.tol) << summary; } @@ -188,10 +178,8 @@ class NNTest : public ::testing::TestWithParam> { raft::device_matrix y; raft::device_vector x_norm; raft::device_vector y_norm; - raft::device_vector out_idx; - raft::device_vector out_dist; - raft::device_vector out_kvp; - raft::device_vector ref_out; + raft::device_vector out; + raft::device_vector ref_out; size_t workspace_size; }; @@ -215,14 +203,9 @@ template const std::vector> input_fp32_fused = [] { auto inputs = input_fp32; #if CUVS_CUTILE_ENABLED - inputs.insert( - inputs.begin() + 6, - NNInputs{512, 1024, 64, DistanceType::InnerProduct, false, uint64_t(31415926), 0.1}); -#endif inputs.push_back( - NNInputs{1000, 8, 32, DistanceType::L2Expanded, false, uint64_t(31415926), 0.1}); - inputs.push_back( - NNInputs{1000, 40, 16, DistanceType::CosineExpanded, false, uint64_t(31415926), 0.1}); + {512, 1024, 64, DistanceType::InnerProduct, false, uint64_t(31415926), 0.1}); +#endif return inputs; }(); @@ -267,260 +250,6 @@ TEST_P(NNTest_fp16_unfused, test) INSTANTIATE_TEST_CASE_P(NNTest, NNTest_fp16_unfused, ::testing::ValuesIn(input_fp16)); -TEST(Fused1nn, ExpandedL2ClampsNegativeRoundoff) -{ - raft::resources handle; - auto stream = raft::resource::get_cuda_stream(handle); - constexpr int k = 64; - - auto x = raft::make_device_matrix(handle, 1, k); - auto y = raft::make_device_matrix(handle, 1, k); - auto x_norm = raft::make_device_vector(handle, 1); - auto y_norm = raft::make_device_vector(handle, 1); - auto out_idx = raft::make_device_vector(handle, 1); - auto out_dist = raft::make_device_vector(handle, 1); - auto out_kvp = raft::make_device_vector, int>(handle, 1); - auto workspace = raft::make_device_vector(handle, 1); - - raft::matrix::fill(handle, x.view(), 1.0006f); - raft::copy(y.data_handle(), x.data_handle(), k, stream); - raft::linalg::rowNorm( - x_norm.data_handle(), x.data_handle(), k, 1, stream); - raft::copy(y_norm.data_handle(), x_norm.data_handle(), 1, stream); - - cuvs::distance::fusedDistanceNNMinReduce(out_idx.data_handle(), - out_dist.data_handle(), - x.data_handle(), - y.data_handle(), - x_norm.data_handle(), - y_norm.data_handle(), - 1, - 1, - k, - workspace.data_handle(), - true, - true, - true, - DistanceType::L2SqrtExpanded, - 0.0f, - out_kvp.data_handle(), - stream); - - int actual_idx; - float actual_dist; - raft::update_host(&actual_idx, out_idx.data_handle(), 1, stream); - raft::update_host(&actual_dist, out_dist.data_handle(), 1, stream); - raft::resource::sync_stream(handle); - EXPECT_EQ(actual_idx, 0); - EXPECT_EQ(actual_dist, 0.0f); -} - -TEST(Fused1nn, CutileAvailabilityRejectsUnsupportedArchitectures) -{ - EXPECT_FALSE(cuvs::detail::jit_lto::cutile_launch_available_for_arch(7, 5, 13000)); - EXPECT_FALSE(cuvs::detail::jit_lto::cutile_launch_available_for_arch(13, 0, 13000)); -} - -#if CUVS_CUTILE_ENABLED -TEST(Fused1nn, ExpectedModuleCompatibilityErrorsAreRecoverable) -{ - using cuvs::detail::jit_lto::is_expected_cutile_unavailable; - EXPECT_TRUE(is_expected_cutile_unavailable(cudaErrorInvalidDeviceFunction)); - EXPECT_TRUE(is_expected_cutile_unavailable(cudaErrorInvalidPtx)); - EXPECT_TRUE(is_expected_cutile_unavailable(cudaErrorNoKernelImageForDevice)); - EXPECT_TRUE(is_expected_cutile_unavailable(cudaErrorSymbolNotFound)); - EXPECT_TRUE(is_expected_cutile_unavailable(cudaErrorUnsupportedPtxVersion)); - EXPECT_TRUE(is_expected_cutile_unavailable(cudaErrorCallRequiresNewerDriver)); - EXPECT_FALSE(is_expected_cutile_unavailable(cudaErrorMemoryAllocation)); - EXPECT_FALSE(is_expected_cutile_unavailable(cudaErrorIllegalAddress)); -} - -template -void run_half_cutile_contract_case(int k) -{ - raft::resources handle; - auto stream = raft::resource::get_cuda_stream(handle); - constexpr IdxT m = 2; - constexpr IdxT n = 2; - - std::vector h_x(static_cast(m) * k, __float2half(0.0f)); - std::vector h_y(static_cast(n) * k, __float2half(0.0f)); - h_x[0] = __float2half(1.0f); - h_x[k + 1] = __float2half(1.0f); - h_y[0] = __float2half(1.0f); - h_y[k + 1] = __float2half(1.0f); - - rmm::device_uvector x(h_x.size(), stream); - rmm::device_uvector y(h_y.size(), stream); - rmm::device_uvector x_norm(m, stream); - rmm::device_uvector y_norm(n, stream); - rmm::device_uvector out_idx(m, stream); - rmm::device_uvector out_dist(m, stream); - rmm::device_uvector workspace(m, stream); - raft::update_device(x.data(), h_x.data(), h_x.size(), stream); - raft::update_device(y.data(), h_y.data(), h_y.size(), stream); - const std::vector h_norms(m, 1.0f); - raft::update_device(x_norm.data(), h_norms.data(), m, stream); - raft::update_device(y_norm.data(), h_norms.data(), n, stream); - - if constexpr (std::is_same_v) { - EXPECT_FALSE((cuvs::distance::detail::try_fused_1nn_tile(out_idx.data(), - out_dist.data(), - x.data(), - y.data(), - x_norm.data(), - y_norm.data(), - m, - n, - static_cast(k), - DistanceType::L2Expanded, - false, - x.data(), - stream))); - } - - for (auto metric : {DistanceType::L2Expanded, - DistanceType::L2SqrtExpanded, - DistanceType::CosineExpanded, - DistanceType::InnerProduct}) { - ASSERT_TRUE(( - cuvs::distance::detail::try_fused_1nn_tile(out_idx.data(), - out_dist.data(), - x.data(), - y.data(), - x_norm.data(), - y_norm.data(), - m, - n, - static_cast(k), - metric, - metric == DistanceType::L2SqrtExpanded, - workspace.data(), - stream))); - std::vector h_idx(m); - raft::update_host(h_idx.data(), out_idx.data(), m, stream); - raft::resource::sync_stream(handle, stream); - EXPECT_EQ(h_idx[0], IdxT{0}); - EXPECT_EQ(h_idx[1], IdxT{1}); - } - - ASSERT_TRUE((cuvs::distance::detail::try_fused_1nn_tile(out_idx.data(), - out_dist.data(), - x.data(), - x.data(), - x_norm.data(), - x_norm.data(), - m, - m, - static_cast(k), - DistanceType::L2Expanded, - false, - workspace.data(), - stream))); - std::vector h_alias_idx(m); - raft::update_host(h_alias_idx.data(), out_idx.data(), m, stream); - raft::resource::sync_stream(handle, stream); - EXPECT_EQ(h_alias_idx[0], IdxT{0}); - EXPECT_EQ(h_alias_idx[1], IdxT{1}); - - raft::copy(y.data() + k, y.data(), k, stream); - ASSERT_TRUE((cuvs::distance::detail::try_fused_1nn_tile(out_idx.data(), - out_dist.data(), - x.data(), - y.data(), - x_norm.data(), - y_norm.data(), - m, - n, - static_cast(k), - DistanceType::L2Expanded, - false, - workspace.data(), - stream))); - IdxT h_tie_idx; - half h_tie_dist; - raft::update_host(&h_tie_idx, out_idx.data(), 1, stream); - raft::update_host(&h_tie_dist, out_dist.data(), 1, stream); - raft::resource::sync_stream(handle, stream); - EXPECT_TRUE(h_tie_idx == IdxT{0} || h_tie_idx == IdxT{1}); - EXPECT_EQ(__half2float(h_tie_dist), 0.0f); -} - -TEST(Fused1nn, HalfUsesFloatNormsAcrossAbisAndIndexTypes) -{ - run_half_cutile_contract_case(8); - run_half_cutile_contract_case(7); - run_half_cutile_contract_case(8); - run_half_cutile_contract_case(7); -} -#endif - -TEST(Fused1nn, Int64IndexWorkspaceUsesLargestChunk) -{ - constexpr int64_t max_batch_m_float = cuvs::distance::detail::fused_1nn_cutile_max_batch_m; - constexpr int64_t max_batch_m_half = cuvs::distance::detail::fused_1nn_cutile_max_batch_m; - EXPECT_EQ(max_batch_m_float, 2147483644); - EXPECT_EQ(max_batch_m_half, 2147483640); - EXPECT_EQ(cuvs::distance::detail::fused_1nn_cutile_index_workspace_rows(1024), 1024); - EXPECT_EQ(cuvs::distance::detail::fused_1nn_cutile_index_workspace_rows( - std::numeric_limits::max()), - static_cast(max_batch_m_float)); - EXPECT_EQ(cuvs::distance::detail::fused_1nn_cutile_index_workspace_rows( - std::numeric_limits::max()), - static_cast(max_batch_m_half)); -} - -#if CUVS_CUTILE_ENABLED -TEST(Fused1nn, PointerAwareProbeRejectsMisalignedArrays) -{ - raft::resources handle; - constexpr int m = 32; - constexpr int n = 32; - constexpr int k = 64; - - auto x = raft::make_device_vector(handle, m * k + 1); - auto y = raft::make_device_vector(handle, n * k); - auto x_norm = raft::make_device_vector(handle, m + 1); - auto y_norm = raft::make_device_vector(handle, n); - auto out_idx = raft::make_device_vector(handle, m); - auto out_dist = raft::make_device_vector(handle, m); - - for (auto metric : - {DistanceType::L2Expanded, DistanceType::CosineExpanded, DistanceType::InnerProduct}) { - EXPECT_FALSE(cuvs::distance::detail::can_launch_fused_1nn_tile(out_idx.data_handle(), - out_dist.data_handle(), - x.data_handle() + 1, - y.data_handle(), - m, - n, - k, - metric)); - - if (cuvs::distance::detail::can_launch_fused_1nn_tile(out_idx.data_handle(), - out_dist.data_handle(), - x.data_handle(), - y.data_handle(), - m, - n, - k, - metric)) { - EXPECT_FALSE(cuvs::distance::detail::can_launch_fused_1nn_tile(out_idx.data_handle(), - out_dist.data_handle(), - x.data_handle(), - y.data_handle(), - x_norm.data_handle() + 1, - y_norm.data_handle(), - m, - n, - k, - metric)); - EXPECT_FALSE(cuvs::distance::detail::can_launch_fused_1nn_tile( - out_idx.data_handle(), x.data_handle(), x.data_handle(), y.data_handle(), m, n, k, metric)); - } - } -} -#endif - template const std::vector> input_int8 = { {4096, 4096, 64, DistanceType::L2Expanded, false, uint64_t(31415926), 0.1}, diff --git a/cpp/tests/neighbors/distance_nn_helper.cuh b/cpp/tests/neighbors/distance_nn_helper.cuh index 51028876ff..5162160363 100644 --- a/cpp/tests/neighbors/distance_nn_helper.cuh +++ b/cpp/tests/neighbors/distance_nn_helper.cuh @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ @@ -66,16 +66,6 @@ __device__ AccT cosine_distance(const DataT* v1, const DataT* v2, IdxT K) } // This is a naive implementation of 1-NN computation -template -__device__ AccT inner_product_score(const DataT* v1, const DataT* v2, IdxT K) -{ - AccT score = AccT(0.0); - for (IdxT i = 0; i < K; i++) { - score += AccT(v1[i]) * AccT(v2[i]); - } - return score; -} - template RAFT_KERNEL ref_nn_kernel( OutT* out, const DataT* A, const DataT* B, IdxT M, IdxT N, IdxT K, bool sqrt, DistanceType metric) @@ -83,47 +73,44 @@ RAFT_KERNEL ref_nn_kernel( IdxT tid = threadIdx.x + blockIdx.x * IdxT(blockDim.x); for (IdxT m = tid; m < M; m += (blockDim.x * gridDim.x)) { - IdxT best_index = N + 1; - AccT best_score = min_val(); - AccT best_dist = max_val(); - - for (IdxT n = 0; n < N; n++) { - if (metric == DistanceType::InnerProduct) { - AccT score = inner_product_score(&A[m * K], &B[n * K], K); - if (score > best_score) { - best_score = score; - best_index = n; + if (metric == DistanceType::InnerProduct) { + IdxT max_index = N + 1; + AccT max_score = min_val(); + for (IdxT n = 0; n < N; n++) { + AccT score = AccT(0.0); + for (IdxT k = 0; k < K; k++) { + score += AccT(A[m * K + k]) * AccT(B[n * K + k]); + } + if (score > max_score) { + max_score = score; + max_index = n; } - continue; } + if constexpr (std::is_fundamental::value) { + out[m] = max_score; + } else { + out[m].key = max_index; + out[m].value = max_score; + } + continue; + } + IdxT min_index = N + 1; + AccT min_dist = max_val(); + + for (IdxT n = 0; n < N; n++) { AccT dist; if (metric == DistanceType::L2SqrtExpanded || metric == DistanceType::L2Expanded) { dist = l2_distance(&A[m * K], &B[n * K], K); } else if (metric == DistanceType::CosineExpanded) { dist = cosine_distance(&A[m * K], &B[n * K], K); - } else { - continue; - } - if (dist < best_dist) { - best_dist = dist; - best_index = n; } - } - - if (metric == DistanceType::InnerProduct) { - if constexpr (std::is_fundamental::value) { - out[m] = AccT(best_score); - } else { - out[m].key = IdxT(best_index); - out[m].value = AccT(best_score); + if (dist < min_dist) { + min_dist = dist; + min_index = n; } - continue; } - IdxT min_index = best_index; - AccT min_dist = best_dist; - if constexpr (std::is_fundamental::value) { static_assert(std::is_same::value, "OutT and AccT are not same type"); out[m] = AccT(min_dist); @@ -209,34 +196,6 @@ class ComparisonSummary { } }; -template -void vector_compare_soa(raft::resources const& handle, - const raft::KeyValuePair* ref, - const IdxT* out_idx, - const AccT* out_dist, - IdxT n, - ComparisonSummary& summary) -{ - auto ref_h = raft::make_host_vector, IdxT>(n); - auto idx_h = raft::make_host_vector(n); - auto dist_h = raft::make_host_vector(n); - - raft::copy(ref_h.data_handle(), ref, n, raft::resource::get_cuda_stream(handle)); - raft::copy(idx_h.data_handle(), out_idx, n, raft::resource::get_cuda_stream(handle)); - raft::copy(dist_h.data_handle(), out_dist, n, raft::resource::get_cuda_stream(handle)); - raft::resource::sync_stream(handle, raft::resource::get_cuda_stream(handle)); - - summary.init(); - - for (IdxT i = 0; i < n; i++) { - const double a_val = double(dist_h(i)); - const double b_val = double(ref_h(i).value); - const bool missed = idx_h(i) != ref_h(i).key; - const double diff = std::abs(a_val - b_val); - summary.update(diff, i, a_val, b_val, missed); - } -} - template void vector_compare( raft::resources const& handle, const OutT* a, const OutT* b, IdxT n, ComparisonSummary& summary) From 8e2a9b7e24ba640b0e05090b3eba92ad37a58bdb Mon Sep 17 00:00:00 2001 From: divyegala Date: Thu, 3 Sep 2026 03:50:53 +0000 Subject: [PATCH 70/82] embeddings --- .../all_cuda-133_arch-aarch64.yaml | 2 + .../all_cuda-133_arch-x86_64.yaml | 2 + .../bench_ann_cuda-133_arch-aarch64.yaml | 2 + .../bench_ann_cuda-133_arch-x86_64.yaml | 2 + conda/recipes/libcuvs/recipe.yaml | 10 + cpp/CMakeLists.txt | 33 ++ .../modules/compute_matrix_product.cmake | 25 +- .../modules/generate_cutile_kernels.cmake | 287 ++++++++++++++++++ .../modules/register_cutile_fragment.cpp.in | 31 ++ .../detail/jit_lto/CutileFragmentEntry.hpp | 119 ++++++++ .../detail/jit_lto/TileAlgorithmPlanner.hpp | 70 +++++ .../cuvs/detail/jit_lto/cutile_arch_tags.hpp | 54 ++++ .../cuvs/detail/jit_lto/cutile_module.hpp | 123 ++++++++ .../detail/jit_lto/cutile_smoke_fragments.hpp | 15 + .../cuvs/detail/jit_lto/tileir_compat.hpp | 111 +++++++ .../detail/jit_lto/TileAlgorithmPlanner.cpp | 144 +++++++++ .../cutile_smoke/cutile_smoke_matrix.json | 16 + .../jit_lto/cutile_smoke/export_smoke.py | 73 +++++ .../jit_lto/cutile_smoke/smoke_kernel.py | 17 ++ cpp/tests/CMakeLists.txt | 7 + cpp/tests/detail/jit_lto/cutile_smoke.cu | 137 +++++++++ dependencies.yaml | 46 +++ python/libcuvs/pyproject.toml | 2 + 23 files changed, 1322 insertions(+), 6 deletions(-) create mode 100644 cpp/cmake/modules/generate_cutile_kernels.cmake create mode 100644 cpp/cmake/modules/register_cutile_fragment.cpp.in create mode 100644 cpp/include/cuvs/detail/jit_lto/CutileFragmentEntry.hpp create mode 100644 cpp/include/cuvs/detail/jit_lto/TileAlgorithmPlanner.hpp create mode 100644 cpp/include/cuvs/detail/jit_lto/cutile_arch_tags.hpp create mode 100644 cpp/include/cuvs/detail/jit_lto/cutile_module.hpp create mode 100644 cpp/include/cuvs/detail/jit_lto/cutile_smoke_fragments.hpp create mode 100644 cpp/include/cuvs/detail/jit_lto/tileir_compat.hpp create mode 100644 cpp/src/detail/jit_lto/TileAlgorithmPlanner.cpp create mode 100644 cpp/src/detail/jit_lto/cutile_smoke/cutile_smoke_matrix.json create mode 100644 cpp/src/detail/jit_lto/cutile_smoke/export_smoke.py create mode 100644 cpp/src/detail/jit_lto/cutile_smoke/smoke_kernel.py create mode 100644 cpp/tests/detail/jit_lto/cutile_smoke.cu diff --git a/conda/environments/all_cuda-133_arch-aarch64.yaml b/conda/environments/all_cuda-133_arch-aarch64.yaml index d95c6854bd..07deffb120 100644 --- a/conda/environments/all_cuda-133_arch-aarch64.yaml +++ b/conda/environments/all_cuda-133_arch-aarch64.yaml @@ -15,8 +15,10 @@ dependencies: - cuda-nvrtc-dev - cuda-nvtx-dev - cuda-profiler-api +- cuda-tileiras - cuda-version=13.3 - cupy>=14.0.1,!=14.1.0 +- cutile-python - cxx-compiler - cython>=3.2.2 - dlpack>=0.8,<1.0 diff --git a/conda/environments/all_cuda-133_arch-x86_64.yaml b/conda/environments/all_cuda-133_arch-x86_64.yaml index 69d8f84605..e6de7ca7a5 100644 --- a/conda/environments/all_cuda-133_arch-x86_64.yaml +++ b/conda/environments/all_cuda-133_arch-x86_64.yaml @@ -15,8 +15,10 @@ dependencies: - cuda-nvrtc-dev - cuda-nvtx-dev - cuda-profiler-api +- cuda-tileiras - cuda-version=13.3 - cupy>=14.0.1,!=14.1.0 +- cutile-python - cxx-compiler - cython>=3.2.2 - dlpack>=0.8,<1.0 diff --git a/conda/environments/bench_ann_cuda-133_arch-aarch64.yaml b/conda/environments/bench_ann_cuda-133_arch-aarch64.yaml index 321a892555..1d5f229c8f 100644 --- a/conda/environments/bench_ann_cuda-133_arch-aarch64.yaml +++ b/conda/environments/bench_ann_cuda-133_arch-aarch64.yaml @@ -15,8 +15,10 @@ dependencies: - cuda-nvrtc-dev - cuda-nvtx-dev - cuda-profiler-api +- cuda-tileiras - cuda-version=13.3 - cupy>=14.0.1,!=14.1.0 +- cutile-python - cuvs==26.10.*,>=0.0.0a0 - cxx-compiler - cython>=3.2.2 diff --git a/conda/environments/bench_ann_cuda-133_arch-x86_64.yaml b/conda/environments/bench_ann_cuda-133_arch-x86_64.yaml index 179b4a4a2f..228d11c3d8 100644 --- a/conda/environments/bench_ann_cuda-133_arch-x86_64.yaml +++ b/conda/environments/bench_ann_cuda-133_arch-x86_64.yaml @@ -15,8 +15,10 @@ dependencies: - cuda-nvrtc-dev - cuda-nvtx-dev - cuda-profiler-api +- cuda-tileiras - cuda-version=13.3 - cupy>=14.0.1,!=14.1.0 +- cutile-python - cuvs==26.10.*,>=0.0.0a0 - cxx-compiler - cython>=3.2.2 diff --git a/conda/recipes/libcuvs/recipe.yaml b/conda/recipes/libcuvs/recipe.yaml index 93a69ea1c0..1baa35ec4d 100644 --- a/conda/recipes/libcuvs/recipe.yaml +++ b/conda/recipes/libcuvs/recipe.yaml @@ -70,6 +70,11 @@ cache: - cuda-version =${{ cuda_version }} - cmake ${{ cmake_version }} - ninja + - python + - if: cuda_major == "13" + then: + - cutile-python + - cuda-tileiras - ${{ stdlib("c") }} host: - libnvjitlink-dev @@ -397,6 +402,11 @@ outputs: - cuda-version =${{ cuda_version }} - cmake ${{ cmake_version }} - ninja + - python + - if: cuda_major == "13" + then: + - cutile-python + - cuda-tileiras - ${{ stdlib("c") }} host: - ${{ pin_subpackage("libcuvs-headers", exact=True) }} diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt index 69ede8404e..caf6b4143a 100644 --- a/cpp/CMakeLists.txt +++ b/cpp/CMakeLists.txt @@ -1153,6 +1153,36 @@ if(NOT BUILD_CPU_ONLY) ) endblock() + include(cmake/modules/generate_cutile_kernels.cmake) + set(cutile_smoke_dir "${CMAKE_CURRENT_SOURCE_DIR}/src/detail/jit_lto/cutile_smoke") + set(cutile_smoke_generated_dir + "${CMAKE_CURRENT_BINARY_DIR}/generated_kernels/detail/jit_lto/cutile_smoke" + ) + generate_cutile_kernels( + cutile_smoke_files + KERNEL_DIR + "${cutile_smoke_dir}" + KERNEL_BASENAME + "cutile_smoke" + KERNEL_PYTHON + "smoke_kernel.py" + EXPORT_SCRIPT + "export_smoke.py" + OUTPUT_DIRECTORY + "${cutile_smoke_generated_dir}" + MATRIX_JSON_FILE + "${cutile_smoke_dir}/cutile_smoke_matrix.json" + FRAGMENT_TAG_FORMAT_CUBIN + "cuvs::detail::jit_lto::fragment_tag_cutile_smoke_add_cubin" + FRAGMENT_TAG_HEADER_FILES + "" + "" + ) + if(NOT DEFINED CUVS_CUTILE_ENABLED) + set(CUVS_CUTILE_ENABLED 0) + endif() + target_compile_definitions(cuvs_cpp_headers INTERFACE CUVS_CUTILE_ENABLED=${CUVS_CUTILE_ENABLED}) + # Note that this matrix contains an `arch_includes` placeholder, since we don't currently have a # way to do an item-wise transform on a list after computing the matrix product and before # configuring the file @@ -1364,6 +1394,7 @@ if(NOT BUILD_CPU_ONLY) src/core/omp_wrapper.cpp src/util/file_io.cpp src/util/host_memory.cpp + src/detail/jit_lto/TileAlgorithmPlanner.cpp src/distance/detail/kernels/gram_matrix.cu src/distance/detail/kernels/kernel_factory.cu src/distance/detail/kernels/kernel_matrices.cu @@ -1460,6 +1491,7 @@ if(NOT BUILD_CPU_ONLY) src/stats/trustworthiness_score.cu ${CUVS_MG_ALGOS} ${jit_lto_files} + ${cutile_smoke_files} ) set_target_properties( @@ -1504,6 +1536,7 @@ if(NOT BUILD_CPU_ONLY) "$" INTERFACE "$" PRIVATE "${CMAKE_CURRENT_SOURCE_DIR}/src" "${CMAKE_CURRENT_BINARY_DIR}/src" + "${cutile_smoke_generated_dir}" ) # Endian detection diff --git a/cpp/cmake/modules/compute_matrix_product.cmake b/cpp/cmake/modules/compute_matrix_product.cmake index 82a34f9242..60b96113f0 100644 --- a/cpp/cmake/modules/compute_matrix_product.cmake +++ b/cpp/cmake/modules/compute_matrix_product.cmake @@ -1,12 +1,23 @@ # ============================================================================= # cmake-format: off -# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # cmake-format: on # ============================================================================= include_guard(GLOBAL) +function(cuvs_find_build_python output_var) + if(DEFINED ENV{BUILD_PREFIX}) + set(Python_ROOT "$ENV{BUILD_PREFIX}") + endif() + find_package(Python REQUIRED COMPONENTS Interpreter) + set(${output_var} + "${Python_EXECUTABLE}" + PARENT_SCOPE + ) +endfunction() + function(compute_matrix_product output_var) set(options) set(one_value MATRIX_JSON_FILE MATRIX_JSON_STRING) @@ -14,19 +25,21 @@ function(compute_matrix_product output_var) cmake_parse_arguments(_JIT_LTO "${options}" "${one_value}" "${multi_value}" ${ARGN}) - find_package(Python3 REQUIRED COMPONENTS Interpreter) + cuvs_find_build_python(_matrix_python_executable) if(_JIT_LTO_MATRIX_JSON_FILE) execute_process( - COMMAND "${Python3_EXECUTABLE}" "${CMAKE_CURRENT_FUNCTION_LIST_DIR}/compute_matrix_product.py" - "${_JIT_LTO_MATRIX_JSON_FILE}" # + COMMAND + "${_matrix_python_executable}" + "${CMAKE_CURRENT_FUNCTION_LIST_DIR}/compute_matrix_product.py" + "${_JIT_LTO_MATRIX_JSON_FILE}" OUTPUT_VARIABLE output COMMAND_ERROR_IS_FATAL ANY ) else() execute_process( COMMAND "${CMAKE_COMMAND}" -E echo "${_JIT_LTO_MATRIX_JSON_STRING}" - COMMAND "${Python3_EXECUTABLE}" "${CMAKE_CURRENT_FUNCTION_LIST_DIR}/compute_matrix_product.py" - - + COMMAND "${_matrix_python_executable}" + "${CMAKE_CURRENT_FUNCTION_LIST_DIR}/compute_matrix_product.py" - OUTPUT_VARIABLE output COMMAND_ERROR_IS_FATAL ANY ) endif() diff --git a/cpp/cmake/modules/generate_cutile_kernels.cmake b/cpp/cmake/modules/generate_cutile_kernels.cmake new file mode 100644 index 0000000000..35fb8c381b --- /dev/null +++ b/cpp/cmake/modules/generate_cutile_kernels.cmake @@ -0,0 +1,287 @@ +# ============================================================================= +# cmake-format: off +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# cmake-format: on +# ============================================================================= + +include_guard(GLOBAL) + +include(${CMAKE_CURRENT_LIST_DIR}/compute_matrix_product.cmake) + +function(generate_cutile_kernels_stub) + set(CUVS_CUTILE_ENABLED + 0 + PARENT_SCOPE + ) +endfunction() + +function(_cutile_fragment_tag_header_files output_var) + set(${output_var} "") + foreach(_header IN LISTS ARGN) + if(NOT _header MATCHES "^(\".*\"|<.*>)$") + set(_header "\"${_header}\"") + endif() + string(APPEND ${output_var} "#include ${_header}\n") + endforeach() + set(${output_var} + "${${output_var}}" + PARENT_SCOPE + ) +endfunction() + +function(_cutile_kernels_setup) + set(options) + set(one_value MATRIX_JSON_FILE OUTPUT_DIRECTORY) + set(multi_value) + cmake_parse_arguments(_CUTILE "${options}" "${one_value}" "${multi_value}" ${ARGN}) + + find_package(CUDAToolkit REQUIRED) + + if(CUDAToolkit_VERSION VERSION_LESS 13.0) + message( + STATUS + "cuTile embedded kernels require CUDA 13.0+; skipping cuTile generation (found ${CUDAToolkit_VERSION})." + ) + set(_CUTILE_SETUP_OK + FALSE + PARENT_SCOPE + ) + return() + endif() + + cuvs_find_build_python(Python3_EXECUTABLE) + + find_program( + CUTILE_BIN2C + NAMES bin2c + PATHS ${CUDAToolkit_BIN_DIR} REQUIRED + ) + + execute_process( + COMMAND "${Python3_EXECUTABLE}" -c "import cuda.tile" + RESULT_VARIABLE _cutile_import_result + ERROR_VARIABLE _cutile_import_error + OUTPUT_QUIET ERROR_STRIP_TRAILING_WHITESPACE + ) + if(NOT _cutile_import_result EQUAL 0) + message( + FATAL_ERROR + "cuda.tile (cuTile Python) is required to build cuTile embedded kernels. " + "Install cutile-python and cuda-tileiras (conda), or cuda-tile[tileiras] (pip).\n" + "Interpreter: ${Python3_EXECUTABLE}\n" + "Import error: ${_cutile_import_error}" + ) + endif() + message(STATUS "Using cuTile Python: ${Python3_EXECUTABLE}") + + set_property( + DIRECTORY + PROPERTY CMAKE_CONFIGURE_DEPENDS "${_CUTILE_MATRIX_JSON_FILE}" + APPEND + ) + + file(MAKE_DIRECTORY "${_CUTILE_OUTPUT_DIRECTORY}") + + set(Python3_EXECUTABLE + "${Python3_EXECUTABLE}" + PARENT_SCOPE + ) + set(CUTILE_BIN2C + "${CUTILE_BIN2C}" + PARENT_SCOPE + ) + set(_CUTILE_SETUP_OK + TRUE + PARENT_SCOPE + ) +endfunction() + +function(_cutile_make_python_args output_var) + set(_python_args + --format + "${output_format}" + --data-type + "${data_type}" + --metric + "${metric}" + --index-type + "${index_type}" + --tile-m + "${tile_m}" + --tile-n + "${tile_n}" + --tile-k + "${tile_k}" + --gpu-code + "${gpu_code}" + ) + if(DEFINED bytecode_version AND NOT "${bytecode_version}" STREQUAL "") + list(APPEND _python_args --bytecode-version "${bytecode_version}") + endif() + if(DEFINED matrix_layout AND NOT "${matrix_layout}" STREQUAL "") + list(APPEND _python_args --matrix-layout "${matrix_layout}") + endif() + if(DEFINED occupancy AND NOT "${occupancy}" STREQUAL "") + list(APPEND _python_args --occupancy "${occupancy}") + endif() + set(${output_var} + "${_python_args}" + PARENT_SCOPE + ) +endfunction() + +function(process_cutile_matrix_entry source_list_var) + set(options) + set(one_value KERNEL_DIR KERNEL_BASENAME KERNEL_PYTHON EXPORT_SCRIPT OUTPUT_DIRECTORY + FRAGMENT_TAG_FORMAT_CUBIN FRAGMENT_TAG_FORMAT_TILEIR MATRIX_JSON_ENTRY + ) + set(multi_value FRAGMENT_TAG_HEADER_FILES) + cmake_parse_arguments(_CUTILE "${options}" "${one_value}" "${multi_value}" ${ARGN}) + + if(NOT Python3_EXECUTABLE) + cuvs_find_build_python(Python3_EXECUTABLE) + endif() + + populate_matrix_variables("${_CUTILE_MATRIX_JSON_ENTRY}") + + if(register STREQUAL "cubin") + string(CONFIGURE "${_CUTILE_FRAGMENT_TAG_FORMAT_CUBIN}" fragment_tag @ONLY) + set(bin2c_symbol embedded_cubin) + set(fragment_entry_type "cuvs::detail::jit_lto::StaticCubinFragmentEntry") + elseif(register STREQUAL "tileir") + string(CONFIGURE "${_CUTILE_FRAGMENT_TAG_FORMAT_TILEIR}" fragment_tag @ONLY) + set(bin2c_symbol embedded_tileir) + set(fragment_entry_type + "cuvs::detail::jit_lto::StaticTileIrBytecodeFragmentEntry" + ) + else() + message(FATAL_ERROR "Unknown cuTile register kind '${register}'") + endif() + + _cutile_fragment_tag_header_files(fragment_tag_header_files ${_CUTILE_FRAGMENT_TAG_HEADER_FILES}) + + string(CONFIGURE "${artifact_basename}" _artifact_basename @ONLY) + set(_artifact_stem "${_CUTILE_KERNEL_BASENAME}_${_artifact_basename}") + set(_artifact_file "${_CUTILE_OUTPUT_DIRECTORY}/${_artifact_stem}.${artifact_ext}") + set(_embedded_header "${_CUTILE_OUTPUT_DIRECTORY}/${_artifact_stem}_${register}.h") + set(_fragment_cpp "${_CUTILE_OUTPUT_DIRECTORY}/${_artifact_stem}_${register}.cpp") + set(embedded_header_file "${_artifact_stem}_${register}.h") + + _cutile_make_python_args(_python_args) + + set(_export_python_executable "${Python3_EXECUTABLE}") + if(DEFINED python_executable AND NOT "${python_executable}" STREQUAL "") + string(CONFIGURE "${python_executable}" _export_python_executable @ONLY) + endif() + + if(DEFINED prebuilt_artifact AND NOT "${prebuilt_artifact}" STREQUAL "") + string(CONFIGURE "${prebuilt_artifact}" _prebuilt_artifact @ONLY) + if(NOT IS_ABSOLUTE "${_prebuilt_artifact}") + set(_prebuilt_artifact "${_CUTILE_KERNEL_DIR}/${_prebuilt_artifact}") + endif() + add_custom_command( + OUTPUT "${_artifact_file}" + COMMAND "${CMAKE_COMMAND}" -E copy_if_different "${_prebuilt_artifact}" "${_artifact_file}" + DEPENDS "${_prebuilt_artifact}" + COMMENT "Copying prebuilt cuTile ${_CUTILE_KERNEL_BASENAME} ${output_format} ${data_type}" + VERBATIM + ) + else() + add_custom_command( + OUTPUT "${_artifact_file}" + COMMAND "${_export_python_executable}" "${_CUTILE_KERNEL_DIR}/${_CUTILE_EXPORT_SCRIPT}" + "${_artifact_file}" ${_python_args} + WORKING_DIRECTORY "${_CUTILE_KERNEL_DIR}" + DEPENDS "${_CUTILE_KERNEL_DIR}/${_CUTILE_EXPORT_SCRIPT}" + "${_CUTILE_KERNEL_DIR}/${_CUTILE_KERNEL_PYTHON}" + COMMENT "Exporting cuTile ${_CUTILE_KERNEL_BASENAME} ${output_format} ${data_type}" + VERBATIM + ) + endif() + + add_custom_command( + OUTPUT "${_embedded_header}" + COMMAND "${CUTILE_BIN2C}" --const --name ${bin2c_symbol} --static "${_artifact_file}" > + "${_embedded_header}" + DEPENDS "${_artifact_file}" + VERBATIM + ) + + configure_file( + "${CMAKE_CURRENT_FUNCTION_LIST_DIR}/register_cutile_fragment.cpp.in" "${_fragment_cpp}" @ONLY + ) + list(APPEND ${source_list_var} "${_embedded_header}" "${_fragment_cpp}") + set(${source_list_var} + "${${source_list_var}}" + PARENT_SCOPE + ) +endfunction() + +function(generate_cutile_kernels source_list_var) + set(options) + set(one_value KERNEL_DIR KERNEL_BASENAME KERNEL_PYTHON EXPORT_SCRIPT OUTPUT_DIRECTORY + MATRIX_JSON_FILE FRAGMENT_TAG_FORMAT_CUBIN FRAGMENT_TAG_FORMAT_TILEIR + ) + set(multi_value FRAGMENT_TAG_HEADER_FILES) + cmake_parse_arguments(_CUTILE "${options}" "${one_value}" "${multi_value}" ${ARGN}) + + if(NOT _CUTILE_KERNEL_BASENAME) + message(FATAL_ERROR "generate_cutile_kernels: KERNEL_BASENAME is required") + endif() + if(NOT _CUTILE_KERNEL_PYTHON) + message(FATAL_ERROR "generate_cutile_kernels: KERNEL_PYTHON is required") + endif() + + _cutile_kernels_setup( + MATRIX_JSON_FILE "${_CUTILE_MATRIX_JSON_FILE}" OUTPUT_DIRECTORY "${_CUTILE_OUTPUT_DIRECTORY}" + ) + if(NOT _CUTILE_SETUP_OK) + generate_cutile_kernels_stub() + set(${source_list_var} + "" + PARENT_SCOPE + ) + return() + endif() + + compute_matrix_product(matrix_product MATRIX_JSON_FILE "${_CUTILE_MATRIX_JSON_FILE}") + + string(JSON len LENGTH "${matrix_product}") + math(EXPR last "${len} - 1") + + # cmake-lint: disable=C0103,E1120 + foreach(i RANGE "${last}") + string(JSON matrix_json_entry GET "${matrix_product}" "${i}") + process_cutile_matrix_entry( + "${source_list_var}" + KERNEL_DIR + "${_CUTILE_KERNEL_DIR}" + KERNEL_BASENAME + "${_CUTILE_KERNEL_BASENAME}" + KERNEL_PYTHON + "${_CUTILE_KERNEL_PYTHON}" + EXPORT_SCRIPT + "${_CUTILE_EXPORT_SCRIPT}" + OUTPUT_DIRECTORY + "${_CUTILE_OUTPUT_DIRECTORY}" + FRAGMENT_TAG_FORMAT_CUBIN + "${_CUTILE_FRAGMENT_TAG_FORMAT_CUBIN}" + FRAGMENT_TAG_FORMAT_TILEIR + "${_CUTILE_FRAGMENT_TAG_FORMAT_TILEIR}" + FRAGMENT_TAG_HEADER_FILES + ${_CUTILE_FRAGMENT_TAG_HEADER_FILES} + MATRIX_JSON_ENTRY + "${matrix_json_entry}" + ) + endforeach() + + set(CUVS_CUTILE_ENABLED + 1 + PARENT_SCOPE + ) + set(${source_list_var} + "${${source_list_var}}" + PARENT_SCOPE + ) +endfunction() diff --git a/cpp/cmake/modules/register_cutile_fragment.cpp.in b/cpp/cmake/modules/register_cutile_fragment.cpp.in new file mode 100644 index 0000000000..7206e88e57 --- /dev/null +++ b/cpp/cmake/modules/register_cutile_fragment.cpp.in @@ -0,0 +1,31 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include "@embedded_header_file@" +#include + +@fragment_tag_header_files@ + + namespace +{ + using fragment_tag = @fragment_tag@; + using fragment_entry = @fragment_entry_type@; + +} // namespace + +template <> +const uint8_t* const fragment_entry::data = @bin2c_symbol@; + +template <> +const size_t fragment_entry::length = sizeof(@bin2c_symbol@); + +template <> +const int fragment_entry::tile_m = @tile_m@; + +template <> +const int fragment_entry::tile_n = @tile_n@; + +template <> +const int fragment_entry::tile_k = @tile_k@; diff --git a/cpp/include/cuvs/detail/jit_lto/CutileFragmentEntry.hpp b/cpp/include/cuvs/detail/jit_lto/CutileFragmentEntry.hpp new file mode 100644 index 0000000000..724662c9dd --- /dev/null +++ b/cpp/include/cuvs/detail/jit_lto/CutileFragmentEntry.hpp @@ -0,0 +1,119 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +#include +#include +#include + +namespace cuvs::detail::jit_lto { + +/** cuTile GEMM-style block geometry embedded in generated static fragment specializations. */ +struct CutileTileConfig { + int tile_m; + int tile_n; + int tile_k; +}; + +/** Embedded CUDA binary module (cubin), loaded directly via cudaLibraryLoadData. */ +struct CubinFragmentEntry { + virtual ~CubinFragmentEntry() = default; + + virtual const uint8_t* get_data() const = 0; + + virtual size_t get_length() const = 0; + + virtual const char* get_key() const = 0; + + virtual int get_cc_major() const = 0; + + virtual int get_cc_minor() const = 0; + + virtual int get_tile_m() const { return 0; } + + virtual int get_tile_n() const { return 0; } + + virtual int get_tile_k() const { return 0; } +}; + +template +struct StaticCubinFragmentEntry final : CubinFragmentEntry { + const uint8_t* get_data() const override { return StaticCubinFragmentEntry::data; } + + size_t get_length() const override { return StaticCubinFragmentEntry::length; } + + const char* get_key() const override + { + return typeid(StaticCubinFragmentEntry).name(); + } + + int get_cc_major() const override { return FragmentTag::cc_major; } + + int get_cc_minor() const override { return FragmentTag::cc_minor; } + + int get_tile_m() const override { return tile_m; } + + int get_tile_n() const override { return tile_n; } + + int get_tile_k() const override { return tile_k; } + + static const int tile_m; + static const int tile_n; + static const int tile_k; + + static const uint8_t* const data; + static const size_t length; +}; + +/** Embedded TileIR bytecode, JIT-compiled by the driver when no matching cubin exists. */ +struct TileIrBytecodeFragmentEntry { + virtual ~TileIrBytecodeFragmentEntry() = default; + + virtual const uint8_t* get_data() const = 0; + + virtual size_t get_length() const = 0; + + virtual const char* get_key() const = 0; + + virtual int get_tile_m() const { return 0; } + + virtual int get_tile_n() const { return 0; } + + virtual int get_tile_k() const { return 0; } +}; + +template +struct StaticTileIrBytecodeFragmentEntry final : TileIrBytecodeFragmentEntry { + const uint8_t* get_data() const override + { + return StaticTileIrBytecodeFragmentEntry::data; + } + + size_t get_length() const override + { + return StaticTileIrBytecodeFragmentEntry::length; + } + + const char* get_key() const override + { + return typeid(StaticTileIrBytecodeFragmentEntry).name(); + } + + int get_tile_m() const override { return tile_m; } + + int get_tile_n() const override { return tile_n; } + + int get_tile_k() const override { return tile_k; } + + static const int tile_m; + static const int tile_n; + static const int tile_k; + + static const uint8_t* const data; + static const size_t length; +}; + +} // namespace cuvs::detail::jit_lto diff --git a/cpp/include/cuvs/detail/jit_lto/TileAlgorithmPlanner.hpp b/cpp/include/cuvs/detail/jit_lto/TileAlgorithmPlanner.hpp new file mode 100644 index 0000000000..c552966e2d --- /dev/null +++ b/cpp/include/cuvs/detail/jit_lto/TileAlgorithmPlanner.hpp @@ -0,0 +1,70 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +#include "CutileFragmentEntry.hpp" + +#include +#include +#include +#include +#include +#include +#include + +#include + +namespace cuvs::detail::jit_lto { + +struct TileLauncherCache { + std::shared_mutex mutex; + std::unordered_map> launchers; + std::unordered_set build_failed; +}; + +/** Loads prebuilt cubins or TileIR bytecode directly through the CUDA library API. */ +struct TileAlgorithmPlanner { + TileAlgorithmPlanner(std::string entrypoint, TileLauncherCache& launcher_cache) + : entrypoint_(std::move(entrypoint)), launcher_cache_(launcher_cache) + { + } + + virtual ~TileAlgorithmPlanner() = default; + + std::shared_ptr get_launcher(); + + /** Returns nullptr when no module can be loaded for the current device (does not RAFT_FAIL). */ + std::shared_ptr try_get_launcher(); + + template + void add_static_fragment() + { + cubin_fragments_.push_back(std::make_unique>()); + } + + template + void add_static_tileir_fragment() + { + tileir_fragment_ = std::make_unique>(); + } + + /** Tile geometry from the cubin or TileIR fragment that would load on this device. */ + CutileTileConfig tile_config() const; + + protected: + std::vector> cubin_fragments_; + std::unique_ptr tileir_fragment_; + + private: + std::string get_planner_key() const; + + std::shared_ptr build(); + + std::string entrypoint_; + TileLauncherCache& launcher_cache_; +}; + +} // namespace cuvs::detail::jit_lto diff --git a/cpp/include/cuvs/detail/jit_lto/cutile_arch_tags.hpp b/cpp/include/cuvs/detail/jit_lto/cutile_arch_tags.hpp new file mode 100644 index 0000000000..2b378dac78 --- /dev/null +++ b/cpp/include/cuvs/detail/jit_lto/cutile_arch_tags.hpp @@ -0,0 +1,54 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +#ifndef CUVS_CUTILE_ENABLED +#define CUVS_CUTILE_ENABLED 0 +#endif + +namespace cuvs::detail::jit_lto { + +#if CUVS_CUTILE_ENABLED + +/** Must stay in sync with cuTile matrix _arch entries and planner add_static_fragment calls. */ +struct cutile_arch_8_0 { + static constexpr int cc_major = 8; + static constexpr int cc_minor = 0; +}; + +struct cutile_arch_8_6 { + static constexpr int cc_major = 8; + static constexpr int cc_minor = 6; +}; + +struct cutile_arch_9_0 { + static constexpr int cc_major = 9; + static constexpr int cc_minor = 0; +}; + +struct cutile_arch_10_0 { + static constexpr int cc_major = 10; + static constexpr int cc_minor = 0; +}; + +struct cutile_arch_12_0 { + static constexpr int cc_major = 12; + static constexpr int cc_minor = 0; +}; + +inline bool is_embedded_cubin_arch(int cc_major, int cc_minor) +{ + if (cc_minor < 0) { return false; } + return cc_major == 8 || cc_major == 9 || cc_major == 10 || cc_major == 12; +} + +#else + +inline bool is_embedded_cubin_arch(int, int) { return false; } + +#endif + +} // namespace cuvs::detail::jit_lto diff --git a/cpp/include/cuvs/detail/jit_lto/cutile_module.hpp b/cpp/include/cuvs/detail/jit_lto/cutile_module.hpp new file mode 100644 index 0000000000..4e39450ba1 --- /dev/null +++ b/cpp/include/cuvs/detail/jit_lto/cutile_module.hpp @@ -0,0 +1,123 @@ +/* + * 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 + +namespace cuvs::detail::jit_lto { + +struct CutileModuleImage { + const uint8_t* data; + size_t size; +}; + +inline bool get_device_compute_capability(int& cc_major, int& cc_minor) +{ + int device = 0; + if (cudaGetDevice(&device) != cudaSuccess) { return false; } + if (cudaDeviceGetAttribute(&cc_major, cudaDevAttrComputeCapabilityMajor, device) != cudaSuccess) { + return false; + } + if (cudaDeviceGetAttribute(&cc_minor, cudaDevAttrComputeCapabilityMinor, device) != cudaSuccess) { + return false; + } + return true; +} + +/** + * Selects the newest compatible cubin in the device's compute-capability major family. + * + * CUDA cubins are forward compatible across minor revisions within a major family, so an SM 8.9 + * device can load SM 8.6 SASS and an SM 12.1 device can load SM 12.0 SASS. + */ +inline const CubinFragmentEntry* find_compatible_cubin_fragment( + int cc_major, + int cc_minor, + const std::vector>& cubin_fragments) +{ + const CubinFragmentEntry* best = nullptr; + for (const auto& fragment : cubin_fragments) { + if (fragment->get_cc_major() != cc_major || fragment->get_cc_minor() > cc_minor) { continue; } + if (best == nullptr || fragment->get_cc_minor() > best->get_cc_minor()) { + best = fragment.get(); + } + } + return best; +} + +/** Selects compatible prebuilt SASS for the device, or TileIR when the driver can JIT it. */ +inline std::optional resolve_cutile_module_image( + int cc_major, + int cc_minor, + int driver_version, + const std::vector>& cubin_fragments, + const TileIrBytecodeFragmentEntry* tileir_fragment) +{ + if (const auto* fragment = find_compatible_cubin_fragment(cc_major, cc_minor, cubin_fragments)) { + return CutileModuleImage{fragment->get_data(), fragment->get_length()}; + } + if (tileir_fragment != nullptr && tileir_fallback_available(driver_version)) { + return CutileModuleImage{tileir_fragment->get_data(), tileir_fragment->get_length()}; + } + return std::nullopt; +} + +inline bool is_expected_cutile_unavailable(cudaError_t status) +{ + switch (status) { + case cudaErrorInvalidDeviceFunction: + case cudaErrorInvalidPtx: + case cudaErrorNoKernelImageForDevice: + case cudaErrorSymbolNotFound: + case cudaErrorUnsupportedPtxVersion: + case cudaErrorCallRequiresNewerDriver: + case cudaErrorSharedObjectSymbolNotFound: + case cudaErrorSharedObjectInitFailed: + case cudaErrorJitCompilerNotFound: return true; + default: return false; + } +} + +/** + * Loads a cuTile launcher, returning null for an expected module/JIT compatibility rejection. + * Unexpected CUDA failures retain the normal RAFT exception behavior. + */ +inline std::shared_ptr try_load_cutile_launcher( + const CutileModuleImage& image, const std::string& kernel_symbol) +{ + cudaLibrary_t library{}; + auto load_status = + cudaLibraryLoadData(&library, image.data, nullptr, nullptr, 0, nullptr, nullptr, 0); + if (load_status != cudaSuccess) { + if (is_expected_cutile_unavailable(load_status)) { return nullptr; } + RAFT_CUDA_TRY(load_status); + } + + cudaKernel_t kernel{}; + load_status = cudaLibraryGetKernel(&kernel, library, kernel_symbol.c_str()); + if (load_status != cudaSuccess) { + RAFT_CUDA_TRY(cudaLibraryUnload(library)); + if (is_expected_cutile_unavailable(load_status)) { return nullptr; } + RAFT_CUDA_TRY(load_status); + } + + return std::make_shared(kernel, library); +} + +} // namespace cuvs::detail::jit_lto diff --git a/cpp/include/cuvs/detail/jit_lto/cutile_smoke_fragments.hpp b/cpp/include/cuvs/detail/jit_lto/cutile_smoke_fragments.hpp new file mode 100644 index 0000000000..3b52f3daf8 --- /dev/null +++ b/cpp/include/cuvs/detail/jit_lto/cutile_smoke_fragments.hpp @@ -0,0 +1,15 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +#pragma once + +namespace cuvs::detail::jit_lto { + +template +struct fragment_tag_cutile_smoke_add_cubin { + static constexpr int cc_major = ArchTag::cc_major; + static constexpr int cc_minor = ArchTag::cc_minor; +}; + +} // namespace cuvs::detail::jit_lto diff --git a/cpp/include/cuvs/detail/jit_lto/tileir_compat.hpp b/cpp/include/cuvs/detail/jit_lto/tileir_compat.hpp new file mode 100644 index 0000000000..a03a7a78fc --- /dev/null +++ b/cpp/include/cuvs/detail/jit_lto/tileir_compat.hpp @@ -0,0 +1,111 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +#ifndef CUVS_CUTILE_ENABLED +#define CUVS_CUTILE_ENABLED 0 +#endif + +#include +#include + +#include + +namespace cuvs::detail::jit_lto { + +/** Minimum CUDA driver version (from cudaDriverGetVersion) for TileIR JIT of embedded bytecode. */ +inline constexpr int kMinTileIrJitDriverVersion = 13010; // CUDA 13.1 / driver >= 590.44 + +/** Minimum CUDA runtime version (from cudaRuntimeGetVersion) for cuTile integration. */ +inline constexpr int kMinCutileRuntimeVersion = 13000; + +inline constexpr bool library_built_with_cutile() +{ +#if CUVS_CUTILE_ENABLED + return true; +#else + return false; +#endif +} + +inline bool runtime_cuda13_or_newer() +{ + int runtime_version = 0; + if (cudaRuntimeGetVersion(&runtime_version) != cudaSuccess) { return false; } + return runtime_version >= kMinCutileRuntimeVersion; +} + +/** True when this build embeds cuTile artifacts and the runtime is CUDA 13+. */ +inline bool cutile_integration_enabled() +{ + return library_built_with_cutile() && runtime_cuda13_or_newer(); +} + +/** True when this build embeds compatible SASS in the device's compute-capability major family. */ +inline bool has_embedded_cubin_for_arch(int cc_major, int cc_minor) +{ + return is_embedded_cubin_arch(cc_major, cc_minor); +} + +/** True when the driver can JIT-compile embedded TileIR bytecode at load time. */ +inline bool tileir_fallback_available(int driver_version) +{ + return driver_version >= kMinTileIrJitDriverVersion; +} + +/** + * True when a cuTile launch may be attempted for the given device: cuTile is enabled, the runtime + * is CUDA 13+, and either compatible same-family SASS exists (no driver JIT required) or the + * driver can JIT the embedded TileIR bytecode fallback. + */ +#if CUVS_CUTILE_ENABLED +inline bool cutile_launch_available_for_arch(int cc_major, int cc_minor, int driver_version) +{ + if (!runtime_cuda13_or_newer()) { return false; } + // The exported fused-1NN kernels require Ampere-or-newer tensor-core semantics, and the current + // integration is validated only through the SM12 family. + if (cc_major < 8 || cc_major > 12) { return false; } + if (has_embedded_cubin_for_arch(cc_major, cc_minor)) { return true; } + return tileir_fallback_available(driver_version); +} +#else +inline constexpr bool cutile_launch_available_for_arch(int, int, int) { return false; } +#endif + +inline bool query_driver_version(int& driver_version) +{ + return cudaDriverGetVersion(&driver_version) == cudaSuccess; +} + +inline bool query_current_device_arch(int& cc_major, int& cc_minor) +{ + int device = 0; + if (cudaGetDevice(&device) != cudaSuccess) { return false; } + if (cudaDeviceGetAttribute(&cc_major, cudaDevAttrComputeCapabilityMajor, device) != cudaSuccess) { + return false; + } + if (cudaDeviceGetAttribute(&cc_minor, cudaDevAttrComputeCapabilityMinor, device) != cudaSuccess) { + return false; + } + return true; +} + +#if CUVS_CUTILE_ENABLED +inline bool cutile_launch_available_on_current_device() +{ + int cc_major = 0; + int cc_minor = 0; + int driver_version = 0; + if (!query_current_device_arch(cc_major, cc_minor)) { return false; } + if (!query_driver_version(driver_version)) { return false; } + return cutile_launch_available_for_arch(cc_major, cc_minor, driver_version); +} +#else +/** Compile-time false when cuTile is not built; use in if constexpr to skip cuTile-only paths. */ +inline constexpr bool cutile_launch_available_on_current_device() { return false; } +#endif + +} // namespace cuvs::detail::jit_lto diff --git a/cpp/src/detail/jit_lto/TileAlgorithmPlanner.cpp b/cpp/src/detail/jit_lto/TileAlgorithmPlanner.cpp new file mode 100644 index 0000000000..d23d337817 --- /dev/null +++ b/cpp/src/detail/jit_lto/TileAlgorithmPlanner.cpp @@ -0,0 +1,144 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include +#include +#include +#include + +#include +#include + +#include +#include + +namespace cuvs::detail::jit_lto { + +namespace { + +template +CutileTileConfig tile_config_from_fragment(const FragmentT* fragment, const std::string& entrypoint) +{ + if (fragment == nullptr) { + RAFT_FAIL("cuTile planner '%s' has no registered fragments", entrypoint.c_str()); + } + const int tile_m = fragment->get_tile_m(); + const int tile_n = fragment->get_tile_n(); + const int tile_k = fragment->get_tile_k(); + if (tile_m <= 0 || tile_n <= 0 || tile_k <= 0) { + RAFT_FAIL( + "cuTile planner '%s' is missing tile geometry in its static fragment (check " + "register_cutile_fragment.cpp generation)", + entrypoint.c_str()); + } + return CutileTileConfig{tile_m, tile_n, tile_k}; +} + +} // namespace + +std::shared_ptr TileAlgorithmPlanner::try_get_launcher() +{ + auto launch_key = this->get_planner_key(); + + { + std::shared_lock read_lock(launcher_cache_.mutex); + if (launcher_cache_.build_failed.count(launch_key)) { return nullptr; } + if (auto it = launcher_cache_.launchers.find(launch_key); + it != launcher_cache_.launchers.end()) { + return it->second; + } + } + + std::unique_lock write_lock(launcher_cache_.mutex); + if (launcher_cache_.build_failed.count(launch_key)) { return nullptr; } + if (auto it = launcher_cache_.launchers.find(launch_key); it != launcher_cache_.launchers.end()) { + return it->second; + } + + auto launcher = this->build(); + if (!launcher) { + launcher_cache_.build_failed.insert(launch_key); + return nullptr; + } + launcher_cache_.launchers[launch_key] = launcher; + return launcher; +} + +std::shared_ptr TileAlgorithmPlanner::get_launcher() +{ + auto launcher = try_get_launcher(); + if (!launcher) { + RAFT_FAIL("Failed to build launcher for kernel entrypoint: %s", entrypoint_.c_str()); + } + return launcher; +} + +std::string TileAlgorithmPlanner::get_planner_key() const +{ + std::string key = entrypoint_; + for (const auto& fragment : cubin_fragments_) { + key += fragment->get_key(); + } + if (tileir_fragment_) { key += tileir_fragment_->get_key(); } + + int device = -1; + int cc_major = -1; + int cc_minor = -1; + int driver_version = -1; + if (cudaGetDevice(&device) == cudaSuccess && + cuvs::detail::jit_lto::get_device_compute_capability(cc_major, cc_minor)) { + key += ":device=" + std::to_string(device); + key += ":cc=" + std::to_string(cc_major) + "." + std::to_string(cc_minor); + if (const auto* fragment = cuvs::detail::jit_lto::find_compatible_cubin_fragment( + cc_major, cc_minor, cubin_fragments_)) { + key += ":cubin=" + std::to_string(fragment->get_cc_major()) + "." + + std::to_string(fragment->get_cc_minor()); + } else { + key += ":tileir"; + } + if (cudaDriverGetVersion(&driver_version) == cudaSuccess) { + key += ":driver=" + std::to_string(driver_version); + } + } + return key; +} + +CutileTileConfig TileAlgorithmPlanner::tile_config() const +{ + int cc_major = 0; + int cc_minor = 0; + if (cuvs::detail::jit_lto::get_device_compute_capability(cc_major, cc_minor)) { + if (const auto* fragment = cuvs::detail::jit_lto::find_compatible_cubin_fragment( + cc_major, cc_minor, cubin_fragments_)) { + return tile_config_from_fragment(fragment, entrypoint_); + } + } + + if (tileir_fragment_) { return tile_config_from_fragment(tileir_fragment_.get(), entrypoint_); } + + if (!cubin_fragments_.empty()) { + return tile_config_from_fragment(cubin_fragments_.front().get(), entrypoint_); + } + + RAFT_FAIL("cuTile planner '%s' has no registered fragments", entrypoint_.c_str()); +} + +std::shared_ptr TileAlgorithmPlanner::build() +{ + int cc_major = 0; + int cc_minor = 0; + if (!cuvs::detail::jit_lto::get_device_compute_capability(cc_major, cc_minor)) { return nullptr; } + + int driver_version = 0; + if (cudaDriverGetVersion(&driver_version) != cudaSuccess) { return nullptr; } + + auto image = cuvs::detail::jit_lto::resolve_cutile_module_image( + cc_major, cc_minor, driver_version, cubin_fragments_, tileir_fragment_.get()); + if (!image) { return nullptr; } + + return cuvs::detail::jit_lto::try_load_cutile_launcher(*image, entrypoint_); +} + +} // namespace cuvs::detail::jit_lto diff --git a/cpp/src/detail/jit_lto/cutile_smoke/cutile_smoke_matrix.json b/cpp/src/detail/jit_lto/cutile_smoke/cutile_smoke_matrix.json new file mode 100644 index 0000000000..3047d111d2 --- /dev/null +++ b/cpp/src/detail/jit_lto/cutile_smoke/cutile_smoke_matrix.json @@ -0,0 +1,16 @@ +[ + { + "_data": [{"data_type": "float", "data_abbrev": "f"}], + "_metric": [{"metric": "add", "metric_abbrev": "add"}], + "_index": [{"index_type": "int32", "index_abbrev": "i32"}], + "_abi": [{"abi": "contiguous", "abi_abbrev": "contiguous"}], + "_tile": [{"tile_m": 256, "tile_n": 1, "tile_k": 1}], + "_export": [ + {"output_format": "cubin", "register": "cubin", "gpu_code": "sm_80", "arch_tag": "cutile_arch_8_0", "artifact_basename": "@gpu_code@", "artifact_ext": "cubin"}, + {"output_format": "cubin", "register": "cubin", "gpu_code": "sm_86", "arch_tag": "cutile_arch_8_6", "artifact_basename": "@gpu_code@", "artifact_ext": "cubin"}, + {"output_format": "cubin", "register": "cubin", "gpu_code": "sm_90", "arch_tag": "cutile_arch_9_0", "artifact_basename": "@gpu_code@", "artifact_ext": "cubin"}, + {"output_format": "cubin", "register": "cubin", "gpu_code": "sm_100", "arch_tag": "cutile_arch_10_0", "artifact_basename": "@gpu_code@", "artifact_ext": "cubin"}, + {"output_format": "cubin", "register": "cubin", "gpu_code": "sm_120", "arch_tag": "cutile_arch_12_0", "artifact_basename": "@gpu_code@", "artifact_ext": "cubin"} + ] + } +] diff --git a/cpp/src/detail/jit_lto/cutile_smoke/export_smoke.py b/cpp/src/detail/jit_lto/cutile_smoke/export_smoke.py new file mode 100644 index 0000000000..64723c0813 --- /dev/null +++ b/cpp/src/detail/jit_lto/cutile_smoke/export_smoke.py @@ -0,0 +1,73 @@ +# ============================================================================= +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# ============================================================================= +"""Export the standalone cuTile embedding smoke kernel.""" + +from __future__ import annotations + +import argparse +from pathlib import Path + +import cuda.tile as ct +from cuda.tile.compilation import ( + ArrayConstraint, + CallingConvention, + KernelSignature, + export_kernel, +) + +from smoke_kernel import TILE_SIZE, cutile_smoke_add + + +def _array_constraint() -> ArrayConstraint: + return ArrayConstraint( + ct.float32, + ndim=1, + index_dtype=ct.int32, + stride_lower_bound_incl=(None,), + alias_groups=(), + may_alias_internally=False, + stride_constant=(1,), + stride_divisible_by=(1,), + shape_divisible_by=(TILE_SIZE,), + base_addr_divisible_by=16, + ) + + +def _signature() -> KernelSignature: + array = _array_constraint() + return KernelSignature( + parameters=[array, array, array], + calling_convention=CallingConvention.cutile_python_v1(), + ).with_symbol("cutile_smoke_add") + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("output_file", type=Path) + parser.add_argument("--format", choices=("cubin",), required=True) + parser.add_argument("--data-type", choices=("float",), required=True) + parser.add_argument("--metric", choices=("add",), required=True) + parser.add_argument("--index-type", choices=("int32",), required=True) + parser.add_argument("--tile-m", type=int, required=True) + parser.add_argument("--tile-n", type=int, required=True) + parser.add_argument("--tile-k", type=int, required=True) + parser.add_argument("--gpu-code", required=True) + args = parser.parse_args() + + if (args.tile_m, args.tile_n, args.tile_k) != (TILE_SIZE, 1, 1): + raise ValueError("cutile smoke kernel requires a 256x1x1 tile") + + export_kernel( + kernel=cutile_smoke_add, + signatures=[_signature()], + output_file=str(args.output_file), + gpu_code=args.gpu_code, + output_format=args.format, + ) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/cpp/src/detail/jit_lto/cutile_smoke/smoke_kernel.py b/cpp/src/detail/jit_lto/cutile_smoke/smoke_kernel.py new file mode 100644 index 0000000000..9b4d6f056a --- /dev/null +++ b/cpp/src/detail/jit_lto/cutile_smoke/smoke_kernel.py @@ -0,0 +1,17 @@ +# ============================================================================= +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# ============================================================================= + +import cuda.tile as ct + + +TILE_SIZE = 256 + + +@ct.kernel +def cutile_smoke_add(lhs, rhs, output): + block = ct.bid(0) + lhs_tile = ct.load(lhs, block, TILE_SIZE) + rhs_tile = ct.load(rhs, block, TILE_SIZE) + ct.store(output, block, lhs_tile + rhs_tile) diff --git a/cpp/tests/CMakeLists.txt b/cpp/tests/CMakeLists.txt index 744bc6a7a2..54f1d9c965 100644 --- a/cpp/tests/CMakeLists.txt +++ b/cpp/tests/CMakeLists.txt @@ -136,6 +136,13 @@ ConfigureTest( PERCENT 100 ) +ConfigureTest( + NAME CUTILE_SMOKE_TEST + PATH detail/jit_lto/cutile_smoke.cu + GPUS 1 + PERCENT 100 +) + ConfigureTest( NAME NEIGHBORS_ANN_IVF_FLAT_UDF_TEST PATH neighbors/ann_ivf_flat/test_udf.cu diff --git a/cpp/tests/detail/jit_lto/cutile_smoke.cu b/cpp/tests/detail/jit_lto/cutile_smoke.cu new file mode 100644 index 0000000000..eb9d4371b5 --- /dev/null +++ b/cpp/tests/detail/jit_lto/cutile_smoke.cu @@ -0,0 +1,137 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include +#include +#include +#include + +#include + +#include + +#include +#include +#include + +namespace cuvs::detail::jit_lto { + +#if !CUVS_CUTILE_ENABLED + +TEST(CutileSmoke, DisabledBuild) +{ + GTEST_SKIP() << "cuTile embedded kernels are disabled in this build"; +} + +#else + +namespace { + +template +using smoke_fragment = StaticCubinFragmentEntry>; + +std::vector> make_smoke_fragments() +{ + std::vector> fragments; + fragments.emplace_back(std::make_unique>()); + fragments.emplace_back(std::make_unique>()); + fragments.emplace_back(std::make_unique>()); + fragments.emplace_back(std::make_unique>()); + fragments.emplace_back(std::make_unique>()); + return fragments; +} + +void add_smoke_fragments(TileAlgorithmPlanner& planner) +{ + planner.add_static_fragment>(); + planner.add_static_fragment>(); + planner.add_static_fragment>(); + planner.add_static_fragment>(); + planner.add_static_fragment>(); +} + +} // namespace + +TEST(CutileSmoke, ResolvesEveryEmbeddedArchitecture) +{ + auto fragments = make_smoke_fragments(); + + EXPECT_EQ(find_compatible_cubin_fragment(8, 0, fragments), fragments[0].get()); + EXPECT_EQ(find_compatible_cubin_fragment(8, 9, fragments), fragments[1].get()); + EXPECT_EQ(find_compatible_cubin_fragment(9, 0, fragments), fragments[2].get()); + EXPECT_EQ(find_compatible_cubin_fragment(10, 0, fragments), fragments[3].get()); + EXPECT_EQ(find_compatible_cubin_fragment(12, 1, fragments), fragments[4].get()); + EXPECT_EQ(find_compatible_cubin_fragment(7, 5, fragments), nullptr); +} + +TEST(CutileSmoke, LaunchesCompatibleCubin) +{ + int cc_major = 0; + int cc_minor = 0; + if (!get_device_compute_capability(cc_major, cc_minor)) { + GTEST_SKIP() << "No CUDA device is available"; + } + + auto fragments = make_smoke_fragments(); + if (find_compatible_cubin_fragment(cc_major, cc_minor, fragments) == nullptr) { + GTEST_SKIP() << "No embedded smoke cubin is compatible with this device"; + } + + TileLauncherCache cache; + TileAlgorithmPlanner planner{"cutile_smoke_add", cache}; + add_smoke_fragments(planner); + auto launcher = planner.try_get_launcher(); + ASSERT_NE(launcher, nullptr); + + cudaStream_t stream = nullptr; + constexpr int count = 256; + std::array host_lhs{}; + std::array host_rhs{}; + std::array host_output{}; + for (int i = 0; i < count; ++i) { + host_lhs[i] = static_cast(i); + host_rhs[i] = static_cast(count - i); + } + + float* lhs = nullptr; + float* rhs = nullptr; + float* output = nullptr; + ASSERT_EQ(cudaMalloc(&lhs, sizeof(host_lhs)), cudaSuccess); + ASSERT_EQ(cudaMalloc(&rhs, sizeof(host_rhs)), cudaSuccess); + ASSERT_EQ(cudaMalloc(&output, sizeof(host_output)), cudaSuccess); + ASSERT_EQ(cudaMemcpy(lhs, host_lhs.data(), sizeof(host_lhs), cudaMemcpyHostToDevice), + cudaSuccess); + ASSERT_EQ(cudaMemcpy(rhs, host_rhs.data(), sizeof(host_rhs), cudaMemcpyHostToDevice), + cudaSuccess); + + using smoke_kernel_t = void(void*, int, int, void*, int, int, void*, int, int); + launcher->template dispatch(stream, + dim3{1, 1, 1}, + dim3{1, 1, 1}, + 0, + static_cast(lhs), + count, + 1, + static_cast(rhs), + count, + 1, + static_cast(output), + count, + 1); + ASSERT_EQ(cudaGetLastError(), cudaSuccess); + ASSERT_EQ(cudaMemcpy(host_output.data(), output, sizeof(host_output), cudaMemcpyDeviceToHost), + cudaSuccess); + ASSERT_EQ(cudaFree(lhs), cudaSuccess); + ASSERT_EQ(cudaFree(rhs), cudaSuccess); + ASSERT_EQ(cudaFree(output), cudaSuccess); + + for (const auto value : host_output) { + EXPECT_FLOAT_EQ(value, static_cast(count)); + } +} + +#endif + +} // namespace cuvs::detail::jit_lto diff --git a/dependencies.yaml b/dependencies.yaml index 86ccd33d06..08b1209411 100644 --- a/dependencies.yaml +++ b/dependencies.yaml @@ -13,6 +13,7 @@ files: - checks - clang - cuda + - cutile_python - cuda_version - depends_on_cuda_python - depends_on_cupy @@ -41,6 +42,7 @@ files: - build_py_cuvs - clang - cuda + - cutile_python - cuda_version - depends_on_cuda_python - depends_on_cupy @@ -78,6 +80,7 @@ files: includes: - clang - cuda + - cutile_python - cuda_version - depends_on_cupy - docs @@ -139,6 +142,7 @@ files: table: tool.rapids-build-backend key: requires includes: + - cutile_python - depends_on_libraft - depends_on_librmm - depends_on_libkvikio @@ -425,6 +429,48 @@ dependencies: - libcusolver-dev - libcusparse-dev - libnvjitlink-dev + cutile_python: + specific: + - output_types: conda + matrices: + - matrix: + cuda: "12.*" + packages: + - matrix: + cuda: "13.3" + packages: + - cutile-python + - cuda-tileiras + - matrix: + cuda: "13.*" + packages: + - cutile-python + - cuda-tileiras + - matrix: + packages: + - cutile-python + - cuda-tileiras + - output_types: [requirements, pyproject] + matrices: + - matrix: + cuda: "12.*" + packages: + - matrix: + cuda: "13.3" + packages: + - cuda-tile + - cuda-toolkit[tileiras]==13.3.* + - matrix: + cuda: "13.*" + packages: + - &cutile_python_cu13 cuda-tile + - &cutile_toolkit_cu13 cuda-toolkit[tileiras]==13.* + # if no matching matrix selectors passed, list the CUDA 13 requirement + # (as a source of documentation in the generated pyproject.toml) + - matrix: + packages: + - *cutile_python_cu13 + - *cutile_toolkit_cu13 cuda_wheels: specific: # cuVS needs 'nvJitLink>={whatever-cuvs-was-built-against}' at runtime, and mixing diff --git a/python/libcuvs/pyproject.toml b/python/libcuvs/pyproject.toml index 9e050be9bc..42b5e8e213 100644 --- a/python/libcuvs/pyproject.toml +++ b/python/libcuvs/pyproject.toml @@ -83,6 +83,8 @@ regex = "(?P.*)" build-backend = "scikit_build_core.build" requires = [ "cmake>=4.0", + "cuda-tile", + "cuda-toolkit[tileiras]==13.*", "libkvikio==26.10.*,>=0.0.0a0", "libraft==26.10.*,>=0.0.0a0", "librmm==26.10.*,>=0.0.0a0", From 1fcea21e500043fba24f9fbab1e9e7f916fd07cc Mon Sep 17 00:00:00 2001 From: divyegala Date: Thu, 3 Sep 2026 04:00:16 +0000 Subject: [PATCH 71/82] style check --- cpp/src/cluster/detail/kmeans_balanced.cuh | 20 ++--- cpp/src/cluster/detail/kmeans_common.cuh | 10 +-- .../detail/minClusterDistanceCompute.cu | 87 +++++++++---------- cpp/src/distance/detail/fused_distance_nn.cuh | 14 +-- .../cutile/export_fused_1nn.py | 24 ++--- cpp/tests/neighbors/distance_nn.cu | 9 +- 6 files changed, 80 insertions(+), 84 deletions(-) diff --git a/cpp/src/cluster/detail/kmeans_balanced.cuh b/cpp/src/cluster/detail/kmeans_balanced.cuh index 1cab6a0ade..b3a6933b48 100644 --- a/cpp/src/cluster/detail/kmeans_balanced.cuh +++ b/cpp/src/cluster/detail/kmeans_balanced.cuh @@ -66,7 +66,7 @@ bool predict_core_min_cluster(const raft::resources& handle, rmm::device_async_resource_ref mr, const MathT* cutile_x_norm) { - auto n_rows = X.extent(0); + auto n_rows = X.extent(0); const auto requirements = get_fused_1nn_requirements(handle, X, centroids, metric); if (requirements.output_layout == Fused1nnOutputLayout::Soa) { @@ -74,11 +74,11 @@ bool predict_core_min_cluster(const raft::resources& handle, // cluster labels are never negative, so no conversion or copy is required. if constexpr (std::is_same_v || std::is_same_v) { if (n_rows <= static_cast(std::numeric_limits::max())) { - using CutileIdxT = int; - const auto cutile_rows = static_cast(n_rows); - const auto cutile_cols = static_cast(X.extent(1)); + using CutileIdxT = int; + const auto cutile_rows = static_cast(n_rows); + const auto cutile_cols = static_cast(X.extent(1)); const auto cutile_clusters = static_cast(centroids.extent(0)); - auto cutile_X = raft::make_device_matrix_view( + auto cutile_X = raft::make_device_matrix_view( X.data_handle(), cutile_rows, cutile_cols); auto cutile_centroids = raft::make_device_matrix_view( centroids.data_handle(), cutile_clusters, cutile_cols); @@ -87,13 +87,13 @@ bool predict_core_min_cluster(const raft::resources& handle, RAFT_EXPECTS(cutile_requirements.output_layout == Fused1nnOutputLayout::Soa, "resolved cuTile plan changed while adapting 32-bit labels"); - auto nearest_dist = - raft::make_device_mdarray(handle, mr, raft::make_extents(cutile_rows)); + auto nearest_dist = raft::make_device_mdarray( + handle, mr, raft::make_extents(cutile_rows)); auto* cutile_labels = reinterpret_cast(labels); auto labels_view = raft::make_device_vector_view(cutile_labels, cutile_rows); - auto cutile_norm = raft::make_device_vector_view( - X_norm.data_handle(), cutile_rows); + auto cutile_norm = + raft::make_device_vector_view(X_norm.data_handle(), cutile_rows); minClusterAndDistanceCompute(handle, cutile_X, cutile_centroids, @@ -117,7 +117,7 @@ bool predict_core_min_cluster(const raft::resources& handle, if constexpr (label_storage_is_cutile_index) { auto* cutile_labels = reinterpret_cast(labels); - auto labels_view = raft::make_device_vector_view(cutile_labels, n_rows); + auto labels_view = raft::make_device_vector_view(cutile_labels, n_rows); minClusterAndDistanceCompute(handle, X, centroids, diff --git a/cpp/src/cluster/detail/kmeans_common.cuh b/cpp/src/cluster/detail/kmeans_common.cuh index 6591e3bf77..0bbd9e2d9d 100644 --- a/cpp/src/cluster/detail/kmeans_common.cuh +++ b/cpp/src/cluster/detail/kmeans_common.cuh @@ -482,7 +482,7 @@ Fused1nnRequirements get_fused_1nn_requirements( raft::device_matrix_view X, raft::device_matrix_view centroids, cuvs::distance::DistanceType metric, - int batch_samples = 0, + int batch_samples = 0, int batch_centroids = 0); template @@ -591,7 +591,7 @@ void countSamplesInCluster(raft::resources const& handle, auto n_clusters = centroids.extent(0); rmm::device_uvector L2NormBuf_OR_DistBuf(0, stream); - auto centroids_const = raft::make_const_mdspan(centroids); + auto centroids_const = raft::make_const_mdspan(centroids); const auto requirements = get_fused_1nn_requirements( handle, X, centroids_const, params.metric, params.batch_samples, params.batch_centroids); @@ -843,11 +843,11 @@ void process_batch(raft::resources const& handle, raft::device_scalar_view clustering_cost, rmm::device_uvector& batch_workspace) { - cudaStream_t stream = raft::resource::get_cuda_stream(handle); - const auto n_samples = batch_data.extent(0); + cudaStream_t stream = raft::resource::get_cuda_stream(handle); + const auto n_samples = batch_data.extent(0); const auto requirements = get_fused_1nn_requirements( handle, batch_data, centroids, metric, batch_samples_param, batch_centroids_param); - auto batch_cost = raft::make_device_scalar(handle, DataT{0}); + auto batch_cost = raft::make_device_scalar(handle, DataT{0}); if (requirements.output_layout == Fused1nnOutputLayout::Soa) { const auto dist_offset = requirements.distance_offset; diff --git a/cpp/src/cluster/detail/minClusterDistanceCompute.cu b/cpp/src/cluster/detail/minClusterDistanceCompute.cu index 11be57e19a..c1da6f6355 100644 --- a/cpp/src/cluster/detail/minClusterDistanceCompute.cu +++ b/cpp/src/cluster/detail/minClusterDistanceCompute.cu @@ -115,35 +115,33 @@ Fused1nnRequirements get_fused_1nn_requirements( int batch_samples, int batch_centroids) { - const auto path = cuvs::distance::detail::resolve_fused_1nn_backend( - handle, - X.data_handle(), - centroids.data_handle(), - X.extent(0), - centroids.extent(0), - X.extent(1), - metric); + const auto path = cuvs::distance::detail::resolve_fused_1nn_backend(handle, + X.data_handle(), + centroids.data_handle(), + X.extent(0), + centroids.extent(0), + X.extent(1), + metric); Fused1nnRequirements requirements{}; - requirements.path = path; - requirements.sample_tile = getDataBatchSize(batch_samples, X.extent(0)); - requirements.centroid_tile = getCentroidsBatchSize(batch_centroids, centroids.extent(0)); + requirements.path = path; + requirements.sample_tile = getDataBatchSize(batch_samples, X.extent(0)); + requirements.centroid_tile = getCentroidsBatchSize(batch_centroids, centroids.extent(0)); requirements.workspace_alignment = alignof(int); if (path == FusedDistancePath::FusedCutile) { - requirements.output_layout = Fused1nnOutputLayout::Soa; - requirements.norm_policy = std::is_same_v - ? Fused1nnNormPolicy::Tf32 - : Fused1nnNormPolicy::Default; + requirements.output_layout = Fused1nnOutputLayout::Soa; + requirements.norm_policy = + std::is_same_v ? Fused1nnNormPolicy::Tf32 : Fused1nnNormPolicy::Default; requirements.result_alignment = 16; requirements.distance_offset = raft::alignTo(sizeof(IndexT) * static_cast(X.extent(0)), size_t{16}); - requirements.result_bytes = requirements.distance_offset + - sizeof(DataT) * static_cast(X.extent(0)); + requirements.result_bytes = + requirements.distance_offset + sizeof(DataT) * static_cast(X.extent(0)); if constexpr (std::is_same_v) { requirements.workspace_bytes = - sizeof(int) * cuvs::distance::detail::fused_1nn_cutile_index_workspace_rows( - X.extent(0)); + sizeof(int) * + cuvs::distance::detail::fused_1nn_cutile_index_workspace_rows(X.extent(0)); } } else { requirements.output_layout = Fused1nnOutputLayout::Kvp; @@ -157,13 +155,15 @@ Fused1nnRequirements get_fused_1nn_requirements( (metric == cuvs::distance::DistanceType::L2Expanded || metric == cuvs::distance::DistanceType::L2SqrtExpanded || metric == cuvs::distance::DistanceType::CosineExpanded)) { - auto sample_tile = requirements.sample_tile; + auto sample_tile = requirements.sample_tile; const auto centroid_tile = requirements.centroid_tile; sample_tile = std::min(sample_tile, std::numeric_limits::max() / centroid_tile); - const size_t distance_bytes = sizeof(DataT) * static_cast(sample_tile) * - static_cast(centroid_tile); - requirements.workspace_alignment = std::max(alignof(DataT), alignof(raft::KeyValuePair)); - requirements.workspace_bytes = raft::alignTo(distance_bytes, alignof(raft::KeyValuePair)); + const size_t distance_bytes = + sizeof(DataT) * static_cast(sample_tile) * static_cast(centroid_tile); + requirements.workspace_alignment = + std::max(alignof(DataT), alignof(raft::KeyValuePair)); + requirements.workspace_bytes = + raft::alignTo(distance_bytes, alignof(raft::KeyValuePair)); if (centroid_tile < centroids.extent(0)) { requirements.workspace_bytes += sizeof(raft::KeyValuePair) * static_cast(sample_tile); @@ -196,7 +196,7 @@ void min_cluster_and_distance_compute_impl(raft::resources const& handle, const bool is_l2_cos = metric == cuvs::distance::DistanceType::L2Expanded || metric == cuvs::distance::DistanceType::L2SqrtExpanded || metric == cuvs::distance::DistanceType::CosineExpanded; - const auto fused_path = requirements.path; + const auto fused_path = requirements.path; const bool cutile_ready = fused_path == FusedDistancePath::FusedCutile; if (workspace.size() < requirements.workspace_bytes) { workspace.resize(requirements.workspace_bytes, stream); @@ -331,13 +331,13 @@ void min_cluster_and_distance_compute_impl(raft::resources const& handle, sizeof(DataT) * static_cast(dataBatchSize) * static_cast(centroidsBatchSize); const size_t batch_min_offset = raft::alignTo(distance_workspace_bytes, alignof(KeyValueT)); const size_t required_workspace_bytes = - batch_min_offset + (tileCentroids ? sizeof(KeyValueT) * static_cast(dataBatchSize) : 0); + batch_min_offset + + (tileCentroids ? sizeof(KeyValueT) * static_cast(dataBatchSize) : 0); if (workspace.size() < required_workspace_bytes) { workspace.resize(required_workspace_bytes, stream); } - auto* batch_min_storage = tileCentroids - ? reinterpret_cast(workspace.data() + batch_min_offset) - : nullptr; + auto* batch_min_storage = + tileCentroids ? reinterpret_cast(workspace.data() + batch_min_offset) : nullptr; for (IndexT dIdx = 0; dIdx < n_samples;) { auto ns = std::min(dataBatchSize, n_samples - dIdx); @@ -346,8 +346,7 @@ void min_cluster_and_distance_compute_impl(raft::resources const& handle, for (IndexT cIdx = 0; cIdx < n_clusters;) { auto nc = std::min(centroidsBatchSize, n_clusters - cIdx); - auto batchMin = tileCentroids ? batch_min_storage - : minClusterAndDistanceView.data_handle(); + auto batchMin = tileCentroids ? batch_min_storage : minClusterAndDistanceView.data_handle(); cuvs::distance::unfusedDistanceNNMinReduce( handle, @@ -469,13 +468,13 @@ void minClusterAndDistanceCompute(raft::resources const& handle, if constexpr (is_cutile_fused_data_type_v) { if (requirements.path == FusedDistancePath::FusedCutile) { RAFT_EXPECTS(cuvs::distance::detail::can_launch_fused_1nn_tile(nearest_idx.data_handle(), - nearest_dist.data_handle(), - X.data_handle(), - centroids.data_handle(), - X.extent(0), - centroids.extent(0), - X.extent(1), - metric), + nearest_dist.data_handle(), + X.data_handle(), + centroids.data_handle(), + X.extent(0), + centroids.extent(0), + X.extent(1), + metric), "resolved cuTile plan has incompatible output storage"); } } @@ -556,13 +555,13 @@ template void computeCutileRowNorms( template void computeCutileRowNorms( raft::resources const&, const float*, float*, int64_t, int64_t, bool); -#define INSTANTIATE_FUSED_1NN_REQUIREMENTS(DataT, IndexT) \ +#define INSTANTIATE_FUSED_1NN_REQUIREMENTS(DataT, IndexT) \ template Fused1nnRequirements get_fused_1nn_requirements( \ - raft::resources const&, \ - raft::device_matrix_view, \ - raft::device_matrix_view, \ - cuvs::distance::DistanceType, \ - int, \ + raft::resources const&, \ + raft::device_matrix_view, \ + raft::device_matrix_view, \ + cuvs::distance::DistanceType, \ + int, \ int); INSTANTIATE_FUSED_1NN_REQUIREMENTS(float, int64_t) diff --git a/cpp/src/distance/detail/fused_distance_nn.cuh b/cpp/src/distance/detail/fused_distance_nn.cuh index e1ec935a15..d25e14b59a 100644 --- a/cpp/src/distance/detail/fused_distance_nn.cuh +++ b/cpp/src/distance/detail/fused_distance_nn.cuh @@ -15,8 +15,8 @@ #include "pairwise_distance_base.cuh" // PairwiseDistances #include #include -#include // raft::KeyValuePair -#include // raft::identity_op +#include // raft::KeyValuePair +#include // raft::identity_op #include #include #include // Policy @@ -40,8 +40,10 @@ enum class Fused1nnBackend : std::uint8_t { }; template -constexpr Fused1nnBackend fused_1nn_legacy_backend( - int cc_major, IdxT m, IdxT n, cuvs::distance::DistanceType metric) +constexpr Fused1nnBackend fused_1nn_legacy_backend(int cc_major, + IdxT m, + IdxT n, + cuvs::distance::DistanceType metric) { if (metric == cuvs::distance::DistanceType::InnerProduct) { return Fused1nnBackend::Unfused; } if (cc_major <= 8 || (cc_major == 9 && (m >= 4096 || n >= 4096))) { @@ -61,9 +63,7 @@ Fused1nnBackend resolve_fused_1nn_backend(const raft::resources& handle, cuvs::distance::DistanceType metric) { if constexpr (is_fused_1nn_cutile_data_v) { - if (can_launch_fused_1nn_tile(x, y, m, n, k, metric)) { - return Fused1nnBackend::FusedCutile; - } + if (can_launch_fused_1nn_tile(x, y, m, n, k, metric)) { return Fused1nnBackend::FusedCutile; } } const auto prop = raft::resource::get_device_properties(handle); return fused_1nn_legacy_backend(prop.major, m, n, metric); diff --git a/cpp/src/distance/detail/fused_distance_nn/cutile/export_fused_1nn.py b/cpp/src/distance/detail/fused_distance_nn/cutile/export_fused_1nn.py index 1fdf313d18..6774f44743 100644 --- a/cpp/src/distance/detail/fused_distance_nn/cutile/export_fused_1nn.py +++ b/cpp/src/distance/detail/fused_distance_nn/cutile/export_fused_1nn.py @@ -281,18 +281,18 @@ def main() -> int: args = parser.parse_args() export_binary( - args.output_file, - output_format=args.format, - data_type=args.data_type, - metric=args.metric, - index_type=args.index_type, - tile_m=args.tile_m, - tile_n=args.tile_n, - tile_k=args.tile_k, - gpu_code=args.gpu_code, - matrix_layout=args.matrix_layout, - occupancy=args.occupancy, - bytecode_version=args.bytecode_version, + args.output_file, + output_format=args.format, + data_type=args.data_type, + metric=args.metric, + index_type=args.index_type, + tile_m=args.tile_m, + tile_n=args.tile_n, + tile_k=args.tile_k, + gpu_code=args.gpu_code, + matrix_layout=args.matrix_layout, + occupancy=args.occupancy, + bytecode_version=args.bytecode_version, ) return 0 diff --git a/cpp/tests/neighbors/distance_nn.cu b/cpp/tests/neighbors/distance_nn.cu index df13981279..0719879024 100644 --- a/cpp/tests/neighbors/distance_nn.cu +++ b/cpp/tests/neighbors/distance_nn.cu @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ @@ -42,9 +42,7 @@ __global__ void fill_int8(int8_t* buff, int len, int seed_offset) template class NNTest : public ::testing::TestWithParam> { public: - using OutT = std::conditional_t>; + using OutT = std::conditional_t>; NNTest() : params_{::testing::TestWithParam>::GetParam()}, m{params_.m}, @@ -203,8 +201,7 @@ template const std::vector> input_fp32_fused = [] { auto inputs = input_fp32; #if CUVS_CUTILE_ENABLED - inputs.push_back( - {512, 1024, 64, DistanceType::InnerProduct, false, uint64_t(31415926), 0.1}); + inputs.push_back({512, 1024, 64, DistanceType::InnerProduct, false, uint64_t(31415926), 0.1}); #endif return inputs; }(); From de3d89abd26b4ebc14ddde7f3cf696d5189fff9e Mon Sep 17 00:00:00 2001 From: divyegala Date: Thu, 3 Sep 2026 04:04:53 +0000 Subject: [PATCH 72/82] c build --- ci/build_standalone_c.sh | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/ci/build_standalone_c.sh b/ci/build_standalone_c.sh index 2b8e0863f9..e1267f71b7 100755 --- a/ci/build_standalone_c.sh +++ b/ci/build_standalone_c.sh @@ -41,6 +41,12 @@ source rapids-configure-sccache source rapids-datetime-string rapids-pip-retry install cmake + +RAPIDS_CUDA_MAJOR="${RAPIDS_CUDA_VERSION%%.*}" +if [[ "${RAPIDS_CUDA_MAJOR}" == "13" ]]; then + rapids-pip-retry install cuda-tile "cuda-toolkit[tileiras]==${RAPIDS_CUDA_VERSION%.*}.*" +fi + pyenv rehash rapids-print-env From e1cfd30d98cf3cb9046f723e2eb80fa0dc1b98bb Mon Sep 17 00:00:00 2001 From: divyegala Date: Thu, 3 Sep 2026 04:12:00 +0000 Subject: [PATCH 73/82] style --- cpp/tests/neighbors/distance_nn_helper.cuh | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/cpp/tests/neighbors/distance_nn_helper.cuh b/cpp/tests/neighbors/distance_nn_helper.cuh index 5162160363..6a41ade123 100644 --- a/cpp/tests/neighbors/distance_nn_helper.cuh +++ b/cpp/tests/neighbors/distance_nn_helper.cuh @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ From ec1c7d7ccfbda43b2cf97fbf74f5a8febb0912f5 Mon Sep 17 00:00:00 2001 From: divyegala Date: Thu, 3 Sep 2026 04:52:46 +0000 Subject: [PATCH 74/82] package --- cpp/src/detail/jit_lto/cutile_smoke/export_smoke.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/cpp/src/detail/jit_lto/cutile_smoke/export_smoke.py b/cpp/src/detail/jit_lto/cutile_smoke/export_smoke.py index 64723c0813..edce04dfde 100644 --- a/cpp/src/detail/jit_lto/cutile_smoke/export_smoke.py +++ b/cpp/src/detail/jit_lto/cutile_smoke/export_smoke.py @@ -8,6 +8,7 @@ import argparse from pathlib import Path +import sys import cuda.tile as ct from cuda.tile.compilation import ( @@ -17,6 +18,8 @@ export_kernel, ) +sys.path.insert(0, str(Path(__file__).resolve().parent)) + from smoke_kernel import TILE_SIZE, cutile_smoke_add From c0bcb1ddc3429aa0ce148313b87e879b1d9253f5 Mon Sep 17 00:00:00 2001 From: divyegala Date: Thu, 3 Sep 2026 04:54:59 +0000 Subject: [PATCH 75/82] style check --- cpp/CMakeLists.txt | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt index 440499a357..2f154cf2ce 100644 --- a/cpp/CMakeLists.txt +++ b/cpp/CMakeLists.txt @@ -1576,8 +1576,7 @@ if(NOT BUILD_CPU_ONLY) "$" INTERFACE "$" PRIVATE "${CMAKE_CURRENT_SOURCE_DIR}/src" "${CMAKE_CURRENT_BINARY_DIR}/src" - "${cutile_fused_1nn_generated_dir}" - "${cutile_smoke_generated_dir}" + "${cutile_fused_1nn_generated_dir}" "${cutile_smoke_generated_dir}" ) # Endian detection From e1358196176532c869cd83b28b0fec7f30dd9cf6 Mon Sep 17 00:00:00 2001 From: divyegala Date: Thu, 3 Sep 2026 18:11:36 +0000 Subject: [PATCH 76/82] linker --- cpp/cmake/config.json | 37 +++++++++++++ .../modules/compute_matrix_product.cmake | 3 ++ .../modules/generate_cutile_kernels.cmake | 41 ++++++-------- .../detail/jit_lto/TileAlgorithmPlanner.hpp | 10 ++-- .../cuvs/detail/jit_lto/cutile_module.hpp | 22 ++------ .../cuvs/detail/jit_lto/tileir_compat.hpp | 52 +++++++++--------- .../detail/jit_lto/TileAlgorithmPlanner.cpp | 53 ++++++++----------- cpp/tests/CMakeLists.txt | 10 ++++ cpp/tests/detail/jit_lto/cutile_smoke.cu | 8 +-- 9 files changed, 131 insertions(+), 105 deletions(-) diff --git a/cpp/cmake/config.json b/cpp/cmake/config.json index cc6e647bae..9dfcda3259 100644 --- a/cpp/cmake/config.json +++ b/cpp/cmake/config.json @@ -20,6 +20,43 @@ "MATRIX_JSON_STRING": "?" } }, + "cuvs_find_build_python": { + "pargs": { + "nargs": 1 + } + }, + "process_cutile_matrix_entry": { + "pargs": { + "nargs": 1 + }, + "kwargs": { + "KERNEL_DIR": 1, + "KERNEL_BASENAME": 1, + "KERNEL_PYTHON": 1, + "EXPORT_SCRIPT": 1, + "OUTPUT_DIRECTORY": 1, + "FRAGMENT_TAG_FORMAT_CUBIN": 1, + "FRAGMENT_TAG_FORMAT_TILEIR": "?", + "FRAGMENT_TAG_HEADER_FILES": "*", + "MATRIX_JSON_ENTRY": 1 + } + }, + "generate_cutile_kernels": { + "pargs": { + "nargs": 1 + }, + "kwargs": { + "KERNEL_DIR": 1, + "KERNEL_BASENAME": 1, + "KERNEL_PYTHON": 1, + "EXPORT_SCRIPT": 1, + "OUTPUT_DIRECTORY": 1, + "MATRIX_JSON_FILE": 1, + "FRAGMENT_TAG_FORMAT_CUBIN": 1, + "FRAGMENT_TAG_FORMAT_TILEIR": "?", + "FRAGMENT_TAG_HEADER_FILES": "*" + } + }, "add_jit_lto_kernel": { "pargs": { "nargs": 1 diff --git a/cpp/cmake/modules/compute_matrix_product.cmake b/cpp/cmake/modules/compute_matrix_product.cmake index 60b96113f0..b4b7afb9b7 100644 --- a/cpp/cmake/modules/compute_matrix_product.cmake +++ b/cpp/cmake/modules/compute_matrix_product.cmake @@ -8,6 +8,9 @@ include_guard(GLOBAL) function(cuvs_find_build_python output_var) + # cuTile is a build dependency. In conda builds, it is installed in BUILD_PREFIX while CMake's + # default search can resolve the host interpreter from PREFIX instead. Use the build prefix so + # configure-time matrix expansion and build-time kernel exports see the cuTile package. if(DEFINED ENV{BUILD_PREFIX}) set(Python_ROOT "$ENV{BUILD_PREFIX}") endif() diff --git a/cpp/cmake/modules/generate_cutile_kernels.cmake b/cpp/cmake/modules/generate_cutile_kernels.cmake index 35fb8c381b..51d58e3eca 100644 --- a/cpp/cmake/modules/generate_cutile_kernels.cmake +++ b/cpp/cmake/modules/generate_cutile_kernels.cmake @@ -9,13 +9,6 @@ include_guard(GLOBAL) include(${CMAKE_CURRENT_LIST_DIR}/compute_matrix_product.cmake) -function(generate_cutile_kernels_stub) - set(CUVS_CUTILE_ENABLED - 0 - PARENT_SCOPE - ) -endfunction() - function(_cutile_fragment_tag_header_files output_var) set(${output_var} "") foreach(_header IN LISTS ARGN) @@ -237,7 +230,12 @@ function(generate_cutile_kernels source_list_var) MATRIX_JSON_FILE "${_CUTILE_MATRIX_JSON_FILE}" OUTPUT_DIRECTORY "${_CUTILE_OUTPUT_DIRECTORY}" ) if(NOT _CUTILE_SETUP_OK) - generate_cutile_kernels_stub() + # This function's parent is cpp/CMakeLists.txt. Propagate the disabled feature state there so + # the compile definition cannot retain a stale value from a previous generator invocation. + set(CUVS_CUTILE_ENABLED + 0 + PARENT_SCOPE + ) set(${source_list_var} "" PARENT_SCOPE @@ -255,24 +253,15 @@ function(generate_cutile_kernels source_list_var) string(JSON matrix_json_entry GET "${matrix_product}" "${i}") process_cutile_matrix_entry( "${source_list_var}" - KERNEL_DIR - "${_CUTILE_KERNEL_DIR}" - KERNEL_BASENAME - "${_CUTILE_KERNEL_BASENAME}" - KERNEL_PYTHON - "${_CUTILE_KERNEL_PYTHON}" - EXPORT_SCRIPT - "${_CUTILE_EXPORT_SCRIPT}" - OUTPUT_DIRECTORY - "${_CUTILE_OUTPUT_DIRECTORY}" - FRAGMENT_TAG_FORMAT_CUBIN - "${_CUTILE_FRAGMENT_TAG_FORMAT_CUBIN}" - FRAGMENT_TAG_FORMAT_TILEIR - "${_CUTILE_FRAGMENT_TAG_FORMAT_TILEIR}" - FRAGMENT_TAG_HEADER_FILES - ${_CUTILE_FRAGMENT_TAG_HEADER_FILES} - MATRIX_JSON_ENTRY - "${matrix_json_entry}" + KERNEL_DIR "${_CUTILE_KERNEL_DIR}" + KERNEL_BASENAME "${_CUTILE_KERNEL_BASENAME}" + KERNEL_PYTHON "${_CUTILE_KERNEL_PYTHON}" + EXPORT_SCRIPT "${_CUTILE_EXPORT_SCRIPT}" + OUTPUT_DIRECTORY "${_CUTILE_OUTPUT_DIRECTORY}" + FRAGMENT_TAG_FORMAT_CUBIN "${_CUTILE_FRAGMENT_TAG_FORMAT_CUBIN}" + FRAGMENT_TAG_FORMAT_TILEIR "${_CUTILE_FRAGMENT_TAG_FORMAT_TILEIR}" + FRAGMENT_TAG_HEADER_FILES ${_CUTILE_FRAGMENT_TAG_HEADER_FILES} + MATRIX_JSON_ENTRY "${matrix_json_entry}" ) endforeach() diff --git a/cpp/include/cuvs/detail/jit_lto/TileAlgorithmPlanner.hpp b/cpp/include/cuvs/detail/jit_lto/TileAlgorithmPlanner.hpp index c552966e2d..fb6025fd64 100644 --- a/cpp/include/cuvs/detail/jit_lto/TileAlgorithmPlanner.hpp +++ b/cpp/include/cuvs/detail/jit_lto/TileAlgorithmPlanner.hpp @@ -19,10 +19,14 @@ namespace cuvs::detail::jit_lto { +struct CutileRuntimeCapabilities; + struct TileLauncherCache { std::shared_mutex mutex; std::unordered_map> launchers; - std::unordered_set build_failed; + // Cache expected compatibility misses so an unsupported module is not loaded on every call. + // Unexpected CUDA errors are raised rather than inserted here. + std::unordered_set unavailable_launchers; }; /** Loads prebuilt cubins or TileIR bytecode directly through the CUDA library API. */ @@ -59,9 +63,9 @@ struct TileAlgorithmPlanner { std::unique_ptr tileir_fragment_; private: - std::string get_planner_key() const; + std::string get_planner_key(const CutileRuntimeCapabilities* capabilities) const; - std::shared_ptr build(); + std::shared_ptr build(const CutileRuntimeCapabilities* capabilities); std::string entrypoint_; TileLauncherCache& launcher_cache_; diff --git a/cpp/include/cuvs/detail/jit_lto/cutile_module.hpp b/cpp/include/cuvs/detail/jit_lto/cutile_module.hpp index 4e39450ba1..bf26e1c9c5 100644 --- a/cpp/include/cuvs/detail/jit_lto/cutile_module.hpp +++ b/cpp/include/cuvs/detail/jit_lto/cutile_module.hpp @@ -27,19 +27,6 @@ struct CutileModuleImage { size_t size; }; -inline bool get_device_compute_capability(int& cc_major, int& cc_minor) -{ - int device = 0; - if (cudaGetDevice(&device) != cudaSuccess) { return false; } - if (cudaDeviceGetAttribute(&cc_major, cudaDevAttrComputeCapabilityMajor, device) != cudaSuccess) { - return false; - } - if (cudaDeviceGetAttribute(&cc_minor, cudaDevAttrComputeCapabilityMinor, device) != cudaSuccess) { - return false; - } - return true; -} - /** * Selects the newest compatible cubin in the device's compute-capability major family. * @@ -63,16 +50,15 @@ inline const CubinFragmentEntry* find_compatible_cubin_fragment( /** Selects compatible prebuilt SASS for the device, or TileIR when the driver can JIT it. */ inline std::optional resolve_cutile_module_image( - int cc_major, - int cc_minor, - int driver_version, + const CutileRuntimeCapabilities& capabilities, const std::vector>& cubin_fragments, const TileIrBytecodeFragmentEntry* tileir_fragment) { - if (const auto* fragment = find_compatible_cubin_fragment(cc_major, cc_minor, cubin_fragments)) { + if (const auto* fragment = find_compatible_cubin_fragment( + capabilities.cc_major, capabilities.cc_minor, cubin_fragments)) { return CutileModuleImage{fragment->get_data(), fragment->get_length()}; } - if (tileir_fragment != nullptr && tileir_fallback_available(driver_version)) { + if (tileir_fragment != nullptr && tileir_fallback_available(capabilities.driver_version)) { return CutileModuleImage{tileir_fragment->get_data(), tileir_fragment->get_length()}; } return std::nullopt; diff --git a/cpp/include/cuvs/detail/jit_lto/tileir_compat.hpp b/cpp/include/cuvs/detail/jit_lto/tileir_compat.hpp index a03a7a78fc..8e5e599069 100644 --- a/cpp/include/cuvs/detail/jit_lto/tileir_compat.hpp +++ b/cpp/include/cuvs/detail/jit_lto/tileir_compat.hpp @@ -16,6 +16,30 @@ namespace cuvs::detail::jit_lto { +/** Runtime/device properties that determine cuTile image selection and launch eligibility. */ +struct CutileRuntimeCapabilities { + int device; + int cc_major; + int cc_minor; + int driver_version; +}; + +inline bool query_current_cutile_runtime_capabilities(CutileRuntimeCapabilities& capabilities) +{ + if (cudaGetDevice(&capabilities.device) != cudaSuccess) { return false; } + if (cudaDeviceGetAttribute(&capabilities.cc_major, + cudaDevAttrComputeCapabilityMajor, + capabilities.device) != cudaSuccess) { + return false; + } + if (cudaDeviceGetAttribute(&capabilities.cc_minor, + cudaDevAttrComputeCapabilityMinor, + capabilities.device) != cudaSuccess) { + return false; + } + return cudaDriverGetVersion(&capabilities.driver_version) == cudaSuccess; +} + /** Minimum CUDA driver version (from cudaDriverGetVersion) for TileIR JIT of embedded bytecode. */ inline constexpr int kMinTileIrJitDriverVersion = 13010; // CUDA 13.1 / driver >= 590.44 @@ -75,33 +99,13 @@ inline bool cutile_launch_available_for_arch(int cc_major, int cc_minor, int dri inline constexpr bool cutile_launch_available_for_arch(int, int, int) { return false; } #endif -inline bool query_driver_version(int& driver_version) -{ - return cudaDriverGetVersion(&driver_version) == cudaSuccess; -} - -inline bool query_current_device_arch(int& cc_major, int& cc_minor) -{ - int device = 0; - if (cudaGetDevice(&device) != cudaSuccess) { return false; } - if (cudaDeviceGetAttribute(&cc_major, cudaDevAttrComputeCapabilityMajor, device) != cudaSuccess) { - return false; - } - if (cudaDeviceGetAttribute(&cc_minor, cudaDevAttrComputeCapabilityMinor, device) != cudaSuccess) { - return false; - } - return true; -} - #if CUVS_CUTILE_ENABLED inline bool cutile_launch_available_on_current_device() { - int cc_major = 0; - int cc_minor = 0; - int driver_version = 0; - if (!query_current_device_arch(cc_major, cc_minor)) { return false; } - if (!query_driver_version(driver_version)) { return false; } - return cutile_launch_available_for_arch(cc_major, cc_minor, driver_version); + CutileRuntimeCapabilities capabilities{}; + if (!query_current_cutile_runtime_capabilities(capabilities)) { return false; } + return cutile_launch_available_for_arch( + capabilities.cc_major, capabilities.cc_minor, capabilities.driver_version); } #else /** Compile-time false when cuTile is not built; use in if constexpr to skip cuTile-only paths. */ diff --git a/cpp/src/detail/jit_lto/TileAlgorithmPlanner.cpp b/cpp/src/detail/jit_lto/TileAlgorithmPlanner.cpp index d23d337817..65363d6fe3 100644 --- a/cpp/src/detail/jit_lto/TileAlgorithmPlanner.cpp +++ b/cpp/src/detail/jit_lto/TileAlgorithmPlanner.cpp @@ -40,11 +40,14 @@ CutileTileConfig tile_config_from_fragment(const FragmentT* fragment, const std: std::shared_ptr TileAlgorithmPlanner::try_get_launcher() { - auto launch_key = this->get_planner_key(); + CutileRuntimeCapabilities capabilities{}; + const auto* current_capabilities = + query_current_cutile_runtime_capabilities(capabilities) ? &capabilities : nullptr; + auto launch_key = this->get_planner_key(current_capabilities); { std::shared_lock read_lock(launcher_cache_.mutex); - if (launcher_cache_.build_failed.count(launch_key)) { return nullptr; } + if (launcher_cache_.unavailable_launchers.count(launch_key)) { return nullptr; } if (auto it = launcher_cache_.launchers.find(launch_key); it != launcher_cache_.launchers.end()) { return it->second; @@ -52,14 +55,14 @@ std::shared_ptr TileAlgorithmPlanner::try_get_launcher } std::unique_lock write_lock(launcher_cache_.mutex); - if (launcher_cache_.build_failed.count(launch_key)) { return nullptr; } + if (launcher_cache_.unavailable_launchers.count(launch_key)) { return nullptr; } if (auto it = launcher_cache_.launchers.find(launch_key); it != launcher_cache_.launchers.end()) { return it->second; } - auto launcher = this->build(); + auto launcher = this->build(current_capabilities); if (!launcher) { - launcher_cache_.build_failed.insert(launch_key); + launcher_cache_.unavailable_launchers.insert(launch_key); return nullptr; } launcher_cache_.launchers[launch_key] = launcher; @@ -75,7 +78,8 @@ std::shared_ptr TileAlgorithmPlanner::get_launcher() return launcher; } -std::string TileAlgorithmPlanner::get_planner_key() const +std::string TileAlgorithmPlanner::get_planner_key( + const CutileRuntimeCapabilities* capabilities) const { std::string key = entrypoint_; for (const auto& fragment : cubin_fragments_) { @@ -83,35 +87,28 @@ std::string TileAlgorithmPlanner::get_planner_key() const } if (tileir_fragment_) { key += tileir_fragment_->get_key(); } - int device = -1; - int cc_major = -1; - int cc_minor = -1; - int driver_version = -1; - if (cudaGetDevice(&device) == cudaSuccess && - cuvs::detail::jit_lto::get_device_compute_capability(cc_major, cc_minor)) { - key += ":device=" + std::to_string(device); - key += ":cc=" + std::to_string(cc_major) + "." + std::to_string(cc_minor); + if (capabilities != nullptr) { + key += ":device=" + std::to_string(capabilities->device); + key += ":cc=" + std::to_string(capabilities->cc_major) + "." + + std::to_string(capabilities->cc_minor); if (const auto* fragment = cuvs::detail::jit_lto::find_compatible_cubin_fragment( - cc_major, cc_minor, cubin_fragments_)) { + capabilities->cc_major, capabilities->cc_minor, cubin_fragments_)) { key += ":cubin=" + std::to_string(fragment->get_cc_major()) + "." + std::to_string(fragment->get_cc_minor()); } else { key += ":tileir"; } - if (cudaDriverGetVersion(&driver_version) == cudaSuccess) { - key += ":driver=" + std::to_string(driver_version); - } + key += ":driver=" + std::to_string(capabilities->driver_version); } return key; } CutileTileConfig TileAlgorithmPlanner::tile_config() const { - int cc_major = 0; - int cc_minor = 0; - if (cuvs::detail::jit_lto::get_device_compute_capability(cc_major, cc_minor)) { + CutileRuntimeCapabilities capabilities{}; + if (query_current_cutile_runtime_capabilities(capabilities)) { if (const auto* fragment = cuvs::detail::jit_lto::find_compatible_cubin_fragment( - cc_major, cc_minor, cubin_fragments_)) { + capabilities.cc_major, capabilities.cc_minor, cubin_fragments_)) { return tile_config_from_fragment(fragment, entrypoint_); } } @@ -125,17 +122,13 @@ CutileTileConfig TileAlgorithmPlanner::tile_config() const RAFT_FAIL("cuTile planner '%s' has no registered fragments", entrypoint_.c_str()); } -std::shared_ptr TileAlgorithmPlanner::build() +std::shared_ptr TileAlgorithmPlanner::build( + const CutileRuntimeCapabilities* capabilities) { - int cc_major = 0; - int cc_minor = 0; - if (!cuvs::detail::jit_lto::get_device_compute_capability(cc_major, cc_minor)) { return nullptr; } - - int driver_version = 0; - if (cudaDriverGetVersion(&driver_version) != cudaSuccess) { return nullptr; } + if (capabilities == nullptr) { return nullptr; } auto image = cuvs::detail::jit_lto::resolve_cutile_module_image( - cc_major, cc_minor, driver_version, cubin_fragments_, tileir_fragment_.get()); + *capabilities, cubin_fragments_, tileir_fragment_.get()); if (!image) { return nullptr; } return cuvs::detail::jit_lto::try_load_cutile_launcher(*image, entrypoint_); diff --git a/cpp/tests/CMakeLists.txt b/cpp/tests/CMakeLists.txt index 54f1d9c965..a56385bfda 100644 --- a/cpp/tests/CMakeLists.txt +++ b/cpp/tests/CMakeLists.txt @@ -142,6 +142,16 @@ ConfigureTest( GPUS 1 PERCENT 100 ) +if(CUVS_CUTILE_ENABLED) + # These are intentionally library-private implementation symbols. Build the smoke executable with + # the generated fragment registrations and planner implementation so it can exercise them without + # exporting the cuTile internals from libcuvs. + target_sources( + CUTILE_SMOKE_TEST PRIVATE ${cutile_smoke_files} + "${CUVS_SOURCE_DIR}/src/detail/jit_lto/TileAlgorithmPlanner.cpp" + ) + target_include_directories(CUTILE_SMOKE_TEST PRIVATE "${cutile_smoke_generated_dir}") +endif() ConfigureTest( NAME NEIGHBORS_ANN_IVF_FLAT_UDF_TEST diff --git a/cpp/tests/detail/jit_lto/cutile_smoke.cu b/cpp/tests/detail/jit_lto/cutile_smoke.cu index eb9d4371b5..843835e22e 100644 --- a/cpp/tests/detail/jit_lto/cutile_smoke.cu +++ b/cpp/tests/detail/jit_lto/cutile_smoke.cu @@ -68,14 +68,14 @@ TEST(CutileSmoke, ResolvesEveryEmbeddedArchitecture) TEST(CutileSmoke, LaunchesCompatibleCubin) { - int cc_major = 0; - int cc_minor = 0; - if (!get_device_compute_capability(cc_major, cc_minor)) { + CutileRuntimeCapabilities capabilities{}; + if (!query_current_cutile_runtime_capabilities(capabilities)) { GTEST_SKIP() << "No CUDA device is available"; } auto fragments = make_smoke_fragments(); - if (find_compatible_cubin_fragment(cc_major, cc_minor, fragments) == nullptr) { + if (find_compatible_cubin_fragment(capabilities.cc_major, capabilities.cc_minor, fragments) == + nullptr) { GTEST_SKIP() << "No embedded smoke cubin is compatible with this device"; } From 63f76fd5e8db66c1c78d3e49a9301a83754ff441 Mon Sep 17 00:00:00 2001 From: divyegala Date: Thu, 3 Sep 2026 18:17:18 +0000 Subject: [PATCH 77/82] style --- cpp/CMakeLists.txt | 25 +++++++++---------------- 1 file changed, 9 insertions(+), 16 deletions(-) diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt index caf6b4143a..ff8ef7f560 100644 --- a/cpp/CMakeLists.txt +++ b/cpp/CMakeLists.txt @@ -1160,23 +1160,16 @@ if(NOT BUILD_CPU_ONLY) ) generate_cutile_kernels( cutile_smoke_files - KERNEL_DIR - "${cutile_smoke_dir}" - KERNEL_BASENAME - "cutile_smoke" - KERNEL_PYTHON - "smoke_kernel.py" - EXPORT_SCRIPT - "export_smoke.py" - OUTPUT_DIRECTORY - "${cutile_smoke_generated_dir}" - MATRIX_JSON_FILE - "${cutile_smoke_dir}/cutile_smoke_matrix.json" + KERNEL_DIR "${cutile_smoke_dir}" + KERNEL_BASENAME "cutile_smoke" + KERNEL_PYTHON "smoke_kernel.py" + EXPORT_SCRIPT "export_smoke.py" + OUTPUT_DIRECTORY "${cutile_smoke_generated_dir}" + MATRIX_JSON_FILE "${cutile_smoke_dir}/cutile_smoke_matrix.json" FRAGMENT_TAG_FORMAT_CUBIN - "cuvs::detail::jit_lto::fragment_tag_cutile_smoke_add_cubin" - FRAGMENT_TAG_HEADER_FILES - "" - "" + "cuvs::detail::jit_lto::fragment_tag_cutile_smoke_add_cubin" + FRAGMENT_TAG_HEADER_FILES "" + "" ) if(NOT DEFINED CUVS_CUTILE_ENABLED) set(CUVS_CUTILE_ENABLED 0) From 6fa9fb101975f427ecf3590c14d5484756b61fd1 Mon Sep 17 00:00:00 2001 From: divyegala Date: Thu, 3 Sep 2026 21:22:48 +0000 Subject: [PATCH 78/82] 1-nn primitive with --- cpp/CMakeLists.txt | 45 +- .../modules/generate_cutile_tile_metadata.py | 54 ++ .../fused_distance_nn/fused_1nn_fragments.hpp | 66 +++ cpp/src/distance/detail/fused_distance_nn.cuh | 12 +- .../cutile/export_fused_1nn.py | 301 ++++++++++ .../cutile/fused_1nn_cutile_matrix.json | 515 ++++++++++++++++++ .../cutile/fused_1nn_kernel.py | 191 +++++++ .../cutile/fused_1nn_planner.hpp | 130 +++++ .../cutile/fused_1nn_tile.cu | 471 ++++++++++++++++ .../cutile/fused_1nn_tile.hpp | 164 ++++++ cpp/src/distance/fused_distance_nn-inl.cuh | 77 ++- 11 files changed, 2022 insertions(+), 4 deletions(-) create mode 100644 cpp/cmake/modules/generate_cutile_tile_metadata.py create mode 100644 cpp/include/cuvs/detail/jit_lto/fused_distance_nn/fused_1nn_fragments.hpp create mode 100644 cpp/src/distance/detail/fused_distance_nn/cutile/export_fused_1nn.py create mode 100644 cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_cutile_matrix.json create mode 100644 cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_kernel.py create mode 100644 cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_planner.hpp create mode 100644 cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_tile.cu create mode 100644 cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_tile.hpp diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt index a633094a88..7b642fa7d4 100644 --- a/cpp/CMakeLists.txt +++ b/cpp/CMakeLists.txt @@ -1176,6 +1176,47 @@ if(NOT BUILD_CPU_ONLY) endif() target_compile_definitions(cuvs_cpp_headers INTERFACE CUVS_CUTILE_ENABLED=${CUVS_CUTILE_ENABLED}) + set(fused_1nn_cutile_dir + "${CMAKE_CURRENT_SOURCE_DIR}/src/distance/detail/fused_distance_nn/cutile" + ) + set(cutile_fused_1nn_generated_dir + "${CMAKE_CURRENT_BINARY_DIR}/generated_kernels/distance/fused_1nn/cutile" + ) + set(cutile_fused_1nn_tiles "${cutile_fused_1nn_generated_dir}/fused_1nn_cutile_tiles.hpp") + generate_cutile_kernels( + cutile_fused_1nn_files + KERNEL_DIR "${fused_1nn_cutile_dir}" + KERNEL_BASENAME "fused_1nn" + KERNEL_PYTHON "fused_1nn_kernel.py" + EXPORT_SCRIPT "export_fused_1nn.py" + OUTPUT_DIRECTORY "${cutile_fused_1nn_generated_dir}" + MATRIX_JSON_FILE "${fused_1nn_cutile_dir}/fused_1nn_cutile_matrix.json" + FRAGMENT_TAG_FORMAT_CUBIN + "cuvs::distance::detail::fragment_tag_fused_1nn_cubin, cuvs::distance::detail::@abi_tag@, cuvs::detail::jit_lto::@arch_tag@>" + FRAGMENT_TAG_FORMAT_TILEIR + "cuvs::distance::detail::fragment_tag_fused_1nn_tileir, cuvs::distance::detail::@abi_tag@>" + FRAGMENT_TAG_HEADER_FILES + "" + "" "" + ) + if(CUVS_CUTILE_ENABLED) + cuvs_find_build_python(cutile_tile_metadata_python) + add_custom_command( + OUTPUT "${cutile_fused_1nn_tiles}" + COMMAND + "${cutile_tile_metadata_python}" + "${CMAKE_CURRENT_SOURCE_DIR}/cmake/modules/generate_cutile_tile_metadata.py" --matrix + "${fused_1nn_cutile_dir}/fused_1nn_cutile_matrix.json" --output "${cutile_fused_1nn_tiles}" + --namespace "cuvs::distance::detail" --include + "" --alias-prefix + fused_1nn_matrix_tile + DEPENDS "${fused_1nn_cutile_dir}/fused_1nn_cutile_matrix.json" + "${CMAKE_CURRENT_SOURCE_DIR}/cmake/modules/generate_cutile_tile_metadata.py" + VERBATIM + ) + list(APPEND cutile_fused_1nn_files "${cutile_fused_1nn_tiles}") + endif() + # Note that this matrix contains an `arch_includes` placeholder, since we don't currently have a # way to do an item-wise transform on a list after computing the matrix product and before # configuring the file @@ -1486,6 +1527,8 @@ if(NOT BUILD_CPU_ONLY) src/stats/trustworthiness_score.cu ${CUVS_MG_ALGOS} ${jit_lto_files} + ${cutile_fused_1nn_files} + $<$:src/distance/detail/fused_distance_nn/cutile/fused_1nn_tile.cu> ${cutile_smoke_files} ) @@ -1531,7 +1574,7 @@ if(NOT BUILD_CPU_ONLY) "$" INTERFACE "$" PRIVATE "${CMAKE_CURRENT_SOURCE_DIR}/src" "${CMAKE_CURRENT_BINARY_DIR}/src" - "${cutile_smoke_generated_dir}" + "${cutile_fused_1nn_generated_dir}" "${cutile_smoke_generated_dir}" ) # Endian detection diff --git a/cpp/cmake/modules/generate_cutile_tile_metadata.py b/cpp/cmake/modules/generate_cutile_tile_metadata.py new file mode 100644 index 0000000000..4415c6a7a8 --- /dev/null +++ b/cpp/cmake/modules/generate_cutile_tile_metadata.py @@ -0,0 +1,54 @@ +#!/usr/bin/env python3 +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import argparse +import json +from pathlib import Path + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--matrix", type=Path, required=True) + parser.add_argument("--output", type=Path, required=True) + parser.add_argument("--namespace", required=True) + parser.add_argument("--include", required=True) + parser.add_argument("--alias-prefix", required=True) + args = parser.parse_args() + aliases = {} + for entry in json.loads(args.matrix.read_text()): + default_tile = entry.get("_tile", [{}])[0] + for data in entry["_data"]: + for abi in entry["_abi"]: + tile = tuple( + abi.get(k, default_tile.get(k)) + for k in ("tile_m", "tile_n", "tile_k") + ) + if any(value is None for value in tile): + raise ValueError("missing cuTile tile geometry") + for exported in entry["_export"]: + suffix = f"{data['data_abbrev']}_{exported.get('arch_tag', 'tileir')}_{abi['abi_abbrev']}" + if suffix in aliases and aliases[suffix] != tile: + raise ValueError( + f"conflicting tile geometry for {suffix}" + ) + aliases[suffix] = tile + lines = [ + "#pragma once", + "", + f"#include {args.include}", + "", + f"namespace {args.namespace} {{", + "", + ] + for suffix, (m, n, k) in sorted(aliases.items()): + lines.append( + f"using {args.alias_prefix}_{suffix} = cutile_tile_config<{m}, {n}, {k}>;" + ) + lines.extend(["", f"}} // namespace {args.namespace}", ""]) + args.output.parent.mkdir(parents=True, exist_ok=True) + args.output.write_text("\n".join(lines)) + + +if __name__ == "__main__": + main() diff --git a/cpp/include/cuvs/detail/jit_lto/fused_distance_nn/fused_1nn_fragments.hpp b/cpp/include/cuvs/detail/jit_lto/fused_distance_nn/fused_1nn_fragments.hpp new file mode 100644 index 0000000000..807dc50e24 --- /dev/null +++ b/cpp/include/cuvs/detail/jit_lto/fused_distance_nn/fused_1nn_fragments.hpp @@ -0,0 +1,66 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +#include + +#include + +#include +namespace cuvs::distance::detail { + +struct cutile_abi_strict {}; +struct cutile_abi_relaxed {}; + +template +struct cutile_tile_config { + static constexpr int tile_m = TileM; + static constexpr int tile_n = TileN; + static constexpr int tile_k = TileK; +}; + +template +struct fused_1nn_data_tag; + +template <> +struct fused_1nn_data_tag { + using type = cuvs::neighbors::detail::tag_f; +}; + +template <> +struct fused_1nn_data_tag { + using type = cuvs::neighbors::detail::tag_h; +}; + +template +using fused_1nn_data_tag_t = typename fused_1nn_data_tag::type; + +template +struct fused_1nn_index_tag; + +template <> +struct fused_1nn_index_tag { + using type = cuvs::neighbors::detail::tag_index_i32; +}; + +template <> +struct fused_1nn_index_tag { + using type = cuvs::neighbors::detail::tag_index_i64; +}; + +template +using fused_1nn_index_tag_t = typename fused_1nn_index_tag::type; + +template +struct fragment_tag_fused_1nn_cubin { + static constexpr int cc_major = ArchTag::cc_major; + static constexpr int cc_minor = ArchTag::cc_minor; +}; + +template +struct fragment_tag_fused_1nn_tileir {}; + +} // namespace cuvs::distance::detail diff --git a/cpp/src/distance/detail/fused_distance_nn.cuh b/cpp/src/distance/detail/fused_distance_nn.cuh index f9dbd968ec..5a65482f13 100644 --- a/cpp/src/distance/detail/fused_distance_nn.cuh +++ b/cpp/src/distance/detail/fused_distance_nn.cuh @@ -1,11 +1,12 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2024, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ #pragma once #include "distance_ops/l2_exp.cuh" // ops::l2_exp_distance_op +#include "fused_distance_nn/cutile/fused_1nn_tile.hpp" #include "fused_distance_nn/cutlass_base.cuh" #include "fused_distance_nn/fused_cosine_nn.cuh" #include "fused_distance_nn/fused_l2_nn.cuh" @@ -20,13 +21,20 @@ #include // raft::ceildiv, raft::shfl #include // size_t -#include // std::numeric_limits +#include +#include // std::numeric_limits namespace cuvs { namespace distance { namespace detail { +/** Explicit implementation selected for the fused 1-NN primitive. */ +enum class Fused1nnBackend : std::uint8_t { + Cutile, + Cutlass, +}; + template str: + return {"half": "h", "float": "f"}[data_type] + + +def _elem_stride_divisible_for_tma(elem_dtype) -> tuple[int, int]: + """Row stride (dim 0) divisible enough for 16-byte TMA access; last dim stride 1.""" + bytes_per_elem = 2 if elem_dtype == ct.float16 else 4 + return (16 // bytes_per_elem, 1) + + +def _elem_shape_divisible_for_ldgsts(elem_dtype) -> tuple[int, int]: + """Matrix extent aligned to the same 16-byte row pitch enforced on strides.""" + bytes_per_elem = 2 if elem_dtype == ct.float16 else 4 + return (1, 16 // bytes_per_elem) + + +def _cuvs_matrix_constraint( + elem_dtype, + *, + index_dtype=ct.int32, + require_tma_friendly_pitch: bool = True, + require_ldgsts_friendly_shape: bool = False, +): + """Row-major device matrices for cuVS KMeans benchmarks. + + Assumes raft/cupy-style contiguous layout: stride[-1]==1, stride[0]==D, + 16-byte base alignment, and row pitch 16-byte aligned (float32 D%4==0, + float16 D%8==0). Applies to both points and centroids matrices. + + SM80/SM86 strict exports also express the row-pitch guarantee as + shape_divisible_by=(1, 4) for float32 or (1, 8) for float16. This + duplicates the stride constraint intentionally so the compiler selects + LDGSTS instead of LDG. Tail tiles remain masked in the kernel. + + Odd D or general layouts need a separate relaxed export profile. + """ + return ArrayConstraint( + elem_dtype, + ndim=2, + index_dtype=index_dtype, + stride_lower_bound_incl=(0, None), + # Dataset and centroid views are read-only and may legally share storage. + alias_groups=("read_only_inputs",), + may_alias_internally=False, + stride_constant=(None, 1), + stride_divisible_by=( + _elem_stride_divisible_for_tma(elem_dtype) + if require_tma_friendly_pitch + else (1, 1) + ), + shape_divisible_by=( + _elem_shape_divisible_for_ldgsts(elem_dtype) + if require_ldgsts_friendly_shape + else (1, 1) + ), + base_addr_divisible_by=16, + ) + + +def _cuvs_vector_constraint( + elem_dtype, *, index_dtype=ct.int32, alias_groups=() +): + """1-D device vectors: contiguous, 16-byte base. Length need not be divisible by 16.""" + return ArrayConstraint( + elem_dtype, + ndim=1, + index_dtype=index_dtype, + stride_lower_bound_incl=(None,), + alias_groups=alias_groups, + may_alias_internally=False, + stride_constant=(1,), + stride_divisible_by=(1,), + shape_divisible_by=(1,), + base_addr_divisible_by=16, + ) + + +def _relaxed_matrix_constraint(elem_dtype): + """Deprecated alias for the arbitrary-row-pitch matrix constraint.""" + return _cuvs_matrix_constraint( + elem_dtype, require_tma_friendly_pitch=False + ) + + +def _relaxed_vector_constraint(elem_dtype, *, tma_friendly: bool = False): + """Deprecated alias; use _cuvs_vector_constraint.""" + del tma_friendly + return _cuvs_vector_constraint(elem_dtype) + + +def _kernel_signature( + data_type: str, + metric: str, + index_type: str, + tile_m: int, + tile_n: int, + tile_k: int, + gpu_code: str, + matrix_layout: str, +) -> KernelSignature: + elem = _dtype_for(data_type) + idx_dtype = _idx_dtype(index_type) + matrix = _cuvs_matrix_constraint( + elem, + index_dtype=idx_dtype, + require_tma_friendly_pitch=matrix_layout == "strict", + require_ldgsts_friendly_shape=( + matrix_layout == "strict" and gpu_code in ("sm_80", "sm_86") + ), + ) + norm_elem = ct.float32 if data_type == "half" else elem + norm_array = _cuvs_vector_constraint( + norm_elem, + index_dtype=idx_dtype, + alias_groups=("read_only_inputs",), + ) + idx_array = _cuvs_vector_constraint(idx_dtype, index_dtype=idx_dtype) + dist_array = _cuvs_vector_constraint(elem, index_dtype=idx_dtype) + + abbrev = _data_abbrev(data_type) + symbol = kernel_symbol( + abbrev, + index_abbrev(index_type), + matrix_layout, + ) + + return KernelSignature( + parameters=[ + matrix, + matrix, + norm_array, + norm_array, + idx_array, + dist_array, + ScalarConstraint(idx_dtype), + ScalarConstraint(idx_dtype), + ScalarConstraint(idx_dtype), + ScalarConstraint(idx_dtype), + ScalarConstraint(idx_dtype), + ScalarConstraint(ct.int32), + ConstantConstraint(tile_m), + ConstantConstraint(tile_n), + ConstantConstraint(tile_k), + ], + calling_convention=CallingConvention.cutile_python_v1(), + ).with_symbol(symbol) + + +def export_binary( + output_file: Path, + *, + output_format: Literal["cubin", "tileir_bytecode"], + data_type: str, + metric: str, + index_type: str, + tile_m: int, + tile_n: int, + tile_k: int, + gpu_code: str, + matrix_layout: str = "strict", + occupancy: int | None = None, + bytecode_version: str | None = None, +) -> str: + kernel = make_kernel( + data_type, + metric, + tile_m, + tile_n, + tile_k, + index_type=index_type, + gpu_code=gpu_code, + matrix_layout=matrix_layout, + occupancy=occupancy, + ) + signature = _kernel_signature( + data_type, + metric, + index_type, + tile_m, + tile_n, + tile_k, + gpu_code, + matrix_layout, + ) + + export_kwargs = { + "kernel": kernel, + "signatures": [signature], + "output_file": str(output_file), + "gpu_code": gpu_code, + "output_format": output_format, + } + if output_format == "tileir_bytecode": + export_kwargs["bytecode_version"] = ( + bytecode_version or DEFAULT_TILEIR_BYTECODE_VERSION + ) + + export_kernel(**export_kwargs) + + return signature.symbol + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("output_file", type=Path) + parser.add_argument( + "--format", choices=("cubin", "tileir_bytecode"), default="cubin" + ) + parser.add_argument( + "--data-type", choices=("half", "float"), required=True + ) + parser.add_argument("--metric", choices=METRICS, required=True) + parser.add_argument("--index-type", choices=INDEX_TYPES, required=True) + parser.add_argument("--tile-m", type=int, required=True) + parser.add_argument("--tile-n", type=int, required=True) + parser.add_argument("--tile-k", type=int, required=True) + parser.add_argument( + "--gpu-code", + default=DEFAULT_TILEIR_EXPORT_GPU_CODE, + help="Target SM for cubin export, or compile hint for TileIR bytecode export", + ) + parser.add_argument( + "--matrix-layout", + choices=("strict", "relaxed"), + default="strict", + ) + parser.add_argument("--occupancy", type=int) + parser.add_argument( + "--bytecode-version", default=DEFAULT_TILEIR_BYTECODE_VERSION + ) + args = parser.parse_args() + + export_binary( + args.output_file, + output_format=args.format, + data_type=args.data_type, + metric=args.metric, + index_type=args.index_type, + tile_m=args.tile_m, + tile_n=args.tile_n, + tile_k=args.tile_k, + gpu_code=args.gpu_code, + matrix_layout=args.matrix_layout, + occupancy=args.occupancy, + bytecode_version=args.bytecode_version, + ) + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_cutile_matrix.json b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_cutile_matrix.json new file mode 100644 index 0000000000..1c3eacc6bf --- /dev/null +++ b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_cutile_matrix.json @@ -0,0 +1,515 @@ +[ + { + "_abi": [ + { + "matrix_layout": "relaxed", + "abi_abbrev": "relaxed", + "abi_tag": "cutile_abi_relaxed", + "tile_m": 64, + "tile_n": 128, + "tile_k": 32 + }, + { + "matrix_layout": "strict", + "abi_abbrev": "strict", + "abi_tag": "cutile_abi_strict", + "tile_m": 64, + "tile_n": 128, + "tile_k": 32, + "occupancy": 2 + } + ], + "_data": [ + { + "data_type": "float", + "data_abbrev": "f" + } + ], + "_metric": [ + { + "metric": "runtime" + } + ], + "_index": [ + { + "index_type": "int32", + "index_abbrev": "i32" + } + ], + "_export": [ + { + "output_format": "cubin", + "artifact_ext": "cubin", + "artifact_basename": "@data_type@_@index_abbrev@_@abi_abbrev@_@gpu_code@", + "register": "cubin", + "gpu_code": "sm_80", + "cc_major": 8, + "cc_minor": 0, + "arch_tag": "cutile_arch_8_0" + }, + { + "output_format": "cubin", + "artifact_ext": "cubin", + "artifact_basename": "@data_type@_@index_abbrev@_@abi_abbrev@_@gpu_code@", + "register": "cubin", + "gpu_code": "sm_86", + "cc_major": 8, + "cc_minor": 6, + "arch_tag": "cutile_arch_8_6" + } + ] + }, + { + "_abi": [ + { + "matrix_layout": "relaxed", + "abi_abbrev": "relaxed", + "abi_tag": "cutile_abi_relaxed", + "tile_m": 128, + "tile_n": 128, + "tile_k": 32 + }, + { + "matrix_layout": "strict", + "abi_abbrev": "strict", + "abi_tag": "cutile_abi_strict", + "tile_m": 128, + "tile_n": 128, + "tile_k": 128, + "occupancy": 2 + } + ], + "_data": [ + { + "data_type": "half", + "data_abbrev": "h" + } + ], + "_metric": [ + { + "metric": "runtime" + } + ], + "_index": [ + { + "index_type": "int32", + "index_abbrev": "i32" + } + ], + "_export": [ + { + "output_format": "cubin", + "artifact_ext": "cubin", + "artifact_basename": "@data_type@_@index_abbrev@_@abi_abbrev@_@gpu_code@", + "register": "cubin", + "gpu_code": "sm_80", + "cc_major": 8, + "cc_minor": 0, + "arch_tag": "cutile_arch_8_0" + } + ] + }, + { + "_abi": [ + { + "matrix_layout": "relaxed", + "abi_abbrev": "relaxed", + "abi_tag": "cutile_abi_relaxed", + "tile_m": 128, + "tile_n": 128, + "tile_k": 32, + "occupancy": 2 + }, + { + "matrix_layout": "strict", + "abi_abbrev": "strict", + "abi_tag": "cutile_abi_strict", + "tile_m": 128, + "tile_n": 128, + "tile_k": 32, + "occupancy": 2 + } + ], + "_data": [ + { + "data_type": "half", + "data_abbrev": "h" + } + ], + "_metric": [ + { + "metric": "runtime" + } + ], + "_index": [ + { + "index_type": "int32", + "index_abbrev": "i32" + } + ], + "_export": [ + { + "output_format": "cubin", + "artifact_ext": "cubin", + "artifact_basename": "@data_type@_@index_abbrev@_@abi_abbrev@_@gpu_code@", + "register": "cubin", + "gpu_code": "sm_86", + "cc_major": 8, + "cc_minor": 6, + "arch_tag": "cutile_arch_8_6" + } + ] + }, + { + "_abi": [ + { + "matrix_layout": "relaxed", + "abi_abbrev": "relaxed", + "abi_tag": "cutile_abi_relaxed", + "tile_m": 128, + "tile_n": 128, + "tile_k": 64 + }, + { + "matrix_layout": "strict", + "abi_abbrev": "strict", + "abi_tag": "cutile_abi_strict", + "tile_m": 64, + "tile_n": 256, + "tile_k": 32 + } + ], + "_data": [ + { + "data_type": "float", + "data_abbrev": "f" + } + ], + "_metric": [ + { + "metric": "runtime" + } + ], + "_index": [ + { + "index_type": "int32", + "index_abbrev": "i32" + } + ], + "_export": [ + { + "output_format": "cubin", + "artifact_ext": "cubin", + "artifact_basename": "@data_type@_@index_abbrev@_@abi_abbrev@_@gpu_code@", + "register": "cubin", + "gpu_code": "sm_90", + "cc_major": 9, + "cc_minor": 0, + "arch_tag": "cutile_arch_9_0" + } + ] + }, + { + "_abi": [ + { + "matrix_layout": "relaxed", + "abi_abbrev": "relaxed", + "abi_tag": "cutile_abi_relaxed" + }, + { + "matrix_layout": "strict", + "abi_abbrev": "strict", + "abi_tag": "cutile_abi_strict" + } + ], + "_data": [ + { + "data_type": "half", + "data_abbrev": "h" + } + ], + "_metric": [ + { + "metric": "runtime" + } + ], + "_index": [ + { + "index_type": "int32", + "index_abbrev": "i32" + } + ], + "_tile": [ + { + "tile_m": 128, + "tile_n": 128, + "tile_k": 128 + } + ], + "_export": [ + { + "output_format": "cubin", + "artifact_ext": "cubin", + "artifact_basename": "@data_type@_@index_abbrev@_@abi_abbrev@_@gpu_code@", + "register": "cubin", + "gpu_code": "sm_90", + "cc_major": 9, + "cc_minor": 0, + "arch_tag": "cutile_arch_9_0" + } + ] + }, + { + "_abi": [ + { + "matrix_layout": "relaxed", + "abi_abbrev": "relaxed", + "abi_tag": "cutile_abi_relaxed", + "tile_m": 128, + "tile_n": 256, + "tile_k": 16 + }, + { + "matrix_layout": "strict", + "abi_abbrev": "strict", + "abi_tag": "cutile_abi_strict", + "tile_m": 128, + "tile_n": 128, + "tile_k": 32 + } + ], + "_data": [ + { + "data_type": "float", + "data_abbrev": "f" + } + ], + "_metric": [ + { + "metric": "runtime" + } + ], + "_index": [ + { + "index_type": "int32", + "index_abbrev": "i32" + } + ], + "_export": [ + { + "output_format": "cubin", + "artifact_ext": "cubin", + "artifact_basename": "@data_type@_@index_abbrev@_@abi_abbrev@_@gpu_code@", + "register": "cubin", + "gpu_code": "sm_100", + "cc_major": 10, + "cc_minor": 0, + "arch_tag": "cutile_arch_10_0" + } + ] + }, + { + "_abi": [ + { + "matrix_layout": "relaxed", + "abi_abbrev": "relaxed", + "abi_tag": "cutile_abi_relaxed", + "tile_m": 128, + "tile_n": 256, + "tile_k": 16, + "occupancy": 2 + }, + { + "matrix_layout": "strict", + "abi_abbrev": "strict", + "abi_tag": "cutile_abi_strict", + "tile_m": 128, + "tile_n": 128, + "tile_k": 128 + } + ], + "_data": [ + { + "data_type": "half", + "data_abbrev": "h" + } + ], + "_metric": [ + { + "metric": "runtime" + } + ], + "_index": [ + { + "index_type": "int32", + "index_abbrev": "i32" + } + ], + "_export": [ + { + "output_format": "cubin", + "artifact_ext": "cubin", + "artifact_basename": "@data_type@_@index_abbrev@_@abi_abbrev@_@gpu_code@", + "register": "cubin", + "gpu_code": "sm_100", + "cc_major": 10, + "cc_minor": 0, + "arch_tag": "cutile_arch_10_0" + } + ] + }, + { + "_abi": [ + { + "matrix_layout": "relaxed", + "abi_abbrev": "relaxed", + "abi_tag": "cutile_abi_relaxed", + "tile_m": 64, + "tile_n": 128, + "tile_k": 64, + "occupancy": 2 + }, + { + "matrix_layout": "strict", + "abi_abbrev": "strict", + "abi_tag": "cutile_abi_strict", + "tile_m": 64, + "tile_n": 128, + "tile_k": 32, + "occupancy": 2 + } + ], + "_data": [ + { + "data_type": "float", + "data_abbrev": "f" + } + ], + "_metric": [ + { + "metric": "runtime" + } + ], + "_index": [ + { + "index_type": "int32", + "index_abbrev": "i32" + } + ], + "_export": [ + { + "output_format": "cubin", + "artifact_ext": "cubin", + "artifact_basename": "@data_type@_@index_abbrev@_@abi_abbrev@_@gpu_code@", + "register": "cubin", + "gpu_code": "sm_120", + "cc_major": 12, + "cc_minor": 0, + "arch_tag": "cutile_arch_12_0" + } + ] + }, + { + "_abi": [ + { + "matrix_layout": "relaxed", + "abi_abbrev": "relaxed", + "abi_tag": "cutile_abi_relaxed", + "tile_m": 64, + "tile_n": 128, + "tile_k": 128, + "occupancy": 2 + }, + { + "matrix_layout": "strict", + "abi_abbrev": "strict", + "abi_tag": "cutile_abi_strict", + "tile_m": 64, + "tile_n": 256, + "tile_k": 64, + "occupancy": 2 + } + ], + "_data": [ + { + "data_type": "half", + "data_abbrev": "h" + } + ], + "_metric": [ + { + "metric": "runtime" + } + ], + "_index": [ + { + "index_type": "int32", + "index_abbrev": "i32" + } + ], + "_export": [ + { + "output_format": "cubin", + "artifact_ext": "cubin", + "artifact_basename": "@data_type@_@index_abbrev@_@abi_abbrev@_@gpu_code@", + "register": "cubin", + "gpu_code": "sm_120", + "cc_major": 12, + "cc_minor": 0, + "arch_tag": "cutile_arch_12_0" + } + ] + }, + { + "_abi": [ + { + "matrix_layout": "relaxed", + "abi_abbrev": "relaxed", + "abi_tag": "cutile_abi_relaxed" + }, + { + "matrix_layout": "strict", + "abi_abbrev": "strict", + "abi_tag": "cutile_abi_strict" + } + ], + "_data": [ + { + "data_type": "half", + "data_abbrev": "h" + }, + { + "data_type": "float", + "data_abbrev": "f" + } + ], + "_metric": [ + { + "metric": "runtime" + } + ], + "_index": [ + { + "index_type": "int32", + "index_abbrev": "i32" + } + ], + "_tile": [ + { + "tile_m": 128, + "tile_n": 128, + "tile_k": 32 + } + ], + "_export": [ + { + "output_format": "tileir_bytecode", + "artifact_ext": "tilebc", + "artifact_basename": "@data_type@_@index_abbrev@_@abi_abbrev@", + "register": "tileir", + "gpu_code": "sm_80", + "bytecode_version": "13.1" + } + ] + } +] diff --git a/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_kernel.py b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_kernel.py new file mode 100644 index 0000000000..ad1ba8fbea --- /dev/null +++ b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_kernel.py @@ -0,0 +1,191 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""cuTile fused GEMM + 1-NN kernel with runtime metric selection.""" + +from __future__ import annotations + +import cuda.tile as ct + +ConstInt = ct.Constant[int] + +# Default tile geometry; overridden per export via make_kernel(..., tile_m, tile_n, tile_k). +DEFAULT_TILE_M = 128 +DEFAULT_TILE_N = 128 +DEFAULT_TILE_K = 32 + +METRICS = ("runtime",) +INDEX_TYPES = ("int32", "int64") +METRIC_L2_EXPANDED = 0 +METRIC_COSINE_EXPANDED = 2 +METRIC_INNER_PRODUCT = 6 + + +def _idx_dtype(index_type: str): + if index_type == "int32": + return ct.int32 + if index_type == "int64": + return ct.int64 + raise ValueError(f"Unsupported index_type {index_type!r}") + + +def make_kernel( + data_type: str, + metric: str, + tile_m: int = DEFAULT_TILE_M, + tile_n: int = DEFAULT_TILE_N, + tile_k: int = DEFAULT_TILE_K, + *, + index_type: str = "int32", + gpu_code: str = "sm_80", + matrix_layout: str = "strict", + occupancy: int | None = None, +): + """Build the flat-reduction runtime-metric cuTile kernel.""" + if data_type not in ("half", "float"): + raise ValueError(f"Unsupported data_type {data_type!r}") + if metric not in METRICS: + raise ValueError(f"Unsupported metric {metric!r}") + if index_type not in INDEX_TYPES: + raise ValueError(f"Unsupported index_type {index_type!r}") + if matrix_layout not in ("strict", "relaxed"): + raise ValueError(f"Unsupported matrix_layout {matrix_layout!r}") + + acc_dtype = ct.float32 + idx_dtype = _idx_dtype(index_type) + out_dist_dtype = ct.float16 if data_type == "half" else ct.float32 + core_shape = (tile_m, tile_n) + best_shape = (tile_m, 1) + kernel_options = {} + if occupancy is not None: + kernel_options["occupancy"] = ct.ByTarget(**{gpu_code: occupancy}) + + @ct.kernel(**kernel_options) + def fused_1nn_kernel( + A, + B, + A_norm, + B_norm, + OutIdx, + OutDist, + M, + N, + K, + apply_sqrt, + store_idx, + metric_code, + tm: ConstInt, + tn: ConstInt, + tk: ConstInt, + ): + bidm = ct.bid(0) + best_dist = ct.full(best_shape, 3.4e38, acc_dtype) + best_idx = ct.zeros(best_shape, idx_dtype) + num_tiles_k = ct.num_tiles(A, axis=1, shape=(tm, tk)) + num_tiles_n = ct.num_tiles(B, axis=0, shape=(tn, tk)) + zero_pad = ct.PaddingMode.ZERO + + def reduce_scores(dists, indices): + def red_op(a_score, a_idx, b_score, b_idx): + cond = (a_score < b_score) | ( + (a_score == b_score) & (a_idx < b_idx) + ) + return ( + ct.where(cond, a_score, b_score), + ct.where(cond, a_idx, b_idx), + ) + + return ct.reduce( + (dists, indices), + 1, + red_op, + (3.4e38, -1), + keepdims=True, + ) + + local_indices = ct.arange(tn, dtype=ct.int16)[None, :] + for n in range(num_tiles_n): + accumulator = ct.full((tm, tn), 0, dtype=acc_dtype) + for k in range(num_tiles_k): + dtype = ct.tfloat32 if A.dtype == ct.float32 else A.dtype + a = ct.load( + A, index=(bidm, k), shape=(tm, tk), padding_mode=zero_pad + ).astype(dtype) + b_T = ct.load( + B, + index=(k, n), + shape=(tk, tn), + padding_mode=zero_pad, + order=(1, 0), + ).astype(dtype) + accumulator = ct.mma(a, b_T, accumulator) + + if metric_code == METRIC_INNER_PRODUCT: + score = -accumulator + else: + b_norm = ct.load( + B_norm, index=(n,), shape=(tn,), padding_mode=zero_pad + ) + if metric_code == METRIC_L2_EXPANDED: + # The A norm is constant across centroids. Excluding it + # avoids cancellation in the score used by argmin. + score = (0.5 * b_norm)[None, :] - accumulator + else: + # Defer the A-norm division until after selecting the + # winning centroid. + score = accumulator / (-b_norm)[None, :] + + if n == num_tiles_n - 1: + col = ct.arange(tn, dtype=ct.int16) + score = ct.where((n * tn + col)[None, :] < N, score, 3.4e38) + + curr_best, curr_idx = reduce_scores( + score.reshape(core_shape), local_indices + ) + update = curr_best < best_dist + best_dist = ct.where(update, curr_best, best_dist) + best_idx = ct.where(update, n * tn + curr_idx, best_idx) + + if metric_code == METRIC_INNER_PRODUCT: + out_dist = -best_dist + else: + a_norm = ct.load( + A_norm, index=(bidm,), shape=(tm,), padding_mode=zero_pad + )[:, None] + if metric_code == METRIC_L2_EXPANDED: + out_dist = a_norm + 2.0 * best_dist + # Separately reduced norms and the MMA can reconstruct a + # slightly negative distance; clamp before an optional sqrt. + out_dist = ct.where(out_dist > 0.0, out_dist, 0.0) + out_dist = ct.where( + apply_sqrt != 0, ct.sqrt(out_dist), out_dist + ) + else: + out_dist = 1.0 + best_dist / a_norm + + if store_idx != 0: + ct.store(OutIdx, index=(bidm,), tile=best_idx.reshape((tm,))) + ct.store( + OutDist, + index=(bidm,), + tile=out_dist.reshape((tm,)).astype(out_dist_dtype), + ) + + return fused_1nn_kernel + + +def kernel_symbol( + data_abbrev: str, + index_abbrev: str, + matrix_layout: str = "strict", +) -> str: + """Must stay in sync with fused_1nn_kernel_entrypoint() in fused_1nn_planner.hpp.""" + base = f"fused_1nn_{data_abbrev}_{index_abbrev}" + if matrix_layout == "strict": + return base + if matrix_layout == "relaxed": + return f"{base}_relaxed" + raise ValueError(f"Unsupported matrix layout {matrix_layout!r}") + + +def index_abbrev(index_type: str) -> str: + return {"int32": "i32", "int64": "i64"}[index_type] diff --git a/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_planner.hpp b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_planner.hpp new file mode 100644 index 0000000000..0755788121 --- /dev/null +++ b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_planner.hpp @@ -0,0 +1,130 @@ +/* + * 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 "fused_1nn_cutile_tiles.hpp" + +namespace cuvs::distance::detail { + +/** Must match kernel_symbol() in fused_1nn_kernel.py (export uses with_symbol). */ +template +inline const char* fused_1nn_kernel_entrypoint() +{ + constexpr bool is_relaxed = std::is_same_v; + static_assert(is_relaxed || std::is_same_v, + "unsupported fused 1-NN cuTile ABI"); + + if constexpr (std::is_same_v) { + return is_relaxed ? "fused_1nn_f_i32_relaxed" : "fused_1nn_f_i32"; + } else if constexpr (std::is_same_v) { + return is_relaxed ? "fused_1nn_h_i32_relaxed" : "fused_1nn_h_i32"; + } else { + static_assert(sizeof(DataTag) == 0, "unsupported fused 1-NN cuTile data type"); + return ""; + } +} + +template +struct Fused1nnTilePlanner : cuvs::detail::jit_lto::TileAlgorithmPlanner { + using DataTag = fused_1nn_data_tag_t; + using IndexTag = cuvs::neighbors::detail::tag_index_i32; + + inline static cuvs::detail::jit_lto::TileLauncherCache launcher_cache{}; + + Fused1nnTilePlanner() + : TileAlgorithmPlanner(fused_1nn_kernel_entrypoint(), launcher_cache) + { + } + + /** Registers embedded cubin modules (one per SM); see register_cutile_fragment.cpp object files. + */ + void add_entrypoint() + { + using cuvs::detail::jit_lto::cutile_arch_10_0; + using cuvs::detail::jit_lto::cutile_arch_12_0; + using cuvs::detail::jit_lto::cutile_arch_8_0; + using cuvs::detail::jit_lto::cutile_arch_8_6; + using cuvs::detail::jit_lto::cutile_arch_9_0; + + constexpr bool is_relaxed = std::is_same_v; + constexpr bool is_float = std::is_same_v; + using Tile80 = + std::conditional_t, + std::conditional_t>; + using Tile86 = + std::conditional_t, + std::conditional_t>; + using Tile90 = + std::conditional_t, + std::conditional_t>; + using Tile100 = + std::conditional_t, + std::conditional_t>; + using Tile120 = + std::conditional_t, + std::conditional_t>; + + this->add_static_fragment< + fragment_tag_fused_1nn_cubin>(); + this->add_static_fragment< + fragment_tag_fused_1nn_cubin>(); + this->add_static_fragment< + fragment_tag_fused_1nn_cubin>(); + this->add_static_fragment< + fragment_tag_fused_1nn_cubin>(); + this->add_static_fragment< + fragment_tag_fused_1nn_cubin>(); + } + + void add_tileir_fallback() + { + constexpr bool is_relaxed = std::is_same_v; + constexpr bool is_float = std::is_same_v; + using TileIr = std::conditional_t, + std::conditional_t>; + this->add_static_tileir_fragment< + fragment_tag_fused_1nn_tileir>(); + } +}; + +} // namespace cuvs::distance::detail diff --git a/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_tile.cu b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_tile.cu new file mode 100644 index 0000000000..870c794789 --- /dev/null +++ b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_tile.cu @@ -0,0 +1,471 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include "fused_1nn_tile.hpp" + +#include "fused_1nn_planner.hpp" + +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include + +namespace cuvs { +namespace distance { +namespace detail { + +namespace { + +bool is_16_byte_aligned(const void* ptr) +{ + return ptr == nullptr || reinterpret_cast(ptr) % 16 == 0; +} + +bool byte_ranges_overlap(const void* lhs, size_t lhs_bytes, const void* rhs, size_t rhs_bytes) +{ + if (lhs == nullptr || rhs == nullptr || lhs_bytes == 0 || rhs_bytes == 0) { return false; } + const auto lhs_begin = reinterpret_cast(lhs); + const auto rhs_begin = reinterpret_cast(rhs); + if (lhs_bytes > std::numeric_limits::max() - lhs_begin || + rhs_bytes > std::numeric_limits::max() - rhs_begin) { + return true; + } + return lhs_begin < rhs_begin + rhs_bytes && rhs_begin < lhs_begin + lhs_bytes; +} + +template +size_t checked_tensor_bytes(IdxT rows, IdxT cols, size_t element_size) +{ + const auto rows_u = static_cast(rows); + const auto cols_u = static_cast(cols); + constexpr auto max_size = std::numeric_limits::max(); + if (cols_u != 0 && rows_u > max_size / cols_u) { return max_size; } + const auto elements = rows_u * cols_u; + if (element_size != 0 && elements > max_size / element_size) { return max_size; } + return static_cast(elements) * element_size; +} + +template +bool has_fused_1nn_tile_launcher() +{ + Fused1nnTilePlanner planner; + planner.add_entrypoint(); + planner.add_tileir_fallback(); + return planner.try_get_launcher() != nullptr; +} + +template +bool launch_fused_1nn_tile(IdxT* nearest_idx, + DataT* nearest_dist, + const DataT* x, + const DataT* y, + const fused_1nn_cutile_norm_t* xn, + const fused_1nn_cutile_norm_t* yn, + IdxT m, + IdxT n, + IdxT k, + cuvs::distance::DistanceType metric, + bool is_sqrt, + cudaStream_t stream) +{ + if constexpr (!std::is_same_v && !std::is_same_v) { return false; } + + if (nearest_dist == nullptr) { return false; } + + Fused1nnTilePlanner planner; + planner.add_entrypoint(); + planner.add_tileir_fallback(); + auto launcher = planner.try_get_launcher(); + if (!launcher) { return false; } + const cuvs::detail::jit_lto::CutileTileConfig tile_cfg = planner.tile_config(); + + int metric_code; + bool apply_sqrt = false; + switch (metric) { + case cuvs::distance::DistanceType::InnerProduct: + metric_code = static_cast(cuvs::distance::DistanceType::InnerProduct); + break; + case cuvs::distance::DistanceType::L2Expanded: + case cuvs::distance::DistanceType::L2SqrtExpanded: + metric_code = static_cast(cuvs::distance::DistanceType::L2Expanded); + apply_sqrt = is_sqrt; + break; + case cuvs::distance::DistanceType::CosineExpanded: + metric_code = static_cast(cuvs::distance::DistanceType::CosineExpanded); + break; + default: return false; + } + + IdxT shape_x[2] = {m, k}; + IdxT stride_x[2] = {k, IdxT{1}}; + IdxT shape_y[2] = {n, k}; + IdxT stride_y[2] = {k, IdxT{1}}; + IdxT shape_xn = m; + IdxT stride_xn = IdxT{1}; + IdxT shape_yn = n; + IdxT stride_yn = IdxT{1}; + IdxT shape_idx = m; + IdxT stride_idx = IdxT{1}; + IdxT shape_dist = m; + IdxT stride_dist = IdxT{1}; + + IdxT M = m; + IdxT N = n; + IdxT K = k; + + void* x_ptr = const_cast(x); + void* y_ptr = const_cast(y); + void* xn_ptr = const_cast*>(xn); + void* yn_ptr = const_cast*>(yn); + const IdxT store_idx = nearest_idx != nullptr ? IdxT{1} : IdxT{0}; + void* idx_ptr = nearest_idx; + void* dist_ptr = nearest_dist; + + const int tile_m = tile_cfg.tile_m; + dim3 grid((static_cast(m) + tile_m - 1) / tile_m, 1, 1); + dim3 block(1, 1, 1); + + using fused_1nn_cutile_kernel_t = void(void*, + IdxT, + IdxT, + IdxT, + IdxT, + void*, + IdxT, + IdxT, + IdxT, + IdxT, + void*, + IdxT, + IdxT, + void*, + IdxT, + IdxT, + void*, + IdxT, + IdxT, + void*, + IdxT, + IdxT, + IdxT, + IdxT, + IdxT, + IdxT, + IdxT, + int); + launcher->template dispatch(stream, + grid, + block, + 0, + x_ptr, + shape_x[0], + shape_x[1], + stride_x[0], + stride_x[1], + y_ptr, + shape_y[0], + shape_y[1], + stride_y[0], + stride_y[1], + xn_ptr, + shape_xn, + stride_xn, + yn_ptr, + shape_yn, + stride_yn, + idx_ptr, + shape_idx, + stride_idx, + dist_ptr, + shape_dist, + stride_dist, + M, + N, + K, + static_cast(apply_sqrt), + store_idx, + metric_code); + RAFT_CUDA_TRY(cudaGetLastError()); + return true; +} + +template +bool try_fused_1nn_tile_dispatch(IdxT* nearest_idx, + DataT* nearest_dist, + const DataT* x, + const DataT* y, + const fused_1nn_cutile_norm_t* xn, + const fused_1nn_cutile_norm_t* yn, + IdxT m, + IdxT n, + IdxT k, + cuvs::distance::DistanceType metric, + bool is_sqrt, + cudaStream_t stream) +{ + return launch_fused_1nn_tile( + nearest_idx, nearest_dist, x, y, xn, yn, m, n, k, metric, is_sqrt, stream); +} + +} // namespace + +template + requires is_fused_1nn_cutile_data_v +bool can_launch_fused_1nn_tile( + const DataT* x, const DataT* y, IdxT m, IdxT n, IdxT k, cuvs::distance::DistanceType metric) +{ + if (!cuvs::detail::jit_lto::cutile_launch_available_on_current_device()) { return false; } + static_assert(std::is_same_v || std::is_same_v); + + if (x == nullptr || y == nullptr || m <= 0 || n <= 0 || k <= 0) { return false; } + if (metric != cuvs::distance::DistanceType::InnerProduct && + metric != cuvs::distance::DistanceType::L2Expanded && + metric != cuvs::distance::DistanceType::L2SqrtExpanded && + metric != cuvs::distance::DistanceType::CosineExpanded) { + return false; + } + + if (!is_16_byte_aligned(x) || !is_16_byte_aligned(y)) { return false; } + if constexpr (std::is_same_v) { + constexpr int64_t max_i32 = std::numeric_limits::max(); + if (n > max_i32 || k > max_i32) { return false; } + } + + constexpr int strict_pitch_elements = 16 / sizeof(DataT); + return k % strict_pitch_elements == 0 ? has_fused_1nn_tile_launcher() + : has_fused_1nn_tile_launcher(); +} + +template + requires is_fused_1nn_cutile_data_v +bool can_launch_fused_1nn_tile(IdxT* nearest_idx, + DataT* nearest_dist, + const DataT* x, + const DataT* y, + IdxT m, + IdxT n, + IdxT k, + cuvs::distance::DistanceType metric) +{ + if (!can_launch_fused_1nn_tile(x, y, m, n, k, metric)) { return false; } + if (nearest_dist == nullptr || !is_16_byte_aligned(nearest_dist)) { return false; } + if constexpr (std::is_same_v) { + if (!is_16_byte_aligned(nearest_idx)) { return false; } + } + const auto x_bytes = checked_tensor_bytes(m, k, sizeof(DataT)); + const auto y_bytes = checked_tensor_bytes(n, k, sizeof(DataT)); + const auto dist_bytes = checked_tensor_bytes(m, IdxT{1}, sizeof(DataT)); + const auto idx_bytes = checked_tensor_bytes(m, IdxT{1}, sizeof(IdxT)); + if (byte_ranges_overlap(nearest_dist, dist_bytes, x, x_bytes) || + byte_ranges_overlap(nearest_dist, dist_bytes, y, y_bytes) || + byte_ranges_overlap(nearest_idx, idx_bytes, x, x_bytes) || + byte_ranges_overlap(nearest_idx, idx_bytes, y, y_bytes) || + byte_ranges_overlap(nearest_idx, idx_bytes, nearest_dist, dist_bytes)) { + return false; + } + return true; +} + +template + requires is_fused_1nn_cutile_data_v +bool can_launch_fused_1nn_tile(IdxT* nearest_idx, + DataT* nearest_dist, + const DataT* x, + const DataT* y, + const fused_1nn_cutile_norm_t* xn, + const fused_1nn_cutile_norm_t* yn, + IdxT m, + IdxT n, + IdxT k, + cuvs::distance::DistanceType metric) +{ + if (!can_launch_fused_1nn_tile(nearest_idx, nearest_dist, x, y, m, n, k, metric)) { + return false; + } + if (metric != cuvs::distance::DistanceType::InnerProduct && (xn == nullptr || yn == nullptr)) { + return false; + } + if (!is_16_byte_aligned(xn) || !is_16_byte_aligned(yn)) { return false; } + const auto xn_bytes = checked_tensor_bytes(m, IdxT{1}, sizeof(*xn)); + const auto yn_bytes = checked_tensor_bytes(n, IdxT{1}, sizeof(*yn)); + const auto dist_bytes = checked_tensor_bytes(m, IdxT{1}, sizeof(DataT)); + const auto idx_bytes = checked_tensor_bytes(m, IdxT{1}, sizeof(IdxT)); + return !byte_ranges_overlap(nearest_dist, dist_bytes, xn, xn_bytes) && + !byte_ranges_overlap(nearest_dist, dist_bytes, yn, yn_bytes) && + !byte_ranges_overlap(nearest_idx, idx_bytes, xn, xn_bytes) && + !byte_ranges_overlap(nearest_idx, idx_bytes, yn, yn_bytes); +} + +template + requires is_fused_1nn_cutile_data_v +bool try_fused_1nn_tile(IdxT* nearest_idx, + DataT* nearest_dist, + const DataT* x, + const DataT* y, + const fused_1nn_cutile_norm_t* xn, + const fused_1nn_cutile_norm_t* yn, + IdxT m, + IdxT n, + IdxT k, + cuvs::distance::DistanceType metric, + bool is_sqrt, + void* index_workspace, + cudaStream_t stream) +{ + if (!can_launch_fused_1nn_tile(nearest_idx, nearest_dist, x, y, xn, yn, m, n, k, metric)) { + return false; + } + + constexpr int strict_pitch_elements = 16 / sizeof(DataT); + const bool use_strict_abi = k % strict_pitch_elements == 0; + + if constexpr (std::is_same_v) { + if (use_strict_abi) { + return try_fused_1nn_tile_dispatch( + nearest_idx, nearest_dist, x, y, xn, yn, m, n, k, metric, is_sqrt, stream); + } + return try_fused_1nn_tile_dispatch( + nearest_idx, nearest_dist, x, y, xn, yn, m, n, k, metric, is_sqrt, stream); + } else { + if (nearest_idx != nullptr && index_workspace == nullptr) { return false; } + if (!is_16_byte_aligned(index_workspace)) { return false; } + const auto workspace_bytes = checked_tensor_bytes(m, IdxT{1}, sizeof(int)); + const auto x_bytes = checked_tensor_bytes(m, k, sizeof(DataT)); + const auto y_bytes = checked_tensor_bytes(n, k, sizeof(DataT)); + const auto norm_x_bytes = checked_tensor_bytes(m, IdxT{1}, sizeof(*xn)); + const auto norm_y_bytes = checked_tensor_bytes(n, IdxT{1}, sizeof(*yn)); + const auto dist_bytes = checked_tensor_bytes(m, IdxT{1}, sizeof(DataT)); + const auto idx_bytes = checked_tensor_bytes(m, IdxT{1}, sizeof(IdxT)); + if (byte_ranges_overlap(index_workspace, workspace_bytes, x, x_bytes) || + byte_ranges_overlap(index_workspace, workspace_bytes, y, y_bytes) || + byte_ranges_overlap(index_workspace, workspace_bytes, xn, norm_x_bytes) || + byte_ranges_overlap(index_workspace, workspace_bytes, yn, norm_y_bytes) || + byte_ranges_overlap(index_workspace, workspace_bytes, nearest_dist, dist_bytes) || + byte_ranges_overlap(index_workspace, workspace_bytes, nearest_idx, idx_bytes)) { + return false; + } + + // Keep every chunk offset 16-byte aligned for x, xn, and nearest_dist. + constexpr int64_t max_batch_m = fused_1nn_cutile_max_batch_m; + auto* tmp_idx = static_cast(index_workspace); + for (int64_t offset = 0; offset < m;) { + const int64_t batch_m64 = std::min(max_batch_m, m - offset); + const int batch_m = static_cast(batch_m64); + const auto* batch_x = x + static_cast(offset) * static_cast(k); + const auto* batch_xn = xn == nullptr ? nullptr : xn + offset; + auto* batch_dist = nearest_dist == nullptr ? nullptr : nearest_dist + offset; + + const bool launched = + use_strict_abi + ? try_fused_1nn_tile_dispatch(tmp_idx, + batch_dist, + batch_x, + y, + batch_xn, + yn, + batch_m, + static_cast(n), + static_cast(k), + metric, + is_sqrt, + stream) + : try_fused_1nn_tile_dispatch(tmp_idx, + batch_dist, + batch_x, + y, + batch_xn, + yn, + batch_m, + static_cast(n), + static_cast(k), + metric, + is_sqrt, + stream); + if (!launched) { return false; } + + if (nearest_idx != nullptr) { + raft::linalg::unaryOp( + nearest_idx + offset, tmp_idx, batch_m, raft::cast_op{}, stream); + } + offset += batch_m64; + } + return true; + } +} + +#define CUVS_INST_CAN_LAUNCH_FUSED_1NN_TILE_INPUTS(DataT, IdxT) \ + template CUVS_EXPORT bool can_launch_fused_1nn_tile( \ + const DataT*, const DataT*, IdxT, IdxT, IdxT, cuvs::distance::DistanceType) + +CUVS_INST_CAN_LAUNCH_FUSED_1NN_TILE_INPUTS(float, int); +CUVS_INST_CAN_LAUNCH_FUSED_1NN_TILE_INPUTS(float, int64_t); +CUVS_INST_CAN_LAUNCH_FUSED_1NN_TILE_INPUTS(half, int); +CUVS_INST_CAN_LAUNCH_FUSED_1NN_TILE_INPUTS(half, int64_t); + +#undef CUVS_INST_CAN_LAUNCH_FUSED_1NN_TILE_INPUTS + +#define CUVS_INST_CAN_LAUNCH_FUSED_1NN_TILE_PREFLIGHT(DataT, IdxT) \ + template CUVS_EXPORT bool can_launch_fused_1nn_tile( \ + IdxT*, DataT*, const DataT*, const DataT*, IdxT, IdxT, IdxT, cuvs::distance::DistanceType) + +CUVS_INST_CAN_LAUNCH_FUSED_1NN_TILE_PREFLIGHT(float, int); +CUVS_INST_CAN_LAUNCH_FUSED_1NN_TILE_PREFLIGHT(float, int64_t); +CUVS_INST_CAN_LAUNCH_FUSED_1NN_TILE_PREFLIGHT(half, int); +CUVS_INST_CAN_LAUNCH_FUSED_1NN_TILE_PREFLIGHT(half, int64_t); + +#undef CUVS_INST_CAN_LAUNCH_FUSED_1NN_TILE_PREFLIGHT + +#define CUVS_INST_CAN_LAUNCH_FUSED_1NN_TILE(DataT, IdxT) \ + template CUVS_EXPORT bool can_launch_fused_1nn_tile( \ + IdxT*, \ + DataT*, \ + const DataT*, \ + const DataT*, \ + const fused_1nn_cutile_norm_t*, \ + const fused_1nn_cutile_norm_t*, \ + IdxT, \ + IdxT, \ + IdxT, \ + cuvs::distance::DistanceType) + +CUVS_INST_CAN_LAUNCH_FUSED_1NN_TILE(float, int); +CUVS_INST_CAN_LAUNCH_FUSED_1NN_TILE(float, int64_t); +CUVS_INST_CAN_LAUNCH_FUSED_1NN_TILE(half, int); +CUVS_INST_CAN_LAUNCH_FUSED_1NN_TILE(half, int64_t); + +#undef CUVS_INST_CAN_LAUNCH_FUSED_1NN_TILE + +#define CUVS_INST_TRY_FUSED_1NN_TILE(DataT, IdxT) \ + template CUVS_EXPORT bool try_fused_1nn_tile(IdxT*, \ + DataT*, \ + const DataT*, \ + const DataT*, \ + const fused_1nn_cutile_norm_t*, \ + const fused_1nn_cutile_norm_t*, \ + IdxT, \ + IdxT, \ + IdxT, \ + cuvs::distance::DistanceType, \ + bool, \ + void*, \ + cudaStream_t) + +CUVS_INST_TRY_FUSED_1NN_TILE(float, int); +CUVS_INST_TRY_FUSED_1NN_TILE(float, int64_t); +CUVS_INST_TRY_FUSED_1NN_TILE(half, int); +CUVS_INST_TRY_FUSED_1NN_TILE(half, int64_t); + +#undef CUVS_INST_TRY_FUSED_1NN_TILE + +} // namespace detail +} // namespace distance +} // namespace cuvs diff --git a/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_tile.hpp b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_tile.hpp new file mode 100644 index 0000000000..add5f7a94e --- /dev/null +++ b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_tile.hpp @@ -0,0 +1,164 @@ +/* + * 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 + +#ifndef CUVS_CUTILE_ENABLED +#define CUVS_CUTILE_ENABLED 0 +#endif + +namespace cuvs { +namespace distance { +namespace detail { + +template +inline constexpr bool is_fused_1nn_cutile_data_v = + std::is_same_v || std::is_same_v; + +// Tensor-core products accumulate in FP32; FP16 norms must remain FP32 through the epilogue. +template +using fused_1nn_cutile_norm_t = std::conditional_t, float, DataT>; + +template +inline constexpr int64_t fused_1nn_cutile_max_batch_m = [] { + constexpr int64_t max_i32 = std::numeric_limits::max(); + constexpr int64_t batch_alignment = 16 / sizeof(DataT); + return max_i32 - max_i32 % batch_alignment; +}(); + +template +constexpr size_t fused_1nn_cutile_index_workspace_rows(IdxT m) +{ + const auto rows = static_cast(m); + if (rows <= 0) { return 0; } + return static_cast( + rows < fused_1nn_cutile_max_batch_m ? rows : fused_1nn_cutile_max_batch_m); +} + +#if CUVS_CUTILE_ENABLED +/** + * Return whether the input problem has a compatible cuTile launcher. + * + * This output-independent probe lets callers select native result storage before allocating it. + */ +template + requires is_fused_1nn_cutile_data_v +bool can_launch_fused_1nn_tile( + const DataT* x, const DataT* y, IdxT m, IdxT n, IdxT k, cuvs::distance::DistanceType metric); + +/** + * Return whether the supplied problem can use cuTile without fallback scratch. + * + * The result includes runtime/device support, exported ABI constraints, and launcher construction. + * A successful probe populates the shared launcher cache used by try_fused_1nn_tile. + * An int64 output index still requires an int32 workspace sized to the largest launch chunk. + */ +template + requires is_fused_1nn_cutile_data_v +bool can_launch_fused_1nn_tile(IdxT* nearest_idx, + DataT* nearest_dist, + const DataT* x, + const DataT* y, + IdxT m, + IdxT n, + IdxT k, + cuvs::distance::DistanceType metric); + +/** + * Return whether the supplied problem and existing norm buffers can use cuTile. + * + * The overload without norm pointers is a preflight probe for callers that allocate aligned norm + * buffers only after the remaining launch requirements have been validated. + */ +template + requires is_fused_1nn_cutile_data_v +bool can_launch_fused_1nn_tile(IdxT* nearest_idx, + DataT* nearest_dist, + const DataT* x, + const DataT* y, + const fused_1nn_cutile_norm_t* xn, + const fused_1nn_cutile_norm_t* yn, + IdxT m, + IdxT n, + IdxT k, + cuvs::distance::DistanceType metric); + +template + requires is_fused_1nn_cutile_data_v +bool try_fused_1nn_tile(IdxT* nearest_idx, + DataT* nearest_dist, + const DataT* x, + const DataT* y, + const fused_1nn_cutile_norm_t* xn, + const fused_1nn_cutile_norm_t* yn, + IdxT m, + IdxT n, + IdxT k, + cuvs::distance::DistanceType metric, + bool is_sqrt, + void* index_workspace, + cudaStream_t stream); +#else +template +bool can_launch_fused_1nn_tile( + const DataT*, const DataT*, IdxT, IdxT, IdxT, cuvs::distance::DistanceType) +{ + return false; +} + +template +bool can_launch_fused_1nn_tile( + IdxT*, DataT*, const DataT*, const DataT*, IdxT, IdxT, IdxT, cuvs::distance::DistanceType) +{ + return false; +} + +template +bool can_launch_fused_1nn_tile(IdxT*, + DataT*, + const DataT*, + const DataT*, + const fused_1nn_cutile_norm_t*, + const fused_1nn_cutile_norm_t*, + IdxT, + IdxT, + IdxT, + cuvs::distance::DistanceType) +{ + return false; +} + +template +bool try_fused_1nn_tile(IdxT*, + DataT*, + const DataT*, + const DataT*, + const fused_1nn_cutile_norm_t*, + const fused_1nn_cutile_norm_t*, + IdxT, + IdxT, + IdxT, + cuvs::distance::DistanceType, + bool, + void*, + cudaStream_t) +{ + return false; +} +#endif + +} // namespace detail +} // namespace distance +} // namespace cuvs diff --git a/cpp/src/distance/fused_distance_nn-inl.cuh b/cpp/src/distance/fused_distance_nn-inl.cuh index 3fa80a9b60..ff7b6575aa 100644 --- a/cpp/src/distance/fused_distance_nn-inl.cuh +++ b/cpp/src/distance/fused_distance_nn-inl.cuh @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ @@ -311,6 +311,81 @@ void fusedDistanceNNMinReduce(OutT* min, stream); } +namespace detail { +template +__global__ void unpack_fused_1nn_kvp(IdxT* nearest_idx, + DataT* nearest_dist, + const raft::KeyValuePair* kvp, + IdxT m) +{ + const auto i = static_cast(blockIdx.x * blockDim.x + threadIdx.x); + if (i >= m) { return; } + if (nearest_idx != nullptr) { nearest_idx[i] = kvp[i].key; } + if (nearest_dist != nullptr) { nearest_dist[i] = kvp[i].value; } +} +} // namespace detail + +template +void fusedDistanceNNMinReduce(IdxT* nearest_idx, + DataT* nearest_dist, + const DataT* x, + const DataT* y, + const NormT* xn, + const NormT* yn, + IdxT m, + IdxT n, + IdxT k, + void* workspace, + bool sqrt, + bool init_out_buffer, + bool is_row_major, + cuvs::distance::DistanceType metric, + float metric_arg, + detail::Fused1nnBackend backend, + raft::KeyValuePair* cutlass_kvp_scratch, + cudaStream_t stream) +{ + RAFT_EXPECTS(is_row_major, "fusedDistanceNN only supports row-major inputs"); + RAFT_EXPECTS(nearest_dist != nullptr, "Explicit fused 1-NN backends require nearest_dist"); + if (backend == detail::Fused1nnBackend::Cutile) { + if constexpr (detail::is_fused_1nn_cutile_data_v && + std::is_same_v>) { + RAFT_EXPECTS( + detail::try_fused_1nn_tile( + nearest_idx, nearest_dist, x, y, xn, yn, m, n, k, metric, sqrt, workspace, stream), + "Requested cuTile fused 1-NN backend is unavailable for this input/device"); + return; + } + RAFT_FAIL("Requested cuTile fused 1-NN backend does not support these data/norm types"); + } + RAFT_EXPECTS(backend == detail::Fused1nnBackend::Cutlass, "Unknown fused 1-NN backend"); + RAFT_EXPECTS(metric != cuvs::distance::DistanceType::InnerProduct, + "CUTLASS fused 1-NN does not support InnerProduct"); + RAFT_EXPECTS(std::is_same_v, "CUTLASS fused 1-NN requires matching norm types"); + RAFT_EXPECTS(cutlass_kvp_scratch != nullptr, + "CUTLASS fused 1-NN with explicit outputs requires KVP scratch storage"); + fusedDistanceNNMinReduce, IdxT>(cutlass_kvp_scratch, + x, + y, + xn, + yn, + m, + n, + k, + workspace, + sqrt, + init_out_buffer, + is_row_major, + metric, + metric_arg, + stream); + constexpr int threads = 256; + detail:: + unpack_fused_1nn_kvp<<((m + threads - 1) / threads), threads, 0, stream>>>( + nearest_idx, nearest_dist, cutlass_kvp_scratch, m); + RAFT_CUDA_TRY(cudaGetLastError()); +} + /** @} */ } // namespace distance From fd30a05e99f3bf14a8f72d99a20d85d2c915752e Mon Sep 17 00:00:00 2001 From: divyegala Date: Thu, 3 Sep 2026 21:34:44 +0000 Subject: [PATCH 79/82] no extra work --- cpp/src/distance/detail/fused_distance_nn.cuh | 25 ++++++ cpp/src/distance/fused_distance_nn-inl.cuh | 30 ++----- cpp/tests/neighbors/distance_nn.cu | 90 +++++++++++++++---- cpp/tests/neighbors/distance_nn_helper.cuh | 26 +++++- 4 files changed, 127 insertions(+), 44 deletions(-) diff --git a/cpp/src/distance/detail/fused_distance_nn.cuh b/cpp/src/distance/detail/fused_distance_nn.cuh index 5a65482f13..d1b0fe0dc3 100644 --- a/cpp/src/distance/detail/fused_distance_nn.cuh +++ b/cpp/src/distance/detail/fused_distance_nn.cuh @@ -35,6 +35,31 @@ enum class Fused1nnBackend : std::uint8_t { Cutlass, }; +/** + * Output-independent backend probe. Call this before allocating backend-native result storage. + * cuTile delegates to its launcher/ABI probe; CUTLASS is available for the legacy L2/cosine + * fused primitive only. + */ +template +bool can_launch_fused_1nn_backend(Fused1nnBackend backend, + const DataT* x, + const DataT* y, + IdxT m, + IdxT n, + IdxT k, + cuvs::distance::DistanceType metric) +{ + if (backend == Fused1nnBackend::Cutile) { + if constexpr (is_fused_1nn_cutile_data_v) { + return can_launch_fused_1nn_tile(x, y, m, n, k, metric); + } + return false; + } + return backend == Fused1nnBackend::Cutlass && + metric != cuvs::distance::DistanceType::InnerProduct && x != nullptr && y != nullptr && + m > 0 && n > 0 && k > 0; +} + template -__global__ void unpack_fused_1nn_kvp(IdxT* nearest_idx, - DataT* nearest_dist, - const raft::KeyValuePair* kvp, - IdxT m) -{ - const auto i = static_cast(blockIdx.x * blockDim.x + threadIdx.x); - if (i >= m) { return; } - if (nearest_idx != nullptr) { nearest_idx[i] = kvp[i].key; } - if (nearest_dist != nullptr) { nearest_dist[i] = kvp[i].value; } -} -} // namespace detail - template void fusedDistanceNNMinReduce(IdxT* nearest_idx, DataT* nearest_dist, @@ -342,11 +328,10 @@ void fusedDistanceNNMinReduce(IdxT* nearest_idx, cuvs::distance::DistanceType metric, float metric_arg, detail::Fused1nnBackend backend, - raft::KeyValuePair* cutlass_kvp_scratch, + raft::KeyValuePair* cutlass_kvp_output, cudaStream_t stream) { RAFT_EXPECTS(is_row_major, "fusedDistanceNN only supports row-major inputs"); - RAFT_EXPECTS(nearest_dist != nullptr, "Explicit fused 1-NN backends require nearest_dist"); if (backend == detail::Fused1nnBackend::Cutile) { if constexpr (detail::is_fused_1nn_cutile_data_v && std::is_same_v>) { @@ -359,12 +344,14 @@ void fusedDistanceNNMinReduce(IdxT* nearest_idx, RAFT_FAIL("Requested cuTile fused 1-NN backend does not support these data/norm types"); } RAFT_EXPECTS(backend == detail::Fused1nnBackend::Cutlass, "Unknown fused 1-NN backend"); + RAFT_EXPECTS(detail::can_launch_fused_1nn_backend(backend, x, y, m, n, k, metric), + "Requested CUTLASS fused 1-NN backend is unavailable for this input"); RAFT_EXPECTS(metric != cuvs::distance::DistanceType::InnerProduct, "CUTLASS fused 1-NN does not support InnerProduct"); RAFT_EXPECTS(std::is_same_v, "CUTLASS fused 1-NN requires matching norm types"); - RAFT_EXPECTS(cutlass_kvp_scratch != nullptr, - "CUTLASS fused 1-NN with explicit outputs requires KVP scratch storage"); - fusedDistanceNNMinReduce, IdxT>(cutlass_kvp_scratch, + RAFT_EXPECTS(cutlass_kvp_output != nullptr, + "CUTLASS fused 1-NN requires its native KVP output buffer"); + fusedDistanceNNMinReduce, IdxT>(cutlass_kvp_output, x, y, xn, @@ -379,11 +366,6 @@ void fusedDistanceNNMinReduce(IdxT* nearest_idx, metric, metric_arg, stream); - constexpr int threads = 256; - detail:: - unpack_fused_1nn_kvp<<((m + threads - 1) / threads), threads, 0, stream>>>( - nearest_idx, nearest_dist, cutlass_kvp_scratch, m); - RAFT_CUDA_TRY(cudaGetLastError()); } /** @} */ diff --git a/cpp/tests/neighbors/distance_nn.cu b/cpp/tests/neighbors/distance_nn.cu index f31f3ebacf..5dcf698f76 100644 --- a/cpp/tests/neighbors/distance_nn.cu +++ b/cpp/tests/neighbors/distance_nn.cu @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ @@ -18,6 +18,14 @@ namespace cuvs::neighbors { enum class ImplType { fused, unfused }; +template +void vector_compare_soa(raft::resources const& handle, + const raft::KeyValuePair* ref, + const IdxT* indices, + const AccT* distances, + IdxT n, + ComparisonSummary& summary); + template struct NNInputs { IdxT m; @@ -27,6 +35,8 @@ struct NNInputs { bool sqrt; uint64_t rng_seed; double tol; + cuvs::distance::detail::Fused1nnBackend backend = + cuvs::distance::detail::Fused1nnBackend::Cutlass; }; __global__ void fill_int8(int8_t* buff, int len, int seed_offset) @@ -50,13 +60,16 @@ class NNTest : public ::testing::TestWithParam> { k{params_.k}, metric{params_.metric}, sqrt{params_.sqrt}, + backend{params_.backend}, stream{raft::resource::get_cuda_stream(handle)}, x{raft::make_device_matrix(handle, m, k)}, y{raft::make_device_matrix(handle, n, k)}, x_norm{raft::make_device_vector(handle, m)}, y_norm{raft::make_device_vector(handle, n)}, out{raft::make_device_vector(handle, m)}, - ref_out{raft::make_device_vector(handle, m)} + ref_out{raft::make_device_vector(handle, m)}, + cutile_idx{raft::make_device_vector(handle, m)}, + cutile_dist{raft::make_device_vector(handle, m)} { } @@ -114,21 +127,32 @@ class NNTest : public ::testing::TestWithParam> { if constexpr (impl == ImplType::fused) { if constexpr (std::is_same_v) { - cuvs::distance::fusedDistanceNNMinReduce(out.data_handle(), - x.data_handle(), - y.data_handle(), - x_norm.data_handle(), - y_norm.data_handle(), - m, - n, - k, - (void*)workspace.data_handle(), - sqrt, - true, - true, - metric, - 0.0, - stream); + if (backend == cuvs::distance::detail::Fused1nnBackend::Cutile && + !cuvs::distance::detail::can_launch_fused_1nn_backend( + backend, x.data_handle(), y.data_handle(), m, n, k, metric)) { + GTEST_SKIP() << "cuTile is not available for this device/input"; + } + cuvs::distance::fusedDistanceNNMinReduce( + backend == cuvs::distance::detail::Fused1nnBackend::Cutile ? cutile_idx.data_handle() + : nullptr, + backend == cuvs::distance::detail::Fused1nnBackend::Cutile ? cutile_dist.data_handle() + : nullptr, + x.data_handle(), + y.data_handle(), + x_norm.data_handle(), + y_norm.data_handle(), + m, + n, + k, + (void*)workspace.data_handle(), + sqrt, + true, + true, + metric, + 0.0, + backend, + backend == cuvs::distance::detail::Fused1nnBackend::Cutlass ? out.data_handle() : nullptr, + stream); } else { static_assert(sizeof(DataT) == 0, "fusedDistanceNNMinReduce is not implemented for datatype other than float"); @@ -156,7 +180,20 @@ class NNTest : public ::testing::TestWithParam> { void compare() { - vector_compare(handle, ref_out.data_handle(), out.data_handle(), m, summary); + if constexpr (impl == ImplType::fused) { + if (backend == cuvs::distance::detail::Fused1nnBackend::Cutile) { + vector_compare_soa(handle, + ref_out.data_handle(), + cutile_idx.data_handle(), + cutile_dist.data_handle(), + m, + summary); + } else { + vector_compare(handle, ref_out.data_handle(), out.data_handle(), m, summary); + } + } else { + vector_compare(handle, ref_out.data_handle(), out.data_handle(), m, summary); + } ASSERT_TRUE(summary.max_diff < params_.tol) << summary; } @@ -170,12 +207,15 @@ class NNTest : public ::testing::TestWithParam> { IdxT k; DistanceType metric; bool sqrt; + cuvs::distance::detail::Fused1nnBackend backend; raft::device_matrix x; raft::device_matrix y; raft::device_vector x_norm; raft::device_vector y_norm; raft::device_vector out; raft::device_vector ref_out; + raft::device_vector cutile_idx; + raft::device_vector cutile_dist; size_t workspace_size; }; @@ -195,6 +235,18 @@ const std::vector> input_fp32 = { // {4096, 8192, 128, DistanceType::CosineExpanded, true, uint64_t(31415926), 0.1}, }; +template +const std::vector> input_fp32_fused = [] { + auto inputs = input_fp32; +#if CUVS_CUTILE_ENABLED + for (auto input : input_fp32) { + input.backend = cuvs::distance::detail::Fused1nnBackend::Cutile; + inputs.push_back(input); + } +#endif + return inputs; +}(); + // Test fused implementation with single-precision typedef NNTest NNTest_fp32_fused; TEST_P(NNTest_fp32_fused, test) @@ -203,7 +255,7 @@ TEST_P(NNTest_fp32_fused, test) this->compare(); } -INSTANTIATE_TEST_CASE_P(NNTest, NNTest_fp32_fused, ::testing::ValuesIn(input_fp32)); +INSTANTIATE_TEST_CASE_P(NNTest, NNTest_fp32_fused, ::testing::ValuesIn(input_fp32_fused)); // Test unfused implementation with single-precision typedef NNTest NNTest_fp32_unfused; diff --git a/cpp/tests/neighbors/distance_nn_helper.cuh b/cpp/tests/neighbors/distance_nn_helper.cuh index fda7b76573..84dbe53c6f 100644 --- a/cpp/tests/neighbors/distance_nn_helper.cuh +++ b/cpp/tests/neighbors/distance_nn_helper.cuh @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ @@ -207,4 +207,28 @@ void vector_compare( } } +template +void vector_compare_soa(raft::resources const& handle, + const raft::KeyValuePair* ref, + const IdxT* indices, + const AccT* distances, + IdxT n, + ComparisonSummary& summary) +{ + auto ref_h = raft::make_host_vector, IdxT>(n); + auto idx_h = raft::make_host_vector(n); + auto dist_h = raft::make_host_vector(n); + auto stream = raft::resource::get_cuda_stream(handle); + raft::copy(ref_h.data_handle(), ref, n, stream); + raft::copy(idx_h.data_handle(), indices, n, stream); + raft::copy(dist_h.data_handle(), distances, n, stream); + raft::resource::sync_stream(handle, stream); + summary.init(); + for (IdxT i = 0; i < n; ++i) { + const auto a = static_cast(ref_h(i).value); + const auto b = static_cast(dist_h(i)); + summary.update(std::abs(a - b), i, a, b, ref_h(i).key != idx_h(i)); + } +} + } // namespace cuvs::neighbors From f012e5c5303f706ee0d4277e0b903cc8224a69c4 Mon Sep 17 00:00:00 2001 From: divyegala Date: Thu, 3 Sep 2026 22:05:14 +0000 Subject: [PATCH 80/82] new prim --- cpp/src/distance/detail/fused_distance_nn.cuh | 13 +- cpp/src/distance/fused_distance_nn-inl.cuh | 120 ++++++++++++++---- cpp/tests/neighbors/distance_nn.cu | 18 ++- 3 files changed, 122 insertions(+), 29 deletions(-) diff --git a/cpp/src/distance/detail/fused_distance_nn.cuh b/cpp/src/distance/detail/fused_distance_nn.cuh index d1b0fe0dc3..78005eee51 100644 --- a/cpp/src/distance/detail/fused_distance_nn.cuh +++ b/cpp/src/distance/detail/fused_distance_nn.cuh @@ -33,6 +33,17 @@ namespace detail { enum class Fused1nnBackend : std::uint8_t { Cutile, Cutlass, + Unfused, +}; + +/** Tuning used only by the bounded-workspace unfused backend. */ +struct UnfusedTop1nnTuning { + std::size_t row_tile = 8192; + std::size_t candidate_tile = 8192; +}; + +struct Top1nnTuning { + UnfusedTop1nnTuning unfused{}; }; /** @@ -55,7 +66,7 @@ bool can_launch_fused_1nn_backend(Fused1nnBackend backend, } return false; } - return backend == Fused1nnBackend::Cutlass && + return (backend == Fused1nnBackend::Cutlass || backend == Fused1nnBackend::Unfused) && metric != cuvs::distance::DistanceType::InnerProduct && x != nullptr && y != nullptr && m > 0 && n > 0 && k > 0; } diff --git a/cpp/src/distance/fused_distance_nn-inl.cuh b/cpp/src/distance/fused_distance_nn-inl.cuh index f49482ff54..fac3ece709 100644 --- a/cpp/src/distance/fused_distance_nn-inl.cuh +++ b/cpp/src/distance/fused_distance_nn-inl.cuh @@ -10,14 +10,19 @@ #include "detail/fused_distance_nn.cuh" #include "fused_distance_nn_helpers.cuh" +#include "unfused_distance_nn.cuh" #include #include +#include #include +#include + #include #include +#include #include #include @@ -312,43 +317,106 @@ void fusedDistanceNNMinReduce(OutT* min, } template -void fusedDistanceNNMinReduce(IdxT* nearest_idx, - DataT* nearest_dist, - const DataT* x, - const DataT* y, - const NormT* xn, - const NormT* yn, - IdxT m, - IdxT n, - IdxT k, - void* workspace, - bool sqrt, - bool init_out_buffer, - bool is_row_major, - cuvs::distance::DistanceType metric, - float metric_arg, - detail::Fused1nnBackend backend, - raft::KeyValuePair* cutlass_kvp_output, - cudaStream_t stream) +void top_1_nn(raft::resources const& handle, + IdxT* nearest_idx, + DataT* nearest_dist, + const DataT* x, + const DataT* y, + const NormT* xn, + const NormT* yn, + IdxT m, + IdxT n, + IdxT k, + const detail::Top1nnTuning& tuning, + void* workspace, + std::size_t workspace_bytes, + bool sqrt, + bool init_out_buffer, + bool is_row_major, + cuvs::distance::DistanceType metric, + float metric_arg, + detail::Fused1nnBackend backend, + raft::KeyValuePair* cutlass_kvp_output, + cudaStream_t stream) { RAFT_EXPECTS(is_row_major, "fusedDistanceNN only supports row-major inputs"); if (backend == detail::Fused1nnBackend::Cutile) { if constexpr (detail::is_fused_1nn_cutile_data_v && std::is_same_v>) { - RAFT_EXPECTS( - detail::try_fused_1nn_tile( - nearest_idx, nearest_dist, x, y, xn, yn, m, n, k, metric, sqrt, workspace, stream), - "Requested cuTile fused 1-NN backend is unavailable for this input/device"); + const bool launched = detail::try_fused_1nn_tile( + nearest_idx, nearest_dist, x, y, xn, yn, m, n, k, metric, sqrt, workspace, stream); + RAFT_EXPECTS(launched, + "Requested cuTile fused 1-NN backend is unavailable for this input/device"); return; } RAFT_FAIL("Requested cuTile fused 1-NN backend does not support these data/norm types"); } - RAFT_EXPECTS(backend == detail::Fused1nnBackend::Cutlass, "Unknown fused 1-NN backend"); RAFT_EXPECTS(detail::can_launch_fused_1nn_backend(backend, x, y, m, n, k, metric), - "Requested CUTLASS fused 1-NN backend is unavailable for this input"); + "Requested fused 1-NN backend is unavailable for this input"); RAFT_EXPECTS(metric != cuvs::distance::DistanceType::InnerProduct, - "CUTLASS fused 1-NN does not support InnerProduct"); - RAFT_EXPECTS(std::is_same_v, "CUTLASS fused 1-NN requires matching norm types"); + "Only cuTile top_1_nn supports InnerProduct (as a maximum reduction)"); + constexpr bool matching_norm_type = std::is_same_v; + RAFT_EXPECTS(matching_norm_type, "CUTLASS and unfused top_1_nn require matching norm types"); + if (backend == detail::Fused1nnBackend::Unfused) { + RAFT_EXPECTS(cutlass_kvp_output != nullptr, + "Unfused top_1_nn requires its native KVP output buffer"); + RAFT_EXPECTS(tuning.unfused.row_tile > 0 && tuning.unfused.candidate_tile > 0, + "Unfused top_1_nn tile dimensions must be positive"); + + const auto max_row_tile = static_cast(m); + const auto max_candidate_tile = static_cast(n); + const auto row_tile = static_cast(std::min(tuning.unfused.row_tile, max_row_tile)); + const auto candidate_tile = + static_cast(std::min(tuning.unfused.candidate_tile, max_candidate_tile)); + const auto required_workspace_bytes = + static_cast(row_tile) * static_cast(candidate_tile) * sizeof(DataT); + RAFT_EXPECTS(workspace != nullptr && workspace_bytes >= required_workspace_bytes, + "Unfused top_1_nn workspace is smaller than its configured tile"); + + using KeyValueT = raft::KeyValuePair; + rmm::device_uvector candidate_min(candidate_tile < n ? row_tile : 0, stream); + for (IdxT row_offset = 0; row_offset < m; row_offset += row_tile) { + const auto rows = std::min(row_tile, static_cast(m - row_offset)); + auto output = + raft::make_device_vector_view(cutlass_kvp_output + row_offset, rows); + for (IdxT candidate_offset = 0; candidate_offset < n; candidate_offset += candidate_tile) { + const auto candidates = std::min(candidate_tile, static_cast(n - candidate_offset)); + auto* tile_output = candidate_offset == 0 ? output.data_handle() : candidate_min.data(); + unfusedDistanceNNMinReduce( + handle, + tile_output, + x + row_offset * k, + y + candidate_offset * k, + xn + row_offset, + yn + candidate_offset, + rows, + candidates, + k, + workspace, + sqrt, + candidate_offset != 0 || init_out_buffer, + is_row_major, + metric, + metric_arg, + stream); + if (candidate_offset != 0) { + auto candidate_output = + raft::make_device_vector_view(candidate_min.data(), rows); + raft::linalg::map( + handle, + output, + [candidate_offset] __device__(KeyValueT current, KeyValueT candidate) { + candidate.key += candidate_offset; + return candidate.value < current.value ? candidate : current; + }, + raft::make_const_mdspan(output), + candidate_output); + } + } + } + return; + } + RAFT_EXPECTS(backend == detail::Fused1nnBackend::Cutlass, "Unknown fused 1-NN backend"); RAFT_EXPECTS(cutlass_kvp_output != nullptr, "CUTLASS fused 1-NN requires its native KVP output buffer"); fusedDistanceNNMinReduce, IdxT>(cutlass_kvp_output, diff --git a/cpp/tests/neighbors/distance_nn.cu b/cpp/tests/neighbors/distance_nn.cu index 5dcf698f76..92ae59111c 100644 --- a/cpp/tests/neighbors/distance_nn.cu +++ b/cpp/tests/neighbors/distance_nn.cu @@ -37,6 +37,7 @@ struct NNInputs { double tol; cuvs::distance::detail::Fused1nnBackend backend = cuvs::distance::detail::Fused1nnBackend::Cutlass; + cuvs::distance::detail::Top1nnTuning tuning{}; }; __global__ void fill_int8(int8_t* buff, int len, int seed_offset) @@ -61,6 +62,7 @@ class NNTest : public ::testing::TestWithParam> { metric{params_.metric}, sqrt{params_.sqrt}, backend{params_.backend}, + tuning{params_.tuning}, stream{raft::resource::get_cuda_stream(handle)}, x{raft::make_device_matrix(handle, m, k)}, y{raft::make_device_matrix(handle, n, k)}, @@ -101,6 +103,10 @@ class NNTest : public ::testing::TestWithParam> { if constexpr (impl == ImplType::fused) { workspace_size = m * sizeof(IdxT); + if (backend == cuvs::distance::detail::Fused1nnBackend::Unfused) { + workspace_size = std::min(m, tuning.unfused.row_tile) * + std::min(n, tuning.unfused.candidate_tile) * sizeof(AccT); + } } else if constexpr (impl == ImplType::unfused) { workspace_size = m * n * sizeof(AccT); } @@ -132,7 +138,8 @@ class NNTest : public ::testing::TestWithParam> { backend, x.data_handle(), y.data_handle(), m, n, k, metric)) { GTEST_SKIP() << "cuTile is not available for this device/input"; } - cuvs::distance::fusedDistanceNNMinReduce( + cuvs::distance::top_1_nn( + handle, backend == cuvs::distance::detail::Fused1nnBackend::Cutile ? cutile_idx.data_handle() : nullptr, backend == cuvs::distance::detail::Fused1nnBackend::Cutile ? cutile_dist.data_handle() @@ -144,14 +151,16 @@ class NNTest : public ::testing::TestWithParam> { m, n, k, + tuning, (void*)workspace.data_handle(), + workspace_size, sqrt, true, true, metric, 0.0, backend, - backend == cuvs::distance::detail::Fused1nnBackend::Cutlass ? out.data_handle() : nullptr, + backend == cuvs::distance::detail::Fused1nnBackend::Cutile ? nullptr : out.data_handle(), stream); } else { static_assert(sizeof(DataT) == 0, @@ -208,6 +217,7 @@ class NNTest : public ::testing::TestWithParam> { DistanceType metric; bool sqrt; cuvs::distance::detail::Fused1nnBackend backend; + cuvs::distance::detail::Top1nnTuning tuning; raft::device_matrix x; raft::device_matrix y; raft::device_vector x_norm; @@ -238,6 +248,10 @@ const std::vector> input_fp32 = { template const std::vector> input_fp32_fused = [] { auto inputs = input_fp32; + for (auto input : input_fp32) { + input.backend = cuvs::distance::detail::Fused1nnBackend::Unfused; + inputs.push_back(input); + } #if CUVS_CUTILE_ENABLED for (auto input : input_fp32) { input.backend = cuvs::distance::detail::Fused1nnBackend::Cutile; From da6bd1d55247e009ba2f07e354bb58a221d50d86 Mon Sep 17 00:00:00 2001 From: divyegala Date: Thu, 3 Sep 2026 22:20:56 +0000 Subject: [PATCH 81/82] no duplicate insts --- cpp/CMakeLists.txt | 1 + cpp/src/distance/fused_distance_nn-inl.cuh | 273 +++++++++++++++------ cpp/src/distance/top_1_nn.cu | 62 +++++ cpp/src/distance/top_1_nn.cuh | 118 +++++++++ 4 files changed, 383 insertions(+), 71 deletions(-) create mode 100644 cpp/src/distance/top_1_nn.cu create mode 100644 cpp/src/distance/top_1_nn.cuh diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt index 7b642fa7d4..2489a35ba7 100644 --- a/cpp/CMakeLists.txt +++ b/cpp/CMakeLists.txt @@ -1434,6 +1434,7 @@ if(NOT BUILD_CPU_ONLY) src/distance/detail/kernels/kernel_matrices.cu ${pairwise_matrix_dispatch_inst_files} src/distance/distance.cu + src/distance/top_1_nn.cu src/distance/kde.cu src/distance/pairwise_distance.cu src/distance/sparse_distance.cu diff --git a/cpp/src/distance/fused_distance_nn-inl.cuh b/cpp/src/distance/fused_distance_nn-inl.cuh index fac3ece709..39ed82688e 100644 --- a/cpp/src/distance/fused_distance_nn-inl.cuh +++ b/cpp/src/distance/fused_distance_nn-inl.cuh @@ -10,6 +10,7 @@ #include "detail/fused_distance_nn.cuh" #include "fused_distance_nn_helpers.cuh" +#include "top_1_nn.cuh" #include "unfused_distance_nn.cuh" #include #include @@ -294,6 +295,54 @@ void fusedDistanceNNMinReduce(OutT* min, float metric_arg, cudaStream_t stream) { + if constexpr (std::is_same_v>) { + detail::Top1nnTuning tuning{}; + top_1_nn(nullptr, + nullptr, + x, + y, + xn, + yn, + m, + n, + k, + tuning, + workspace, + 0, + sqrt, + initOutBuffer, + isRowMajor, + metric, + metric_arg, + detail::Fused1nnBackend::Cutlass, + min, + stream); + return; + } else if constexpr (std::is_same_v) { + detail::Top1nnTuning tuning{}; + top_1_nn(nullptr, + min, + x, + y, + xn, + yn, + m, + n, + k, + tuning, + workspace, + 0, + sqrt, + initOutBuffer, + isRowMajor, + metric, + metric_arg, + detail::Fused1nnBackend::Cutlass, + nullptr, + stream); + return; + } + MinAndDistanceReduceOp redOp; KVPMinReduce pairRedOp; @@ -316,7 +365,7 @@ void fusedDistanceNNMinReduce(OutT* min, stream); } -template +template void top_1_nn(raft::resources const& handle, IdxT* nearest_idx, DataT* nearest_dist, @@ -357,83 +406,165 @@ void top_1_nn(raft::resources const& handle, "Only cuTile top_1_nn supports InnerProduct (as a maximum reduction)"); constexpr bool matching_norm_type = std::is_same_v; RAFT_EXPECTS(matching_norm_type, "CUTLASS and unfused top_1_nn require matching norm types"); - if (backend == detail::Fused1nnBackend::Unfused) { - RAFT_EXPECTS(cutlass_kvp_output != nullptr, - "Unfused top_1_nn requires its native KVP output buffer"); - RAFT_EXPECTS(tuning.unfused.row_tile > 0 && tuning.unfused.candidate_tile > 0, - "Unfused top_1_nn tile dimensions must be positive"); + if constexpr (matching_norm_type) { + if (backend == detail::Fused1nnBackend::Unfused) { + RAFT_EXPECTS(cutlass_kvp_output != nullptr, + "Unfused top_1_nn requires its native KVP output buffer"); + RAFT_EXPECTS(tuning.unfused.row_tile > 0 && tuning.unfused.candidate_tile > 0, + "Unfused top_1_nn tile dimensions must be positive"); - const auto max_row_tile = static_cast(m); - const auto max_candidate_tile = static_cast(n); - const auto row_tile = static_cast(std::min(tuning.unfused.row_tile, max_row_tile)); - const auto candidate_tile = - static_cast(std::min(tuning.unfused.candidate_tile, max_candidate_tile)); - const auto required_workspace_bytes = - static_cast(row_tile) * static_cast(candidate_tile) * sizeof(DataT); - RAFT_EXPECTS(workspace != nullptr && workspace_bytes >= required_workspace_bytes, - "Unfused top_1_nn workspace is smaller than its configured tile"); + const auto max_row_tile = static_cast(m); + const auto max_candidate_tile = static_cast(n); + const auto row_tile = static_cast(std::min(tuning.unfused.row_tile, max_row_tile)); + const auto candidate_tile = + static_cast(std::min(tuning.unfused.candidate_tile, max_candidate_tile)); + const auto required_workspace_bytes = static_cast(row_tile) * + static_cast(candidate_tile) * + sizeof(DataT); + RAFT_EXPECTS(workspace != nullptr && workspace_bytes >= required_workspace_bytes, + "Unfused top_1_nn workspace is smaller than its configured tile"); - using KeyValueT = raft::KeyValuePair; - rmm::device_uvector candidate_min(candidate_tile < n ? row_tile : 0, stream); - for (IdxT row_offset = 0; row_offset < m; row_offset += row_tile) { - const auto rows = std::min(row_tile, static_cast(m - row_offset)); - auto output = - raft::make_device_vector_view(cutlass_kvp_output + row_offset, rows); - for (IdxT candidate_offset = 0; candidate_offset < n; candidate_offset += candidate_tile) { - const auto candidates = std::min(candidate_tile, static_cast(n - candidate_offset)); - auto* tile_output = candidate_offset == 0 ? output.data_handle() : candidate_min.data(); - unfusedDistanceNNMinReduce( - handle, - tile_output, - x + row_offset * k, - y + candidate_offset * k, - xn + row_offset, - yn + candidate_offset, - rows, - candidates, - k, - workspace, - sqrt, - candidate_offset != 0 || init_out_buffer, - is_row_major, - metric, - metric_arg, - stream); - if (candidate_offset != 0) { - auto candidate_output = - raft::make_device_vector_view(candidate_min.data(), rows); - raft::linalg::map( + using KeyValueT = raft::KeyValuePair; + rmm::device_uvector candidate_min(candidate_tile < n ? row_tile : 0, stream); + for (IdxT row_offset = 0; row_offset < m; row_offset += row_tile) { + const auto rows = std::min(row_tile, static_cast(m - row_offset)); + auto output = + raft::make_device_vector_view(cutlass_kvp_output + row_offset, rows); + for (IdxT candidate_offset = 0; candidate_offset < n; candidate_offset += candidate_tile) { + const auto candidates = std::min(candidate_tile, static_cast(n - candidate_offset)); + auto* tile_output = candidate_offset == 0 ? output.data_handle() : candidate_min.data(); + unfusedDistanceNNMinReduce( handle, - output, - [candidate_offset] __device__(KeyValueT current, KeyValueT candidate) { - candidate.key += candidate_offset; - return candidate.value < current.value ? candidate : current; - }, - raft::make_const_mdspan(output), - candidate_output); + tile_output, + x + row_offset * k, + y + candidate_offset * k, + xn + row_offset, + yn + candidate_offset, + rows, + candidates, + k, + workspace, + sqrt, + candidate_offset != 0 || init_out_buffer, + is_row_major, + metric, + metric_arg, + stream); + if (candidate_offset != 0) { + auto candidate_output = + raft::make_device_vector_view(candidate_min.data(), rows); + raft::linalg::map( + handle, + output, + [candidate_offset] __device__(KeyValueT current, KeyValueT candidate) { + candidate.key += candidate_offset; + return candidate.value < current.value ? candidate : current; + }, + raft::make_const_mdspan(output), + candidate_output); + } } } + return; } - return; + RAFT_EXPECTS(backend == detail::Fused1nnBackend::Cutlass, "Unknown fused 1-NN backend"); + RAFT_EXPECTS(cutlass_kvp_output != nullptr, + "CUTLASS fused 1-NN requires its native KVP output buffer"); + MinAndDistanceReduceOp red_op; + KVPMinReduce pair_red_op; + fusedDistanceNN, IdxT>(cutlass_kvp_output, + x, + y, + xn, + yn, + m, + n, + k, + workspace, + red_op, + pair_red_op, + sqrt, + init_out_buffer, + is_row_major, + metric, + metric_arg, + stream); + } +} + +template +void top_1_nn(IdxT* nearest_idx, + DataT* nearest_dist, + const DataT* x, + const DataT* y, + const NormT* xn, + const NormT* yn, + IdxT m, + IdxT n, + IdxT k, + const detail::Top1nnTuning&, + void* workspace, + std::size_t, + bool sqrt, + bool init_out_buffer, + bool is_row_major, + cuvs::distance::DistanceType metric, + float metric_arg, + detail::Fused1nnBackend backend, + raft::KeyValuePair* cutlass_kvp_output, + cudaStream_t stream) +{ + RAFT_EXPECTS(backend == detail::Fused1nnBackend::Cutlass, + "The no-handle top_1_nn compatibility path supports CUTLASS only"); + RAFT_EXPECTS(is_row_major, "fusedDistanceNN only supports row-major inputs"); + RAFT_EXPECTS(detail::can_launch_fused_1nn_backend(backend, x, y, m, n, k, metric), + "Requested CUTLASS fused 1-NN backend is unavailable for this input"); + constexpr bool matching_norm_type = std::is_same_v; + RAFT_EXPECTS(matching_norm_type, "CUTLASS top_1_nn requires matching norm types"); + + if constexpr (std::is_same_v) { + MinAndDistanceReduceOp red_op; + KVPMinReduce pair_red_op; + if (cutlass_kvp_output != nullptr) { + fusedDistanceNN, IdxT>(cutlass_kvp_output, + x, + y, + xn, + yn, + m, + n, + k, + workspace, + red_op, + pair_red_op, + sqrt, + init_out_buffer, + is_row_major, + metric, + metric_arg, + stream); + return; + } + RAFT_EXPECTS(nearest_idx == nullptr && nearest_dist != nullptr, + "CUTLASS scalar top_1_nn requires a distance output and no index output"); + fusedDistanceNN(nearest_dist, + x, + y, + xn, + yn, + m, + n, + k, + workspace, + red_op, + pair_red_op, + sqrt, + init_out_buffer, + is_row_major, + metric, + metric_arg, + stream); } - RAFT_EXPECTS(backend == detail::Fused1nnBackend::Cutlass, "Unknown fused 1-NN backend"); - RAFT_EXPECTS(cutlass_kvp_output != nullptr, - "CUTLASS fused 1-NN requires its native KVP output buffer"); - fusedDistanceNNMinReduce, IdxT>(cutlass_kvp_output, - x, - y, - xn, - yn, - m, - n, - k, - workspace, - sqrt, - init_out_buffer, - is_row_major, - metric, - metric_arg, - stream); } /** @} */ diff --git a/cpp/src/distance/top_1_nn.cu b/cpp/src/distance/top_1_nn.cu new file mode 100644 index 0000000000..5e4daef89f --- /dev/null +++ b/cpp/src/distance/top_1_nn.cu @@ -0,0 +1,62 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include "fused_distance_nn.cuh" + +namespace cuvs::distance { + +#define CUVS_INSTANTIATE_TOP_1_NN(DataT, IdxT, NormT) \ + template CUVS_EXPORT void top_1_nn(raft::resources const&, \ + IdxT*, \ + DataT*, \ + const DataT*, \ + const DataT*, \ + const NormT*, \ + const NormT*, \ + IdxT, \ + IdxT, \ + IdxT, \ + const detail::Top1nnTuning&, \ + void*, \ + std::size_t, \ + bool, \ + bool, \ + bool, \ + DistanceType, \ + float, \ + detail::Fused1nnBackend, \ + raft::KeyValuePair*, \ + cudaStream_t); \ + template CUVS_EXPORT void top_1_nn(IdxT*, \ + DataT*, \ + const DataT*, \ + const DataT*, \ + const NormT*, \ + const NormT*, \ + IdxT, \ + IdxT, \ + IdxT, \ + const detail::Top1nnTuning&, \ + void*, \ + std::size_t, \ + bool, \ + bool, \ + bool, \ + DistanceType, \ + float, \ + detail::Fused1nnBackend, \ + raft::KeyValuePair*, \ + cudaStream_t) + +CUVS_INSTANTIATE_TOP_1_NN(float, int, float); +CUVS_INSTANTIATE_TOP_1_NN(float, int64_t, float); +CUVS_INSTANTIATE_TOP_1_NN(double, int, double); +CUVS_INSTANTIATE_TOP_1_NN(double, int64_t, double); +CUVS_INSTANTIATE_TOP_1_NN(half, int, float); +CUVS_INSTANTIATE_TOP_1_NN(half, int64_t, float); + +#undef CUVS_INSTANTIATE_TOP_1_NN + +} // namespace cuvs::distance diff --git a/cpp/src/distance/top_1_nn.cuh b/cpp/src/distance/top_1_nn.cuh new file mode 100644 index 0000000000..3c7bdbe3a3 --- /dev/null +++ b/cpp/src/distance/top_1_nn.cuh @@ -0,0 +1,118 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +#include "detail/fused_distance_nn.cuh" + +#include + +#include +#include + +#include + +namespace cuvs::distance { + +/** Dispatch 1-NN to a selected backend using backend-native output storage. */ +template +CUVS_EXPORT void top_1_nn(raft::resources const& handle, + IdxT* nearest_idx, + DataT* nearest_dist, + const DataT* x, + const DataT* y, + const NormT* xn, + const NormT* yn, + IdxT m, + IdxT n, + IdxT k, + const detail::Top1nnTuning& tuning, + void* workspace, + std::size_t workspace_bytes, + bool sqrt, + bool init_out_buffer, + bool is_row_major, + DistanceType metric, + float metric_arg, + detail::Fused1nnBackend backend, + raft::KeyValuePair* cutlass_kvp_output, + cudaStream_t stream); + +/** CUTLASS-only overload used by the no-handle legacy compatibility wrapper. */ +template +CUVS_EXPORT void top_1_nn(IdxT* nearest_idx, + DataT* nearest_dist, + const DataT* x, + const DataT* y, + const NormT* xn, + const NormT* yn, + IdxT m, + IdxT n, + IdxT k, + const detail::Top1nnTuning& tuning, + void* workspace, + std::size_t workspace_bytes, + bool sqrt, + bool init_out_buffer, + bool is_row_major, + DistanceType metric, + float metric_arg, + detail::Fused1nnBackend backend, + raft::KeyValuePair* cutlass_kvp_output, + cudaStream_t stream); + +#define CUVS_EXTERN_TOP_1_NN(DataT, IdxT, NormT) \ + extern template void top_1_nn(raft::resources const&, \ + IdxT*, \ + DataT*, \ + const DataT*, \ + const DataT*, \ + const NormT*, \ + const NormT*, \ + IdxT, \ + IdxT, \ + IdxT, \ + const detail::Top1nnTuning&, \ + void*, \ + std::size_t, \ + bool, \ + bool, \ + bool, \ + DistanceType, \ + float, \ + detail::Fused1nnBackend, \ + raft::KeyValuePair*, \ + cudaStream_t); \ + extern template void top_1_nn(IdxT*, \ + DataT*, \ + const DataT*, \ + const DataT*, \ + const NormT*, \ + const NormT*, \ + IdxT, \ + IdxT, \ + IdxT, \ + const detail::Top1nnTuning&, \ + void*, \ + std::size_t, \ + bool, \ + bool, \ + bool, \ + DistanceType, \ + float, \ + detail::Fused1nnBackend, \ + raft::KeyValuePair*, \ + cudaStream_t) + +CUVS_EXTERN_TOP_1_NN(float, int, float); +CUVS_EXTERN_TOP_1_NN(float, int64_t, float); +CUVS_EXTERN_TOP_1_NN(double, int, double); +CUVS_EXTERN_TOP_1_NN(double, int64_t, double); +CUVS_EXTERN_TOP_1_NN(half, int, float); +CUVS_EXTERN_TOP_1_NN(half, int64_t, float); + +#undef CUVS_EXTERN_TOP_1_NN + +} // namespace cuvs::distance From 97beec73ced454b027a53afe18fca178197d30fe Mon Sep 17 00:00:00 2001 From: divyegala Date: Thu, 3 Sep 2026 23:44:35 +0000 Subject: [PATCH 82/82] corrct fallback, remove legacy --- cpp/src/cluster/detail/kmeans_common.cuh | 9 +- .../detail/minClusterDistanceCompute.cu | 178 ++---------- cpp/src/distance/detail/fused_distance_nn.cuh | 16 +- .../cutile/fused_1nn_tile.cu | 5 +- .../fused_distance_nn/fused_cosine_nn.cuh | 23 +- .../detail/fused_distance_nn/fused_l2_nn.cuh | 25 +- .../fused_distance_nn/helper_structs.cuh | 103 +------ cpp/src/distance/fused_distance_nn-inl.cuh | 258 +++++++++++++----- cpp/tests/neighbors/distance_nn.cu | 11 +- 9 files changed, 282 insertions(+), 346 deletions(-) diff --git a/cpp/src/cluster/detail/kmeans_common.cuh b/cpp/src/cluster/detail/kmeans_common.cuh index fec28c526e..b7c286edee 100644 --- a/cpp/src/cluster/detail/kmeans_common.cuh +++ b/cpp/src/cluster/detail/kmeans_common.cuh @@ -131,10 +131,8 @@ FusedDistancePath use_legacy_fused(const raft::resources& handle, /** * @brief Selects the fused-distance assignment path for KMeans. * - * Float/half: cuTile when the build and device support it. Otherwise L2/L2Sqrt/Cosine may use - * legacy CUTLASS fused on Ampere/Hopper (large enough problems). InnerProduct without cuTile uses - * Unfused. Double never uses cuTile; keeps historical CUTLASS/unfused heuristics on pre-Blackwell - * GPUs. + * With CUDA 13, float/half use cuTile whenever the build and device support it. CUDA 12 and a + * failed CUDA 13 cuTile probe use the historical CUTLASS/unfused heuristic. */ template FusedDistancePath use_fused( @@ -142,6 +140,7 @@ FusedDistancePath use_fused( { (void)k; +#if CUDART_VERSION >= 13000 if constexpr (is_cutile_fused_data_type_v) { if constexpr (cuvs::detail::jit_lto::library_built_with_cutile()) { const bool dimensions_fit_i32 = n <= static_cast(std::numeric_limits::max()) && @@ -151,8 +150,8 @@ FusedDistancePath use_fused( return FusedDistancePath::Cutile; } } - return use_legacy_fused(handle, m, n, metric); } +#endif return use_legacy_fused(handle, m, n, metric); } diff --git a/cpp/src/cluster/detail/minClusterDistanceCompute.cu b/cpp/src/cluster/detail/minClusterDistanceCompute.cu index 32e1a434d7..4d9403a6d3 100644 --- a/cpp/src/cluster/detail/minClusterDistanceCompute.cu +++ b/cpp/src/cluster/detail/minClusterDistanceCompute.cu @@ -67,32 +67,6 @@ void compute_tf32_row_norms(raft::resources const& handle, } } -template -__global__ void unpack_kvp_to_soa(IndexT* nearest_idx, - DataT* nearest_dist, - const raft::KeyValuePair* kvp, - IndexT n) -{ - IndexT i = blockIdx.x * blockDim.x + threadIdx.x; - if (i < n) { - if (nearest_idx != nullptr) { nearest_idx[i] = kvp[i].key; } - if (nearest_dist != nullptr) { nearest_dist[i] = kvp[i].value; } - } -} - -template -void unpack_kvp(raft::resources const& handle, - IndexT* nearest_idx, - DataT* nearest_dist, - const raft::KeyValuePair* kvp, - IndexT n) -{ - auto stream = raft::resource::get_cuda_stream(handle); - int blks = static_cast((n + 255) / 256); - unpack_kvp_to_soa<<>>(nearest_idx, nearest_dist, kvp, n); - RAFT_CUDA_TRY(cudaGetLastError()); -} - } // namespace template @@ -124,9 +98,16 @@ Fused1nnRequirements get_fused_1nn_requirements( metric); Fused1nnRequirements requirements{}; - requirements.path = path; - requirements.sample_tile = getDataBatchSize(batch_samples, X.extent(0)); - requirements.centroid_tile = getCentroidsBatchSize(batch_centroids, centroids.extent(0)); + requirements.path = path; + const cuvs::distance::detail::Top1nnTuning default_tuning{}; + requirements.sample_tile = std::min( + getDataBatchSize(batch_samples, X.extent(0)), + static_cast(std::min(default_tuning.unfused.row_tile, + static_cast(std::numeric_limits::max())))); + requirements.centroid_tile = std::min( + getCentroidsBatchSize(batch_centroids, centroids.extent(0)), + static_cast(std::min(default_tuning.unfused.candidate_tile, + static_cast(std::numeric_limits::max())))); requirements.workspace_alignment = alignof(int); if (path == FusedDistancePath::Cutile) { @@ -201,10 +182,15 @@ void min_cluster_and_distance_compute_impl(raft::resources const& handle, if (workspace.size() < requirements.workspace_bytes) { workspace.resize(requirements.workspace_bytes, stream); } - RAFT_EXPECTS(!cutile_ready || native_kvp == nullptr, - "cuTile fused 1-NN requires separate index and distance outputs"); + if (cutile_ready) { + RAFT_EXPECTS(native_kvp == nullptr && nearest_idx != nullptr && nearest_dist != nullptr, + "cuTile fused 1-NN requires native separate index and distance outputs"); + } else { + RAFT_EXPECTS(native_kvp != nullptr && nearest_idx == nullptr && nearest_dist == nullptr, + "CUTLASS and unfused 1-NN require their native KVP output"); + } - if (uses_fused_distance_nn(fused_path)) { + if (is_l2_cos || cutile_ready) { const DataT* x_norm_ptr = nullptr; const DataT* centroids_norm_ptr = nullptr; if (is_l2_cos) { @@ -257,28 +243,10 @@ void min_cluster_and_distance_compute_impl(raft::resources const& handle, } } - raft::KeyValuePair* cutlass_kvp_scratch = nullptr; - rmm::device_uvector> temp_kvp(0, stream); const bool needs_index_workspace = cutile_ready && std::is_same_v; - if (!cutile_ready) { - if (native_kvp != nullptr) { - cutlass_kvp_scratch = native_kvp; - } else { - temp_kvp.resize(n_samples, stream); - cutlass_kvp_scratch = temp_kvp.data(); - } - if (workspace.size() < sizeof(int) * static_cast(n_samples)) { - workspace.resize(sizeof(int) * static_cast(n_samples), stream); - } - } else if (needs_index_workspace) { - const auto workspace_rows = - cuvs::distance::detail::fused_1nn_cutile_index_workspace_rows(n_samples); - if (workspace.size() < sizeof(int) * workspace_rows) { - workspace.resize(sizeof(int) * workspace_rows, stream); - } - } - cuvs::distance::detail::Top1nnTuning tuning{}; + tuning.unfused.row_tile = static_cast(requirements.sample_tile); + tuning.unfused.candidate_tile = static_cast(requirements.centroid_tile); cuvs::distance::top_1_nn( handle, cutile_ready ? nearest_idx : nullptr, @@ -299,99 +267,8 @@ void min_cluster_and_distance_compute_impl(raft::resources const& handle, metric, 0.0f, fused_path, - cutlass_kvp_scratch, + native_kvp, stream); - } else if (is_l2_cos) { - L2NormBuf_OR_DistBuf.resize(n_clusters, stream); - auto centroidsNorm = - raft::make_device_vector_view(L2NormBuf_OR_DistBuf.data(), n_clusters); - - if (metric == cuvs::distance::DistanceType::CosineExpanded) { - raft::linalg::norm( - handle, centroids, centroidsNorm, raft::sqrt_op{}); - } else { - raft::linalg::norm( - handle, centroids, centroidsNorm); - } - - auto centroidsNormConst = - raft::make_device_vector_view(L2NormBuf_OR_DistBuf.data(), n_clusters); - - auto dataBatchSize = getDataBatchSize(batch_samples, n_samples); - auto centroidsBatchSize = getCentroidsBatchSize(batch_centroids, n_clusters); - - // The unfused reduction indexes its distance matrix with IndexT. - dataBatchSize = - std::min(dataBatchSize, std::numeric_limits::max() / centroidsBatchSize); - - using KeyValueT = raft::KeyValuePair; - rmm::device_uvector temp_kvp(native_kvp == nullptr ? n_samples : 0, stream); - auto* kvp_output = native_kvp == nullptr ? temp_kvp.data() : native_kvp; - auto kvp_output_view = raft::make_device_vector_view(kvp_output, n_samples); - KeyValueT initial_value(0, std::numeric_limits::max()); - raft::matrix::fill(handle, kvp_output_view, initial_value); - - const bool tileCentroids = centroidsBatchSize < n_clusters; - const size_t distance_workspace_bytes = - sizeof(DataT) * static_cast(dataBatchSize) * static_cast(centroidsBatchSize); - const size_t batch_min_offset = raft::alignTo(distance_workspace_bytes, alignof(KeyValueT)); - const size_t required_workspace_bytes = - batch_min_offset + - (tileCentroids ? sizeof(KeyValueT) * static_cast(dataBatchSize) : 0); - if (workspace.size() < required_workspace_bytes) { - workspace.resize(required_workspace_bytes, stream); - } - auto* batch_min_storage = - tileCentroids ? reinterpret_cast(workspace.data() + batch_min_offset) : nullptr; - - for (IndexT dIdx = 0; dIdx < n_samples;) { - auto ns = std::min(dataBatchSize, n_samples - dIdx); - auto minClusterAndDistanceView = - raft::make_device_vector_view(kvp_output + dIdx, ns); - - for (IndexT cIdx = 0; cIdx < n_clusters;) { - auto nc = std::min(centroidsBatchSize, n_clusters - cIdx); - auto batchMin = tileCentroids ? batch_min_storage : minClusterAndDistanceView.data_handle(); - - cuvs::distance::unfusedDistanceNNMinReduce( - handle, - batchMin, - X.data_handle() + dIdx * n_features, - centroids.data_handle() + cIdx * n_features, - L2NormX.data_handle() + dIdx, - centroidsNormConst.data_handle() + cIdx, - ns, - nc, - n_features, - (void*)workspace.data(), - metric != cuvs::distance::DistanceType::L2Expanded, - tileCentroids, - true, - metric, - 0.0f, - stream); - - if (tileCentroids) { - // Convert tile-local centroid indices and merge the tile minima. - auto batchMinView = raft::make_device_vector_view(batchMin, ns); - raft::linalg::map( - handle, - minClusterAndDistanceView, - [cIdx] __device__(KeyValueT current, KeyValueT batch) { - batch.key += cIdx; - return batch.value < current.value ? batch : current; - }, - raft::make_const_mdspan(minClusterAndDistanceView), - batchMinView); - } - cIdx += nc; - } - dIdx += ns; - } - - if (native_kvp == nullptr) { - unpack_kvp(handle, nearest_idx, nearest_dist, kvp_output, n_samples); - } } else { auto dataBatchSize = getDataBatchSize(batch_samples, n_samples); auto centroidsBatchSize = getCentroidsBatchSize(batch_centroids, n_clusters); @@ -401,9 +278,8 @@ void min_cluster_and_distance_compute_impl(raft::resources const& handle, auto pairwiseDistance = raft::make_device_matrix_view( L2NormBuf_OR_DistBuf.data(), dataBatchSize, centroidsBatchSize); - using KeyValueT = raft::KeyValuePair; - rmm::device_uvector temp_kvp(native_kvp == nullptr ? n_samples : 0, stream); - auto* kvp_output = native_kvp == nullptr ? temp_kvp.data() : native_kvp; + using KeyValueT = raft::KeyValuePair; + auto* kvp_output = native_kvp; auto kvp_output_view = raft::make_device_vector_view(kvp_output, n_samples); KeyValueT initial_value(0, std::numeric_limits::max()); raft::matrix::fill(handle, kvp_output_view, initial_value); @@ -446,10 +322,6 @@ void min_cluster_and_distance_compute_impl(raft::resources const& handle, raft::identity_op{}); } } - - if (native_kvp == nullptr) { - unpack_kvp(handle, nearest_idx, nearest_dist, kvp_output, n_samples); - } } } @@ -469,7 +341,7 @@ void minClusterAndDistanceCompute(raft::resources const& handle, const DataT* cutile_x_norm) { RAFT_EXPECTS(requirements.output_layout == Fused1nnOutputLayout::Soa, - "resolved fused 1-NN plan requires KVP output"); + "resolved fused 1-NN plan requires separate output arrays"); if constexpr (is_cutile_fused_data_type_v) { if (requirements.path == FusedDistancePath::Cutile) { RAFT_EXPECTS(cuvs::distance::detail::can_launch_fused_1nn_tile(nearest_idx.data_handle(), @@ -515,7 +387,7 @@ void minClusterAndDistanceComputeKvp( const Fused1nnRequirements& requirements) { RAFT_EXPECTS(requirements.output_layout == Fused1nnOutputLayout::Kvp, - "resolved fused 1-NN plan requires separate output arrays"); + "resolved fused 1-NN plan requires KVP output"); min_cluster_and_distance_compute_impl(handle, X, centroids, diff --git a/cpp/src/distance/detail/fused_distance_nn.cuh b/cpp/src/distance/detail/fused_distance_nn.cuh index d799003b97..d8cfde6e0c 100644 --- a/cpp/src/distance/detail/fused_distance_nn.cuh +++ b/cpp/src/distance/detail/fused_distance_nn.cuh @@ -48,20 +48,24 @@ struct Top1nnTuning { UnfusedTop1nnTuning unfused{}; }; -/** Legacy policy used only by algorithm-level AUTO selection. */ +/** Select the pre-cuTile implementation used by algorithm-level AUTO dispatch. */ template constexpr Fused1nnBackend fused_1nn_legacy_backend(int cc_major, IdxT m, IdxT n, cuvs::distance::DistanceType metric) { - if (metric == cuvs::distance::DistanceType::InnerProduct) { return Fused1nnBackend::Unfused; } + const bool legacy_fused_metric = metric == cuvs::distance::DistanceType::L2Expanded || + metric == cuvs::distance::DistanceType::L2SqrtExpanded || + metric == cuvs::distance::DistanceType::CosineExpanded; + if (!legacy_fused_metric) { return Fused1nnBackend::Unfused; } if (cc_major <= 8 || (cc_major == 9 && (m >= 4096 || n >= 4096))) { return Fused1nnBackend::Cutlass; } return Fused1nnBackend::Unfused; } +/** Resolve AUTO before allocating backend-native output and workspace storage. */ template Fused1nnBackend resolve_fused_1nn_backend(const raft::resources& handle, const DataT* x, @@ -71,9 +75,11 @@ Fused1nnBackend resolve_fused_1nn_backend(const raft::resources& handle, IdxT k, cuvs::distance::DistanceType metric) { +#if CUDART_VERSION >= 13000 if constexpr (is_fused_1nn_cutile_data_v) { if (can_launch_fused_1nn_tile(x, y, m, n, k, metric)) { return Fused1nnBackend::Cutile; } } +#endif const auto prop = raft::resource::get_device_properties(handle); return fused_1nn_legacy_backend(prop.major, m, n, metric); } @@ -98,9 +104,11 @@ bool can_launch_fused_1nn_backend(Fused1nnBackend backend, } return false; } + const bool supported_metric = metric == cuvs::distance::DistanceType::L2Expanded || + metric == cuvs::distance::DistanceType::L2SqrtExpanded || + metric == cuvs::distance::DistanceType::CosineExpanded; return (backend == Fused1nnBackend::Cutlass || backend == Fused1nnBackend::Unfused) && - metric != cuvs::distance::DistanceType::InnerProduct && x != nullptr && y != nullptr && - m > 0 && n > 0 && k > 0; + supported_metric && x != nullptr && y != nullptr && m > 0 && n > 0 && k > 0; } template -void fusedCosineNN(IdxT* nearest_idx, - DataT* nearest_dist, +void fusedCosineNN(OutT* min, const DataT* x, const DataT* y, const DataT* xn, @@ -43,18 +42,14 @@ void fusedCosineNN(IdxT* nearest_idx, ReduceOpT redOp, KVPReduceOpT pairRedOp, bool sqrt, - OutT* cutlass_out, cudaStream_t stream) { + // The kernel policy is determined by fusedL2NN. typedef Policy P; dim3 blk(P::Nthreads); constexpr auto maxVal = std::numeric_limits::max(); - - if (cutlass_out == nullptr) { - initFused1nnOutput(nearest_idx, nearest_dist, m, maxVal, stream); - RAFT_CUDA_TRY(cudaGetLastError()); - } + typedef raft::KeyValuePair KVPair; namespace arch = raft::util::arch; using AccT = DataT; @@ -71,13 +66,18 @@ void fusedCosineNN(IdxT* nearest_idx, decltype(distance_op), decltype(fin_op)>; + // Get pointer to fp32 SIMT kernel to determine the runtime architecture of the + // current system. Other methods to determine the architecture (that do not + // require a pointer) can be error prone. See: + // https://github.com/NVIDIA/cub/issues/545 void* kernel_ptr = reinterpret_cast(kernel); auto runtime_arch = arch::kernel_virtual_arch(kernel_ptr); auto cutlass_range = arch::SM_range(arch::SM_80(), arch::SM_future()); if (cutlass_range.contains(runtime_arch)) { + // If device is SM_80 or later, use CUTLASS-based kernel. using cosineOp = cuvs::distance::detail::ops::cosine_cutlass_op; - using kvp_cg_min_reduce_op_ = kvp_cg_min_reduce_op; + using kvp_cg_min_reduce_op_ = kvp_cg_min_reduce_op; kvp_cg_min_reduce_op_ cg_reduce_op; cosineOp cosine_dist_op; @@ -102,7 +102,7 @@ void fusedCosineNN(IdxT* nearest_idx, lda, ldb, ldd, - cutlass_out, + min, workspace, cg_reduce_op, cosine_dist_op, @@ -110,11 +110,12 @@ void fusedCosineNN(IdxT* nearest_idx, pairRedOp, stream); } else { + // If device less than SM_80, use fp32 SIMT kernel. constexpr size_t shmemSize = P::SmemSize + ((P::Mblk + P::Nblk) * sizeof(DataT)); dim3 grid = launchConfigGenerator

(m, n, shmemSize, kernel); kernel<<>>( - cutlass_out, x, y, xn, yn, m, n, k, maxVal, workspace, redOp, pairRedOp, distance_op, fin_op); + min, x, y, xn, yn, m, n, k, maxVal, workspace, redOp, pairRedOp, distance_op, fin_op); RAFT_CUDA_TRY(cudaGetLastError()); } } diff --git a/cpp/src/distance/detail/fused_distance_nn/fused_l2_nn.cuh b/cpp/src/distance/detail/fused_distance_nn/fused_l2_nn.cuh index 8e945a508c..142bc57909 100644 --- a/cpp/src/distance/detail/fused_distance_nn/fused_l2_nn.cuh +++ b/cpp/src/distance/detail/fused_distance_nn/fused_l2_nn.cuh @@ -30,8 +30,7 @@ template -void fusedL2NNImpl(IdxT* nearest_idx, - DataT* nearest_dist, +void fusedL2NNImpl(OutT* min, const DataT* x, const DataT* y, const DataT* xn, @@ -44,16 +43,19 @@ void fusedL2NNImpl(IdxT* nearest_idx, KVPReduceOpT pairRedOp, bool sqrt, bool initOutBuffer, - OutT* cutlass_out, cudaStream_t stream) { + // The kernel policy is determined by fusedL2NN. typedef Policy P; dim3 blk(P::Nthreads); + auto nblks = raft::ceildiv(m, P::Nthreads); constexpr auto maxVal = std::numeric_limits::max(); + typedef raft::KeyValuePair KVPair; - if (initOutBuffer && cutlass_out == nullptr) { - initFused1nnOutput(nearest_idx, nearest_dist, m, maxVal, stream); + if (initOutBuffer) { + initKernel + <<>>(min, m, maxVal, redOp); RAFT_CUDA_TRY(cudaGetLastError()); } @@ -72,13 +74,19 @@ void fusedL2NNImpl(IdxT* nearest_idx, decltype(distance_op), decltype(fin_op)>; + // Get pointer to fp32 SIMT kernel to determine the best compute architecture + // out of all for which the kernel was compiled for that matches closely + // to the current device. Other methods to determine the architecture (that do not + // require a pointer) can be error prone. See: + // https://github.com/NVIDIA/cub/issues/545 void* kernel_ptr = reinterpret_cast(kernel); auto runtime_arch = arch::kernel_virtual_arch(kernel_ptr); auto cutlass_range = arch::SM_range(arch::SM_80(), arch::SM_future()); if (cutlass_range.contains(runtime_arch)) { + // If device is SM_80 or later, use CUTLASS-based kernel. using L2Op = cuvs::distance::detail::ops::l2_exp_cutlass_op; - using kvp_cg_min_reduce_op_ = kvp_cg_min_reduce_op; + using kvp_cg_min_reduce_op_ = kvp_cg_min_reduce_op; kvp_cg_min_reduce_op_ cg_reduce_op; L2Op L2_dist_op(sqrt); @@ -103,7 +111,7 @@ void fusedL2NNImpl(IdxT* nearest_idx, lda, ldb, ldd, - cutlass_out, + min, workspace, cg_reduce_op, L2_dist_op, @@ -111,11 +119,12 @@ void fusedL2NNImpl(IdxT* nearest_idx, pairRedOp, stream); } else { + // If device less than SM_80, use fp32 SIMT kernel. constexpr size_t shmemSize = P::SmemSize + ((P::Mblk + P::Nblk) * sizeof(DataT)); dim3 grid = launchConfigGenerator

(m, n, shmemSize, kernel); kernel<<>>( - cutlass_out, x, y, xn, yn, m, n, k, maxVal, workspace, redOp, pairRedOp, distance_op, fin_op); + min, x, y, xn, yn, m, n, k, maxVal, workspace, redOp, pairRedOp, distance_op, fin_op); RAFT_CUDA_TRY(cudaGetLastError()); } } diff --git a/cpp/src/distance/detail/fused_distance_nn/helper_structs.cuh b/cpp/src/distance/detail/fused_distance_nn/helper_structs.cuh index 3f4e839f35..19a358a27a 100644 --- a/cpp/src/distance/detail/fused_distance_nn/helper_structs.cuh +++ b/cpp/src/distance/detail/fused_distance_nn/helper_structs.cuh @@ -32,43 +32,20 @@ struct KVPMinReduceImpl { }; // KVPMinReduce -/** Writes fused 1-NN results to separate idx/dist arrays (dist may be null). */ template struct MinAndDistanceReduceOpImpl { typedef typename raft::KeyValuePair KVP; - LabelT* out_idx{nullptr}; - DataT* out_dist{nullptr}; - /** When set, CUTLASS/SIMT global merge writes here instead of SoA (caller unpacks). */ - KVP* out_kvp{nullptr}; - - DI void merge(LabelT rid, const KVP& other) const - { - if (out_kvp != nullptr) { - if (other.value < out_kvp[rid].value) { out_kvp[rid] = other; } - } else if (out_dist != nullptr) { - if (other.value < out_dist[rid]) { - out_dist[rid] = other.value; - if (out_idx != nullptr) { out_idx[rid] = other.key; } - } - } else if (out_idx != nullptr) { - // Idx-only output: dist must still be tracked for multi-tile merge; caller must provide - // out_dist or use a single-pass backend (cuTile). KMeans always passes both buffers. - out_idx[rid] = other.key; - } - } - DI void operator()(LabelT rid, KVP* out, const KVP& other) const { - if (out != nullptr && other.value < out->value) { + if (other.value < out->value) { out->key = other.key; out->value = other.value; } } - DI void operator()(LabelT rid, volatile KVP* out, const KVP& other) const { - if (out != nullptr && other.value < out->value) { + if (other.value < out->value) { out->key = other.key; out->value = other.value; } @@ -76,41 +53,35 @@ struct MinAndDistanceReduceOpImpl { DI void operator()(LabelT rid, DataT* out, const KVP& other) const { - if (out != nullptr && other.value < *out) { *out = other.value; } + if (other.value < *out) { *out = other.value; } } DI void operator()(LabelT rid, volatile DataT* out, const KVP& other) const { - if (out != nullptr && other.value < *out) { *out = other.value; } + if (other.value < *out) { *out = other.value; } } DI void operator()(LabelT rid, DataT* out, const DataT& other) const { - if (out != nullptr && other < *out) { *out = other; } + if (other < *out) { *out = other; } } DI void operator()(LabelT rid, volatile DataT* out, const DataT& other) const { - if (out != nullptr && other < *out) { *out = other; } - } - - DI void init(DataT* out, DataT maxVal) const - { - if (out != nullptr) { *out = maxVal; } + if (other < *out) { *out = other; } } + DI void init(DataT* out, DataT maxVal) const { *out = maxVal; } DI void init(KVP* out, DataT maxVal) const { out->value = maxVal; - out->key = LabelT(0); + out->key = 0xfffffff0; } - DI void init_key(DataT& /*out*/, LabelT /*idx*/) const {} - + DI void init_key(DataT& out, LabelT idx) const { return; } DI void init_key(KVP& out, LabelT idx) const { out.key = idx; } DI DataT get_value(KVP& out) const { return out.value; } - DI DataT get_value(DataT& out) const { return out; } }; @@ -125,53 +96,6 @@ struct MinReduceOpImpl { DI void init(DataT* out, DataT maxVal) { *out = maxVal; } }; -template -RAFT_KERNEL initFused1nnOutputKernel(IdxT* nearest_idx, DataT* nearest_dist, IdxT m, DataT maxVal) -{ - IdxT tid = IdxT(blockIdx.x) * blockDim.x + threadIdx.x; - if (tid < m) { - if (nearest_idx != nullptr) { nearest_idx[tid] = IdxT(0); } - if (nearest_dist != nullptr) { nearest_dist[tid] = maxVal; } - } -} - -template -void initFused1nnOutput( - IdxT* nearest_idx, DataT* nearest_dist, IdxT m, DataT maxVal, cudaStream_t stream) -{ - if (nearest_idx == nullptr && nearest_dist == nullptr) { return; } - auto blks = raft::ceildiv(m, 256); - initFused1nnOutputKernel - <<>>(nearest_idx, nearest_dist, m, maxVal); -} - -template -RAFT_KERNEL unpackFused1nnKvpToSoaKernel(IdxT* nearest_idx, - DataT* nearest_dist, - const raft::KeyValuePair* kvp, - IdxT n) -{ - IdxT i = IdxT(blockIdx.x) * blockDim.x + threadIdx.x; - if (i < n) { - if (nearest_idx != nullptr) { nearest_idx[i] = kvp[i].key; } - if (nearest_dist != nullptr) { nearest_dist[i] = kvp[i].value; } - } -} - -template -void unpackFused1nnKvpToSoa(IdxT* nearest_idx, - DataT* nearest_dist, - const raft::KeyValuePair* kvp, - IdxT m, - cudaStream_t stream) -{ - if (nearest_idx == nullptr && nearest_dist == nullptr) { return; } - auto blks = raft::ceildiv(m, 256); - unpackFused1nnKvpToSoaKernel - <<>>(nearest_idx, nearest_dist, kvp, m); - RAFT_CUDA_TRY(cudaGetLastError()); -} - template RAFT_KERNEL initKernel(OutT* min, IdxT m, DataT maxVal, ReduceOpT redOp) { @@ -182,13 +106,15 @@ RAFT_KERNEL initKernel(OutT* min, IdxT m, DataT maxVal, ReduceOpT redOp) template void initialize(OutT* min, IdxT m, DataT maxVal, ReduceOpT redOp, cudaStream_t stream) { - auto blks = raft::ceildiv(m, 256); - initKernel<<>>(min, m, maxVal, redOp); + auto blks = raft::ceildiv(m, 256); + initKernel<<>>(min, m, maxVal, redOp); } // cg::reduce functor for FusedDistanceNN used in its cutlass version // to output the min distance value & key(loc id). -template +// This is used in fused_distance_nn/predicated_tile_iterator_reduced_vec.h +// store_with_byte_offset() passed to cg::reduce() & select_reduce. +template struct kvp_cg_min_reduce_op { typedef typename raft::KeyValuePair KVP; @@ -196,6 +122,7 @@ struct kvp_cg_min_reduce_op { using AccTypeT = AccType; using IndexT = Index; + // functor signature. __host__ __device__ KVP operator()(KVP a, KVP b) const { return a.value < b.value ? a : b; } __host__ __device__ AccType operator()(AccType a, AccType b) const { return min(a, b); } diff --git a/cpp/src/distance/fused_distance_nn-inl.cuh b/cpp/src/distance/fused_distance_nn-inl.cuh index e4f70d2c23..f2a50a22fc 100644 --- a/cpp/src/distance/fused_distance_nn-inl.cuh +++ b/cpp/src/distance/fused_distance_nn-inl.cuh @@ -17,8 +17,6 @@ #include #include -#include - #include #include @@ -34,14 +32,52 @@ namespace distance { * \ingroup fused_l2_nn * @{ */ - -template -void fusedDistanceNN(IdxT* nearest_idx, - DataT* nearest_dist, +/** + * @brief Fused L2 distance and 1-nearest-neighbor computation in a single call. + * + * The benefits of such a call are 2-fold: 1) eliminate the need for an + * intermediate buffer to store the output of gemm 2) reduce the memory read + * traffic on this intermediate buffer, otherwise needed during the reduction + * phase for 1-NN. + * + * @tparam DataT data type + * @tparam OutT output type to either store 1-NN indices and their minimum + * distances or store only the min distances. Accordingly, one + * has to pass an appropriate `ReduceOpT` + * @tparam IdxT indexing arithmetic type + * @tparam ReduceOpT A struct to perform the final needed reduction operation + * and also to initialize the output array elements with the + * appropriate initial value needed for reduction. + * @tparam KVPReduceOpT A struct providing functions for key-value pair comparison. + * + * @param[out] min will contain the reduced output (Length = `m`) + * (on device) + * @param[in] x first matrix. Row major. Dim = `m x k`. + * (on device). + * @param[in] y second matrix. Row major. Dim = `n x k`. + * (on device). + * @param[in] xn L2 squared norm of `x`. Length = `m`. (on device). + * @param[in] yn L2 squared norm of `y`. Length = `n`. (on device) + * @param[in] m gemm m + * @param[in] n gemm n + * @param[in] k gemm k + * @param[in] workspace temp workspace. Size = sizeof(int)*m. (on device) + * @param[in] redOp reduction operator in the epilogue + * @param[in] pairRedOp reduction operation on key value pairs + * @param[in] sqrt Whether the output `minDist` should contain L2-sqrt + * @param[in] initOutBuffer whether to initialize the output buffer before the + * main kernel launch + * @param[in] isRowMajor whether the input/output is row or column major. + * @param[in] metric Distance metric to be used (supports L2, cosine) + * @param[in] metric_arg power argument for distances like Minkowski (not supported for now) + * @param[in] stream cuda stream + */ +template +void fusedDistanceNN(OutT* min, const DataT* x, const DataT* y, - const NormT* xn, - const NormT* yn, + const DataT* xn, + const DataT* yn, IdxT m, IdxT n, IdxT k, @@ -53,10 +89,12 @@ void fusedDistanceNN(IdxT* nearest_idx, bool isRowMajor, cuvs::distance::DistanceType metric, float metric_arg, - raft::KeyValuePair* cutlass_kvp_scratch, cudaStream_t stream) { ASSERT(isRowMajor, "fusedDistanceNN only supports row major inputs"); + // When k is smaller than 32, the Policy4x4 results in redundant calculations + // as it uses tiles that have k=32. Therefore, use a "skinny" policy instead + // that uses tiles with a smaller value of k. bool is_skinny = k < 32; size_t bytes = sizeof(DataT) * k; @@ -66,11 +104,10 @@ void fusedDistanceNN(IdxT* nearest_idx, if (is_skinny) { detail::fusedDistanceNNImpl< DataT, - NormT, + OutT, IdxT, typename raft::linalg::Policy4x4Skinny::Policy, - ReduceOpT>(nearest_idx, - nearest_dist, + ReduceOpT>(min, x, y, xn, @@ -86,16 +123,14 @@ void fusedDistanceNN(IdxT* nearest_idx, isRowMajor, metric, metric_arg, - cutlass_kvp_scratch, stream); } else { detail::fusedDistanceNNImpl< DataT, - NormT, + OutT, IdxT, typename raft::linalg::Policy4x4::Policy, - ReduceOpT>(nearest_idx, - nearest_dist, + ReduceOpT>(min, x, y, xn, @@ -111,18 +146,16 @@ void fusedDistanceNN(IdxT* nearest_idx, isRowMajor, metric, metric_arg, - cutlass_kvp_scratch, stream); } } else if (8 % sizeof(DataT) == 0 && bytes % 8 == 0 && px % 8 == 0 && py % 8 == 0) { if (is_skinny) { detail::fusedDistanceNNImpl< DataT, - NormT, + OutT, IdxT, typename raft::linalg::Policy4x4Skinny::Policy, - ReduceOpT>(nearest_idx, - nearest_dist, + ReduceOpT>(min, x, y, xn, @@ -138,16 +171,14 @@ void fusedDistanceNN(IdxT* nearest_idx, isRowMajor, metric, metric_arg, - cutlass_kvp_scratch, stream); } else { detail::fusedDistanceNNImpl< DataT, - NormT, + OutT, IdxT, typename raft::linalg::Policy4x4::Policy, - ReduceOpT>(nearest_idx, - nearest_dist, + ReduceOpT>(min, x, y, xn, @@ -163,17 +194,15 @@ void fusedDistanceNN(IdxT* nearest_idx, isRowMajor, metric, metric_arg, - cutlass_kvp_scratch, stream); } } else { if (is_skinny) { detail::fusedDistanceNNImpl::Policy, - ReduceOpT>(nearest_idx, - nearest_dist, + ReduceOpT>(min, x, y, xn, @@ -189,15 +218,13 @@ void fusedDistanceNN(IdxT* nearest_idx, isRowMajor, metric, metric_arg, - cutlass_kvp_scratch, stream); } else { detail::fusedDistanceNNImpl::Policy, - ReduceOpT>(nearest_idx, - nearest_dist, + ReduceOpT>(min, x, y, xn, @@ -213,28 +240,48 @@ void fusedDistanceNN(IdxT* nearest_idx, isRowMajor, metric, metric_arg, - cutlass_kvp_scratch, stream); } } } /** - * @brief Fused GEMM + 1-NN minimum reduction. + * @brief Wrapper around fusedDistanceNN with minimum reduction operators. + * + * fusedDistanceNN cannot be compiled in the distance library due to the lambda + * operators, so this wrapper covers the most common case (minimum). * - * @param[out] nearest_idx Nearest neighbor index per row, length `m` (optional). - * @param[out] nearest_dist Minimum distance per row, length `m` (optional, may be null). - * @param[in] cutlass_kvp_scratch Temp KVP buffer, length `m`, for index-bearing CUTLASS/SIMT - * output. Distance-only output may pass null and write directly - * to `nearest_dist`. + * @tparam DataT data type + * @tparam OutT output type to either store 1-NN indices and their minimum + * distances (e.g. raft::KeyValuePair) or store only the min + * distances. + * @tparam IdxT indexing arithmetic type + * @param[out] min will contain the reduced output (Length = `m`) + * (on device) + * @param[in] x first matrix. Row major. Dim = `m x k`. + * (on device). + * @param[in] y second matrix. Row major. Dim = `n x k`. + * (on device). + * @param[in] xn L2 squared norm of `x`. Length = `m`. (on device). + * @param[in] yn L2 squared norm of `y`. Length = `n`. (on device) + * @param[in] m gemm m + * @param[in] n gemm n + * @param[in] k gemm k + * @param[in] workspace temp workspace. Size = sizeof(int)*m. (on device) + * @param[in] sqrt Whether the output `minDist` should contain L2-sqrt + * @param[in] initOutBuffer whether to initialize the output buffer before the + * main kernel launch + * @param[in] isRowMajor whether the input/output is row or column major. + * @param[in] metric Distance metric to be used (supports L2, cosine) + * @param[in] metric_arg power argument for distances like Minkowski (not supported for now) + * @param[in] stream cuda stream */ -template -void fusedDistanceNNMinReduce(IdxT* nearest_idx, - DataT* nearest_dist, +template +void fusedDistanceNNMinReduce(OutT* min, const DataT* x, const DataT* y, - const NormT* xn, - const NormT* yn, + const DataT* xn, + const DataT* yn, IdxT m, IdxT n, IdxT k, @@ -244,33 +291,76 @@ void fusedDistanceNNMinReduce(IdxT* nearest_idx, bool isRowMajor, cuvs::distance::DistanceType metric, float metric_arg, - raft::KeyValuePair* cutlass_kvp_scratch, cudaStream_t stream) { + if constexpr (std::is_same_v>) { + detail::Top1nnTuning tuning{}; + top_1_nn(nullptr, + nullptr, + x, + y, + xn, + yn, + m, + n, + k, + tuning, + workspace, + 0, + sqrt, + initOutBuffer, + isRowMajor, + metric, + metric_arg, + detail::Fused1nnBackend::Cutlass, + min, + stream); + return; + } else if constexpr (std::is_same_v) { + detail::Top1nnTuning tuning{}; + top_1_nn(nullptr, + min, + x, + y, + xn, + yn, + m, + n, + k, + tuning, + workspace, + 0, + sqrt, + initOutBuffer, + isRowMajor, + metric, + metric_arg, + detail::Fused1nnBackend::Cutlass, + nullptr, + stream); + return; + } + MinAndDistanceReduceOp redOp; - redOp.out_idx = nearest_idx; - redOp.out_dist = nearest_dist; KVPMinReduce pairRedOp; - fusedDistanceNN(nearest_idx, - nearest_dist, - x, - y, - xn, - yn, - m, - n, - k, - workspace, - redOp, - pairRedOp, - sqrt, - initOutBuffer, - isRowMajor, - metric, - metric_arg, - cutlass_kvp_scratch, - stream); + fusedDistanceNN(min, + x, + y, + xn, + yn, + m, + n, + k, + workspace, + redOp, + pairRedOp, + sqrt, + initOutBuffer, + isRowMajor, + metric, + metric_arg, + stream); } template @@ -297,21 +387,34 @@ void top_1_nn(raft::resources const& handle, cudaStream_t stream) { RAFT_EXPECTS(is_row_major, "fusedDistanceNN only supports row-major inputs"); + RAFT_EXPECTS(backend != detail::Fused1nnBackend::Auto, + "top_1_nn requires AUTO to be resolved before dispatch"); if (backend == detail::Fused1nnBackend::Cutile) { + RAFT_EXPECTS(cutlass_kvp_output == nullptr, + "cuTile top_1_nn requires its native separate output arrays"); if constexpr (detail::is_fused_1nn_cutile_data_v && std::is_same_v>) { + if constexpr (std::is_same_v) { + const auto required_workspace_bytes = + sizeof(int) * detail::fused_1nn_cutile_index_workspace_rows(m); + RAFT_EXPECTS(workspace != nullptr && workspace_bytes >= required_workspace_bytes, + "cuTile top_1_nn workspace is too small for int64 output batching"); + } const bool launched = detail::try_fused_1nn_tile( nearest_idx, nearest_dist, x, y, xn, yn, m, n, k, metric, sqrt, workspace, stream); RAFT_EXPECTS(launched, "Requested cuTile fused 1-NN backend is unavailable for this input/device"); return; + } else { + RAFT_FAIL("Requested cuTile fused 1-NN backend does not support these data/norm types"); } - RAFT_FAIL("Requested cuTile fused 1-NN backend does not support these data/norm types"); } RAFT_EXPECTS(detail::can_launch_fused_1nn_backend(backend, x, y, m, n, k, metric), "Requested fused 1-NN backend is unavailable for this input"); RAFT_EXPECTS(metric != cuvs::distance::DistanceType::InnerProduct, "Only cuTile top_1_nn supports InnerProduct (as a maximum reduction)"); + RAFT_EXPECTS(nearest_idx == nullptr && nearest_dist == nullptr, + "CUTLASS and unfused top_1_nn require their native KVP output buffer"); constexpr bool matching_norm_type = std::is_same_v; RAFT_EXPECTS(matching_norm_type, "CUTLASS and unfused top_1_nn require matching norm types"); if constexpr (matching_norm_type) { @@ -326,21 +429,31 @@ void top_1_nn(raft::resources const& handle, const auto row_tile = static_cast(std::min(tuning.unfused.row_tile, max_row_tile)); const auto candidate_tile = static_cast(std::min(tuning.unfused.candidate_tile, max_candidate_tile)); - const auto required_workspace_bytes = static_cast(row_tile) * + const auto distance_workspace_bytes = static_cast(row_tile) * static_cast(candidate_tile) * sizeof(DataT); + const auto candidate_min_offset = + raft::alignTo(distance_workspace_bytes, alignof(raft::KeyValuePair)); + const auto candidate_min_bytes = + candidate_tile < n + ? static_cast(row_tile) * sizeof(raft::KeyValuePair) + : std::size_t{0}; + const auto required_workspace_bytes = candidate_min_offset + candidate_min_bytes; RAFT_EXPECTS(workspace != nullptr && workspace_bytes >= required_workspace_bytes, "Unfused top_1_nn workspace is smaller than its configured tile"); using KeyValueT = raft::KeyValuePair; - rmm::device_uvector candidate_min(candidate_tile < n ? row_tile : 0, stream); + auto* candidate_min = + candidate_tile < n + ? reinterpret_cast(static_cast(workspace) + candidate_min_offset) + : nullptr; for (IdxT row_offset = 0; row_offset < m; row_offset += row_tile) { const auto rows = std::min(row_tile, static_cast(m - row_offset)); auto output = raft::make_device_vector_view(cutlass_kvp_output + row_offset, rows); for (IdxT candidate_offset = 0; candidate_offset < n; candidate_offset += candidate_tile) { const auto candidates = std::min(candidate_tile, static_cast(n - candidate_offset)); - auto* tile_output = candidate_offset == 0 ? output.data_handle() : candidate_min.data(); + auto* tile_output = candidate_offset == 0 ? output.data_handle() : candidate_min; unfusedDistanceNNMinReduce( handle, tile_output, @@ -360,7 +473,7 @@ void top_1_nn(raft::resources const& handle, stream); if (candidate_offset != 0) { auto candidate_output = - raft::make_device_vector_view(candidate_min.data(), rows); + raft::make_device_vector_view(candidate_min, rows); raft::linalg::map( handle, output, @@ -378,6 +491,9 @@ void top_1_nn(raft::resources const& handle, RAFT_EXPECTS(backend == detail::Fused1nnBackend::Cutlass, "Unknown fused 1-NN backend"); RAFT_EXPECTS(cutlass_kvp_output != nullptr, "CUTLASS fused 1-NN requires its native KVP output buffer"); + RAFT_EXPECTS( + workspace != nullptr && workspace_bytes >= sizeof(int) * static_cast(m), + "CUTLASS top_1_nn workspace is too small"); MinAndDistanceReduceOp red_op; KVPMinReduce pair_red_op; fusedDistanceNN, IdxT>(cutlass_kvp_output, diff --git a/cpp/tests/neighbors/distance_nn.cu b/cpp/tests/neighbors/distance_nn.cu index 92ae59111c..42d59b6af4 100644 --- a/cpp/tests/neighbors/distance_nn.cu +++ b/cpp/tests/neighbors/distance_nn.cu @@ -102,10 +102,15 @@ class NNTest : public ::testing::TestWithParam> { } if constexpr (impl == ImplType::fused) { - workspace_size = m * sizeof(IdxT); + workspace_size = m * sizeof(int); if (backend == cuvs::distance::detail::Fused1nnBackend::Unfused) { - workspace_size = std::min(m, tuning.unfused.row_tile) * - std::min(n, tuning.unfused.candidate_tile) * sizeof(AccT); + using KvpT = raft::KeyValuePair; + const auto row_tile = std::min(m, tuning.unfused.row_tile); + const auto candidate_tile = std::min(n, tuning.unfused.candidate_tile); + const auto distance_bytes = row_tile * candidate_tile * sizeof(AccT); + workspace_size = + raft::alignTo(distance_bytes, alignof(KvpT)) + + (candidate_tile < static_cast(n) ? row_tile * sizeof(KvpT) : std::size_t{0}); } } else if constexpr (impl == ImplType::unfused) { workspace_size = m * n * sizeof(AccT);