diff --git a/c/src/cluster/kmeans.cpp b/c/src/cluster/kmeans.cpp index 9d371c182b..06e21fcda7 100644 --- a/c/src/cluster/kmeans.cpp +++ b/c/src/cluster/kmeans.cpp @@ -4,6 +4,7 @@ */ #include +#include #include @@ -41,6 +42,56 @@ cuvs::cluster::kmeans::balanced_params convert_balanced_params(const cuvsKMeansP return kmeans_params; } +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 cuvsKMeansParams& params, @@ -54,8 +105,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"); @@ -66,24 +119,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, @@ -92,7 +145,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; @@ -109,10 +162,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; } @@ -164,13 +227,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 { @@ -196,7 +268,7 @@ void _predict(cuvsResources_t res, } } -template +template void _cluster_cost(cuvsResources_t res, DLManagedTensor* X_tensor, DLManagedTensor* centroids_tensor, @@ -259,10 +331,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, @@ -282,10 +368,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, @@ -301,10 +401,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/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/conda/environments/all_cuda-133_arch-aarch64.yaml b/conda/environments/all_cuda-133_arch-aarch64.yaml index 545dcf5326..aa50303efa 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 7b06a9c20e..47d0226234 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 29aebfd06a..9004819b23 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 @@ -392,6 +397,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 4d51f17836..a6d11fb19d 100644 --- a/cpp/CMakeLists.txt +++ b/cpp/CMakeLists.txt @@ -1172,6 +1172,42 @@ 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_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(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" @@ -1363,6 +1399,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 @@ -1459,6 +1496,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( @@ -1482,7 +1521,9 @@ if(NOT BUILD_CPU_ONLY) target_compile_definitions( cuvs_objs PRIVATE $<$:CUVS_BUILD_CAGRA_HNSWLIB> - $<$:CUVS_BUILD_MG_ALGOS> $<$:NVTX_ENABLED> + $<$:CUVS_BUILD_MG_ALGOS> + $<$:NVTX_ENABLED> + CUVS_CUTILE_ENABLED=${CUVS_CUTILE_ENABLED} ) target_link_libraries( @@ -1502,6 +1543,7 @@ if(NOT BUILD_CPU_ONLY) "$" INTERFACE "$" PRIVATE "${CMAKE_CURRENT_SOURCE_DIR}/src" "${CMAKE_CURRENT_BINARY_DIR}/src" + "${cutile_fused_1nn_generated_dir}" ) # Endian detection diff --git a/cpp/cmake/modules/compute_matrix_product.cmake b/cpp/cmake/modules/compute_matrix_product.cmake index 82a34f9242..b5f3d06f86 100644 --- a/cpp/cmake/modules/compute_matrix_product.cmake +++ b/cpp/cmake/modules/compute_matrix_product.cmake @@ -1,12 +1,25 @@ # ============================================================================= # 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() + 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 + ) +endfunction() + function(compute_matrix_product output_var) set(options) set(one_value MATRIX_JSON_FILE MATRIX_JSON_STRING) @@ -14,19 +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}" # + 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..1d81963be8 --- /dev/null +++ b/cpp/cmake/modules/generate_cutile_kernels.cmake @@ -0,0 +1,372 @@ +# ============================================================================= +# 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() + +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) + 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 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() + + _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() + file( + WRITE "${header_path}" + "/* + * Generated from ${matrix_json_file} by generate_cutile_kernels.cmake — do not edit. + */ +#pragma once + +#include + +namespace cuvs::distance::detail { + +${_tile_aliases} + +} // 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 + 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") + + 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() + + 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) + 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}") + + 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") + + # 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..ae46b523e1 --- /dev/null +++ b/cpp/include/cuvs/detail/jit_lto/cutile_module.hpp @@ -0,0 +1,94 @@ +/* + * 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 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..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/include/cuvs/detail/jit_lto/tileir_compat.hpp b/cpp/include/cuvs/detail/jit_lto/tileir_compat.hpp new file mode 100644 index 0000000000..bde7dab302 --- /dev/null +++ b/cpp/include/cuvs/detail/jit_lto/tileir_compat.hpp @@ -0,0 +1,108 @@ +/* + * 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; } + 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/cluster/detail/kmeans.cuh b/cpp/src/cluster/detail/kmeans.cuh index e3ffb4a439..abb6b1a648 100644 --- a/cpp/src/cluster/detail/kmeans.cuh +++ b/cpp/src/cluster/detail/kmeans.cuh @@ -684,8 +684,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, device_buffer_samples); + auto nearest_idx = raft::make_device_vector(handle, device_buffer_samples); + auto nearest_dist = raft::make_device_vector(handle, device_buffer_samples); 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); @@ -855,8 +855,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) { @@ -885,7 +887,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, @@ -1073,8 +1076,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 @@ -1084,50 +1086,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 ac0430b430..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 @@ -99,55 +100,115 @@ 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>()); + 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()); + } break; } case cuvs::distance::DistanceType::InnerProduct: { - // TODO: pass buffer - rmm::device_uvector distances(n_rows * n_clusters, stream, mr); + 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); - 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: { @@ -186,14 +247,24 @@ 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; + case distance::DistanceType::L2SqrtExpanded: + case distance::DistanceType::InnerProduct: { + switch (use_fused(handle, n_rows, n_clusters, dim, metric)) { + 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); + 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. @@ -212,6 +283,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/cluster/detail/kmeans_common.cuh b/cpp/src/cluster/detail/kmeans_common.cuh index ab3ef0a05a..3ece6b0e19 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,29 +58,70 @@ namespace cuvs::cluster::kmeans::detail { +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. + * @brief Selects the fused-distance assignment path for KMeans. * - * 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. + * 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. */ 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; 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()) { + 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; + } + } + 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 false; + + 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 @@ -339,8 +381,19 @@ 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) { + 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 or L2SqrtExpanded distance, have %i", + RAFT_FAIL("kmeans requires L2Expanded, L2SqrtExpanded, or L2Unexpanded distance, have %i", static_cast(metric)); } } @@ -378,33 +431,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) @@ -463,22 +515,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, @@ -486,12 +532,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, @@ -676,7 +718,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 @@ -685,29 +728,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, @@ -715,36 +759,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 dbe2c23039..c1f51026eb 100644 --- a/cpp/src/cluster/detail/kmeans_mg.cuh +++ b/cpp/src/cluster/detail/kmeans_mg.cuh @@ -214,8 +214,8 @@ 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 minClusterAndDistance = - raft::make_device_vector, IndexT>(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, data_on_device ? IndexT{0} : alloc_batch_size); rmm::device_uvector L2NormBuf_OR_DistBuf(0, stream); @@ -448,9 +448,10 @@ void mnmg_fit( L2NormBatch_const = raft::make_const_mdspan(norm_slice); } - 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, @@ -460,7 +461,8 @@ void mnmg_fit( metric, params.batch_samples, params.batch_centroids, - minClusterAndDistance_view, + nearest_idx_view, + nearest_dist_view, L2NormBatch_const, L2NormBuf_OR_DistBuf, workspace, @@ -555,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/detail/minClusterDistanceCompute.cu b/cpp/src/cluster/detail/minClusterDistanceCompute.cu index ee3cc3cdfd..b230b8af16 100644 --- a/cpp/src/cluster/detail/minClusterDistanceCompute.cu +++ b/cpp/src/cluster/detail/minClusterDistanceCompute.cu @@ -7,35 +7,190 @@ #include "../../distance/unfused_distance_nn.cuh" #include "kmeans_common.cuh" +#include #include +#include +#include + 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 { + +__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, + 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); - const bool can_use_fused_path = metric == cuvs::distance::DistanceType::L2Expanded || - metric == cuvs::distance::DistanceType::L2SqrtExpanded || - metric == cuvs::distance::DistanceType::CosineExpanded; - - if (can_use_fused_path) { + 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)) { + 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 (!(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); + } + } + + 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(), + x_norm_ptr, + centroids_norm_ptr, + 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); @@ -48,149 +203,110 @@ void minClusterAndDistanceCompute( handle, centroids, centroidsNorm); } - raft::KeyValuePair initial_value(0, std::numeric_limits::max()); - raft::matrix::fill(handle, minClusterAndDistance, initial_value); - - const bool use_fused_path = - use_fused(handle, n_samples, n_clusters, n_features); - - if (use_fused_path) { - workspace.resize((sizeof(int)) * n_samples, stream); - - cuvs::distance::fusedDistanceNNMinReduce, IndexT>( - minClusterAndDistance.data_handle(), - X.data_handle(), - centroids.data_handle(), - L2NormX.data_handle(), - centroidsNorm.data_handle(), - n_samples, - n_clusters, - n_features, - (void*)workspace.data(), - metric != cuvs::distance::DistanceType::L2Expanded, - false, - true, - metric, - 0.0f, - stream); - } else { - auto dataBatchSize = getDataBatchSize(batch_samples, n_samples); - auto centroidsBatchSize = getCentroidsBatchSize(batch_centroids, n_clusters); + auto centroidsNormConst = + raft::make_device_vector_view(L2NormBuf_OR_DistBuf.data(), 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); + auto dataBatchSize = getDataBatchSize(batch_samples, n_samples); + auto centroidsBatchSize = getCentroidsBatchSize(batch_centroids, n_clusters); - using KeyValueT = raft::KeyValuePair; - const bool tileCentroids = centroidsBatchSize < n_clusters; - rmm::device_uvector batchMinClusterAndDistance(tileCentroids ? dataBatchSize : 0, - stream); + // 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); - for (IndexT dIdx = 0; dIdx < n_samples;) { - auto ns = std::min(dataBatchSize, n_samples - dIdx); - auto minClusterAndDistanceView = raft::make_device_vector_view( - minClusterAndDistance.data_handle() + dIdx, ns); + 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); - for (IndexT cIdx = 0; cIdx < n_clusters;) { - auto nc = std::min(centroidsBatchSize, n_clusters - cIdx); - auto batchMin = tileCentroids ? batchMinClusterAndDistance.data() - : minClusterAndDistanceView.data_handle(); + const bool tileCentroids = centroidsBatchSize < n_clusters; + rmm::device_uvector batchMinClusterAndDistance(tileCentroids ? dataBatchSize : 0, + stream); - cuvs::distance::unfusedDistanceNNMinReduce( + 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, - batchMin, - X.data_handle() + dIdx * n_features, - centroids.data_handle() + cIdx * n_features, - L2NormX.data_handle() + dIdx, - centroidsNorm.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; + minClusterAndDistanceView, + [cIdx] __device__(KeyValueT current, KeyValueT batch) { + batch.key += cIdx; + return batch.value < current.value ? batch : current; + }, + raft::make_const_mdspan(minClusterAndDistanceView), + batchMinView); } - dIdx += ns; + 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); - // 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), @@ -207,20 +323,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) @@ -247,49 +366,87 @@ 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) { - 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 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)) { + 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(); } - workspace.resize(sizeof(int) * n_samples, stream); + 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); + } + } - cuvs::distance::fusedDistanceNNMinReduce( + 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(), + x_norm_ptr, + centroids_norm_ptr, 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); @@ -300,8 +457,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); @@ -311,7 +466,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 f6e2c7d819..d635426229 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 { @@ -331,6 +335,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 +343,61 @@ 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"); + + 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); - 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, @@ -435,22 +478,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/detail/jit_lto/TileAlgorithmPlanner.cpp b/cpp/src/detail/jit_lto/TileAlgorithmPlanner.cpp new file mode 100644 index 0000000000..d3fee7e9dc --- /dev/null +++ b/cpp/src/detail/jit_lto/TileAlgorithmPlanner.cpp @@ -0,0 +1,125 @@ +/* + * 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; + } + + 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); + 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(); } + 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::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 f9dbd968ec..dbc87f468d 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" @@ -13,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 @@ -27,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, @@ -49,36 +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; - - dim3 blk(P::Nthreads); - auto nblks = raft::ceildiv(m, P::Nthreads); + typedef raft::KeyValuePair KVP; constexpr auto maxVal = std::numeric_limits::max(); - typedef raft::KeyValuePair KVPair; - RAFT_CUDA_TRY(cudaMemsetAsync(workspace, 0, sizeof(int) * m, stream)); + 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, workspace, stream)) { + return; + } + } + } + + RAFT_EXPECTS(cutlass_kvp_scratch != nullptr, "CUTLASS fused 1-NN requires a scratch KVP buffer"); + 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: - // 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); + 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 new file mode 100644 index 0000000000..2ddc717b60 --- /dev/null +++ b/cpp/src/distance/detail/fused_distance_nn/cutile/export_fused_1nn.py @@ -0,0 +1,271 @@ +# 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.""" + +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, +) + +# 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, + index_abbrev, + kernel_symbol, + make_kernel, +) + +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 _data_abbrev(data_type: str) -> 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 _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, + 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=index_dtype, + stride_lower_bound_incl=(0, None), + alias_groups=(), + 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=(1, 1), + base_addr_divisible_by=16, + ) + + +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=index_dtype, + 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=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, + 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", + ) + 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, + 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", + 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, + ) + signature = _kernel_signature( + data_type, + metric, + index_type, + tile_m, + tile_n, + tile_k, + 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( + "--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, + 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, + 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..b870f70597 --- /dev/null +++ b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_cutile_matrix.json @@ -0,0 +1,272 @@ +[ + { + "_abi": [ + { + "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": 64, + "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_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": "strict", + "abi_abbrev": "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", + "tile_m": 128, + "tile_n": 256, + "tile_k": 16 + } + ], + "_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_100", + "cc_major": 10, + "cc_minor": 0, + "arch_tag": "cutile_arch_10_0" + } + ] + }, + { + "_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": "tileir_bytecode", + "artifact_ext": "tilebc", + "artifact_basename": "@data_type@_@index_abbrev@_@abi_abbrev@", + "register": "tileir", + "gpu_code": "sm_80", + "bytecode_version": "13.1" + } + ] + }, + { + "_abi": [ + { + "matrix_layout": "strict", + "abi_abbrev": "strict", + "abi_tag": "cutile_abi_strict", + "tile_m": 64, + "tile_n": 256, + "tile_k": 32 + }, + { + "matrix_layout": "relaxed", + "abi_abbrev": "relaxed", + "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": [ + { + "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_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_kernel.py b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_kernel.py new file mode 100644 index 0000000000..2e14cc0d97 --- /dev/null +++ b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_kernel.py @@ -0,0 +1,185 @@ +# 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", +): + """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) + + @ct.kernel(occupancy=ct.ByTarget(sm_120=2)) + 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: + # 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, :] + + 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 + 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..029fb87746 --- /dev/null +++ b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_planner.hpp @@ -0,0 +1,97 @@ +/* + * 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; + using Tile90 = std::conditional_t; + using Tile100 = std::conditional_t; + using Tile120 = 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>(); + } + + 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>(); + } +}; + +} // 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..3e1a46cdf1 --- /dev/null +++ b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_tile.cu @@ -0,0 +1,288 @@ +/* + * 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 { + +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 cuvs::detail::jit_lto::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); + 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 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_tile.hpp b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_tile.hpp new file mode 100644 index 0000000000..7cdbabd411 --- /dev/null +++ b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_tile.hpp @@ -0,0 +1,65 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +#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; + +#if CUVS_CUTILE_ENABLED +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); +#else +template +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) +{ + return false; +} +#endif + +} // 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..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,11 +1,11 @@ /* - * 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/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..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-2025, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2021-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * 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..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-2024, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2021-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * 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..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 @@ -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/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 */ diff --git a/cpp/src/distance/fused_distance_nn-inl.cuh b/cpp/src/distance/fused_distance_nn-inl.cuh index 3fa80a9b60..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 */ @@ -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/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 66a13a6cba..d28cbb6773 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. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ @@ -179,7 +179,12 @@ __device__ __forceinline__ void interleaved_scan_impl(const uint32_t query_smem_ } if constexpr (kManageLocalTopK) { - queue.add(val, sample_offset + vec_id); + // 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); } else { if (vec_id < list_length) distances[sample_offset + vec_id] = val; } @@ -202,6 +207,16 @@ __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_rabitq/gpu_index/quantizer_gpu.cu b/cpp/src/neighbors/ivf_rabitq/gpu_index/quantizer_gpu.cu index af0876acca..c7701ae6ef 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. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ @@ -519,6 +519,8 @@ 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); @@ -635,6 +637,8 @@ 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(), @@ -712,6 +716,8 @@ 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); @@ -755,6 +761,8 @@ 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_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]; diff --git a/cpp/tests/neighbors/distance_nn.cu b/cpp/tests/neighbors/distance_nn.cu index f31f3ebacf..376a60a36b 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 */ @@ -14,6 +14,8 @@ #include #include +#include + namespace cuvs::neighbors { enum class ImplType { fused, unfused }; @@ -42,7 +44,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 +57,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 +96,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 +109,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 +158,16 @@ 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); + // 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); + } ASSERT_TRUE(summary.max_diff < params_.tol) << summary; } @@ -174,8 +185,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; }; @@ -195,6 +208,19 @@ 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; + inputs.insert( + inputs.begin() + 6, + NNInputs{512, 1024, 64, DistanceType::InnerProduct, false, uint64_t(31415926), 0.1}); + 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}); + return inputs; +}(); + // Test fused implementation with single-precision typedef NNTest NNTest_fp32_fused; TEST_P(NNTest_fp32_fused, test) @@ -203,7 +229,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..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 */ @@ -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); @@ -174,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) diff --git a/dependencies.yaml b/dependencies.yaml index a58e00cb58..33dc4fbf1e 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 @@ -421,6 +425,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 6e8789349c..795d91fbc8 100644 --- a/python/libcuvs/pyproject.toml +++ b/python/libcuvs/pyproject.toml @@ -81,6 +81,8 @@ regex = "(?P.*)" build-backend = "scikit_build_core.build" requires = [ "cmake>=4.0", + "cuda-tile", + "cuda-toolkit[tileiras]==13.*", "libraft==26.10.*,>=0.0.0a0", "librmm==26.10.*,>=0.0.0a0", "ninja",