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 a633094a88..6e7e53165a 100644 --- a/cpp/CMakeLists.txt +++ b/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::distance::detail::@abi_tag@, cuvs::detail::jit_lto::@arch_tag@>" + FRAGMENT_TAG_FORMAT_TILEIR + "cuvs::distance::detail::fragment_tag_fused_1nn_tileir, cuvs::distance::detail::@abi_tag@>" + FRAGMENT_TAG_HEADER_FILES + "" + "" "" + ) + if(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 + "" --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 $<$: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( @@ -1531,7 +1580,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/modules/generate_cutile_tile_metadata.py b/cpp/cmake/modules/generate_cutile_tile_metadata.py new file mode 100644 index 0000000000..3d2b224b96 --- /dev/null +++ b/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() 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..1f8e2e91e9 --- /dev/null +++ b/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 +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; +}; + +template +struct fragment_tag_fused_1nn_tileir {}; + +} // namespace cuvs::distance::detail diff --git a/cpp/src/cluster/detail/minClusterDistanceCompute.cu b/cpp/src/cluster/detail/minClusterDistanceCompute.cu index ee3cc3cdfd..30f419ea75 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/distance/detail/fused_distance_nn.cuh b/cpp/src/distance/detail/fused_distance_nn.cuh index f9dbd968ec..2cd81a8351 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,77 @@ #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{}; +}; + +/** + * 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 (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 && + metric != cuvs::distance::DistanceType::InnerProduct && 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 +101,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; 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..6774f44743 --- /dev/null +++ b/cpp/src/distance/detail/fused_distance_nn/cutile/export_fused_1nn.py @@ -0,0 +1,301 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Export fused 1-NN cuTile kernels to cubin or TileIR bytecode.""" + +from __future__ import annotations + +import argparse +import sys +from pathlib import Path +from typing import Literal + +import cuda.tile as ct +from cuda.tile.compilation import ( + ArrayConstraint, + CallingConvention, + ConstantConstraint, + KernelSignature, + ScalarConstraint, + export_kernel, +) + +# CI enables Python safe-path mode, so the script directory is not guaranteed +# to be importable even when this file is executed directly. +SCRIPT_DIR = Path(__file__).resolve().parent +if str(SCRIPT_DIR) not in sys.path: + sys.path.insert(0, str(SCRIPT_DIR)) + +from fused_1nn_kernel import ( # noqa: E402 + INDEX_TYPES, + METRICS, + _idx_dtype, + index_abbrev, + kernel_symbol, + make_kernel, +) + +DEFAULT_TILEIR_BYTECODE_VERSION = "13.1" +# cuTile requires a gpu_code even for TileIR bytecode export: it selects the compilation +# target / feature set for lowering, not the runtime architecture (the driver JITs at load). +DEFAULT_TILEIR_EXPORT_GPU_CODE = "sm_80" + + +def _dtype_for(data_type: str): + if data_type == "half": + return ct.float16 + if data_type == "float": + return ct.float32 + raise ValueError(f"Unsupported data_type {data_type!r}") + + +def _data_abbrev(data_type: str) -> str: + return {"half": "h", "float": "f"}[data_type] + + +def _elem_stride_divisible_for_tma(elem_dtype) -> tuple[int, int]: + """Row stride (dim 0) divisible enough for 16-byte TMA access; last dim stride 1.""" + bytes_per_elem = 2 if elem_dtype == ct.float16 else 4 + return (16 // bytes_per_elem, 1) + + +def _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(elem, index_dtype=idx_dtype) + + abbrev = _data_abbrev(data_type) + symbol = kernel_symbol( + abbrev, + index_abbrev(index_type), + matrix_layout, + ) + + return KernelSignature( + parameters=[ + matrix, + matrix, + norm_array, + norm_array, + idx_array, + dist_array, + ScalarConstraint(idx_dtype), + ScalarConstraint(idx_dtype), + ScalarConstraint(idx_dtype), + ScalarConstraint(idx_dtype), + ScalarConstraint(idx_dtype), + ScalarConstraint(ct.int32), + ConstantConstraint(tile_m), + ConstantConstraint(tile_n), + ConstantConstraint(tile_k), + ], + calling_convention=CallingConvention.cutile_python_v1(), + ).with_symbol(symbol) + + +def export_binary( + output_file: Path, + *, + output_format: Literal["cubin", "tileir_bytecode"], + data_type: str, + metric: str, + index_type: str, + tile_m: int, + tile_n: int, + tile_k: int, + gpu_code: str, + matrix_layout: str = "strict", + occupancy: int | None = None, + bytecode_version: str | None = None, +) -> str: + kernel = make_kernel( + data_type, + metric, + tile_m, + tile_n, + tile_k, + index_type=index_type, + gpu_code=gpu_code, + matrix_layout=matrix_layout, + occupancy=occupancy, + ) + signature = _kernel_signature( + data_type, + metric, + index_type, + tile_m, + tile_n, + tile_k, + gpu_code, + matrix_layout, + ) + + export_kwargs = { + "kernel": kernel, + "signatures": [signature], + "output_file": str(output_file), + "gpu_code": gpu_code, + "output_format": output_format, + } + if output_format == "tileir_bytecode": + export_kwargs["bytecode_version"] = ( + bytecode_version or DEFAULT_TILEIR_BYTECODE_VERSION + ) + + export_kernel(**export_kwargs) + + return signature.symbol + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("output_file", type=Path) + parser.add_argument( + "--format", choices=("cubin", "tileir_bytecode"), default="cubin" + ) + parser.add_argument( + "--data-type", choices=("half", "float"), required=True + ) + parser.add_argument("--metric", choices=METRICS, required=True) + parser.add_argument("--index-type", choices=INDEX_TYPES, required=True) + parser.add_argument("--tile-m", type=int, required=True) + parser.add_argument("--tile-n", type=int, required=True) + parser.add_argument("--tile-k", type=int, required=True) + parser.add_argument( + "--gpu-code", + default=DEFAULT_TILEIR_EXPORT_GPU_CODE, + help="Target SM for cubin export, or compile hint for TileIR bytecode export", + ) + parser.add_argument( + "--matrix-layout", + choices=("strict", "relaxed"), + default="strict", + ) + parser.add_argument("--occupancy", type=int) + parser.add_argument( + "--bytecode-version", default=DEFAULT_TILEIR_BYTECODE_VERSION + ) + args = parser.parse_args() + + export_binary( + args.output_file, + output_format=args.format, + data_type=args.data_type, + metric=args.metric, + index_type=args.index_type, + tile_m=args.tile_m, + tile_n=args.tile_n, + tile_k=args.tile_k, + gpu_code=args.gpu_code, + matrix_layout=args.matrix_layout, + occupancy=args.occupancy, + bytecode_version=args.bytecode_version, + ) + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_cutile_matrix.json b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_cutile_matrix.json new file mode 100644 index 0000000000..eb578d0f8e --- /dev/null +++ b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_cutile_matrix.json @@ -0,0 +1,298 @@ +{ + "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" + } + ] + }, + { + "output_format": "tileir_bytecode", + "artifact_ext": "tilebc", + "artifact_basename": "@data_type@_@index_abbrev@_@abi_abbrev@", + "register": "tileir", + "gpu_code": "sm_80", + "bytecode_version": "13.1", + "tile_m": 128, + "tile_n": 128, + "tile_k": 32, + "_data": [ + { + "data_type": "half", + "data_abbrev": "h" + }, + { + "data_type": "float", + "data_abbrev": "f" + } + ], + "_abi": [ + { + "matrix_layout": "relaxed", + "abi_abbrev": "relaxed", + "abi_tag": "cutile_abi_relaxed" + }, + { + "matrix_layout": "strict", + "abi_abbrev": "strict", + "abi_tag": "cutile_abi_strict" + } + ] + } + ] +} 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..ad1ba8fbea --- /dev/null +++ b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_kernel.py @@ -0,0 +1,191 @@ +# 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) + out_dist_dtype = ct.float16 if data_type == "half" else ct.float32 + 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: + # Defer the A-norm division until after selecting the + # winning centroid. + score = accumulator / (-b_norm)[None, :] + + if n == num_tiles_n - 1: + col = ct.arange(tn, dtype=ct.int16) + score = ct.where((n * tn + col)[None, :] < N, score, 3.4e38) + + curr_best, curr_idx = reduce_scores( + score.reshape(core_shape), local_indices + ) + update = curr_best < best_dist + best_dist = ct.where(update, curr_best, best_dist) + best_idx = ct.where(update, n * tn + curr_idx, best_idx) + + if metric_code == METRIC_INNER_PRODUCT: + out_dist = -best_dist + else: + a_norm = ct.load( + A_norm, index=(bidm,), shape=(tm,), padding_mode=zero_pad + )[:, None] + if metric_code == METRIC_L2_EXPANDED: + out_dist = a_norm + 2.0 * best_dist + # 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: + out_dist = 1.0 + best_dist / a_norm + + if store_idx != 0: + ct.store(OutIdx, index=(bidm,), tile=best_idx.reshape((tm,))) + ct.store( + OutDist, + index=(bidm,), + tile=out_dist.reshape((tm,)).astype(out_dist_dtype), + ) + + return fused_1nn_kernel + + +def kernel_symbol( + data_abbrev: str, + index_abbrev: str, + matrix_layout: str = "strict", +) -> str: + """Must stay in sync with fused_1nn_kernel_entrypoint() in fused_1nn_planner.hpp.""" + base = f"fused_1nn_{data_abbrev}_{index_abbrev}" + if matrix_layout == "strict": + return base + if matrix_layout == "relaxed": + return f"{base}_relaxed" + raise ValueError(f"Unsupported matrix layout {matrix_layout!r}") + + +def index_abbrev(index_type: str) -> str: + return {"int32": "i32", "int64": "i64"}[index_type] diff --git a/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_planner.hpp b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_planner.hpp new file mode 100644 index 0000000000..22b425cff6 --- /dev/null +++ b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_planner.hpp @@ -0,0 +1,132 @@ +/* + * 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>(); + } + + void add_tileir_fallback() + { + constexpr bool is_relaxed = std::is_same_v; + constexpr bool is_float = std::is_same_v; + using TileIr = std::conditional_t, + std::conditional_t>; + this->add_static_tileir_fragment< + fragment_tag_fused_1nn_tileir>(); + } +}; + +} // namespace cuvs::distance::detail diff --git a/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_tile.cu b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_tile.cu new file mode 100644 index 0000000000..432230d961 --- /dev/null +++ b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_tile.cu @@ -0,0 +1,403 @@ +/* + * 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); + if (lhs_bytes > std::numeric_limits::max() - lhs_begin || + rhs_bytes > std::numeric_limits::max() - rhs_begin) { + return true; + } + return lhs_begin < rhs_begin + rhs_bytes && rhs_begin < lhs_begin + lhs_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(); + planner.add_tileir_fallback(); + return planner.try_get_launcher() != nullptr; +} + +template +void launch_fused_1nn_tile_impl(raft::resources const& handle, + IdxT* nearest_idx, + DataT* 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(); + planner.add_tileir_fallback(); + 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, + DataT* 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(DataT)); + 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, + DataT* 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*, \ + DataT*, \ + 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..1c3e2d9271 --- /dev/null +++ b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_tile.hpp @@ -0,0 +1,82 @@ +/* + * 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 + +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; + +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 / 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, + DataT* 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..a71d8537dd 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,394 @@ 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) % 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; row_offset += row_tile) { + 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; candidate_offset += candidate_tile) { + 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); + } + } + } + } 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"); + 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..4c14d7990c --- /dev/null +++ b/cpp/src/distance/top_1_nn.cuh @@ -0,0 +1,118 @@ +/* + * 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 + +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 +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..78e9032f84 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 { @@ -131,6 +132,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 +140,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 +172,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 +255,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 +274,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 +290,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 +303,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/neighbors/distance_nn.cu b/cpp/tests/neighbors/distance_nn.cu index f31f3ebacf..4b7df9b7ef 100644 --- a/cpp/tests/neighbors/distance_nn.cu +++ b/cpp/tests/neighbors/distance_nn.cu @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ @@ -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)} { } @@ -87,7 +96,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); } @@ -96,10 +105,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); } @@ -110,25 +120,41 @@ 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); + 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); + 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) { + run_top_1_nn(cuvs::distance::Top1nnOutput{cutile_idx.data_handle(), + cutile_dist.data_handle()}); + } else { + run_top_1_nn(out.data_handle()); + } } else { static_assert(sizeof(DataT) == 0, "fusedDistanceNNMinReduce is not implemented for datatype other than float"); @@ -149,14 +175,45 @@ class NNTest : public ::testing::TestWithParam> { true, true, metric, - 0.0, - stream); + 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) { + // FP32 cuTile MMA uses TF32-rounded inputs, so nearly tied candidates can produce a + // different index from the scalar FP32 reference. Validate that the returned index selects + // a candidate within the same numerical tolerance of the true minimum. + 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; } @@ -170,12 +227,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; }; @@ -195,6 +258,22 @@ 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); + } +#endif + return inputs; +}(); + // Test fused implementation with single-precision typedef NNTest NNTest_fp32_fused; TEST_P(NNTest_fp32_fused, test) @@ -203,7 +282,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; diff --git a/cpp/tests/neighbors/distance_nn_helper.cuh b/cpp/tests/neighbors/distance_nn_helper.cuh index fda7b76573..2d20e037b8 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 */ @@ -107,16 +107,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 +125,53 @@ 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 { + dist = cosine_distance(&A[m * K], &B[n * K], K); + } + out[m] = sqrt ? 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: