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/distance/detail/fused_distance_nn.cuh b/cpp/src/distance/detail/fused_distance_nn.cuh index f9dbd968ec..91b42b1087 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" @@ -20,13 +23,60 @@ #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 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..0984b08a0f --- /dev/null +++ b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_tile.cu @@ -0,0 +1,400 @@ +/* + * 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 + +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(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, + cudaStream_t stream) +{ + 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(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, + cudaStream_t stream) +{ + 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( + nearest_idx, nearest_dist, x, y, xn, yn, m, n, k, metric, is_sqrt, stream); + } else { + launch_fused_1nn_tile_impl( + nearest_idx, nearest_dist, x, y, xn, yn, m, n, k, metric, is_sqrt, stream); + } + } else { + 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(tmp_idx, + batch_dist, + batch_x, + y, + batch_xn, + yn, + batch_m, + static_cast(n), + static_cast(k), + metric, + is_sqrt, + stream); + } else { + launch_fused_1nn_tile_impl(tmp_idx, + batch_dist, + batch_x, + y, + batch_xn, + yn, + batch_m, + static_cast(n), + static_cast(k), + metric, + is_sqrt, + stream); + } + + 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( \ + 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*, \ + cudaStream_t) + +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..fb83746448 --- /dev/null +++ b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_tile.hpp @@ -0,0 +1,81 @@ +/* + * 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 + +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(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, + cudaStream_t stream); +} // 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..fb5e29709f 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,19 @@ #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 #include @@ -289,26 +294,395 @@ void fusedDistanceNNMinReduce(OutT* min, float metric_arg, cudaStream_t stream) { - 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"); + raft::resources handle; + raft::resource::set_cuda_stream(handle, stream); + 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, + stream); +} + +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(OutputT output, + const DataT* x, + const DataT* y, + const NormT* xn, + const NormT* yn, + IdxT m, + IdxT n, + IdxT k, + void* workspace, + bool sqrt, + cuvs::distance::DistanceType metric, + cudaStream_t stream) +{ + 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(output.nearest_idx, + output.nearest_dist, + x, + y, + xn, + yn, + m, + n, + k, + metric, + sqrt, + workspace, + stream); + } 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(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, + cudaStream_t stream) +{ + 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(output, + x, + y, + xn, + yn, + m, + n, + k, + workspace, + red_op, + pair_red_op, + sqrt, + init_out_buffer, + is_row_major, + metric, + metric_arg, + stream); + } 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, + cudaStream_t stream) +{ + 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 + row_offset * k, + y + candidate_offset * k, + xn + row_offset, + yn + candidate_offset, + rows, + candidates, + k, + workspace, + sqrt, + candidate_offset != 0 || init_out_buffer, + is_row_major, + metric, + metric_arg, + stream); + 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, + cudaStream_t stream) +{ + 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(output, x, y, xn, yn, m, n, k, workspace, sqrt, metric, stream); +#else + RAFT_FAIL("Requested cuTile fused 1-NN backend was not built"); +#endif + return; + case detail::Top1nnBackend::Cutlass: + detail::top_1_nn_legacy_fused(output, + x, + y, + xn, + yn, + m, + n, + k, + workspace, + sqrt, + init_out_buffer, + is_row_major, + metric, + metric_arg, + stream); + 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, + stream); + 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..b5d78bb348 --- /dev/null +++ b/cpp/src/distance/top_1_nn.cu @@ -0,0 +1,61 @@ +/* + * 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, \ + cudaStream_t) + +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..1a68aa2d18 --- /dev/null +++ b/cpp/src/distance/top_1_nn.cuh @@ -0,0 +1,122 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +#include "detail/fused_distance_nn.cuh" + +#include + +#include +#include + +#include + +namespace cuvs::distance { + +/** Separate index and distance arrays used by backends with structure-of-arrays output. */ +template +struct Top1nnOutput { + IdxT* nearest_idx; + DistT* nearest_dist; +}; + +namespace detail { + +template +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, + cudaStream_t stream); + +#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, \ + cudaStream_t) + +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/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..993b1cda7a 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,18 @@ 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)}, + cutile_idx{raft::make_device_vector(handle, m)}, + cutile_dist{raft::make_device_vector(handle, m)} { } @@ -87,7 +95,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 +104,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); } @@ -114,21 +123,38 @@ class NNTest : public ::testing::TestWithParam> { 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, + stream); + }; + 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"); @@ -156,7 +182,24 @@ class NNTest : public ::testing::TestWithParam> { 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 valid index from the scalar FP32 reference. Compare the resulting minimum + // distance using the test's existing numerical tolerance. + raft::linalg::unaryOp( + ref_dist.data_handle(), ref_out.data_handle(), m, raft::value_op{}, 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 +213,17 @@ 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 cutile_idx; + raft::device_vector cutile_dist; size_t workspace_size; }; @@ -195,6 +243,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 +267,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..e7931fb86d 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 */