New 1-nn primitive using cutile

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+=(

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::neighbors::detail::tag_@data_abbrev@, cuvs::neighbors::detail::tag_index_@index_abbrev@, cuvs::distance::detail::cutile_tile_config<@tile_m@, @tile_n@, @tile_k@>, cuvs::distance::detail::@abi_tag@, cuvs::detail::jit_lto::@arch_tag@>"
+    FRAGMENT_TAG_HEADER_FILES
+      "<cuvs/detail/jit_lto/fused_distance_nn/fused_1nn_fragments.hpp>"
+      "<cuvs/detail/jit_lto/cutile_arch_tags.hpp>" "<cuvs/detail/jit_lto/common_fragments.hpp>"
+  )
+  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 <cuvs/detail/jit_lto/fused_distance_nn/fused_1nn_fragments.hpp>"
+    )
+    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 $<$<BOOL:${BUILD_CAGRA_HNSWLIB}>:CUVS_BUILD_CAGRA_HNSWLIB>
-            $<$<BOOL:${BUILD_MG_ALGOS}>:CUVS_BUILD_MG_ALGOS> $<$<BOOL:${CUVS_NVTX}>:NVTX_ENABLED>
+    PRIVATE CUVS_CUTILE_ENABLED=${CUVS_CUTILE_ENABLED}
+            $<$<BOOL:${BUILD_CAGRA_HNSWLIB}>:CUVS_BUILD_CAGRA_HNSWLIB>
+            $<$<BOOL:${BUILD_MG_ALGOS}>:CUVS_BUILD_MG_ALGOS>
+            $<$<BOOL:${CUVS_NVTX}>:NVTX_ENABLED>
   )
 
   target_link_libraries(
@@ -1541,7 +1592,7 @@ if(NOT BUILD_CPU_ONLY)
            "$<BUILD_INTERFACE:${CMAKE_CURRENT_SOURCE_DIR}/../c/include>"
     INTERFACE "$<INSTALL_INTERFACE:include>"
     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

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": "*"
         }
       },

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

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<fragment_tag>")
-  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<fragment_tag>"
-    )
-  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<fragment_tag>")
 
   _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}"
     )

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 <typename FragmentTag>
-struct StaticTileIrBytecodeFragmentEntry final : TileIrBytecodeFragmentEntry {
-  const uint8_t* get_data() const override
-  {
-    return StaticTileIrBytecodeFragmentEntry<FragmentTag>::data;
-  }
-
-  size_t get_length() const override
-  {
-    return StaticTileIrBytecodeFragmentEntry<FragmentTag>::length;
-  }
-
-  const char* get_key() const override
-  {
-    return typeid(StaticTileIrBytecodeFragmentEntry<FragmentTag>).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

cpp/include/cuvs/detail/jit_lto/TileAlgorithmPlanner.hpp

@@ -29,7 +29,7 @@ struct TileLauncherCache {
   std::unordered_set<std::string> 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<StaticCubinFragmentEntry<FragmentTag>>());
   }
 
-  template <typename FragmentTag>
-  void add_static_tileir_fragment()
-  {
-    tileir_fragment_ = std::make_unique<StaticTileIrBytecodeFragmentEntry<FragmentTag>>();
-  }
-
-  /** 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<std::unique_ptr<CubinFragmentEntry>> cubin_fragments_;
-  std::unique_ptr<TileIrBytecodeFragmentEntry> tileir_fragment_;
 
  private:
   std::string get_planner_key(const CutileRuntimeCapabilities* capabilities) const;

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

cpp/include/cuvs/detail/jit_lto/tileir_compat.hpp → cpp/include/cuvs/detail/jit_lto/cutile_compat.hpp

@@ -9,7 +9,6 @@
 #define CUVS_CUTILE_ENABLED 0
 #endif
 
-#include <cuda.h>
 #include <cuda_runtime.h>
 
 #include <cuvs/detail/jit_lto/cutile_arch_tags.hpp>
@@ -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,44 +63,32 @@ 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
 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. */