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

@@ -1174,8 +1174,47 @@ 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_FORMAT_TILEIR
+      "cuvs::distance::detail::fragment_tag_fused_1nn_tileir<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@>"
+    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)
+    cuvs_find_build_python(cutile_tile_metadata_python)
+    add_custom_command(
+      OUTPUT "${cutile_fused_1nn_tiles}"
+      COMMAND
+        "${cutile_tile_metadata_python}"
+        "${CMAKE_CURRENT_SOURCE_DIR}/cmake/modules/generate_cutile_tile_metadata.py" --matrix
+        "${fused_1nn_cutile_dir}/fused_1nn_cutile_matrix.json" --output "${cutile_fused_1nn_tiles}"
+        --namespace "cuvs::distance::detail" --include
+        "<cuvs/detail/jit_lto/fused_distance_nn/fused_1nn_fragments.hpp>" --alias-prefix
+        fused_1nn_matrix_tile
+      DEPENDS "${fused_1nn_cutile_dir}/fused_1nn_cutile_matrix.json"
+              "${CMAKE_CURRENT_SOURCE_DIR}/cmake/modules/generate_cutile_tile_metadata.py"
+      VERBATIM
+    )
+    list(APPEND cutile_fused_1nn_files "${cutile_fused_1nn_tiles}")
+  endif()
   # Note that this matrix contains an `arch_includes` placeholder, since we don't currently have a
   # way to do an item-wise transform on a list after computing the matrix product and before
   # configuring the file
@@ -1393,6 +1432,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
@@ -1486,9 +1526,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
@@ -1509,8 +1556,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(
@@ -1531,7 +1580,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/modules/generate_cutile_tile_metadata.py

@@ -0,0 +1,55 @@
+#!/usr/bin/env python3
+# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+import argparse
+import json
+import runpy
+from pathlib import Path
+
+iterate_matrix_product = runpy.run_path(
+    str(Path(__file__).with_name("compute_matrix_product.py"))
+)["iterate_matrix_product"]
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--matrix", type=Path, required=True)
+    parser.add_argument("--output", type=Path, required=True)
+    parser.add_argument("--namespace", required=True)
+    parser.add_argument("--include", required=True)
+    parser.add_argument("--alias-prefix", required=True)
+    args = parser.parse_args()
+    aliases = {}
+    matrix = json.loads(args.matrix.read_text())
+    for entry in iterate_matrix_product(matrix=matrix):
+        tile = tuple(entry.get(key) for key in ("tile_m", "tile_n", "tile_k"))
+        if any(value is None for value in tile):
+            raise ValueError("missing cuTile tile geometry")
+        suffix = (
+            f"{entry['data_abbrev']}_"
+            f"{entry.get('arch_tag', 'tileir')}_"
+            f"{entry['abi_abbrev']}"
+        )
+        if suffix in aliases and aliases[suffix] != tile:
+            raise ValueError(f"conflicting tile geometry for {suffix}")
+        aliases[suffix] = tile
+    lines = [
+        "#pragma once",
+        "",
+        f"#include {args.include}",
+        "",
+        f"namespace {args.namespace} {{",
+        "",
+    ]
+    for suffix, (m, n, k) in sorted(aliases.items()):
+        lines.append(
+            f"using {args.alias_prefix}_{suffix} = cutile_tile_config<{m}, {n}, {k}>;"
+        )
+    lines.extend(["", f"}}  // namespace {args.namespace}", ""])
+    args.output.parent.mkdir(parents=True, exist_ok=True)
+    args.output.write_text("\n".join(lines))
+
+
+if __name__ == "__main__":
+    main()

cpp/include/cuvs/detail/jit_lto/fused_distance_nn/fused_1nn_fragments.hpp

@@ -0,0 +1,29 @@
+/*
+ * 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 <int TileM, int TileN, int TileK>
+struct cutile_tile_config {
+  static constexpr int tile_m = TileM;
+  static constexpr int tile_n = TileN;
+  static constexpr int tile_k = TileK;
+};
+
+template <typename DataTag, typename IndexTag, typename TileTag, typename AbiTag, typename ArchTag>
+struct fragment_tag_fused_1nn_cubin {
+  static constexpr int cc_major = ArchTag::cc_major;
+  static constexpr int cc_minor = ArchTag::cc_minor;
+};
+
+template <typename DataTag, typename IndexTag, typename TileTag, typename AbiTag>
+struct fragment_tag_fused_1nn_tileir {};
+
+}  // namespace cuvs::distance::detail

cpp/src/cluster/detail/minClusterDistanceCompute.cu

@@ -58,6 +58,7 @@ void minClusterAndDistanceCompute(
       workspace.resize((sizeof(int)) * n_samples, stream);
 
       cuvs::distance::fusedDistanceNNMinReduce<DataT, raft::KeyValuePair<IndexT, DataT>, 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<DataT, DataT, KeyValueT, IndexT>(
             handle,
             batchMin,
-            X.data_handle() + dIdx * n_features,
-            centroids.data_handle() + cIdx * n_features,
+            X.data_handle() + static_cast<std::size_t>(dIdx) * static_cast<std::size_t>(n_features),
+            centroids.data_handle() +
+              static_cast<std::size_t>(cIdx) * static_cast<std::size_t>(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<const DataT, IndexT>(
-        X.data_handle() + (dIdx * n_features), ns, n_features);
+        X.data_handle() + static_cast<std::size_t>(dIdx) * static_cast<std::size_t>(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<const DataT, IndexT>(
-          centroids.data_handle() + (cIdx * n_features), nc, n_features);
+          centroids.data_handle() +
+            static_cast<std::size_t>(cIdx) * static_cast<std::size_t>(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<DataT, DataT, IndexT>(
+      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<const DataT, IndexT>(
-        X.data_handle() + dIdx * n_features, ns, n_features);
+        X.data_handle() + static_cast<std::size_t>(dIdx) * static_cast<std::size_t>(n_features),
+        ns,
+        n_features);
 
       auto minClusterDistanceView =
         raft::make_device_vector_view<DataT, IndexT>(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<DataT, IndexT>(
-          centroids.data_handle() + cIdx * n_features, nc, n_features);
+          centroids.data_handle() +
+            static_cast<std::size_t>(cIdx) * static_cast<std::size_t>(n_features),
+          nc,
+          n_features);
 
         auto pairwiseDistanceView =
           raft::make_device_matrix_view<DataT, IndexT>(pairwiseDistance.data_handle(), ns, nc);