diff --git a/ci/validate_wheel.sh b/ci/validate_wheel.sh index bf43e9e170..8ed1b8cd8b 100755 --- a/ci/validate_wheel.sh +++ b/ci/validate_wheel.sh @@ -21,7 +21,7 @@ PYDISTCHECK_ARGS=( if [[ "${package_dir}" == "python/libcuvs" ]]; then if [[ "${RAPIDS_CUDA_MAJOR}" == "12" ]]; then PYDISTCHECK_ARGS+=( - --max-allowed-size-compressed '360Mi' + --max-allowed-size-compressed '365Mi' ) else PYDISTCHECK_ARGS+=( diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt index 66f139c090..faab3482f4 100644 --- a/cpp/CMakeLists.txt +++ b/cpp/CMakeLists.txt @@ -1184,8 +1184,49 @@ if(NOT BUILD_CPU_ONLY) 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}) + set(fused_1nn_cutile_dir + "${CMAKE_CURRENT_SOURCE_DIR}/src/distance/detail/fused_distance_nn/cutile" + ) + set(cutile_fused_1nn_generated_dir + "${CMAKE_CURRENT_BINARY_DIR}/generated_kernels/distance/fused_1nn/cutile" + ) + set(cutile_fused_1nn_tiles "${cutile_fused_1nn_generated_dir}/fused_1nn_cutile_tiles.hpp") + generate_cutile_kernels( + cutile_fused_1nn_files + KERNEL_DIR "${fused_1nn_cutile_dir}" + KERNEL_BASENAME "fused_1nn" + KERNEL_PYTHON "fused_1nn_kernel.py" + EXPORT_SCRIPT "export_fused_1nn.py" + OUTPUT_DIRECTORY "${cutile_fused_1nn_generated_dir}" + MATRIX_JSON_FILE "${fused_1nn_cutile_dir}/fused_1nn_cutile_matrix.json" + FRAGMENT_TAG_FORMAT_CUBIN + "cuvs::distance::detail::fragment_tag_fused_1nn_cubin, cuvs::distance::detail::@abi_tag@, cuvs::detail::jit_lto::@arch_tag@>" + FRAGMENT_TAG_HEADER_FILES + "" + "" "" + ) + if(CUVS_CUTILE_ENABLED) + generate_string_matrix( + cutile_fused_1nn_tile_aliases + MATRIX_JSON_FILE + "${fused_1nn_cutile_dir}/fused_1nn_cutile_matrix.json" + ITEM_FORMAT + "using fused_1nn_matrix_tile_@data_abbrev@_@arch_tag@_@abi_abbrev@ = cutile_tile_config<@tile_m@, @tile_n@, @tile_k@>@semicolon@" + GLUE + "\n" + ) + set(CUTILE_TILE_METADATA_INCLUDE + "#include " + ) + set(CUTILE_TILE_METADATA_NAMESPACE "cuvs::distance::detail") + set(CUTILE_TILE_METADATA_ALIASES "${cutile_fused_1nn_tile_aliases}") + configure_file( + "${CMAKE_CURRENT_SOURCE_DIR}/cmake/modules/cutile_tile_metadata.hpp.in" + "${cutile_fused_1nn_tiles}" @ONLY + ) + list(APPEND cutile_fused_1nn_files "${cutile_fused_1nn_tiles}") + endif() # Note that this matrix contains an `arch_includes` placeholder, since we don't currently have a # way to do an item-wise transform on a list after computing the matrix product and before # configuring the file @@ -1403,6 +1444,7 @@ if(NOT BUILD_CPU_ONLY) src/distance/detail/kernels/kernel_matrices.cu ${pairwise_matrix_dispatch_inst_files} src/distance/distance.cu + src/distance/top_1_nn.cu src/distance/kde.cu src/distance/pairwise_distance.cu src/distance/sparse_distance.cu @@ -1496,9 +1538,16 @@ if(NOT BUILD_CPU_ONLY) src/stats/trustworthiness_score.cu ${CUVS_MG_ALGOS} ${jit_lto_files} - ${cutile_smoke_files} ) + if(CUVS_CUTILE_ENABLED) + target_sources( + cuvs_objs + PRIVATE ${cutile_fused_1nn_files} + src/distance/detail/fused_distance_nn/cutile/fused_1nn_tile.cu ${cutile_smoke_files} + ) + endif() + set_target_properties( cuvs_objs PROPERTIES CXX_STANDARD 20 @@ -1519,8 +1568,10 @@ if(NOT BUILD_CPU_ONLY) target_compile_definitions( cuvs_objs - PRIVATE $<$:CUVS_BUILD_CAGRA_HNSWLIB> - $<$:CUVS_BUILD_MG_ALGOS> $<$:NVTX_ENABLED> + PRIVATE CUVS_CUTILE_ENABLED=${CUVS_CUTILE_ENABLED} + $<$:CUVS_BUILD_CAGRA_HNSWLIB> + $<$:CUVS_BUILD_MG_ALGOS> + $<$:NVTX_ENABLED> ) target_link_libraries( @@ -1541,7 +1592,7 @@ if(NOT BUILD_CPU_ONLY) "$" INTERFACE "$" PRIVATE "${CMAKE_CURRENT_SOURCE_DIR}/src" "${CMAKE_CURRENT_BINARY_DIR}/src" - "${cutile_smoke_generated_dir}" + "${cutile_fused_1nn_generated_dir}" "${cutile_smoke_generated_dir}" ) # Endian detection diff --git a/cpp/cmake/config.json b/cpp/cmake/config.json index 9dfcda3259..5f50b5d5bb 100644 --- a/cpp/cmake/config.json +++ b/cpp/cmake/config.json @@ -36,7 +36,6 @@ "EXPORT_SCRIPT": 1, "OUTPUT_DIRECTORY": 1, "FRAGMENT_TAG_FORMAT_CUBIN": 1, - "FRAGMENT_TAG_FORMAT_TILEIR": "?", "FRAGMENT_TAG_HEADER_FILES": "*", "MATRIX_JSON_ENTRY": 1 } @@ -53,7 +52,6 @@ "OUTPUT_DIRECTORY": 1, "MATRIX_JSON_FILE": 1, "FRAGMENT_TAG_FORMAT_CUBIN": 1, - "FRAGMENT_TAG_FORMAT_TILEIR": "?", "FRAGMENT_TAG_HEADER_FILES": "*" } }, diff --git a/cpp/cmake/modules/cutile_tile_metadata.hpp.in b/cpp/cmake/modules/cutile_tile_metadata.hpp.in new file mode 100644 index 0000000000..3aeed43403 --- /dev/null +++ b/cpp/cmake/modules/cutile_tile_metadata.hpp.in @@ -0,0 +1,16 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +// clang-format off +@CUTILE_TILE_METADATA_INCLUDE@ + +namespace @CUTILE_TILE_METADATA_NAMESPACE@ { + +@CUTILE_TILE_METADATA_ALIASES@ + +} // namespace +// clang-format on diff --git a/cpp/cmake/modules/generate_cutile_kernels.cmake b/cpp/cmake/modules/generate_cutile_kernels.cmake index 51d58e3eca..d5803953f5 100644 --- a/cpp/cmake/modules/generate_cutile_kernels.cmake +++ b/cpp/cmake/modules/generate_cutile_kernels.cmake @@ -109,9 +109,6 @@ function(_cutile_make_python_args output_var) --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() @@ -127,7 +124,7 @@ 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 + FRAGMENT_TAG_FORMAT_CUBIN MATRIX_JSON_ENTRY ) set(multi_value FRAGMENT_TAG_HEADER_FILES) cmake_parse_arguments(_CUTILE "${options}" "${one_value}" "${multi_value}" ${ARGN}) @@ -138,19 +135,12 @@ function(process_cutile_matrix_entry source_list_var) 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() + if(NOT register STREQUAL "cubin") message(FATAL_ERROR "Unknown cuTile register kind '${register}'") endif() + string(CONFIGURE "${_CUTILE_FRAGMENT_TAG_FORMAT_CUBIN}" fragment_tag @ONLY) + set(bin2c_symbol embedded_cubin) + set(fragment_entry_type "cuvs::detail::jit_lto::StaticCubinFragmentEntry") _cutile_fragment_tag_header_files(fragment_tag_header_files ${_CUTILE_FRAGMENT_TAG_HEADER_FILES}) @@ -214,7 +204,7 @@ 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 + MATRIX_JSON_FILE FRAGMENT_TAG_FORMAT_CUBIN ) set(multi_value FRAGMENT_TAG_HEADER_FILES) cmake_parse_arguments(_CUTILE "${options}" "${one_value}" "${multi_value}" ${ARGN}) @@ -259,7 +249,6 @@ function(generate_cutile_kernels source_list_var) 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}" ) diff --git a/cpp/include/cuvs/detail/jit_lto/CutileFragmentEntry.hpp b/cpp/include/cuvs/detail/jit_lto/CutileFragmentEntry.hpp index 724662c9dd..632b953e98 100644 --- a/cpp/include/cuvs/detail/jit_lto/CutileFragmentEntry.hpp +++ b/cpp/include/cuvs/detail/jit_lto/CutileFragmentEntry.hpp @@ -68,52 +68,4 @@ struct StaticCubinFragmentEntry final : CubinFragmentEntry { 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 index fb6025fd64..6fec6e69d1 100644 --- a/cpp/include/cuvs/detail/jit_lto/TileAlgorithmPlanner.hpp +++ b/cpp/include/cuvs/detail/jit_lto/TileAlgorithmPlanner.hpp @@ -29,7 +29,7 @@ struct TileLauncherCache { std::unordered_set unavailable_launchers; }; -/** Loads prebuilt cubins or TileIR bytecode directly through the CUDA library API. */ +/** Loads compatible prebuilt cubins directly through the CUDA library API. */ struct TileAlgorithmPlanner { TileAlgorithmPlanner(std::string entrypoint, TileLauncherCache& launcher_cache) : entrypoint_(std::move(entrypoint)), launcher_cache_(launcher_cache) @@ -49,18 +49,11 @@ struct TileAlgorithmPlanner { 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. */ + /** Tile geometry from the cubin 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 CutileRuntimeCapabilities* capabilities) const; diff --git a/cpp/include/cuvs/detail/jit_lto/cutile_arch_tags.hpp b/cpp/include/cuvs/detail/jit_lto/cutile_arch_tags.hpp index 2b378dac78..ca1b62d2d9 100644 --- a/cpp/include/cuvs/detail/jit_lto/cutile_arch_tags.hpp +++ b/cpp/include/cuvs/detail/jit_lto/cutile_arch_tags.hpp @@ -41,8 +41,8 @@ struct cutile_arch_12_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; + return (cc_major == 8 && (cc_minor == 0 || cc_minor == 6)) || (cc_major == 9 && cc_minor == 0) || + (cc_major == 10 && cc_minor == 0) || (cc_major == 12 && cc_minor == 0); } #else @@ -51,4 +51,10 @@ inline bool is_embedded_cubin_arch(int, int) { return false; } #endif +/** SM89 is binary compatible with the embedded SM86 SASS. */ +inline constexpr bool can_use_sm86_compat_cubin(int cc_major, int cc_minor) +{ + return cc_major == 8 && cc_minor == 9; +} + } // namespace cuvs::detail::jit_lto diff --git a/cpp/include/cuvs/detail/jit_lto/tileir_compat.hpp b/cpp/include/cuvs/detail/jit_lto/cutile_compat.hpp similarity index 61% rename from cpp/include/cuvs/detail/jit_lto/tileir_compat.hpp rename to cpp/include/cuvs/detail/jit_lto/cutile_compat.hpp index 8e5e599069..5954533a61 100644 --- a/cpp/include/cuvs/detail/jit_lto/tileir_compat.hpp +++ b/cpp/include/cuvs/detail/jit_lto/cutile_compat.hpp @@ -9,7 +9,6 @@ #define CUVS_CUTILE_ENABLED 0 #endif -#include #include #include @@ -21,7 +20,6 @@ struct CutileRuntimeCapabilities { int device; int cc_major; int cc_minor; - int driver_version; }; inline bool query_current_cutile_runtime_capabilities(CutileRuntimeCapabilities& capabilities) @@ -37,12 +35,9 @@ inline bool query_current_cutile_runtime_capabilities(CutileRuntimeCapabilities& capabilities.device) != cudaSuccess) { return false; } - return cudaDriverGetVersion(&capabilities.driver_version) == cudaSuccess; + return true; } -/** 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; @@ -68,35 +63,24 @@ 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) +/** True when this build embeds compatible SASS for the device. */ +inline bool has_compatible_embedded_cubin_for_arch(int cc_major, int cc_minor) { - return driver_version >= kMinTileIrJitDriverVersion; + return is_embedded_cubin_arch(cc_major, cc_minor) || + (can_use_sm86_compat_cubin(cc_major, cc_minor) && is_embedded_cubin_arch(8, 6)); } /** * 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. + * is CUDA 13+, and an exact cubin exists, except that SM89 may use the embedded SM86 cubin. */ #if CUVS_CUTILE_ENABLED -inline bool cutile_launch_available_for_arch(int cc_major, int cc_minor, int driver_version) +inline bool cutile_launch_available_for_arch(int cc_major, int cc_minor) { - if (!runtime_cuda13_or_newer()) { return false; } - // The exported fused-1NN kernels require Ampere-or-newer tensor-core semantics, and the current - // integration is validated only through the SM12 family. - if (cc_major < 8 || cc_major > 12) { return false; } - if (has_embedded_cubin_for_arch(cc_major, cc_minor)) { return true; } - return tileir_fallback_available(driver_version); + return runtime_cuda13_or_newer() && has_compatible_embedded_cubin_for_arch(cc_major, cc_minor); } #else -inline constexpr bool cutile_launch_available_for_arch(int, int, int) { return false; } +inline constexpr bool cutile_launch_available_for_arch(int, int) { return false; } #endif #if CUVS_CUTILE_ENABLED @@ -104,8 +88,7 @@ inline bool cutile_launch_available_on_current_device() { CutileRuntimeCapabilities capabilities{}; if (!query_current_cutile_runtime_capabilities(capabilities)) { return false; } - return cutile_launch_available_for_arch( - capabilities.cc_major, capabilities.cc_minor, capabilities.driver_version); + return cutile_launch_available_for_arch(capabilities.cc_major, capabilities.cc_minor); } #else /** Compile-time false when cuTile is not built; use in if constexpr to skip cuTile-only paths. */ diff --git a/cpp/include/cuvs/detail/jit_lto/cutile_module.hpp b/cpp/include/cuvs/detail/jit_lto/cutile_module.hpp index bf26e1c9c5..04117d104c 100644 --- a/cpp/include/cuvs/detail/jit_lto/cutile_module.hpp +++ b/cpp/include/cuvs/detail/jit_lto/cutile_module.hpp @@ -15,7 +15,7 @@ #include #include -#include +#include #include #include @@ -27,40 +27,33 @@ struct CutileModuleImage { size_t size; }; -/** - * 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. - */ +/** Selects an exact architecture-specific cubin, with SM86 accepted only for SM89. */ inline const CubinFragmentEntry* find_compatible_cubin_fragment( int cc_major, int cc_minor, const std::vector>& cubin_fragments) { - const CubinFragmentEntry* best = nullptr; + const CubinFragmentEntry* sm86_fallback = 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(); + if (fragment->get_cc_major() == cc_major && fragment->get_cc_minor() == cc_minor) { + return fragment.get(); + } + if (fragment->get_cc_major() == 8 && fragment->get_cc_minor() == 6) { + sm86_fallback = fragment.get(); } } - return best; + return can_use_sm86_compat_cubin(cc_major, cc_minor) ? sm86_fallback : nullptr; } -/** Selects compatible prebuilt SASS for the device, or TileIR when the driver can JIT it. */ +/** Selects compatible prebuilt SASS for the device. */ inline std::optional resolve_cutile_module_image( const CutileRuntimeCapabilities& capabilities, - const std::vector>& cubin_fragments, - const TileIrBytecodeFragmentEntry* tileir_fragment) + const std::vector>& cubin_fragments) { if (const auto* fragment = find_compatible_cubin_fragment( capabilities.cc_major, capabilities.cc_minor, cubin_fragments)) { return CutileModuleImage{fragment->get_data(), fragment->get_length()}; } - if (tileir_fragment != nullptr && tileir_fallback_available(capabilities.driver_version)) { - return CutileModuleImage{tileir_fragment->get_data(), tileir_fragment->get_length()}; - } return std::nullopt; } diff --git a/cpp/include/cuvs/detail/jit_lto/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..24f05a13bf --- /dev/null +++ b/cpp/include/cuvs/detail/jit_lto/fused_distance_nn/fused_1nn_fragments.hpp @@ -0,0 +1,26 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +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 fragment_tag_fused_1nn_cubin { + static constexpr int cc_major = ArchTag::cc_major; + static constexpr int cc_minor = ArchTag::cc_minor; +}; + +} // namespace cuvs::distance::detail diff --git a/cpp/src/cluster/detail/minClusterDistanceCompute.cu b/cpp/src/cluster/detail/minClusterDistanceCompute.cu index 344f0bf5e5..60bd54d53c 100644 --- a/cpp/src/cluster/detail/minClusterDistanceCompute.cu +++ b/cpp/src/cluster/detail/minClusterDistanceCompute.cu @@ -58,6 +58,7 @@ void minClusterAndDistanceCompute( workspace.resize((sizeof(int)) * n_samples, stream); cuvs::distance::fusedDistanceNNMinReduce, IndexT>( + handle, minClusterAndDistance.data_handle(), X.data_handle(), centroids.data_handle(), @@ -71,8 +72,7 @@ void minClusterAndDistanceCompute( false, true, metric, - 0.0f, - stream); + 0.0f); } else { auto dataBatchSize = getDataBatchSize(batch_samples, n_samples); auto centroidsBatchSize = getCentroidsBatchSize(batch_centroids, n_clusters); @@ -101,8 +101,9 @@ void minClusterAndDistanceCompute( cuvs::distance::unfusedDistanceNNMinReduce( handle, batchMin, - X.data_handle() + dIdx * n_features, - centroids.data_handle() + cIdx * n_features, + X.data_handle() + static_cast(dIdx) * static_cast(n_features), + centroids.data_handle() + + static_cast(cIdx) * static_cast(n_features), L2NormX.data_handle() + dIdx, centroidsNorm.data_handle() + cIdx, ns, @@ -113,8 +114,7 @@ void minClusterAndDistanceCompute( tileCentroids, true, metric, - 0.0f, - stream); + 0.0f); if (tileCentroids) { // Convert tile-local centroid indices and merge the tile minima. @@ -159,7 +159,9 @@ void minClusterAndDistanceCompute( // 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); + X.data_handle() + static_cast(dIdx) * static_cast(n_features), + ns, + n_features); // minClusterAndDistanceView [ns x n_clusters] auto minClusterAndDistanceView = @@ -174,7 +176,10 @@ void minClusterAndDistanceCompute( // 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); + centroids.data_handle() + + static_cast(cIdx) * static_cast(n_features), + nc, + n_features); // pairwiseDistanceView [ns x nc] - view representing the pairwise // distance for current batch @@ -276,6 +281,7 @@ void minClusterDistanceCompute(raft::resources const& handle, workspace.resize(sizeof(int) * n_samples, stream); cuvs::distance::fusedDistanceNNMinReduce( + handle, minClusterDistance.data_handle(), X.data_handle(), centroids.data_handle(), @@ -289,8 +295,7 @@ void minClusterDistanceCompute(raft::resources const& handle, false, true, metric, - 0.0f, - stream); + 0.0f); } else { auto dataBatchSize = getDataBatchSize(batch_samples, n_samples); auto centroidsBatchSize = getCentroidsBatchSize(batch_centroids, n_clusters); @@ -306,7 +311,9 @@ void minClusterDistanceCompute(raft::resources const& handle, auto ns = std::min((IndexT)dataBatchSize, n_samples - dIdx); auto datasetView = raft::make_device_matrix_view( - X.data_handle() + dIdx * n_features, ns, n_features); + X.data_handle() + static_cast(dIdx) * static_cast(n_features), + ns, + n_features); auto minClusterDistanceView = raft::make_device_vector_view(minClusterDistance.data_handle() + dIdx, ns); @@ -316,7 +323,10 @@ void minClusterDistanceCompute(raft::resources const& handle, auto nc = std::min((IndexT)centroidsBatchSize, n_clusters - cIdx); auto centroidsView = raft::make_device_matrix_view( - centroids.data_handle() + cIdx * n_features, nc, n_features); + centroids.data_handle() + + static_cast(cIdx) * static_cast(n_features), + nc, + n_features); auto pairwiseDistanceView = raft::make_device_matrix_view(pairwiseDistance.data_handle(), ns, nc); diff --git a/cpp/src/detail/jit_lto/TileAlgorithmPlanner.cpp b/cpp/src/detail/jit_lto/TileAlgorithmPlanner.cpp index 65363d6fe3..5415015e80 100644 --- a/cpp/src/detail/jit_lto/TileAlgorithmPlanner.cpp +++ b/cpp/src/detail/jit_lto/TileAlgorithmPlanner.cpp @@ -85,8 +85,6 @@ std::string TileAlgorithmPlanner::get_planner_key( for (const auto& fragment : cubin_fragments_) { key += fragment->get_key(); } - if (tileir_fragment_) { key += tileir_fragment_->get_key(); } - if (capabilities != nullptr) { key += ":device=" + std::to_string(capabilities->device); key += ":cc=" + std::to_string(capabilities->cc_major) + "." + @@ -96,9 +94,8 @@ std::string TileAlgorithmPlanner::get_planner_key( key += ":cubin=" + std::to_string(fragment->get_cc_major()) + "." + std::to_string(fragment->get_cc_minor()); } else { - key += ":tileir"; + key += ":unavailable"; } - key += ":driver=" + std::to_string(capabilities->driver_version); } return key; } @@ -113,13 +110,7 @@ CutileTileConfig TileAlgorithmPlanner::tile_config() const } } - 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()); + RAFT_FAIL("cuTile planner '%s' has no compatible cubin", entrypoint_.c_str()); } std::shared_ptr TileAlgorithmPlanner::build( @@ -127,8 +118,7 @@ std::shared_ptr TileAlgorithmPlanner::build( { if (capabilities == nullptr) { return nullptr; } - auto image = cuvs::detail::jit_lto::resolve_cutile_module_image( - *capabilities, cubin_fragments_, tileir_fragment_.get()); + auto image = cuvs::detail::jit_lto::resolve_cutile_module_image(*capabilities, cubin_fragments_); if (!image) { return nullptr; } return cuvs::detail::jit_lto::try_load_cutile_launcher(*image, entrypoint_); diff --git a/cpp/src/distance/detail/fused_distance_nn.cuh b/cpp/src/distance/detail/fused_distance_nn.cuh index f9dbd968ec..7836ff698f 100644 --- a/cpp/src/distance/detail/fused_distance_nn.cuh +++ b/cpp/src/distance/detail/fused_distance_nn.cuh @@ -1,11 +1,14 @@ /* - * 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 +#if CUVS_CUTILE_ENABLED +#include "fused_distance_nn/cutile/fused_1nn_tile.hpp" +#endif #include "fused_distance_nn/cutlass_base.cuh" #include "fused_distance_nn/fused_cosine_nn.cuh" #include "fused_distance_nn/fused_l2_nn.cuh" @@ -13,27 +16,91 @@ #include "fused_distance_nn/simt_kernel.cuh" #include "pairwise_distance_base.cuh" // PairwiseDistances #include -#include // raft::KeyValuePair -#include // raft::identity_op +#include // raft::KeyValuePair +#include // raft::identity_op +#include +#include #include // Policy #include // raft::util::arch::SM_* #include // raft::ceildiv, raft::shfl #include // size_t -#include // std::numeric_limits +#include +#include // std::numeric_limits namespace cuvs { namespace distance { namespace detail { +/** Explicit implementation selected for the top-1 nearest-neighbor primitive. */ +enum class Top1nnBackend : std::uint8_t { + Cutile, + /** Legacy fused dispatcher: CUTLASS on SM80+, with its existing SIMT path before SM80. */ + Cutlass, + Unfused, +}; + +/** Tuning used only by the bounded-workspace unfused backend. */ +struct UnfusedTop1nnTuning { + std::size_t row_tile = 8192; + std::size_t candidate_tile = 8192; +}; + +struct Top1nnTuning { + UnfusedTop1nnTuning unfused{}; +}; + +inline constexpr bool is_top_1_nn_metric_supported(Top1nnBackend backend, DistanceType metric) +{ + switch (backend) { + case Top1nnBackend::Cutile: + return metric == DistanceType::InnerProduct || metric == DistanceType::L2Expanded || + metric == DistanceType::L2SqrtExpanded || metric == DistanceType::CosineExpanded; + case Top1nnBackend::Cutlass: + case Top1nnBackend::Unfused: + return metric == DistanceType::L2Expanded || metric == DistanceType::L2SqrtExpanded || + metric == DistanceType::CosineExpanded; + } + return false; +} + +/** + * Output-independent backend probe. Call this before allocating backend-native result storage. + * cuTile delegates to its launcher/ABI probe. The unfused implementation is always built; + * backend-specific input validation remains the responsibility of top_1_nn. + */ +template +bool is_top_1_nn_backend_available(Top1nnBackend backend, + const DataT* x, + const DataT* y, + IdxT m, + IdxT n, + IdxT k, + cuvs::distance::DistanceType metric) +{ + if (!is_top_1_nn_metric_supported(backend, metric)) { return false; } + if (backend == Top1nnBackend::Cutile) { +#if CUVS_CUTILE_ENABLED + if constexpr (is_fused_1nn_cutile_data_v) { + return is_fused_1nn_tile_available(x, y, m, n, k, metric); + } +#endif + return false; + } + if (backend == Top1nnBackend::Unfused) { return true; } + return backend == Top1nnBackend::Cutlass && x != nullptr && y != nullptr && m > 0 && n > 0 && + k > 0; +} + template -void fusedDistanceNNImpl(OutT* min, +void fusedDistanceNNImpl(raft::resources const& handle, + OutT* min, const DataT* x, const DataT* y, const DataT* xn, @@ -48,9 +115,9 @@ void fusedDistanceNNImpl(OutT* min, bool initOutBuffer, bool isRowMajor, cuvs::distance::DistanceType metric, - float metric_arg, - cudaStream_t stream) + float metric_arg) { + const auto stream = raft::resource::get_cuda_stream(handle); // The kernel policy is determined by fusedDistanceNN. typedef Policy P; @@ -77,7 +144,7 @@ void fusedDistanceNNImpl(OutT* min, fusedL2NNImpl( min, x, y, xn, yn, m, n, k, workspace, redOp, pairRedOp, sqrt, false, stream); break; - default: assert("only cosine/l2 metric is supported with fusedDistanceNN\n"); break; + default: RAFT_FAIL("Only cosine and L2 metrics are supported by fusedDistanceNN"); } } 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..091603f32c --- /dev/null +++ b/cpp/src/distance/detail/fused_distance_nn/cutile/export_fused_1nn.py @@ -0,0 +1,278 @@ +# 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.""" + +from __future__ import annotations + +import argparse +import sys +from pathlib import Path + +import cuda.tile as ct +from cuda.tile.compilation import ( + ArrayConstraint, + CallingConvention, + ConstantConstraint, + KernelSignature, + 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, +) + + +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 _elem_shape_divisible_for_ldgsts(elem_dtype) -> tuple[int, int]: + """Matrix extent aligned to the same 16-byte row pitch enforced on strides.""" + bytes_per_elem = 2 if elem_dtype == ct.float16 else 4 + return (1, 16 // bytes_per_elem) + + +def _cuvs_matrix_constraint( + elem_dtype, + *, + index_dtype=ct.int32, + require_tma_friendly_pitch: bool = True, + require_ldgsts_friendly_shape: bool = False, +): + """Row-major device matrices for cuVS KMeans benchmarks. + + Assumes raft/cupy-style contiguous layout: stride[-1]==1, stride[0]==D, + 16-byte base alignment, and row pitch 16-byte aligned (float32 D%4==0, + float16 D%8==0). Applies to both points and centroids matrices. + + SM80/SM86 strict exports also express the row-pitch guarantee as + shape_divisible_by=(1, 4) for float32 or (1, 8) for float16. This + duplicates the stride constraint intentionally so the compiler selects + LDGSTS instead of LDG. Tail tiles remain masked in the kernel. + + Odd D or general layouts need a separate relaxed export profile. + """ + return ArrayConstraint( + elem_dtype, + ndim=2, + index_dtype=index_dtype, + stride_lower_bound_incl=(0, None), + # Dataset and centroid views are read-only and may legally share storage. + alias_groups=("read_only_inputs",), + may_alias_internally=False, + stride_constant=(None, 1), + stride_divisible_by=( + _elem_stride_divisible_for_tma(elem_dtype) + if require_tma_friendly_pitch + else (1, 1) + ), + shape_divisible_by=( + _elem_shape_divisible_for_ldgsts(elem_dtype) + if require_ldgsts_friendly_shape + else (1, 1) + ), + base_addr_divisible_by=16, + ) + + +def _cuvs_vector_constraint( + elem_dtype, *, index_dtype=ct.int32, alias_groups=() +): + """1-D device vectors: contiguous, 16-byte base. Length need not be divisible by 16.""" + return ArrayConstraint( + elem_dtype, + ndim=1, + index_dtype=index_dtype, + stride_lower_bound_incl=(None,), + alias_groups=alias_groups, + may_alias_internally=False, + stride_constant=(1,), + stride_divisible_by=(1,), + shape_divisible_by=(1,), + base_addr_divisible_by=16, + ) + + +def _relaxed_matrix_constraint(elem_dtype): + """Deprecated alias for the arbitrary-row-pitch matrix constraint.""" + return _cuvs_matrix_constraint( + elem_dtype, require_tma_friendly_pitch=False + ) + + +def _relaxed_vector_constraint(elem_dtype, *, tma_friendly: bool = False): + """Deprecated alias; use _cuvs_vector_constraint.""" + del tma_friendly + return _cuvs_vector_constraint(elem_dtype) + + +def _kernel_signature( + data_type: str, + metric: str, + index_type: str, + tile_m: int, + tile_n: int, + tile_k: int, + gpu_code: str, + matrix_layout: str, +) -> KernelSignature: + elem = _dtype_for(data_type) + idx_dtype = _idx_dtype(index_type) + matrix = _cuvs_matrix_constraint( + elem, + index_dtype=idx_dtype, + require_tma_friendly_pitch=matrix_layout == "strict", + require_ldgsts_friendly_shape=( + matrix_layout == "strict" and gpu_code in ("sm_80", "sm_86") + ), + ) + norm_elem = ct.float32 if data_type == "half" else elem + norm_array = _cuvs_vector_constraint( + norm_elem, + index_dtype=idx_dtype, + alias_groups=("read_only_inputs",), + ) + idx_array = _cuvs_vector_constraint(idx_dtype, index_dtype=idx_dtype) + dist_array = _cuvs_vector_constraint(ct.float32, 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, + *, + data_type: str, + metric: str, + index_type: str, + tile_m: int, + tile_n: int, + tile_k: int, + gpu_code: str, + matrix_layout: str = "strict", + occupancy: int | None = None, +) -> str: + kernel = make_kernel( + data_type, + metric, + tile_m, + tile_n, + tile_k, + index_type=index_type, + gpu_code=gpu_code, + matrix_layout=matrix_layout, + occupancy=occupancy, + ) + signature = _kernel_signature( + data_type, + metric, + index_type, + tile_m, + tile_n, + tile_k, + gpu_code, + matrix_layout, + ) + + export_kernel( + kernel=kernel, + signatures=[signature], + output_file=str(output_file), + gpu_code=gpu_code, + output_format="cubin", + ) + + return signature.symbol + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("output_file", type=Path) + parser.add_argument("--format", choices=("cubin",), 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", required=True, help="Target SM for cubin export" + ) + parser.add_argument( + "--matrix-layout", + choices=("strict", "relaxed"), + default="strict", + ) + parser.add_argument("--occupancy", type=int) + args = parser.parse_args() + + export_binary( + args.output_file, + data_type=args.data_type, + metric=args.metric, + index_type=args.index_type, + tile_m=args.tile_m, + tile_n=args.tile_n, + tile_k=args.tile_k, + gpu_code=args.gpu_code, + matrix_layout=args.matrix_layout, + occupancy=args.occupancy, + ) + 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..5b3d007eec --- /dev/null +++ b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_cutile_matrix.json @@ -0,0 +1,265 @@ +{ + "metric": "runtime", + "index_type": "int32", + "index_abbrev": "i32", + "_format": [ + { + "output_format": "cubin", + "artifact_ext": "cubin", + "artifact_basename": "@data_type@_@index_abbrev@_@abi_abbrev@_@gpu_code@", + "register": "cubin", + "_specialization": [ + { + "data_type": "float", + "data_abbrev": "f", + "_abi": [ + { + "matrix_layout": "relaxed", + "abi_abbrev": "relaxed", + "abi_tag": "cutile_abi_relaxed", + "tile_m": 64, + "tile_n": 128, + "tile_k": 32 + }, + { + "matrix_layout": "strict", + "abi_abbrev": "strict", + "abi_tag": "cutile_abi_strict", + "tile_m": 64, + "tile_n": 128, + "tile_k": 32, + "occupancy": 2 + } + ], + "_architecture": [ + { + "gpu_code": "sm_80", + "cc_major": 8, + "cc_minor": 0, + "arch_tag": "cutile_arch_8_0" + }, + { + "gpu_code": "sm_86", + "cc_major": 8, + "cc_minor": 6, + "arch_tag": "cutile_arch_8_6" + } + ] + }, + { + "data_type": "half", + "data_abbrev": "h", + "_abi": [ + { + "matrix_layout": "relaxed", + "abi_abbrev": "relaxed", + "abi_tag": "cutile_abi_relaxed", + "tile_m": 128, + "tile_n": 128, + "tile_k": 32 + }, + { + "matrix_layout": "strict", + "abi_abbrev": "strict", + "abi_tag": "cutile_abi_strict", + "tile_m": 128, + "tile_n": 128, + "tile_k": 128, + "occupancy": 2 + } + ], + "gpu_code": "sm_80", + "cc_major": 8, + "cc_minor": 0, + "arch_tag": "cutile_arch_8_0" + }, + { + "data_type": "half", + "data_abbrev": "h", + "_abi": [ + { + "matrix_layout": "relaxed", + "abi_abbrev": "relaxed", + "abi_tag": "cutile_abi_relaxed", + "tile_m": 128, + "tile_n": 128, + "tile_k": 32, + "occupancy": 2 + }, + { + "matrix_layout": "strict", + "abi_abbrev": "strict", + "abi_tag": "cutile_abi_strict", + "tile_m": 128, + "tile_n": 128, + "tile_k": 32, + "occupancy": 2 + } + ], + "gpu_code": "sm_86", + "cc_major": 8, + "cc_minor": 6, + "arch_tag": "cutile_arch_8_6" + }, + { + "data_type": "float", + "data_abbrev": "f", + "_abi": [ + { + "matrix_layout": "relaxed", + "abi_abbrev": "relaxed", + "abi_tag": "cutile_abi_relaxed", + "tile_m": 128, + "tile_n": 128, + "tile_k": 64 + }, + { + "matrix_layout": "strict", + "abi_abbrev": "strict", + "abi_tag": "cutile_abi_strict", + "tile_m": 64, + "tile_n": 256, + "tile_k": 32 + } + ], + "gpu_code": "sm_90", + "cc_major": 9, + "cc_minor": 0, + "arch_tag": "cutile_arch_9_0" + }, + { + "data_type": "half", + "data_abbrev": "h", + "tile_m": 128, + "tile_n": 128, + "tile_k": 128, + "_abi": [ + { + "matrix_layout": "relaxed", + "abi_abbrev": "relaxed", + "abi_tag": "cutile_abi_relaxed" + }, + { + "matrix_layout": "strict", + "abi_abbrev": "strict", + "abi_tag": "cutile_abi_strict" + } + ], + "gpu_code": "sm_90", + "cc_major": 9, + "cc_minor": 0, + "arch_tag": "cutile_arch_9_0" + }, + { + "data_type": "float", + "data_abbrev": "f", + "_abi": [ + { + "matrix_layout": "relaxed", + "abi_abbrev": "relaxed", + "abi_tag": "cutile_abi_relaxed", + "tile_m": 128, + "tile_n": 256, + "tile_k": 16 + }, + { + "matrix_layout": "strict", + "abi_abbrev": "strict", + "abi_tag": "cutile_abi_strict", + "tile_m": 128, + "tile_n": 128, + "tile_k": 32 + } + ], + "gpu_code": "sm_100", + "cc_major": 10, + "cc_minor": 0, + "arch_tag": "cutile_arch_10_0" + }, + { + "data_type": "half", + "data_abbrev": "h", + "_abi": [ + { + "matrix_layout": "relaxed", + "abi_abbrev": "relaxed", + "abi_tag": "cutile_abi_relaxed", + "tile_m": 128, + "tile_n": 256, + "tile_k": 16, + "occupancy": 2 + }, + { + "matrix_layout": "strict", + "abi_abbrev": "strict", + "abi_tag": "cutile_abi_strict", + "tile_m": 128, + "tile_n": 128, + "tile_k": 128 + } + ], + "gpu_code": "sm_100", + "cc_major": 10, + "cc_minor": 0, + "arch_tag": "cutile_arch_10_0" + }, + { + "data_type": "float", + "data_abbrev": "f", + "_abi": [ + { + "matrix_layout": "relaxed", + "abi_abbrev": "relaxed", + "abi_tag": "cutile_abi_relaxed", + "tile_m": 64, + "tile_n": 128, + "tile_k": 64, + "occupancy": 2 + }, + { + "matrix_layout": "strict", + "abi_abbrev": "strict", + "abi_tag": "cutile_abi_strict", + "tile_m": 64, + "tile_n": 128, + "tile_k": 32, + "occupancy": 2 + } + ], + "gpu_code": "sm_120", + "cc_major": 12, + "cc_minor": 0, + "arch_tag": "cutile_arch_12_0" + }, + { + "data_type": "half", + "data_abbrev": "h", + "_abi": [ + { + "matrix_layout": "relaxed", + "abi_abbrev": "relaxed", + "abi_tag": "cutile_abi_relaxed", + "tile_m": 64, + "tile_n": 128, + "tile_k": 128, + "occupancy": 2 + }, + { + "matrix_layout": "strict", + "abi_abbrev": "strict", + "abi_tag": "cutile_abi_strict", + "tile_m": 64, + "tile_n": 256, + "tile_k": 64, + "occupancy": 2 + } + ], + "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..1ad341beff --- /dev/null +++ b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_kernel.py @@ -0,0 +1,196 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""cuTile fused GEMM + 1-NN kernel with runtime metric selection.""" + +from __future__ import annotations + +import cuda.tile as ct + +ConstInt = ct.Constant[int] + +# Default tile geometry; overridden per export via make_kernel(..., tile_m, tile_n, tile_k). +DEFAULT_TILE_M = 128 +DEFAULT_TILE_N = 128 +DEFAULT_TILE_K = 32 + +METRICS = ("runtime",) +INDEX_TYPES = ("int32", "int64") +METRIC_L2_EXPANDED = 0 +METRIC_COSINE_EXPANDED = 2 +METRIC_INNER_PRODUCT = 6 + + +def _idx_dtype(index_type: str): + if index_type == "int32": + return ct.int32 + if index_type == "int64": + return ct.int64 + raise ValueError(f"Unsupported index_type {index_type!r}") + + +def make_kernel( + data_type: str, + metric: str, + tile_m: int = DEFAULT_TILE_M, + tile_n: int = DEFAULT_TILE_N, + tile_k: int = DEFAULT_TILE_K, + *, + index_type: str = "int32", + gpu_code: str = "sm_80", + matrix_layout: str = "strict", + occupancy: int | None = None, +): + """Build the flat-reduction runtime-metric cuTile kernel.""" + if data_type not in ("half", "float"): + raise ValueError(f"Unsupported data_type {data_type!r}") + if metric not in METRICS: + raise ValueError(f"Unsupported metric {metric!r}") + if index_type not in INDEX_TYPES: + raise ValueError(f"Unsupported index_type {index_type!r}") + if matrix_layout not in ("strict", "relaxed"): + raise ValueError(f"Unsupported matrix_layout {matrix_layout!r}") + + acc_dtype = ct.float32 + idx_dtype = _idx_dtype(index_type) + core_shape = (tile_m, tile_n) + best_shape = (tile_m, 1) + kernel_options = {} + if occupancy is not None: + kernel_options["occupancy"] = ct.ByTarget(**{gpu_code: occupancy}) + + @ct.kernel(**kernel_options) + def fused_1nn_kernel( + A, + B, + A_norm, + B_norm, + OutIdx, + OutDist, + M, + N, + K, + apply_sqrt, + store_idx, + metric_code, + tm: ConstInt, + tn: ConstInt, + tk: ConstInt, + ): + bidm = ct.bid(0) + best_dist = ct.full(best_shape, 3.4e38, acc_dtype) + best_idx = ct.zeros(best_shape, idx_dtype) + num_tiles_k = ct.num_tiles(A, axis=1, shape=(tm, tk)) + num_tiles_n = ct.num_tiles(B, axis=0, shape=(tn, tk)) + zero_pad = ct.PaddingMode.ZERO + + def reduce_scores(dists, indices): + def red_op(a_score, a_idx, b_score, b_idx): + cond = (a_score < b_score) | ( + (a_score == b_score) & (a_idx < b_idx) + ) + return ( + ct.where(cond, a_score, b_score), + ct.where(cond, a_idx, b_idx), + ) + + return ct.reduce( + (dists, indices), + 1, + red_op, + (3.4e38, -1), + keepdims=True, + ) + + local_indices = ct.arange(tn, dtype=ct.int16)[None, :] + for n in range(num_tiles_n): + accumulator = ct.full((tm, tn), 0, dtype=acc_dtype) + for k in range(num_tiles_k): + dtype = ct.tfloat32 if A.dtype == ct.float32 else A.dtype + a = ct.load( + A, index=(bidm, k), shape=(tm, tk), padding_mode=zero_pad + ).astype(dtype) + b_T = ct.load( + B, + index=(k, n), + shape=(tk, tn), + padding_mode=zero_pad, + order=(1, 0), + ).astype(dtype) + accumulator = ct.mma(a, b_T, accumulator) + + if metric_code == METRIC_INNER_PRODUCT: + score = -accumulator + else: + b_norm = ct.load( + B_norm, index=(n,), shape=(tn,), padding_mode=zero_pad + ) + if metric_code == METRIC_L2_EXPANDED: + # The A norm is constant across centroids. Excluding it + # avoids cancellation in the score used by argmin. + score = (0.5 * b_norm)[None, :] - accumulator + else: + # Cosine distance involving any zero-norm vector is defined as 1. + # A zero B norm therefore contributes a normalized dot score of 0. + positive_b_norm = b_norm > 0.0 + safe_b_norm = ct.where(positive_b_norm, b_norm, 1.0) + score = accumulator / (-safe_b_norm)[None, :] + score = ct.where(positive_b_norm[None, :], score, 0.0) + + if n == num_tiles_n - 1: + col = ct.arange(tn, dtype=ct.int16) + score = ct.where((n * tn + col)[None, :] < N, score, 3.4e38) + + curr_best, curr_idx = reduce_scores( + score.reshape(core_shape), local_indices + ) + update = curr_best < best_dist + best_dist = ct.where(update, curr_best, best_dist) + best_idx = ct.where(update, n * tn + curr_idx, best_idx) + + if metric_code == METRIC_INNER_PRODUCT: + out_dist = -best_dist + else: + a_norm = ct.load( + A_norm, index=(bidm,), shape=(tm,), padding_mode=zero_pad + )[:, None] + if metric_code == METRIC_L2_EXPANDED: + out_dist = a_norm + 2.0 * best_dist + # Separately reduced norms and the MMA can reconstruct a + # slightly negative distance; clamp before an optional sqrt. + out_dist = ct.where(out_dist > 0.0, out_dist, 0.0) + out_dist = ct.where( + apply_sqrt != 0, ct.sqrt(out_dist), out_dist + ) + else: + positive_a_norm = a_norm > 0.0 + safe_a_norm = ct.where(positive_a_norm, a_norm, 1.0) + out_dist = 1.0 + best_dist / safe_a_norm + out_dist = ct.where(positive_a_norm, out_dist, 1.0) + + if store_idx != 0: + ct.store(OutIdx, index=(bidm,), tile=best_idx.reshape((tm,))) + ct.store( + OutDist, + index=(bidm,), + tile=out_dist.reshape((tm,)), + ) + + 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..eb879353ee --- /dev/null +++ b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_planner.hpp @@ -0,0 +1,117 @@ +/* + * 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 = std::conditional_t, + cuvs::neighbors::detail::tag_f, + cuvs::neighbors::detail::tag_h>; + using IndexTag = cuvs::neighbors::detail::tag_index_i32; + + inline static cuvs::detail::jit_lto::TileLauncherCache launcher_cache{}; + + Fused1nnTilePlanner() + : TileAlgorithmPlanner(fused_1nn_kernel_entrypoint(), launcher_cache) + { + } + + /** Registers embedded cubin modules (one per SM); see register_cutile_fragment.cpp object files. + */ + void add_entrypoint() + { + using cuvs::detail::jit_lto::cutile_arch_10_0; + using cuvs::detail::jit_lto::cutile_arch_12_0; + using cuvs::detail::jit_lto::cutile_arch_8_0; + using cuvs::detail::jit_lto::cutile_arch_8_6; + using cuvs::detail::jit_lto::cutile_arch_9_0; + + constexpr bool is_relaxed = std::is_same_v; + constexpr bool is_float = std::is_same_v; + using Tile80 = + std::conditional_t, + std::conditional_t>; + using Tile86 = + std::conditional_t, + std::conditional_t>; + using Tile90 = + std::conditional_t, + std::conditional_t>; + using Tile100 = + std::conditional_t, + std::conditional_t>; + using Tile120 = + std::conditional_t, + std::conditional_t>; + + this->add_static_fragment< + fragment_tag_fused_1nn_cubin>(); + this->add_static_fragment< + fragment_tag_fused_1nn_cubin>(); + this->add_static_fragment< + fragment_tag_fused_1nn_cubin>(); + this->add_static_fragment< + fragment_tag_fused_1nn_cubin>(); + this->add_static_fragment< + fragment_tag_fused_1nn_cubin>(); + } +}; + +} // 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..85b450c21f --- /dev/null +++ b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_tile.cu @@ -0,0 +1,398 @@ +/* + * 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 +#include + +namespace cuvs { +namespace distance { +namespace detail { + +namespace { + +bool is_16_byte_aligned(const void* ptr) +{ + return ptr == nullptr || reinterpret_cast(ptr) % 16 == 0; +} + +bool byte_ranges_overlap(const void* lhs, size_t lhs_bytes, const void* rhs, size_t rhs_bytes) +{ + if (lhs == nullptr || rhs == nullptr || lhs_bytes == 0 || rhs_bytes == 0) { return false; } + const auto lhs_begin = reinterpret_cast(lhs); + const auto rhs_begin = reinterpret_cast(rhs); + return lhs_begin <= rhs_begin ? rhs_begin - lhs_begin < lhs_bytes + : lhs_begin - rhs_begin < rhs_bytes; +} + +template +size_t checked_tensor_bytes(IdxT rows, IdxT cols, size_t element_size) +{ + const auto rows_u = static_cast(rows); + const auto cols_u = static_cast(cols); + constexpr auto max_size = std::numeric_limits::max(); + if (cols_u != 0 && rows_u > max_size / cols_u) { return max_size; } + const auto elements = rows_u * cols_u; + if (element_size != 0 && elements > max_size / element_size) { return max_size; } + return static_cast(elements) * element_size; +} + +template +bool has_fused_1nn_tile_launcher() +{ + Fused1nnTilePlanner planner; + planner.add_entrypoint(); + return planner.try_get_launcher() != nullptr; +} + +template +void launch_fused_1nn_tile_impl(raft::resources const& handle, + IdxT* nearest_idx, + fused_1nn_cutile_distance_t* nearest_dist, + const DataT* x, + const DataT* y, + const fused_1nn_cutile_norm_t* xn, + const fused_1nn_cutile_norm_t* yn, + IdxT m, + IdxT n, + IdxT k, + cuvs::distance::DistanceType metric, + bool is_sqrt) +{ + const auto stream = raft::resource::get_cuda_stream(handle); + Fused1nnTilePlanner planner; + planner.add_entrypoint(); + auto launcher = planner.try_get_launcher(); + RAFT_EXPECTS(launcher != nullptr, "Requested cuTile fused 1-NN launcher is unavailable"); + const cuvs::detail::jit_lto::CutileTileConfig tile_cfg = planner.tile_config(); + + int metric_code; + bool apply_sqrt = false; + switch (metric) { + case cuvs::distance::DistanceType::InnerProduct: + metric_code = static_cast(cuvs::distance::DistanceType::InnerProduct); + break; + case cuvs::distance::DistanceType::L2Expanded: + case cuvs::distance::DistanceType::L2SqrtExpanded: + metric_code = static_cast(cuvs::distance::DistanceType::L2Expanded); + apply_sqrt = is_sqrt; + break; + case cuvs::distance::DistanceType::CosineExpanded: + metric_code = static_cast(cuvs::distance::DistanceType::CosineExpanded); + break; + default: RAFT_FAIL("Unsupported cuTile fused 1-NN metric"); + } + + 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()); +} + +template +void validate_fused_1nn_tile_launch(IdxT* nearest_idx, + fused_1nn_cutile_distance_t* nearest_dist, + const DataT* x, + const DataT* y, + const fused_1nn_cutile_norm_t* xn, + const fused_1nn_cutile_norm_t* yn, + IdxT m, + IdxT n, + IdxT k, + cuvs::distance::DistanceType metric, + void* index_workspace) +{ + RAFT_EXPECTS(is_fused_1nn_tile_available(x, y, m, n, k, metric), + "Requested cuTile fused 1-NN backend is unavailable for this input/device"); + RAFT_EXPECTS(nearest_dist != nullptr && is_16_byte_aligned(nearest_dist), + "cuTile fused 1-NN requires a 16-byte-aligned distance output"); + if constexpr (std::is_same_v) { + RAFT_EXPECTS(is_16_byte_aligned(nearest_idx), + "cuTile fused 1-NN requires a 16-byte-aligned int32 index output"); + } + RAFT_EXPECTS( + metric == cuvs::distance::DistanceType::InnerProduct || (xn != nullptr && yn != nullptr), + "cuTile fused 1-NN requires norm buffers for this metric"); + RAFT_EXPECTS(is_16_byte_aligned(xn) && is_16_byte_aligned(yn), + "cuTile fused 1-NN requires 16-byte-aligned norm buffers"); + + const auto x_bytes = checked_tensor_bytes(m, k, sizeof(DataT)); + const auto y_bytes = checked_tensor_bytes(n, k, sizeof(DataT)); + const auto dist_bytes = checked_tensor_bytes(m, IdxT{1}, sizeof(*nearest_dist)); + const auto idx_bytes = checked_tensor_bytes(m, IdxT{1}, sizeof(IdxT)); + const auto xn_bytes = checked_tensor_bytes(m, IdxT{1}, sizeof(*xn)); + const auto yn_bytes = checked_tensor_bytes(n, IdxT{1}, sizeof(*yn)); + RAFT_EXPECTS(!byte_ranges_overlap(nearest_dist, dist_bytes, x, x_bytes) && + !byte_ranges_overlap(nearest_dist, dist_bytes, y, y_bytes) && + !byte_ranges_overlap(nearest_idx, idx_bytes, x, x_bytes) && + !byte_ranges_overlap(nearest_idx, idx_bytes, y, y_bytes) && + !byte_ranges_overlap(nearest_idx, idx_bytes, nearest_dist, dist_bytes) && + !byte_ranges_overlap(nearest_dist, dist_bytes, xn, xn_bytes) && + !byte_ranges_overlap(nearest_dist, dist_bytes, yn, yn_bytes) && + !byte_ranges_overlap(nearest_idx, idx_bytes, xn, xn_bytes) && + !byte_ranges_overlap(nearest_idx, idx_bytes, yn, yn_bytes), + "cuTile fused 1-NN input, norm, and output buffers must not overlap"); + + if constexpr (std::is_same_v) { + RAFT_EXPECTS(nearest_idx == nullptr || index_workspace != nullptr, + "cuTile fused 1-NN requires int32 workspace for int64 index output"); + RAFT_EXPECTS(is_16_byte_aligned(index_workspace), + "cuTile fused 1-NN requires 16-byte-aligned index workspace"); + const auto workspace_rows = static_cast(fused_1nn_cutile_index_workspace_rows(m)); + const auto workspace_bytes = checked_tensor_bytes(workspace_rows, IdxT{1}, sizeof(int)); + RAFT_EXPECTS( + !byte_ranges_overlap(index_workspace, workspace_bytes, x, x_bytes) && + !byte_ranges_overlap(index_workspace, workspace_bytes, y, y_bytes) && + !byte_ranges_overlap(index_workspace, workspace_bytes, xn, xn_bytes) && + !byte_ranges_overlap(index_workspace, workspace_bytes, yn, yn_bytes) && + !byte_ranges_overlap(index_workspace, workspace_bytes, nearest_dist, dist_bytes) && + !byte_ranges_overlap(index_workspace, workspace_bytes, nearest_idx, idx_bytes), + "cuTile fused 1-NN index workspace must not overlap input or output buffers"); + } +} + +} // namespace + +template + requires is_fused_1nn_cutile_data_v +bool is_fused_1nn_tile_available( + const DataT* x, const DataT* y, IdxT m, IdxT n, IdxT k, cuvs::distance::DistanceType metric) +{ + if (!cuvs::detail::jit_lto::cutile_launch_available_on_current_device()) { return false; } + static_assert(std::is_same_v || std::is_same_v); + + if (x == nullptr || y == nullptr || m <= 0 || n <= 0 || k <= 0) { return false; } + if (metric != cuvs::distance::DistanceType::InnerProduct && + metric != cuvs::distance::DistanceType::L2Expanded && + metric != cuvs::distance::DistanceType::L2SqrtExpanded && + metric != cuvs::distance::DistanceType::CosineExpanded) { + return false; + } + + if (!is_16_byte_aligned(x) || !is_16_byte_aligned(y)) { return false; } + if constexpr (std::is_same_v) { + constexpr int64_t max_i32 = std::numeric_limits::max(); + if (n > max_i32 || k > max_i32) { return false; } + } + + constexpr int strict_pitch_elements = 16 / sizeof(DataT); + return k % strict_pitch_elements == 0 ? has_fused_1nn_tile_launcher() + : has_fused_1nn_tile_launcher(); +} + +template + requires is_fused_1nn_cutile_data_v +void launch_fused_1nn_tile(raft::resources const& handle, + IdxT* nearest_idx, + fused_1nn_cutile_distance_t* nearest_dist, + const DataT* x, + const DataT* y, + const fused_1nn_cutile_norm_t* xn, + const fused_1nn_cutile_norm_t* yn, + IdxT m, + IdxT n, + IdxT k, + cuvs::distance::DistanceType metric, + bool is_sqrt, + void* index_workspace) +{ + validate_fused_1nn_tile_launch( + nearest_idx, nearest_dist, x, y, xn, yn, m, n, k, metric, index_workspace); + + constexpr int strict_pitch_elements = 16 / sizeof(DataT); + const bool use_strict_abi = k % strict_pitch_elements == 0; + + if constexpr (std::is_same_v) { + if (use_strict_abi) { + launch_fused_1nn_tile_impl( + handle, nearest_idx, nearest_dist, x, y, xn, yn, m, n, k, metric, is_sqrt); + } else { + launch_fused_1nn_tile_impl( + handle, nearest_idx, nearest_dist, x, y, xn, yn, m, n, k, metric, is_sqrt); + } + } else { + const auto stream = raft::resource::get_cuda_stream(handle); + constexpr int64_t max_batch_m = fused_1nn_cutile_max_batch_m; + auto* tmp_idx = static_cast(index_workspace); + for (int64_t offset = 0; offset < m;) { + const int64_t batch_m64 = std::min(max_batch_m, m - offset); + const int batch_m = static_cast(batch_m64); + const auto* batch_x = x + static_cast(offset) * static_cast(k); + const auto* batch_xn = xn == nullptr ? nullptr : xn + offset; + auto* batch_dist = nearest_dist + offset; + + if (use_strict_abi) { + launch_fused_1nn_tile_impl(handle, + tmp_idx, + batch_dist, + batch_x, + y, + batch_xn, + yn, + batch_m, + static_cast(n), + static_cast(k), + metric, + is_sqrt); + } else { + launch_fused_1nn_tile_impl(handle, + tmp_idx, + batch_dist, + batch_x, + y, + batch_xn, + yn, + batch_m, + static_cast(n), + static_cast(k), + metric, + is_sqrt); + } + + if (nearest_idx != nullptr) { + raft::linalg::unaryOp( + nearest_idx + offset, tmp_idx, batch_m, raft::cast_op{}, stream); + } + offset += batch_m64; + } + } +} + +#define CUVS_INST_IS_FUSED_1NN_TILE_AVAILABLE(DataT, IdxT) \ + template CUVS_EXPORT bool is_fused_1nn_tile_available( \ + const DataT*, const DataT*, IdxT, IdxT, IdxT, cuvs::distance::DistanceType) + +CUVS_INST_IS_FUSED_1NN_TILE_AVAILABLE(float, int); +CUVS_INST_IS_FUSED_1NN_TILE_AVAILABLE(float, int64_t); +CUVS_INST_IS_FUSED_1NN_TILE_AVAILABLE(half, int); +CUVS_INST_IS_FUSED_1NN_TILE_AVAILABLE(half, int64_t); + +#undef CUVS_INST_IS_FUSED_1NN_TILE_AVAILABLE + +#define CUVS_INST_LAUNCH_FUSED_1NN_TILE(DataT, IdxT) \ + template CUVS_EXPORT void launch_fused_1nn_tile( \ + raft::resources const&, \ + IdxT*, \ + fused_1nn_cutile_distance_t*, \ + const DataT*, \ + const DataT*, \ + const fused_1nn_cutile_norm_t*, \ + const fused_1nn_cutile_norm_t*, \ + IdxT, \ + IdxT, \ + IdxT, \ + cuvs::distance::DistanceType, \ + bool, \ + void*) + +CUVS_INST_LAUNCH_FUSED_1NN_TILE(float, int); +CUVS_INST_LAUNCH_FUSED_1NN_TILE(float, int64_t); +CUVS_INST_LAUNCH_FUSED_1NN_TILE(half, int); +CUVS_INST_LAUNCH_FUSED_1NN_TILE(half, int64_t); + +#undef CUVS_INST_LAUNCH_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..91fb9d4120 --- /dev/null +++ b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_tile.hpp @@ -0,0 +1,87 @@ +/* + * 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 + +namespace cuvs { +namespace distance { +namespace detail { + +template +inline constexpr bool is_fused_1nn_cutile_data_v = + std::is_same_v || std::is_same_v; + +// Norm buffers use FP32 storage. Accumulate FP16 norms in FP32; for FP32 inputs, accumulate +// the squares of TF32-rounded values in FP32 to match the cuTile MMA operands. +template +using fused_1nn_cutile_norm_t = float; + +// Both FP16 and FP32 MMA paths accumulate and reconstruct distances in FP32. +template +using fused_1nn_cutile_distance_t = float; + +template +inline constexpr int64_t fused_1nn_cutile_max_batch_m = [] { + constexpr int64_t max_i32 = std::numeric_limits::max(); + constexpr int64_t batch_alignment = 16 / raft::gcd(16, sizeof(DataT)); + return max_i32 - max_i32 % batch_alignment; +}(); + +template +constexpr size_t fused_1nn_cutile_index_workspace_rows(IdxT m) +{ + const auto rows = static_cast(m); + if (rows <= 0) { return 0; } + return static_cast( + rows < fused_1nn_cutile_max_batch_m ? rows : fused_1nn_cutile_max_batch_m); +} + +/** + * Return whether the input problem has a compatible cuTile launcher. + * + * This output-independent probe lets callers select native result storage before allocating it. + */ +template + requires is_fused_1nn_cutile_data_v +bool is_fused_1nn_tile_available( + const DataT* x, const DataT* y, IdxT m, IdxT n, IdxT k, cuvs::distance::DistanceType metric); + +/** + * Launch fused 1-NN with cuTile. + * + * All launch arguments are validated. An int64 output index requires an int32 workspace sized to + * fused_1nn_cutile_index_workspace_rows(m). This function throws instead of falling back + * when the explicitly requested cuTile backend is unavailable. + */ +template + requires is_fused_1nn_cutile_data_v +void launch_fused_1nn_tile(raft::resources const& handle, + IdxT* nearest_idx, + fused_1nn_cutile_distance_t* nearest_dist, + const DataT* x, + const DataT* y, + const fused_1nn_cutile_norm_t* xn, + const fused_1nn_cutile_norm_t* yn, + IdxT m, + IdxT n, + IdxT k, + cuvs::distance::DistanceType metric, + bool is_sqrt, + void* index_workspace); +} // namespace detail +} // namespace distance +} // namespace cuvs diff --git a/cpp/src/distance/fused_distance_nn-inl.cuh b/cpp/src/distance/fused_distance_nn-inl.cuh index 3fa80a9b60..87da85dbe9 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 */ @@ -10,14 +10,18 @@ #include "detail/fused_distance_nn.cuh" #include "fused_distance_nn_helpers.cuh" +#include "top_1_nn.cuh" +#include "unfused_distance_nn.cuh" #include #include +#include #include #include #include +#include #include #include @@ -66,10 +70,11 @@ namespace distance { * @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[in] handle RAFT resources containing the caller-provided CUDA stream */ template -void fusedDistanceNN(OutT* min, +void fusedDistanceNN(raft::resources const& handle, + OutT* min, const DataT* x, const DataT* y, const DataT* xn, @@ -84,8 +89,7 @@ void fusedDistanceNN(OutT* min, bool initOutBuffer, bool isRowMajor, cuvs::distance::DistanceType metric, - float metric_arg, - cudaStream_t stream) + float metric_arg) { ASSERT(isRowMajor, "fusedDistanceNN only supports row major inputs"); // When k is smaller than 32, the Policy4x4 results in redundant calculations @@ -103,7 +107,8 @@ void fusedDistanceNN(OutT* min, OutT, IdxT, typename raft::linalg::Policy4x4Skinny::Policy, - ReduceOpT>(min, + ReduceOpT>(handle, + min, x, y, xn, @@ -118,15 +123,15 @@ void fusedDistanceNN(OutT* min, initOutBuffer, isRowMajor, metric, - metric_arg, - stream); + metric_arg); } else { detail::fusedDistanceNNImpl< DataT, OutT, IdxT, typename raft::linalg::Policy4x4::Policy, - ReduceOpT>(min, + ReduceOpT>(handle, + min, x, y, xn, @@ -141,8 +146,7 @@ void fusedDistanceNN(OutT* min, initOutBuffer, isRowMajor, metric, - metric_arg, - stream); + metric_arg); } } else if (8 % sizeof(DataT) == 0 && bytes % 8 == 0 && px % 8 == 0 && py % 8 == 0) { if (is_skinny) { @@ -151,7 +155,8 @@ void fusedDistanceNN(OutT* min, OutT, IdxT, typename raft::linalg::Policy4x4Skinny::Policy, - ReduceOpT>(min, + ReduceOpT>(handle, + min, x, y, xn, @@ -166,15 +171,15 @@ void fusedDistanceNN(OutT* min, initOutBuffer, isRowMajor, metric, - metric_arg, - stream); + metric_arg); } else { detail::fusedDistanceNNImpl< DataT, OutT, IdxT, typename raft::linalg::Policy4x4::Policy, - ReduceOpT>(min, + ReduceOpT>(handle, + min, x, y, xn, @@ -189,8 +194,7 @@ void fusedDistanceNN(OutT* min, initOutBuffer, isRowMajor, metric, - metric_arg, - stream); + metric_arg); } } else { if (is_skinny) { @@ -198,7 +202,8 @@ void fusedDistanceNN(OutT* min, OutT, IdxT, typename raft::linalg::Policy4x4Skinny::Policy, - ReduceOpT>(min, + ReduceOpT>(handle, + min, x, y, xn, @@ -213,14 +218,14 @@ void fusedDistanceNN(OutT* min, initOutBuffer, isRowMajor, metric, - metric_arg, - stream); + metric_arg); } else { detail::fusedDistanceNNImpl::Policy, - ReduceOpT>(min, + ReduceOpT>(handle, + min, x, y, xn, @@ -235,8 +240,7 @@ void fusedDistanceNN(OutT* min, initOutBuffer, isRowMajor, metric, - metric_arg, - stream); + metric_arg); } } } @@ -270,10 +274,11 @@ void fusedDistanceNN(OutT* min, * @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[in] handle RAFT resources containing the caller-provided CUDA stream */ template -void fusedDistanceNNMinReduce(OutT* min, +void fusedDistanceNNMinReduce(raft::resources const& handle, + OutT* min, const DataT* x, const DataT* y, const DataT* xn, @@ -286,29 +291,403 @@ void fusedDistanceNNMinReduce(OutT* min, bool initOutBuffer, bool isRowMajor, cuvs::distance::DistanceType metric, - float metric_arg, - cudaStream_t stream) + float metric_arg) { - MinAndDistanceReduceOp redOp; - KVPMinReduce pairRedOp; - - fusedDistanceNN(min, - x, - y, - xn, - yn, - m, - n, - k, - workspace, - redOp, - pairRedOp, - sqrt, - initOutBuffer, - isRowMajor, - metric, - metric_arg, - stream); + static_assert( + std::is_same_v> || std::is_same_v, + "fusedDistanceNNMinReduce supports KVP or scalar distance output"); + detail::Top1nnTuning tuning{}; + const auto workspace_bytes = + top_1_nn_workspace_size(m, n, tuning, detail::Top1nnBackend::Cutlass); + top_1_nn(handle, + min, + x, + y, + xn, + yn, + m, + n, + k, + tuning, + workspace, + workspace_bytes, + sqrt, + initOutBuffer, + isRowMajor, + metric, + metric_arg, + detail::Top1nnBackend::Cutlass); +} + +namespace detail { + +inline std::size_t checked_top_1_nn_workspace_multiply(std::size_t lhs, std::size_t rhs) +{ + RAFT_EXPECTS(rhs == 0 || lhs <= std::numeric_limits::max() / rhs, + "top_1_nn workspace size overflowed"); + return lhs * rhs; +} + +inline std::size_t checked_top_1_nn_workspace_add(std::size_t lhs, std::size_t rhs) +{ + RAFT_EXPECTS(lhs <= std::numeric_limits::max() - rhs, + "top_1_nn workspace size overflowed"); + return lhs + rhs; +} + +template +std::size_t checked_top_1_nn_extent(IdxT value) +{ + static_assert(std::is_integral_v); + if constexpr (std::is_signed_v) { + RAFT_EXPECTS(value >= 0, "top_1_nn dimensions must be non-negative"); + } + using UnsignedIdxT = std::make_unsigned_t; + RAFT_EXPECTS(static_cast(value) <= std::numeric_limits::max(), + "top_1_nn dimension does not fit in size_t"); + return static_cast(value); +} + +template +struct UnfusedTop1nnWorkspaceLayout { + IdxT row_tile; + IdxT candidate_tile; + std::size_t candidate_offset; + std::size_t candidate_bytes; + std::size_t total_bytes; +}; + +template +UnfusedTop1nnWorkspaceLayout make_unfused_top_1_nn_workspace_layout( + IdxT m, IdxT n, const Top1nnTuning& tuning) +{ + RAFT_EXPECTS(tuning.unfused.row_tile > 0 && tuning.unfused.candidate_tile > 0, + "Unfused top_1_nn tile dimensions must be positive"); + + const auto rows = checked_top_1_nn_extent(m); + const auto candidates = checked_top_1_nn_extent(n); + const auto row_tile = std::min(tuning.unfused.row_tile, rows); + const auto candidate_tile = std::min(tuning.unfused.candidate_tile, candidates); + const auto distance_bytes = checked_top_1_nn_workspace_multiply( + checked_top_1_nn_workspace_multiply(row_tile, candidate_tile), sizeof(DataT)); + + using KeyValueT = raft::KeyValuePair; + const auto candidate_bytes = candidate_tile < candidates + ? checked_top_1_nn_workspace_multiply(row_tile, sizeof(KeyValueT)) + : 0; + auto candidate_offset = distance_bytes; + if (candidate_bytes != 0) { + constexpr auto alignment = alignof(KeyValueT); + const auto padding = alignment - distance_bytes % alignment; + candidate_offset = checked_top_1_nn_workspace_add(distance_bytes, padding); + } + const auto total_bytes = checked_top_1_nn_workspace_add(candidate_offset, candidate_bytes); + + return {static_cast(row_tile), + static_cast(candidate_tile), + candidate_offset, + candidate_bytes, + total_bytes}; +} + +#if CUVS_CUTILE_ENABLED +template +void top_1_nn_cutile(raft::resources const& handle, + OutputT output, + const DataT* x, + const DataT* y, + const NormT* xn, + const NormT* yn, + IdxT m, + IdxT n, + IdxT k, + void* workspace, + bool sqrt, + bool init_out_buffer, + cuvs::distance::DistanceType metric) +{ + RAFT_EXPECTS(init_out_buffer, + "cuTile top_1_nn does not support accumulating into an initialized output"); + using OutputTypes = Top1nnOutputTypes; + constexpr bool is_separate_output = std::is_same_v; + if constexpr (is_fused_1nn_cutile_data_v && + std::is_same_v> && is_separate_output) { + launch_fused_1nn_tile(handle, + output.nearest_idx, + output.nearest_dist, + x, + y, + xn, + yn, + m, + n, + k, + metric, + sqrt, + workspace); + } else { + RAFT_FAIL( + "Requested cuTile fused 1-NN backend does not support these data, norm, or output types"); + } +} + +#endif + +template +void top_1_nn_legacy_fused(raft::resources const& handle, + OutputT output, + const DataT* x, + const DataT* y, + const NormT* xn, + const NormT* yn, + IdxT m, + IdxT n, + IdxT k, + void* workspace, + bool sqrt, + bool init_out_buffer, + bool is_row_major, + cuvs::distance::DistanceType metric, + float metric_arg) +{ + using OutputTypes = Top1nnOutputTypes; + using NativeOutputT = std::remove_pointer_t; + constexpr bool is_native_output = std::is_same_v || + std::is_same_v; + constexpr bool matching_norm_type = std::is_same_v; + + RAFT_EXPECTS(metric != cuvs::distance::DistanceType::InnerProduct, + "Legacy fused top_1_nn does not support InnerProduct"); + RAFT_EXPECTS(is_top_1_nn_backend_available(Top1nnBackend::Cutlass, x, y, m, n, k, metric), + "Requested legacy fused 1-NN backend is unavailable for this input"); + RAFT_EXPECTS(matching_norm_type, "Legacy fused top_1_nn requires matching norm types"); + + MinAndDistanceReduceOp red_op; + KVPMinReduce pair_red_op; + if constexpr (matching_norm_type && is_native_output) { + RAFT_EXPECTS(output != nullptr, "Legacy fused 1-NN requires a native output buffer"); + fusedDistanceNN(handle, + output, + x, + y, + xn, + yn, + m, + n, + k, + workspace, + red_op, + pair_red_op, + sqrt, + init_out_buffer, + is_row_major, + metric, + metric_arg); + } else { + RAFT_FAIL("Legacy fused top_1_nn requires matching norm types and native KVP or scalar output"); + } +} + +template +void top_1_nn_unfused(raft::resources const& handle, + OutputT output, + const DataT* x, + const DataT* y, + const NormT* xn, + const NormT* yn, + IdxT m, + IdxT n, + IdxT k, + const Top1nnTuning& tuning, + void* workspace, + std::size_t workspace_bytes, + bool sqrt, + bool init_out_buffer, + bool is_row_major, + cuvs::distance::DistanceType metric, + float metric_arg) +{ + using OutputTypes = Top1nnOutputTypes; + using NativeOutputT = std::remove_pointer_t; + using KeyValueT = raft::KeyValuePair; + constexpr bool is_native_output = std::is_same_v || + std::is_same_v; + constexpr bool matching_norm_type = std::is_same_v; + + RAFT_EXPECTS(metric != cuvs::distance::DistanceType::InnerProduct, + "Unfused top_1_nn does not support InnerProduct"); + RAFT_EXPECTS(matching_norm_type, "Unfused top_1_nn requires matching norm types"); + + if constexpr (matching_norm_type && is_native_output) { + RAFT_EXPECTS(output != nullptr, "Unfused top_1_nn requires a native output buffer"); + const auto layout = make_unfused_top_1_nn_workspace_layout(m, n, tuning); + RAFT_EXPECTS(layout.total_bytes == 0 || workspace != nullptr, + "Unfused top_1_nn requires a workspace buffer"); + RAFT_EXPECTS(workspace_bytes >= layout.total_bytes, + "Unfused top_1_nn workspace is smaller than its configured tile"); + + const auto row_tile = layout.row_tile; + const auto candidate_tile = layout.candidate_tile; + auto* candidate_min = + layout.candidate_bytes == 0 + ? nullptr + : reinterpret_cast(static_cast(workspace) + layout.candidate_offset); + for (IdxT row_offset = 0; row_offset < m;) { + const auto rows = std::min(row_tile, static_cast(m - row_offset)); + auto row_output = + raft::make_device_vector_view(output + row_offset, rows); + for (IdxT candidate_offset = 0; candidate_offset < n;) { + const auto candidates = std::min(candidate_tile, static_cast(n - candidate_offset)); + auto* tile_output = candidate_offset == 0 ? row_output.data_handle() : candidate_min; + unfusedDistanceNNMinReduce( + handle, + tile_output, + x + static_cast(row_offset) * static_cast(k), + y + static_cast(candidate_offset) * static_cast(k), + xn + row_offset, + yn + candidate_offset, + rows, + candidates, + k, + workspace, + sqrt, + candidate_offset != 0 || init_out_buffer, + is_row_major, + metric, + metric_arg); + if (candidate_offset != 0) { + auto candidate_output = + raft::make_device_vector_view(candidate_min, rows); + raft::linalg::map( + handle, + row_output, + [candidate_offset] __device__(NativeOutputT current, NativeOutputT candidate) { + if constexpr (std::is_same_v) { + candidate.key += candidate_offset; + return candidate.value < current.value ? candidate : current; + } else { + return candidate < current ? candidate : current; + } + }, + raft::make_const_mdspan(row_output), + candidate_output); + } + candidate_offset += candidates; + } + row_offset += rows; + } + } else { + RAFT_FAIL("Unfused top_1_nn requires matching norm types and native KVP or scalar output"); + } +} + +} // namespace detail + +template +std::size_t top_1_nn_workspace_size(IdxT m, + IdxT n, + const detail::Top1nnTuning& tuning, + detail::Top1nnBackend backend) +{ + const auto rows = detail::checked_top_1_nn_extent(m); + detail::checked_top_1_nn_extent(n); + switch (backend) { + case detail::Top1nnBackend::Cutile: +#if CUVS_CUTILE_ENABLED + if constexpr (std::is_same_v) { + return detail::checked_top_1_nn_workspace_multiply( + detail::fused_1nn_cutile_index_workspace_rows(m), sizeof(int)); + } +#endif + return 0; + case detail::Top1nnBackend::Cutlass: + return detail::checked_top_1_nn_workspace_multiply(rows, sizeof(int)); + case detail::Top1nnBackend::Unfused: + return detail::make_unfused_top_1_nn_workspace_layout(m, n, tuning).total_bytes; + } + RAFT_FAIL("Unknown top_1_nn backend"); +} + +template +void top_1_nn(raft::resources const& handle, + OutputT output, + const DataT* x, + const DataT* y, + const NormT* xn, + const NormT* yn, + IdxT m, + IdxT n, + IdxT k, + const detail::Top1nnTuning& tuning, + void* workspace, + std::size_t workspace_bytes, + bool sqrt, + bool init_out_buffer, + bool is_row_major, + cuvs::distance::DistanceType metric, + float metric_arg, + detail::Top1nnBackend backend) +{ + RAFT_EXPECTS(is_row_major, "top_1_nn only supports row-major inputs"); + RAFT_EXPECTS(m > 0 && n > 0 && k > 0, "top_1_nn requires positive m, n, and k"); + RAFT_EXPECTS(detail::is_top_1_nn_metric_supported(backend, metric), + "Selected top_1_nn backend does not support the requested metric"); + RAFT_EXPECTS(x != nullptr && y != nullptr, "top_1_nn requires non-null input buffers"); + RAFT_EXPECTS( + metric == cuvs::distance::DistanceType::InnerProduct || (xn != nullptr && yn != nullptr), + "top_1_nn requires non-null norm buffers for the requested metric"); + const auto required_workspace_bytes = top_1_nn_workspace_size(m, n, tuning, backend); + RAFT_EXPECTS(required_workspace_bytes == 0 || workspace != nullptr, + "top_1_nn requires a workspace buffer for the selected backend"); + RAFT_EXPECTS(workspace_bytes >= required_workspace_bytes, + "top_1_nn workspace is too small for the selected backend"); + switch (backend) { + case detail::Top1nnBackend::Cutile: +#if CUVS_CUTILE_ENABLED + detail::top_1_nn_cutile( + handle, output, x, y, xn, yn, m, n, k, workspace, sqrt, init_out_buffer, metric); +#else + RAFT_FAIL("Requested cuTile fused 1-NN backend was not built"); +#endif + return; + case detail::Top1nnBackend::Cutlass: + detail::top_1_nn_legacy_fused(handle, + output, + x, + y, + xn, + yn, + m, + n, + k, + workspace, + sqrt, + init_out_buffer, + is_row_major, + metric, + metric_arg); + return; + case detail::Top1nnBackend::Unfused: + detail::top_1_nn_unfused(handle, + output, + x, + y, + xn, + yn, + m, + n, + k, + tuning, + workspace, + workspace_bytes, + sqrt, + init_out_buffer, + is_row_major, + metric, + metric_arg); + return; + } + RAFT_FAIL("Unknown top_1_nn backend"); } /** @} */ diff --git a/cpp/src/distance/top_1_nn.cu b/cpp/src/distance/top_1_nn.cu new file mode 100644 index 0000000000..b9573d358d --- /dev/null +++ b/cpp/src/distance/top_1_nn.cu @@ -0,0 +1,60 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include "fused_distance_nn.cuh" + +namespace cuvs::distance { + +#define CUVS_INSTANTIATE_TOP_1_NN_WORKSPACE_SIZE(DataT, IdxT) \ + template CUVS_EXPORT std::size_t top_1_nn_workspace_size( \ + IdxT, IdxT, const detail::Top1nnTuning&, detail::Top1nnBackend) + +CUVS_INSTANTIATE_TOP_1_NN_WORKSPACE_SIZE(float, int); +CUVS_INSTANTIATE_TOP_1_NN_WORKSPACE_SIZE(float, int64_t); +CUVS_INSTANTIATE_TOP_1_NN_WORKSPACE_SIZE(double, int); +CUVS_INSTANTIATE_TOP_1_NN_WORKSPACE_SIZE(double, int64_t); +CUVS_INSTANTIATE_TOP_1_NN_WORKSPACE_SIZE(half, int); +CUVS_INSTANTIATE_TOP_1_NN_WORKSPACE_SIZE(half, int64_t); + +#undef CUVS_INSTANTIATE_TOP_1_NN_WORKSPACE_SIZE + +#define CUVS_INSTANTIATE_TOP_1_NN(DataT, IdxT, NormT, OutputKind) \ + template CUVS_EXPORT void \ + top_1_nn::OutputKind, NormT>( \ + raft::resources const&, \ + typename detail::Top1nnOutputTypes::OutputKind, \ + const DataT*, \ + const DataT*, \ + const NormT*, \ + const NormT*, \ + IdxT, \ + IdxT, \ + IdxT, \ + const detail::Top1nnTuning&, \ + void*, \ + std::size_t, \ + bool, \ + bool, \ + bool, \ + DistanceType, \ + float, \ + detail::Top1nnBackend) + +CUVS_INSTANTIATE_TOP_1_NN(float, int, float, kvp); +CUVS_INSTANTIATE_TOP_1_NN(float, int, float, scalar); +CUVS_INSTANTIATE_TOP_1_NN(float, int, float, separate); +CUVS_INSTANTIATE_TOP_1_NN(float, int64_t, float, kvp); +CUVS_INSTANTIATE_TOP_1_NN(float, int64_t, float, scalar); +CUVS_INSTANTIATE_TOP_1_NN(float, int64_t, float, separate); +CUVS_INSTANTIATE_TOP_1_NN(double, int, double, kvp); +CUVS_INSTANTIATE_TOP_1_NN(double, int, double, scalar); +CUVS_INSTANTIATE_TOP_1_NN(double, int64_t, double, kvp); +CUVS_INSTANTIATE_TOP_1_NN(double, int64_t, double, scalar); +CUVS_INSTANTIATE_TOP_1_NN(half, int, float, separate); +CUVS_INSTANTIATE_TOP_1_NN(half, int64_t, float, separate); + +#undef CUVS_INSTANTIATE_TOP_1_NN + +} // namespace cuvs::distance diff --git a/cpp/src/distance/top_1_nn.cuh b/cpp/src/distance/top_1_nn.cuh new file mode 100644 index 0000000000..b43450468c --- /dev/null +++ b/cpp/src/distance/top_1_nn.cuh @@ -0,0 +1,123 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +#include "detail/fused_distance_nn.cuh" + +#include + +#include +#include + +#include + +namespace cuvs::distance { + +/** Separate index and distance arrays used by backends with structure-of-arrays output. */ +template +struct Top1nnOutput { + IdxT* nearest_idx; + DistT* nearest_dist; +}; + +namespace detail { + +template +using top_1_nn_distance_t = std::conditional_t, float, DataT>; + +template +struct Top1nnOutputTypes { + using kvp = raft::KeyValuePair*; + using scalar = DataT*; + using separate = Top1nnOutput>; +}; + +} // namespace detail + +/** + * Return the workspace bytes required for one top-1 NN call. + * + * Callers that batch a larger problem should pass their maximum batch dimensions and reuse one + * allocation across calls. + */ +template +CUVS_EXPORT std::size_t top_1_nn_workspace_size(IdxT m, + IdxT n, + const detail::Top1nnTuning& tuning, + detail::Top1nnBackend backend); + +/** Dispatch 1-NN to a selected backend using its native output representation. */ +template +CUVS_EXPORT void top_1_nn(raft::resources const& handle, + OutputT output, + const DataT* x, + const DataT* y, + const NormT* xn, + const NormT* yn, + IdxT m, + IdxT n, + IdxT k, + const detail::Top1nnTuning& tuning, + void* workspace, + std::size_t workspace_bytes, + bool sqrt, + bool init_out_buffer, + bool is_row_major, + DistanceType metric, + float metric_arg, + detail::Top1nnBackend backend); + +#define CUVS_EXTERN_TOP_1_NN_WORKSPACE_SIZE(DataT, IdxT) \ + extern template std::size_t top_1_nn_workspace_size( \ + IdxT, IdxT, const detail::Top1nnTuning&, detail::Top1nnBackend) + +CUVS_EXTERN_TOP_1_NN_WORKSPACE_SIZE(float, int); +CUVS_EXTERN_TOP_1_NN_WORKSPACE_SIZE(float, int64_t); +CUVS_EXTERN_TOP_1_NN_WORKSPACE_SIZE(double, int); +CUVS_EXTERN_TOP_1_NN_WORKSPACE_SIZE(double, int64_t); +CUVS_EXTERN_TOP_1_NN_WORKSPACE_SIZE(half, int); +CUVS_EXTERN_TOP_1_NN_WORKSPACE_SIZE(half, int64_t); + +#undef CUVS_EXTERN_TOP_1_NN_WORKSPACE_SIZE + +#define CUVS_EXTERN_TOP_1_NN(DataT, IdxT, NormT, OutputKind) \ + extern template void \ + top_1_nn::OutputKind, NormT>( \ + raft::resources const&, \ + typename detail::Top1nnOutputTypes::OutputKind, \ + const DataT*, \ + const DataT*, \ + const NormT*, \ + const NormT*, \ + IdxT, \ + IdxT, \ + IdxT, \ + const detail::Top1nnTuning&, \ + void*, \ + std::size_t, \ + bool, \ + bool, \ + bool, \ + DistanceType, \ + float, \ + detail::Top1nnBackend) + +CUVS_EXTERN_TOP_1_NN(float, int, float, kvp); +CUVS_EXTERN_TOP_1_NN(float, int, float, scalar); +CUVS_EXTERN_TOP_1_NN(float, int, float, separate); +CUVS_EXTERN_TOP_1_NN(float, int64_t, float, kvp); +CUVS_EXTERN_TOP_1_NN(float, int64_t, float, scalar); +CUVS_EXTERN_TOP_1_NN(float, int64_t, float, separate); +CUVS_EXTERN_TOP_1_NN(double, int, double, kvp); +CUVS_EXTERN_TOP_1_NN(double, int, double, scalar); +CUVS_EXTERN_TOP_1_NN(double, int64_t, double, kvp); +CUVS_EXTERN_TOP_1_NN(double, int64_t, double, scalar); +CUVS_EXTERN_TOP_1_NN(half, int, float, separate); +CUVS_EXTERN_TOP_1_NN(half, int64_t, float, separate); + +#undef CUVS_EXTERN_TOP_1_NN + +} // namespace cuvs::distance diff --git a/cpp/src/distance/unfused_distance_nn.cuh b/cpp/src/distance/unfused_distance_nn.cuh index f85de31937..a688c30869 100644 --- a/cpp/src/distance/unfused_distance_nn.cuh +++ b/cpp/src/distance/unfused_distance_nn.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 */ @@ -11,6 +11,7 @@ #include #include +#include #include namespace cuvs { @@ -81,10 +82,9 @@ __global__ void reduce_min_kernel(OutT* out, // GEMM round-off can produce slightly negative expanded distances; clamp to zero. dist = (dist > AccT(0)) ? dist : AccT(0); } else if constexpr (metric == DistanceType::CosineExpanded) { - // Guard against zero-norm vectors to avoid inf/NaN from division by zero. - AccT denom = x_norm_row * y_norm[col]; - denom = (denom > AccT(0)) ? denom : AccT(1); - dist = AccT(1.0) - (z[row * n + col] / denom); + // Cosine distance involving any zero-norm vector is defined as 1. + const AccT denom = x_norm_row * y_norm[col]; + dist = denom > AccT(0) ? AccT(1.0) - (z[row * n + col] / denom) : AccT(1.0); } if (dist < thread_min.value) { thread_min.value = dist; @@ -131,6 +131,7 @@ __global__ void reduce_min_kernel(OutT* out, * @tparam IdxT Index type * @tparam metric Distance metric type (L2Expanded, L2SqrtExpanded, or CosineExpanded) * + * @param[in] handle RAFT resources containing the caller-provided CUDA stream * @param[out] out Output array containing minimum distances (and optionally indices) * per row. Length = `m`. (on device) * @param[in] z GEMM output matrix (x * y^T). Dim = `m x n`. (on device) @@ -138,22 +139,22 @@ __global__ void reduce_min_kernel(OutT* out, * @param[in] y_norm Norms of rows in y. Length = `n`. (on device) * @param[in] m Number of rows in x (and output) * @param[in] n Number of rows in y (columns to reduce over) - * @param[in] stream CUDA stream for kernel launch * @param[in] is_sqrt Whether to apply square root to the final distance * @param[in] initOutBuffer Whether to initialize the output buffer or merge with existing values */ template -void reduce_min(OutT* out, +void reduce_min(raft::resources const& handle, + OutT* out, const AccT* z, const AccT* x_norm, const AccT* y_norm, IdxT m, IdxT n, - cudaStream_t stream, bool is_sqrt, bool initOutBuffer) { - const int TPB = 128; + const auto stream = raft::resource::get_cuda_stream(handle); + const int TPB = 128; int blocks = m; reduce_min_kernel @@ -170,9 +171,9 @@ void pairwise_distance_gemm(raft::resources const& handle, IdxT N, IdxT K, const AccT* x_norm, - const AccT* y_norm, - cudaStream_t stream) + const AccT* y_norm) { + const auto stream = raft::resource::get_cuda_stream(handle); cudaDataType_t xyType, zType; cublasComputeType_t computeType; @@ -253,6 +254,7 @@ void pairwise_distance_gemm(raft::resources const& handle, * distances or store only the min distances. * @tparam IdxT indexing arithmetic type * + * @param[in] handle RAFT resources containing the caller-provided CUDA stream * @param[out] min will contain the reduced output (Length = `m`) * (on device) * @param[in] x first matrix. Row major. Dim = `m x k`. @@ -271,7 +273,6 @@ void pairwise_distance_gemm(raft::resources const& handle, * @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 unfusedDistanceNNMinReduce(raft::resources const& handle, @@ -288,14 +289,12 @@ void unfusedDistanceNNMinReduce(raft::resources const& handle, bool initOutBuffer, bool isRowMajor, DistanceType metric, - float metric_arg, - cudaStream_t stream) + float metric_arg) { ASSERT(isRowMajor, "unfusedDistanceNN only supports row major inputs"); ASSERT(m > 0 && n > 0 && k > 0, "unfusedDistanceNN requires non-zero m, n, and k"); - pairwise_distance_gemm( - handle, (AccT*)workspace, x, y, m, n, k, xn, yn, stream); + pairwise_distance_gemm(handle, (AccT*)workspace, x, y, m, n, k, xn, yn); ASSERT((metric == DistanceType::CosineExpanded) || (metric == DistanceType::L2Expanded) || (metric == DistanceType::L2SqrtExpanded), @@ -303,13 +302,13 @@ void unfusedDistanceNNMinReduce(raft::resources const& handle, if (metric == DistanceType::L2Expanded) { reduce_min( - min, (AccT*)workspace, xn, yn, m, n, stream, is_sqrt, initOutBuffer); + handle, min, (AccT*)workspace, xn, yn, m, n, is_sqrt, initOutBuffer); } else if (metric == DistanceType::L2SqrtExpanded) { reduce_min( - min, (AccT*)workspace, xn, yn, m, n, stream, is_sqrt, initOutBuffer); + handle, min, (AccT*)workspace, xn, yn, m, n, is_sqrt, initOutBuffer); } else if (metric == DistanceType::CosineExpanded) { reduce_min( - min, (AccT*)workspace, xn, yn, m, n, stream, is_sqrt, initOutBuffer); + handle, min, (AccT*)workspace, xn, yn, m, n, is_sqrt, initOutBuffer); } } diff --git a/cpp/tests/CMakeLists.txt b/cpp/tests/CMakeLists.txt index 7d3720be08..6a78207855 100644 --- a/cpp/tests/CMakeLists.txt +++ b/cpp/tests/CMakeLists.txt @@ -114,6 +114,7 @@ ConfigureTest( GPUS 1 PERCENT 100 ) +target_compile_definitions(NEIGHBORS_TEST PRIVATE CUVS_CUTILE_ENABLED=${CUVS_CUTILE_ENABLED}) ConfigureTest( NAME NEIGHBORS_TIERED_INDEX_TEST @@ -142,6 +143,7 @@ ConfigureTest( GPUS 1 PERCENT 100 ) +target_compile_definitions(CUTILE_SMOKE_TEST PRIVATE CUVS_CUTILE_ENABLED=${CUVS_CUTILE_ENABLED}) if(CUVS_CUTILE_ENABLED) # These are intentionally library-private implementation symbols. Build the smoke executable with # the generated fragment registrations and planner implementation so it can exercise them without diff --git a/cpp/tests/detail/jit_lto/cutile_smoke.cu b/cpp/tests/detail/jit_lto/cutile_smoke.cu index 843835e22e..27b8116dc8 100644 --- a/cpp/tests/detail/jit_lto/cutile_smoke.cu +++ b/cpp/tests/detail/jit_lto/cutile_smoke.cu @@ -54,15 +54,20 @@ void add_smoke_fragments(TileAlgorithmPlanner& planner) } // namespace -TEST(CutileSmoke, ResolvesEveryEmbeddedArchitecture) +TEST(CutileSmoke, ResolvesExactArchitectureOrSm89Compatibility) { auto fragments = make_smoke_fragments(); EXPECT_EQ(find_compatible_cubin_fragment(8, 0, fragments), fragments[0].get()); + EXPECT_EQ(find_compatible_cubin_fragment(8, 6, fragments), fragments[1].get()); EXPECT_EQ(find_compatible_cubin_fragment(8, 9, fragments), fragments[1].get()); EXPECT_EQ(find_compatible_cubin_fragment(9, 0, fragments), fragments[2].get()); + EXPECT_EQ(find_compatible_cubin_fragment(9, 1, fragments), nullptr); EXPECT_EQ(find_compatible_cubin_fragment(10, 0, fragments), fragments[3].get()); - EXPECT_EQ(find_compatible_cubin_fragment(12, 1, fragments), fragments[4].get()); + EXPECT_EQ(find_compatible_cubin_fragment(10, 1, fragments), nullptr); + EXPECT_EQ(find_compatible_cubin_fragment(12, 0, fragments), fragments[4].get()); + EXPECT_EQ(find_compatible_cubin_fragment(12, 1, fragments), nullptr); + EXPECT_EQ(find_compatible_cubin_fragment(13, 0, fragments), nullptr); EXPECT_EQ(find_compatible_cubin_fragment(7, 5, fragments), nullptr); } diff --git a/cpp/tests/neighbors/distance_nn.cu b/cpp/tests/neighbors/distance_nn.cu index 7c90c515d1..fccbf6cf60 100644 --- a/cpp/tests/neighbors/distance_nn.cu +++ b/cpp/tests/neighbors/distance_nn.cu @@ -9,6 +9,7 @@ #include "../../src/distance/fused_distance_nn.cuh" #include "../../src/distance/unfused_distance_nn.cuh" +#include #include #include #include @@ -27,6 +28,8 @@ struct NNInputs { bool sqrt; uint64_t rng_seed; double tol; + cuvs::distance::detail::Top1nnBackend backend = cuvs::distance::detail::Top1nnBackend::Cutlass; + cuvs::distance::detail::Top1nnTuning tuning{}; }; __global__ void fill_int8(int8_t* buff, int len, int seed_offset) @@ -50,13 +53,19 @@ class NNTest : public ::testing::TestWithParam> { k{params_.k}, metric{params_.metric}, sqrt{params_.sqrt}, + backend{params_.backend}, + tuning{params_.tuning}, stream{raft::resource::get_cuda_stream(handle)}, x{raft::make_device_matrix(handle, m, k)}, y{raft::make_device_matrix(handle, n, k)}, x_norm{raft::make_device_vector(handle, m)}, y_norm{raft::make_device_vector(handle, n)}, out{raft::make_device_vector(handle, m)}, - ref_out{raft::make_device_vector(handle, m)} + ref_out{raft::make_device_vector(handle, m)}, + ref_dist{raft::make_device_vector(handle, m)}, + selected_dist{raft::make_device_vector(handle, m)}, + cutile_idx{raft::make_device_vector(handle, m)}, + cutile_dist{raft::make_device_vector(handle, m)} { } @@ -89,7 +98,7 @@ class NNTest : public ::testing::TestWithParam> { } if constexpr (impl == ImplType::fused) { - workspace_size = m * sizeof(IdxT); + workspace_size = cuvs::distance::top_1_nn_workspace_size(m, n, tuning, backend); } else if constexpr (impl == ImplType::unfused) { workspace_size = m * n * sizeof(AccT); } @@ -98,10 +107,11 @@ class NNTest : public ::testing::TestWithParam> { 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}); + handle, raft::make_device_matrix_view(out.data_handle(), m, IdxT{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.data_handle(), m, IdxT{1}), OutT{0}); } raft::resource::sync_stream(handle, stream); } @@ -112,28 +122,45 @@ class NNTest : public ::testing::TestWithParam> { raft::make_device_vector(handle, workspace_size); ref_nn( - ref_out.data_handle(), x.data_handle(), y.data_handle(), m, n, k, sqrt, metric, stream.get()); + handle, ref_out.data_handle(), x.data_handle(), y.data_handle(), m, n, k, sqrt, metric); 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.get()); + if (backend == cuvs::distance::detail::Top1nnBackend::Cutile && + !cuvs::distance::detail::is_top_1_nn_backend_available( + backend, x.data_handle(), y.data_handle(), m, n, k, metric)) { + GTEST_SKIP() << "cuTile is not available for this device/input"; + } + auto run_top_1_nn = [&](auto output) { + cuvs::distance::top_1_nn(handle, + output, + x.data_handle(), + y.data_handle(), + x_norm.data_handle(), + y_norm.data_handle(), + m, + n, + k, + tuning, + (void*)workspace.data_handle(), + workspace_size, + sqrt, + true, + true, + metric, + 0.0, + backend); + }; + if (backend == cuvs::distance::detail::Top1nnBackend::Cutile) { + if constexpr (std::is_same_v || std::is_same_v) { + run_top_1_nn(cuvs::distance::Top1nnOutput{cutile_idx.data_handle(), + cutile_dist.data_handle()}); + } else { + RAFT_FAIL("cuTile top_1_nn test requires FP16 or FP32 data"); + } + } else if constexpr (std::is_same_v) { + run_top_1_nn(out.data_handle()); } else { - static_assert(sizeof(DataT) == 0, - "fusedDistanceNNMinReduce is not implemented for datatype other than float"); + RAFT_FAIL("Legacy fused top_1_nn test requires FP32 data"); } } else if constexpr (impl == ImplType::unfused) { cuvs::distance::unfusedDistanceNNMinReduce( @@ -151,14 +178,45 @@ class NNTest : public ::testing::TestWithParam> { true, true, metric, - 0.0, - stream.get()); + 0.0); } } void compare() { - vector_compare(handle, ref_out.data_handle(), out.data_handle(), m, summary); + if constexpr (impl == ImplType::fused) { + if (backend == cuvs::distance::detail::Top1nnBackend::Cutile) { + // cuTile MMA arithmetic can produce a different index for nearly tied candidates. + // Validate that the returned index selects a candidate within the same numerical tolerance + // of the true optimum. + raft::linalg::unaryOp( + ref_dist.data_handle(), ref_out.data_handle(), m, raft::value_op{}, stream); + ref_nn_selected(handle, + selected_dist.data_handle(), + cutile_idx.data_handle(), + x.data_handle(), + y.data_handle(), + m, + n, + k, + sqrt, + metric); + ASSERT_TRUE(cuvs::devArrMatch(ref_dist.data_handle(), + selected_dist.data_handle(), + m, + cuvs::CompareApproxNoScaling{AccT(params_.tol)}, + stream)); + ASSERT_TRUE(cuvs::devArrMatch(ref_dist.data_handle(), + cutile_dist.data_handle(), + m, + cuvs::CompareApproxNoScaling{AccT(params_.tol)}, + stream)); + return; + } + vector_compare(handle, ref_out.data_handle(), out.data_handle(), m, summary); + } else { + vector_compare(handle, ref_out.data_handle(), out.data_handle(), m, summary); + } ASSERT_TRUE(summary.max_diff < params_.tol) << summary; } @@ -172,12 +230,18 @@ class NNTest : public ::testing::TestWithParam> { IdxT k; DistanceType metric; bool sqrt; + cuvs::distance::detail::Top1nnBackend backend; + cuvs::distance::detail::Top1nnTuning tuning; raft::device_matrix x; raft::device_matrix y; raft::device_vector x_norm; raft::device_vector y_norm; raft::device_vector out; raft::device_vector ref_out; + raft::device_vector ref_dist; + raft::device_vector selected_dist; + raft::device_vector cutile_idx; + raft::device_vector cutile_dist; size_t workspace_size; }; @@ -197,6 +261,31 @@ 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; + for (auto input : input_fp32) { + input.backend = cuvs::distance::detail::Top1nnBackend::Unfused; + inputs.push_back(input); + } +#if CUVS_CUTILE_ENABLED + for (auto input : input_fp32) { + input.backend = cuvs::distance::detail::Top1nnBackend::Cutile; + inputs.push_back(input); + } + // Non-vector-aligned k selects the relaxed ABI; InnerProduct exercises its argmax path. + inputs.push_back({257, + 263, + 65, + DistanceType::InnerProduct, + false, + uint64_t(31415926), + 0.1, + cuvs::distance::detail::Top1nnBackend::Cutile}); +#endif + return inputs; +}(); + // Test fused implementation with single-precision typedef NNTest NNTest_fp32_fused; TEST_P(NNTest_fp32_fused, test) @@ -205,7 +294,24 @@ 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)); + +#if CUVS_CUTILE_ENABLED +const std::vector> input_fp32_cutile_i64 = [] { + auto input = input_fp32.front(); + input.backend = cuvs::distance::detail::Top1nnBackend::Cutile; + return std::vector>{input}; +}(); + +using NNTest_fp32_fused_i64 = NNTest; +TEST_P(NNTest_fp32_fused_i64, test) +{ + this->compute_1nn(); + this->compare(); +} + +INSTANTIATE_TEST_CASE_P(NNTest, NNTest_fp32_fused_i64, ::testing::ValuesIn(input_fp32_cutile_i64)); +#endif // Test unfused implementation with single-precision typedef NNTest NNTest_fp32_unfused; @@ -227,8 +333,40 @@ const std::vector> input_fp16 = { {4096, 16384, 128, DistanceType::CosineExpanded, true, uint64_t(31415926), 0.1}, }; +#if CUVS_CUTILE_ENABLED +template +// k=64 and k=65 select the strict and relaxed FP16 ABI variants, respectively. +const std::vector> input_fp16_cutile = { + {257, + 263, + 64, + DistanceType::L2Expanded, + false, + uint64_t(31415926), + 0.1, + cuvs::distance::detail::Top1nnBackend::Cutile}, + {257, + 263, + 65, + DistanceType::CosineExpanded, + false, + uint64_t(31415926), + 0.1, + cuvs::distance::detail::Top1nnBackend::Cutile}, +}; + +using NNTest_fp16_fused = NNTest; +TEST_P(NNTest_fp16_fused, test) +{ + this->compute_1nn(); + this->compare(); +} + +INSTANTIATE_TEST_CASE_P(NNTest, NNTest_fp16_fused, ::testing::ValuesIn(input_fp16_cutile)); +#endif + // Test unfused implementation with fp16, int8 -// Fused implementation has no support for fp16, int8 so no test for it +// Legacy fused implementation has no support for fp16, int8 typedef NNTest NNTest_fp16_unfused; TEST_P(NNTest_fp16_unfused, test) { diff --git a/cpp/tests/neighbors/distance_nn_helper.cuh b/cpp/tests/neighbors/distance_nn_helper.cuh index fda7b76573..c9d1a19de5 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 */ @@ -47,6 +47,16 @@ __device__ AccT l2_distance(const DataT* v1, const DataT* v2, IdxT K) return th_dist; } +template +__device__ AccT inner_product(const DataT* v1, const DataT* v2, IdxT K) +{ + AccT result = AccT(0.0); + for (IdxT k = 0; k < K; k++) { + result += AccT(v1[k]) * AccT(v2[k]); + } + return result; +} + template __device__ AccT cosine_distance(const DataT* v1, const DataT* v2, IdxT K) { @@ -75,6 +85,7 @@ RAFT_KERNEL ref_nn_kernel( for (IdxT m = tid; m < M; m += (blockDim.x * gridDim.x)) { IdxT min_index = N + 1; AccT min_dist = max_val(); + AccT min_score = max_val(); for (IdxT n = 0; n < N; n++) { AccT dist; @@ -82,8 +93,12 @@ RAFT_KERNEL ref_nn_kernel( 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 { + dist = inner_product(&A[m * K], &B[n * K], K); } - if (dist < min_dist) { + const auto score = metric == DistanceType::InnerProduct ? -dist : dist; + if (score < min_score) { + min_score = score; min_dist = dist; min_index = n; } @@ -97,7 +112,7 @@ RAFT_KERNEL ref_nn_kernel( static_assert(std::is_same>::value, "OutT is not raft::KeyValuePair<> type"); out[m].key = IdxT(min_index); - if (sqrt) { + if (sqrt && metric != DistanceType::InnerProduct) { out[m].value = raft::sqrt(AccT(min_dist)); } else { out[m].value = AccT(min_dist); @@ -107,16 +122,17 @@ RAFT_KERNEL ref_nn_kernel( } template -void ref_nn(OutT* out, +void ref_nn(raft::resources const& handle, + OutT* out, const DataT* A, const DataT* B, IdxT m, IdxT n, IdxT k, bool sqrt, - DistanceType metric, - cudaStream_t stream) + DistanceType metric) { + const auto stream = raft::resource::get_cuda_stream(handle); ref_nn_kernel <<<(m + 127) / 128, 128, 0, stream>>>(out, A, B, m, n, k, sqrt, metric); @@ -124,6 +140,55 @@ void ref_nn(OutT* out, return; } +template +RAFT_KERNEL ref_nn_selected_kernel(AccT* out, + const IdxT* selected_indices, + const DataT* A, + const DataT* B, + IdxT M, + IdxT N, + IdxT K, + bool sqrt, + DistanceType metric) +{ + IdxT tid = threadIdx.x + blockIdx.x * IdxT(blockDim.x); + for (IdxT m = tid; m < M; m += blockDim.x * gridDim.x) { + const auto n = selected_indices[m]; + if (n < 0 || n >= N) { + out[m] = max_val(); + 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 { + dist = inner_product(&A[m * K], &B[n * K], K); + } + out[m] = sqrt && metric != DistanceType::InnerProduct ? raft::sqrt(dist) : dist; + } +} + +template +void ref_nn_selected(raft::resources const& handle, + AccT* out, + const IdxT* selected_indices, + const DataT* A, + const DataT* B, + IdxT m, + IdxT n, + IdxT k, + bool sqrt, + DistanceType metric) +{ + const auto stream = raft::resource::get_cuda_stream(handle); + ref_nn_selected_kernel + <<<(m + 127) / 128, 128, 0, stream>>>(out, selected_indices, A, B, m, n, k, sqrt, metric); + RAFT_CUDA_TRY(cudaGetLastError()); +} + // Structure to track comparison failures class ComparisonSummary { public: