Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion ci/validate_wheel.sh
Original file line number Diff line number Diff line change
Expand Up @@ -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+=(
Expand Down
59 changes: 54 additions & 5 deletions cpp/CMakeLists.txt
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand All @@ -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(
Expand All @@ -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
Expand Down
55 changes: 55 additions & 0 deletions cpp/cmake/modules/generate_cutile_tile_metadata.py
Original file line number Diff line number Diff line change
@@ -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()
Original file line number Diff line number Diff line change
@@ -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
54 changes: 52 additions & 2 deletions cpp/src/distance/detail/fused_distance_nn.cuh
Original file line number Diff line number Diff line change
@@ -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"
Expand All @@ -20,13 +23,60 @@
#include <raft/util/cuda_utils.cuh> // raft::ceildiv, raft::shfl

#include <cstddef> // size_t
#include <limits> // std::numeric_limits
#include <cstdint>
#include <limits> // 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{};
};

/**
* 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 <typename DataT, typename IdxT>
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 (backend == Top1nnBackend::Cutile) {
#if CUVS_CUTILE_ENABLED
if constexpr (is_fused_1nn_cutile_data_v<DataT>) {
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 &&
metric != cuvs::distance::DistanceType::InnerProduct && x != nullptr && y != nullptr &&
m > 0 && n > 0 && k > 0;
}

template <typename DataT,
typename OutT,
typename IdxT,
Expand Down
Loading
Loading