diff --git a/c/src/cluster/kmeans.cpp b/c/src/cluster/kmeans.cpp index 9d371c182b..cb5f045df0 100644 --- a/c/src/cluster/kmeans.cpp +++ b/c/src/cluster/kmeans.cpp @@ -4,6 +4,7 @@ */ #include +#include #include @@ -41,6 +42,50 @@ cuvs::cluster::kmeans::balanced_params convert_balanced_params(const cuvsKMeansP return kmeans_params; } +constexpr int64_t kKMeansInt32IndexMax = std::numeric_limits::max(); + +bool dlpack_tensor_size_exceeds_int32_index(const DLTensor& tensor) +{ + int64_t size = 1; + for (int i = 0; i < tensor.ndim; ++i) { + const int64_t extent = tensor.shape[i]; + if (extent <= 0) { return false; } + if (extent > kKMeansInt32IndexMax / size) { return true; } + size *= extent; + } + return false; +} + +bool kmeans_tensor_shapes_use_int64_index(DLManagedTensor* X, DLManagedTensor* centroids) +{ + if (dlpack_tensor_size_exceeds_int32_index(X->dl_tensor)) { return true; } + if (dlpack_tensor_size_exceeds_int32_index(centroids->dl_tensor)) { return true; } + return false; +} + +bool kmeans_fit_uses_int64_index(DLManagedTensor* X, + DLManagedTensor* centroids, + int n_clusters) +{ + if (static_cast(n_clusters) > kKMeansInt32IndexMax) { return true; } + return kmeans_tensor_shapes_use_int64_index(X, centroids); +} + +bool kmeans_labels_are_int64(const DLTensor& labels) +{ + return labels.dtype.code == kDLInt && labels.dtype.bits == 64; +} + +void validate_kmeans_labels_dtype(const DLTensor& labels) +{ + if (labels.dtype.code == kDLInt && (labels.dtype.bits == 32 || labels.dtype.bits == 64)) { + return; + } + RAFT_FAIL("Unsupported labels DLtensor dtype: %d and bits: %d", + labels.dtype.code, + labels.dtype.bits); +} + template void _fit(cuvsResources_t res, const cuvsKMeansParams& params, @@ -54,8 +99,10 @@ void _fit(cuvsResources_t res, auto res_ptr = reinterpret_cast(res); if (!cuvs::core::is_dlpack_device_compatible(X)) { - auto n_samples = static_cast(X.shape[0]); - auto n_features = static_cast(X.shape[1]); + // Host fit overloads are only exposed with int64_t index types. + using HostIdxT = int64_t; + auto n_samples = static_cast(X.shape[0]); + auto n_features = static_cast(X.shape[1]); if (params.hierarchical) { RAFT_FAIL("hierarchical kmeans is not supported with host data"); @@ -66,24 +113,24 @@ void _fit(cuvsResources_t res, RAFT_FAIL("centroids must be on device memory"); } - auto X_view = raft::make_host_matrix_view( + auto X_view = raft::make_host_matrix_view( reinterpret_cast(X.data), n_samples, n_features); auto centroids_view = - cuvs::core::from_dlpack>( + cuvs::core::from_dlpack>( centroids_tensor); - std::optional> sample_weight; + std::optional> sample_weight; if (sample_weight_tensor != NULL) { auto sw = sample_weight_tensor->dl_tensor; if (!cuvs::core::is_dlpack_host_compatible(sw)) { RAFT_FAIL("sample_weight must be host accessible when X is on host"); } - sample_weight = raft::make_host_vector_view( + sample_weight = raft::make_host_vector_view( reinterpret_cast(sw.data), n_samples); } T inertia_temp; - IdxT n_iter_temp; + HostIdxT n_iter_temp; auto kmeans_params = convert_params(params); cuvs::cluster::kmeans::fit(*res_ptr, @@ -92,7 +139,7 @@ void _fit(cuvsResources_t res, sample_weight, centroids_view, raft::make_host_scalar_view(&inertia_temp), - raft::make_host_scalar_view(&n_iter_temp)); + raft::make_host_scalar_view(&n_iter_temp)); *inertia = inertia_temp; *n_iter = n_iter_temp; @@ -109,10 +156,20 @@ void _fit(cuvsResources_t res, if constexpr (std::is_same_v) { RAFT_FAIL("float64 is an unsupported dtype for hierarchical kmeans"); } else { - auto kmeans_params = convert_balanced_params(params); + // Balanced fit overloads are only exposed with int64_t index types. + using BalancedIdxT = int64_t; + using balanced_const_mdspan_type = + raft::device_matrix_view; + using balanced_mdspan_type = raft::device_matrix_view; + auto kmeans_params = convert_balanced_params(params); T inertia_temp; auto inertia_view = raft::make_host_scalar_view(&inertia_temp); - cuvs::cluster::kmeans::fit(*res_ptr, kmeans_params, cuvs::core::from_dlpack(X_tensor), cuvs::core::from_dlpack(centroids_tensor), std::make_optional(inertia_view)); + cuvs::cluster::kmeans::fit( + *res_ptr, + kmeans_params, + cuvs::core::from_dlpack(X_tensor), + cuvs::core::from_dlpack(centroids_tensor), + std::make_optional(inertia_view)); *inertia = inertia_temp; *n_iter = params.hierarchical_n_iters; } @@ -164,13 +221,22 @@ void _predict(cuvsResources_t res, if constexpr (std::is_same_v) { RAFT_FAIL("float64 is an unsupported dtype for hierarchical kmeans"); + } else if constexpr (!std::is_same_v) { + RAFT_FAIL("int64 labels are unsupported for hierarchical kmeans"); } else { + // Balanced predict overloads are only exposed with int64_t index types and int32 labels. + using BalancedIdxT = int64_t; + using balanced_const_mdspan_type = + raft::device_matrix_view; + using balanced_labels_mdspan_type = + raft::device_vector_view; auto kmeans_params = convert_balanced_params(params); - cuvs::cluster::kmeans::predict(*res_ptr, - kmeans_params, - cuvs::core::from_dlpack(X_tensor), - cuvs::core::from_dlpack(centroids_tensor), - cuvs::core::from_dlpack(labels_tensor)); + cuvs::cluster::kmeans::predict( + *res_ptr, + kmeans_params, + cuvs::core::from_dlpack(X_tensor), + cuvs::core::from_dlpack(centroids_tensor), + cuvs::core::from_dlpack(labels_tensor)); *inertia = 0; } } else { @@ -196,7 +262,7 @@ void _predict(cuvsResources_t res, } } -template +template void _cluster_cost(cuvsResources_t res, DLManagedTensor* X_tensor, DLManagedTensor* centroids_tensor, @@ -259,10 +325,24 @@ extern "C" cuvsError_t cuvsKMeansFit(cuvsResources_t res, { return cuvs::core::translate_exceptions([=] { auto dataset = X->dl_tensor; + const bool use_int64_index = + kmeans_fit_uses_int64_index(X, centroids, params->n_clusters); if (dataset.dtype.code == kDLFloat && dataset.dtype.bits == 32) { - _fit(res, *params, X, sample_weight, centroids, inertia, n_iter); + if (use_int64_index) { + _fit( + res, *params, X, sample_weight, centroids, inertia, n_iter); + } else { + _fit( + res, *params, X, sample_weight, centroids, inertia, n_iter); + } } else if (dataset.dtype.code == kDLFloat && dataset.dtype.bits == 64) { - _fit(res, *params, X, sample_weight, centroids, inertia, n_iter); + if (use_int64_index) { + _fit( + res, *params, X, sample_weight, centroids, inertia, n_iter); + } else { + _fit( + res, *params, X, sample_weight, centroids, inertia, n_iter); + } } else { RAFT_FAIL("Unsupported dataset DLtensor dtype: %d and bits: %d", dataset.dtype.code, @@ -282,10 +362,45 @@ extern "C" cuvsError_t cuvsKMeansPredict(cuvsResources_t res, { return cuvs::core::translate_exceptions([=] { auto dataset = X->dl_tensor; + validate_kmeans_labels_dtype(labels->dl_tensor); + const bool use_int64_index = kmeans_fit_uses_int64_index(X, centroids, params->n_clusters); + const bool use_int64_labels = kmeans_labels_are_int64(labels->dl_tensor); if (dataset.dtype.code == kDLFloat && dataset.dtype.bits == 32) { - _predict(res, *params, X, sample_weight, centroids, labels, normalize_weight, inertia); + if (use_int64_index) { + if (use_int64_labels) { + _predict( + res, *params, X, sample_weight, centroids, labels, normalize_weight, inertia); + } else { + _predict( + res, *params, X, sample_weight, centroids, labels, normalize_weight, inertia); + } + } else { + if (use_int64_labels) { + _predict( + res, *params, X, sample_weight, centroids, labels, normalize_weight, inertia); + } else { + _predict( + res, *params, X, sample_weight, centroids, labels, normalize_weight, inertia); + } + } } else if (dataset.dtype.code == kDLFloat && dataset.dtype.bits == 64) { - _predict(res, *params, X, sample_weight, centroids, labels, normalize_weight, inertia); + if (use_int64_index) { + if (use_int64_labels) { + _predict( + res, *params, X, sample_weight, centroids, labels, normalize_weight, inertia); + } else { + _predict( + res, *params, X, sample_weight, centroids, labels, normalize_weight, inertia); + } + } else { + if (use_int64_labels) { + _predict( + res, *params, X, sample_weight, centroids, labels, normalize_weight, inertia); + } else { + _predict( + res, *params, X, sample_weight, centroids, labels, normalize_weight, inertia); + } + } } else { RAFT_FAIL("Unsupported dataset DLtensor dtype: %d and bits: %d", dataset.dtype.code, @@ -301,10 +416,19 @@ extern "C" cuvsError_t cuvsKMeansClusterCost(cuvsResources_t res, { return cuvs::core::translate_exceptions([=] { auto dataset = X->dl_tensor; + const bool use_int64_index = kmeans_tensor_shapes_use_int64_index(X, centroids); if (dataset.dtype.code == kDLFloat && dataset.dtype.bits == 32) { - _cluster_cost(res, X, centroids, cost); + if (use_int64_index) { + _cluster_cost(res, X, centroids, cost); + } else { + _cluster_cost(res, X, centroids, cost); + } } else if (dataset.dtype.code == kDLFloat && dataset.dtype.bits == 64) { - _cluster_cost(res, X, centroids, cost); + if (use_int64_index) { + _cluster_cost(res, X, centroids, cost); + } else { + _cluster_cost(res, X, centroids, cost); + } } else { RAFT_FAIL("Unsupported dataset DLtensor dtype: %d and bits: %d", dataset.dtype.code, diff --git a/c/tests/cluster/kmeans_c.cu b/c/tests/cluster/kmeans_c.cu index 5caac1ba98..a1d3b1f64d 100644 --- a/c/tests/cluster/kmeans_c.cu +++ b/c/tests/cluster/kmeans_c.cu @@ -44,6 +44,7 @@ float kInitCentroids[kNClusters][kNFeatures] = { float kExpectedCentroids[kNClusters * kNFeatures] = {1.5f, 1.5f, 10.5f, 10.5f}; int32_t kExpectedLabels[kNSamples] = {0, 0, 0, 0, 1, 1, 1, 1}; +int64_t kExpectedLabels64[kNSamples] = {0, 0, 0, 0, 1, 1, 1, 1}; // 8 points, each at squared distance 0.5 from its cluster mean -> 4.0. constexpr double kExpectedInertia = 4.0; @@ -56,6 +57,7 @@ void test_fit_predict() rmm::device_uvector dataset_d(kNSamples * kNFeatures, stream); rmm::device_uvector centroids_d(kNClusters * kNFeatures, stream); rmm::device_uvector labels_d(kNSamples, stream); + rmm::device_uvector labels64_d(kNSamples, stream); raft::copy(dataset_d.data(), reinterpret_cast(kDataset), @@ -90,6 +92,9 @@ void test_fit_predict() DLManagedTensor labels_t{}; cuvs::core::to_dlpack( raft::make_device_vector_view(labels_d.data(), kNSamples), &labels_t); + DLManagedTensor labels64_t{}; + cuvs::core::to_dlpack( + raft::make_device_vector_view(labels64_d.data(), kNSamples), &labels64_t); double inertia = -1.0; int n_iter = -1; @@ -101,6 +106,9 @@ void test_fit_predict() ASSERT_EQ(cuvsKMeansPredict( res, params, &dataset_t, NULL, ¢roids_t, &labels_t, false, &predict_inertia), CUVS_SUCCESS); + ASSERT_EQ(cuvsKMeansPredict( + res, params, &dataset_t, NULL, ¢roids_t, &labels64_t, false, &predict_inertia), + CUVS_SUCCESS); ASSERT_EQ(cuvsKMeansClusterCost(res, &dataset_t, ¢roids_t, &cluster_cost), CUVS_SUCCESS); ASSERT_TRUE(cuvs::devArrMatchHost(kExpectedCentroids, @@ -109,12 +117,15 @@ void test_fit_predict() cuvs::CompareApprox(1e-4f))); ASSERT_TRUE(cuvs::devArrMatchHost( kExpectedLabels, labels_d.data(), kNSamples, cuvs::Compare())); + ASSERT_TRUE(cuvs::devArrMatchHost( + kExpectedLabels64, labels64_d.data(), kNSamples, cuvs::Compare())); EXPECT_GT(n_iter, 0); EXPECT_NEAR(inertia, kExpectedInertia, 1e-4); EXPECT_NEAR(predict_inertia, kExpectedInertia, 1e-4); EXPECT_NEAR(cluster_cost, kExpectedInertia, 1e-4); + labels64_t.deleter(&labels64_t); labels_t.deleter(&labels_t); centroids_t.deleter(¢roids_t); dataset_t.deleter(&dataset_t); diff --git a/ci/build_standalone_c.sh b/ci/build_standalone_c.sh index 2b8e0863f9..e1267f71b7 100755 --- a/ci/build_standalone_c.sh +++ b/ci/build_standalone_c.sh @@ -41,6 +41,12 @@ source rapids-configure-sccache source rapids-datetime-string rapids-pip-retry install cmake + +RAPIDS_CUDA_MAJOR="${RAPIDS_CUDA_VERSION%%.*}" +if [[ "${RAPIDS_CUDA_MAJOR}" == "13" ]]; then + rapids-pip-retry install cuda-tile "cuda-toolkit[tileiras]==${RAPIDS_CUDA_VERSION%.*}.*" +fi + pyenv rehash rapids-print-env diff --git a/conda/environments/all_cuda-133_arch-aarch64.yaml b/conda/environments/all_cuda-133_arch-aarch64.yaml index d95c6854bd..07deffb120 100644 --- a/conda/environments/all_cuda-133_arch-aarch64.yaml +++ b/conda/environments/all_cuda-133_arch-aarch64.yaml @@ -15,8 +15,10 @@ dependencies: - cuda-nvrtc-dev - cuda-nvtx-dev - cuda-profiler-api +- cuda-tileiras - cuda-version=13.3 - cupy>=14.0.1,!=14.1.0 +- cutile-python - cxx-compiler - cython>=3.2.2 - dlpack>=0.8,<1.0 diff --git a/conda/environments/all_cuda-133_arch-x86_64.yaml b/conda/environments/all_cuda-133_arch-x86_64.yaml index 69d8f84605..e6de7ca7a5 100644 --- a/conda/environments/all_cuda-133_arch-x86_64.yaml +++ b/conda/environments/all_cuda-133_arch-x86_64.yaml @@ -15,8 +15,10 @@ dependencies: - cuda-nvrtc-dev - cuda-nvtx-dev - cuda-profiler-api +- cuda-tileiras - cuda-version=13.3 - cupy>=14.0.1,!=14.1.0 +- cutile-python - cxx-compiler - cython>=3.2.2 - dlpack>=0.8,<1.0 diff --git a/conda/environments/bench_ann_cuda-133_arch-aarch64.yaml b/conda/environments/bench_ann_cuda-133_arch-aarch64.yaml index 321a892555..1d5f229c8f 100644 --- a/conda/environments/bench_ann_cuda-133_arch-aarch64.yaml +++ b/conda/environments/bench_ann_cuda-133_arch-aarch64.yaml @@ -15,8 +15,10 @@ dependencies: - cuda-nvrtc-dev - cuda-nvtx-dev - cuda-profiler-api +- cuda-tileiras - cuda-version=13.3 - cupy>=14.0.1,!=14.1.0 +- cutile-python - cuvs==26.10.*,>=0.0.0a0 - cxx-compiler - cython>=3.2.2 diff --git a/conda/environments/bench_ann_cuda-133_arch-x86_64.yaml b/conda/environments/bench_ann_cuda-133_arch-x86_64.yaml index 179b4a4a2f..228d11c3d8 100644 --- a/conda/environments/bench_ann_cuda-133_arch-x86_64.yaml +++ b/conda/environments/bench_ann_cuda-133_arch-x86_64.yaml @@ -15,8 +15,10 @@ dependencies: - cuda-nvrtc-dev - cuda-nvtx-dev - cuda-profiler-api +- cuda-tileiras - cuda-version=13.3 - cupy>=14.0.1,!=14.1.0 +- cutile-python - cuvs==26.10.*,>=0.0.0a0 - cxx-compiler - cython>=3.2.2 diff --git a/conda/recipes/libcuvs/recipe.yaml b/conda/recipes/libcuvs/recipe.yaml index 93a69ea1c0..1baa35ec4d 100644 --- a/conda/recipes/libcuvs/recipe.yaml +++ b/conda/recipes/libcuvs/recipe.yaml @@ -70,6 +70,11 @@ cache: - cuda-version =${{ cuda_version }} - cmake ${{ cmake_version }} - ninja + - python + - if: cuda_major == "13" + then: + - cutile-python + - cuda-tileiras - ${{ stdlib("c") }} host: - libnvjitlink-dev @@ -397,6 +402,11 @@ outputs: - cuda-version =${{ cuda_version }} - cmake ${{ cmake_version }} - ninja + - python + - if: cuda_major == "13" + then: + - cutile-python + - cuda-tileiras - ${{ stdlib("c") }} host: - ${{ pin_subpackage("libcuvs-headers", exact=True) }} diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt index 847cf543d6..c10202bf5e 100644 --- a/cpp/CMakeLists.txt +++ b/cpp/CMakeLists.txt @@ -1153,6 +1153,70 @@ if(NOT BUILD_CPU_ONLY) ) endblock() + include(cmake/modules/generate_cutile_kernels.cmake) + set(cutile_smoke_dir "${CMAKE_CURRENT_SOURCE_DIR}/src/detail/jit_lto/cutile_smoke") + set(cutile_smoke_generated_dir + "${CMAKE_CURRENT_BINARY_DIR}/generated_kernels/detail/jit_lto/cutile_smoke" + ) + generate_cutile_kernels( + cutile_smoke_files + KERNEL_DIR "${cutile_smoke_dir}" + KERNEL_BASENAME "cutile_smoke" + KERNEL_PYTHON "smoke_kernel.py" + EXPORT_SCRIPT "export_smoke.py" + OUTPUT_DIRECTORY "${cutile_smoke_generated_dir}" + MATRIX_JSON_FILE "${cutile_smoke_dir}/cutile_smoke_matrix.json" + FRAGMENT_TAG_FORMAT_CUBIN + "cuvs::detail::jit_lto::fragment_tag_cutile_smoke_add_cubin" + FRAGMENT_TAG_HEADER_FILES "" + "" + ) + 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 @@ -1173,6 +1237,35 @@ if(NOT BUILD_CPU_ONLY) OUTPUT_FILE_FORMAT "${CMAKE_CURRENT_BINARY_DIR}/src/distance/detail/pairwise_matrix/dispatch_rbf_inst_data_@data_abbrev@_acc_@acc_abbrev@_out_@out_abbrev@_index_@index_abbrev@_op_@op_abbrev@.cu" ) + + include(cmake/modules/generate_cutile_kernels.cmake) + 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" + ) + 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(NOT DEFINED CUVS_CUTILE_ENABLED) + set(CUVS_CUTILE_ENABLED 0) + endif() + target_compile_definitions(cuvs_cpp_headers INTERFACE CUVS_CUTILE_ENABLED=${CUVS_CUTILE_ENABLED}) + generate_inst_matrix( cagra_build_inst_files MATRIX_JSON_FILE "${CMAKE_CURRENT_SOURCE_DIR}/src/neighbors/cagra_build_matrix.json" @@ -1364,11 +1457,13 @@ if(NOT BUILD_CPU_ONLY) src/core/omp_wrapper.cpp src/util/file_io.cpp src/util/host_memory.cpp + src/detail/jit_lto/TileAlgorithmPlanner.cpp src/distance/detail/kernels/gram_matrix.cu src/distance/detail/kernels/kernel_factory.cu 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 @@ -1462,6 +1557,9 @@ if(NOT BUILD_CPU_ONLY) src/stats/trustworthiness_score.cu ${CUVS_MG_ALGOS} ${jit_lto_files} + ${cutile_fused_1nn_files} + $<$:src/distance/detail/fused_distance_nn/cutile/fused_1nn_tile.cu> + ${cutile_smoke_files} ) set_target_properties( @@ -1485,7 +1583,9 @@ if(NOT BUILD_CPU_ONLY) target_compile_definitions( cuvs_objs PRIVATE $<$:CUVS_BUILD_CAGRA_HNSWLIB> - $<$:CUVS_BUILD_MG_ALGOS> $<$:NVTX_ENABLED> + $<$:CUVS_BUILD_MG_ALGOS> + $<$:NVTX_ENABLED> + CUVS_CUTILE_ENABLED=${CUVS_CUTILE_ENABLED} ) target_link_libraries( @@ -1506,6 +1606,7 @@ if(NOT BUILD_CPU_ONLY) "$" INTERFACE "$" PRIVATE "${CMAKE_CURRENT_SOURCE_DIR}/src" "${CMAKE_CURRENT_BINARY_DIR}/src" + "${cutile_fused_1nn_generated_dir}" "${cutile_smoke_generated_dir}" ) # Endian detection diff --git a/cpp/cmake/config.json b/cpp/cmake/config.json index cc6e647bae..9dfcda3259 100644 --- a/cpp/cmake/config.json +++ b/cpp/cmake/config.json @@ -20,6 +20,43 @@ "MATRIX_JSON_STRING": "?" } }, + "cuvs_find_build_python": { + "pargs": { + "nargs": 1 + } + }, + "process_cutile_matrix_entry": { + "pargs": { + "nargs": 1 + }, + "kwargs": { + "KERNEL_DIR": 1, + "KERNEL_BASENAME": 1, + "KERNEL_PYTHON": 1, + "EXPORT_SCRIPT": 1, + "OUTPUT_DIRECTORY": 1, + "FRAGMENT_TAG_FORMAT_CUBIN": 1, + "FRAGMENT_TAG_FORMAT_TILEIR": "?", + "FRAGMENT_TAG_HEADER_FILES": "*", + "MATRIX_JSON_ENTRY": 1 + } + }, + "generate_cutile_kernels": { + "pargs": { + "nargs": 1 + }, + "kwargs": { + "KERNEL_DIR": 1, + "KERNEL_BASENAME": 1, + "KERNEL_PYTHON": 1, + "EXPORT_SCRIPT": 1, + "OUTPUT_DIRECTORY": 1, + "MATRIX_JSON_FILE": 1, + "FRAGMENT_TAG_FORMAT_CUBIN": 1, + "FRAGMENT_TAG_FORMAT_TILEIR": "?", + "FRAGMENT_TAG_HEADER_FILES": "*" + } + }, "add_jit_lto_kernel": { "pargs": { "nargs": 1 diff --git a/cpp/cmake/modules/compute_matrix_product.cmake b/cpp/cmake/modules/compute_matrix_product.cmake index 82a34f9242..b4b7afb9b7 100644 --- a/cpp/cmake/modules/compute_matrix_product.cmake +++ b/cpp/cmake/modules/compute_matrix_product.cmake @@ -1,12 +1,26 @@ # ============================================================================= # cmake-format: off -# 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 # cmake-format: on # ============================================================================= include_guard(GLOBAL) +function(cuvs_find_build_python output_var) + # cuTile is a build dependency. In conda builds, it is installed in BUILD_PREFIX while CMake's + # default search can resolve the host interpreter from PREFIX instead. Use the build prefix so + # configure-time matrix expansion and build-time kernel exports see the cuTile package. + if(DEFINED ENV{BUILD_PREFIX}) + set(Python_ROOT "$ENV{BUILD_PREFIX}") + endif() + find_package(Python REQUIRED COMPONENTS Interpreter) + set(${output_var} + "${Python_EXECUTABLE}" + PARENT_SCOPE + ) +endfunction() + function(compute_matrix_product output_var) set(options) set(one_value MATRIX_JSON_FILE MATRIX_JSON_STRING) @@ -14,19 +28,21 @@ function(compute_matrix_product output_var) cmake_parse_arguments(_JIT_LTO "${options}" "${one_value}" "${multi_value}" ${ARGN}) - find_package(Python3 REQUIRED COMPONENTS Interpreter) + cuvs_find_build_python(_matrix_python_executable) if(_JIT_LTO_MATRIX_JSON_FILE) execute_process( - COMMAND "${Python3_EXECUTABLE}" "${CMAKE_CURRENT_FUNCTION_LIST_DIR}/compute_matrix_product.py" - "${_JIT_LTO_MATRIX_JSON_FILE}" # + COMMAND + "${_matrix_python_executable}" + "${CMAKE_CURRENT_FUNCTION_LIST_DIR}/compute_matrix_product.py" + "${_JIT_LTO_MATRIX_JSON_FILE}" OUTPUT_VARIABLE output COMMAND_ERROR_IS_FATAL ANY ) else() execute_process( COMMAND "${CMAKE_COMMAND}" -E echo "${_JIT_LTO_MATRIX_JSON_STRING}" - COMMAND "${Python3_EXECUTABLE}" "${CMAKE_CURRENT_FUNCTION_LIST_DIR}/compute_matrix_product.py" - - + COMMAND "${_matrix_python_executable}" + "${CMAKE_CURRENT_FUNCTION_LIST_DIR}/compute_matrix_product.py" - OUTPUT_VARIABLE output COMMAND_ERROR_IS_FATAL ANY ) endif() diff --git a/cpp/cmake/modules/generate_cutile_kernels.cmake b/cpp/cmake/modules/generate_cutile_kernels.cmake new file mode 100644 index 0000000000..22489a951c --- /dev/null +++ b/cpp/cmake/modules/generate_cutile_kernels.cmake @@ -0,0 +1,385 @@ +# ============================================================================= +# cmake-format: off +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# cmake-format: on +# ============================================================================= + +include_guard(GLOBAL) + +include(${CMAKE_CURRENT_LIST_DIR}/compute_matrix_product.cmake) + +function(_cutile_fragment_tag_header_files output_var) + set(${output_var} "") + foreach(_header IN LISTS ARGN) + if(NOT _header MATCHES "^(\".*\"|<.*>)$") + set(_header "\"${_header}\"") + endif() + string(APPEND ${output_var} "#include ${_header}\n") + endforeach() + set(${output_var} + "${${output_var}}" + PARENT_SCOPE + ) +endfunction() + +function(_cutile_kernels_setup) + set(options) + set(one_value MATRIX_JSON_FILE OUTPUT_DIRECTORY) + set(multi_value) + cmake_parse_arguments(_CUTILE "${options}" "${one_value}" "${multi_value}" ${ARGN}) + + find_package(CUDAToolkit REQUIRED) + + if(CUDAToolkit_VERSION VERSION_LESS 13.0) + message( + STATUS + "cuTile embedded kernels require CUDA 13.0+; skipping cuTile generation (found ${CUDAToolkit_VERSION})." + ) + set(_CUTILE_SETUP_OK + FALSE + PARENT_SCOPE + ) + return() + endif() + + cuvs_find_build_python(Python3_EXECUTABLE) + + find_program( + CUTILE_BIN2C + NAMES bin2c + PATHS ${CUDAToolkit_BIN_DIR} REQUIRED + ) + + execute_process( + COMMAND "${Python3_EXECUTABLE}" -c "import cuda.tile" + RESULT_VARIABLE _cutile_import_result + ERROR_VARIABLE _cutile_import_error + OUTPUT_QUIET ERROR_STRIP_TRAILING_WHITESPACE + ) + if(NOT _cutile_import_result EQUAL 0) + message( + FATAL_ERROR + "cuda.tile (cuTile Python) is required to build cuTile embedded kernels. " + "Install cutile-python and cuda-tileiras (conda), or cuda-tile[tileiras] (pip).\n" + "Interpreter: ${Python3_EXECUTABLE}\n" + "Import error: ${_cutile_import_error}" + ) + endif() + message(STATUS "Using cuTile Python: ${Python3_EXECUTABLE}") + + set_property( + DIRECTORY + PROPERTY CMAKE_CONFIGURE_DEPENDS "${_CUTILE_MATRIX_JSON_FILE}" + APPEND + ) + + file(MAKE_DIRECTORY "${_CUTILE_OUTPUT_DIRECTORY}") + + set(Python3_EXECUTABLE + "${Python3_EXECUTABLE}" + PARENT_SCOPE + ) + set(CUTILE_BIN2C + "${CUTILE_BIN2C}" + PARENT_SCOPE + ) + set(_CUTILE_SETUP_OK + TRUE + PARENT_SCOPE + ) +endfunction() + +macro(_cutile_append_matrix_tile_aliases entry data_abbrev abi_abbrev tile_geometry) + set(_cutile_tile_geometry "${tile_geometry}") + list(GET _cutile_tile_geometry 0 tile_m) + list(GET _cutile_tile_geometry 1 tile_n) + list(GET _cutile_tile_geometry 2 tile_k) + string(JSON _cutile_export_len LENGTH "${entry}" "_export") + set(_cutile_export_idx 0) + while(_cutile_export_idx LESS _cutile_export_len) + string(JSON _cutile_export_entry GET "${entry}" "_export" "${_cutile_export_idx}") + string(JSON _cutile_register GET "${_cutile_export_entry}" "register") + if(_cutile_register STREQUAL "cubin") + string(JSON _cutile_arch_tag GET "${_cutile_export_entry}" "arch_tag") + set(_cutile_alias_suffix "${data_abbrev}_${_cutile_arch_tag}_${abi_abbrev}") + elseif(_cutile_register STREQUAL "tileir") + set(_cutile_alias_suffix "${data_abbrev}_tileir_${abi_abbrev}") + else() + message(FATAL_ERROR "Unknown cuTile register kind '${_cutile_register}'") + endif() + + set(_cutile_tile_value "${tile_m},${tile_n},${tile_k}") + if(DEFINED _tile_alias_value_${_cutile_alias_suffix}) + if(NOT "${_tile_alias_value_${_cutile_alias_suffix}}" STREQUAL "${_cutile_tile_value}") + message(FATAL_ERROR "Conflicting cuTile tile geometry for ${_cutile_alias_suffix}: " + "${_tile_alias_value_${_cutile_alias_suffix}} vs ${_cutile_tile_value}" + ) + endif() + else() + set(_tile_alias_value_${_cutile_alias_suffix} "${_cutile_tile_value}") + string( + APPEND + _tile_aliases + "using fused_1nn_matrix_tile_${_cutile_alias_suffix} = cutile_tile_config<${tile_m}, ${tile_n}, ${tile_k}>;\n" + ) + endif() + math(EXPR _cutile_export_idx "${_cutile_export_idx} + 1") + endwhile() +endmacro() + +function(_cutile_generate_matrix_tiles_header header_path matrix_json_file) + file(READ "${matrix_json_file}" _matrix_json) + set(_tile_aliases "") + string(JSON _entry_len LENGTH "${_matrix_json}") + set(_entry_idx 0) + while(_entry_idx LESS _entry_len) + string(JSON _entry GET "${_matrix_json}" "${_entry_idx}") + string(JSON _entry_tile ERROR_VARIABLE _entry_tile_error GET "${_entry}" "_tile" 0) + if(NOT _entry_tile_error) + string(JSON _default_tile_m GET "${_entry_tile}" "tile_m") + string(JSON _default_tile_n GET "${_entry_tile}" "tile_n") + string(JSON _default_tile_k GET "${_entry_tile}" "tile_k") + endif() + + string(JSON _data_len LENGTH "${_entry}" "_data") + set(_data_idx 0) + while(_data_idx LESS _data_len) + string(JSON _data_entry GET "${_entry}" "_data" "${_data_idx}") + string(JSON _data_abbrev GET "${_data_entry}" "data_abbrev") + + string(JSON _abi_len LENGTH "${_entry}" "_abi") + set(_abi_idx 0) + while(_abi_idx LESS _abi_len) + string(JSON _abi_entry GET "${_entry}" "_abi" "${_abi_idx}") + string(JSON _abi_abbrev GET "${_abi_entry}" "abi_abbrev") + string(JSON _tile_m ERROR_VARIABLE _tile_m_error GET "${_abi_entry}" "tile_m") + string(JSON _tile_n ERROR_VARIABLE _tile_n_error GET "${_abi_entry}" "tile_n") + string(JSON _tile_k ERROR_VARIABLE _tile_k_error GET "${_abi_entry}" "tile_k") + if(_tile_m_error + OR _tile_n_error + OR _tile_k_error + ) + if(_entry_tile_error) + message(FATAL_ERROR "Missing cuTile geometry for ${_data_abbrev}/${_abi_abbrev}") + endif() + set(_tile_m "${_default_tile_m}") + set(_tile_n "${_default_tile_n}") + set(_tile_k "${_default_tile_k}") + endif() + + set(_tile_geometry "${_tile_m};${_tile_n};${_tile_k}") + _cutile_append_matrix_tile_aliases( + "${_entry}" "${_data_abbrev}" "${_abi_abbrev}" "${_tile_geometry}" + ) + math(EXPR _abi_idx "${_abi_idx} + 1") + endwhile() + math(EXPR _data_idx "${_data_idx} + 1") + endwhile() + math(EXPR _entry_idx "${_entry_idx} + 1") + endwhile() + file( + WRITE "${header_path}" + "/* + * Generated from ${matrix_json_file} by generate_cutile_kernels.cmake — do not edit. + */ +#pragma once + +#include + +namespace cuvs::distance::detail { + +${_tile_aliases} + +} // namespace cuvs::distance::detail +" + ) +endfunction() + +function(_cutile_make_python_args output_var) + set(_python_args + --format + "${output_format}" + --data-type + "${data_type}" + --metric + "${metric}" + --index-type + "${index_type}" + --tile-m + "${tile_m}" + --tile-n + "${tile_n}" + --tile-k + "${tile_k}" + --gpu-code + "${gpu_code}" + ) + if(DEFINED bytecode_version AND NOT "${bytecode_version}" STREQUAL "") + list(APPEND _python_args --bytecode-version "${bytecode_version}") + endif() + if(DEFINED matrix_layout AND NOT "${matrix_layout}" STREQUAL "") + list(APPEND _python_args --matrix-layout "${matrix_layout}") + endif() + if(DEFINED occupancy AND NOT "${occupancy}" STREQUAL "") + list(APPEND _python_args --occupancy "${occupancy}") + endif() + set(${output_var} + "${_python_args}" + PARENT_SCOPE + ) +endfunction() + +function(process_cutile_matrix_entry source_list_var) + set(options) + set(one_value KERNEL_DIR KERNEL_BASENAME KERNEL_PYTHON EXPORT_SCRIPT OUTPUT_DIRECTORY + FRAGMENT_TAG_FORMAT_CUBIN FRAGMENT_TAG_FORMAT_TILEIR MATRIX_JSON_ENTRY + ) + set(multi_value FRAGMENT_TAG_HEADER_FILES) + cmake_parse_arguments(_CUTILE "${options}" "${one_value}" "${multi_value}" ${ARGN}) + + if(NOT Python3_EXECUTABLE) + cuvs_find_build_python(Python3_EXECUTABLE) + endif() + + populate_matrix_variables("${_CUTILE_MATRIX_JSON_ENTRY}") + + if(register STREQUAL "cubin") + string(CONFIGURE "${_CUTILE_FRAGMENT_TAG_FORMAT_CUBIN}" fragment_tag @ONLY) + set(bin2c_symbol embedded_cubin) + set(fragment_entry_type "cuvs::detail::jit_lto::StaticCubinFragmentEntry") + elseif(register STREQUAL "tileir") + string(CONFIGURE "${_CUTILE_FRAGMENT_TAG_FORMAT_TILEIR}" fragment_tag @ONLY) + set(bin2c_symbol embedded_tileir) + set(fragment_entry_type + "cuvs::detail::jit_lto::StaticTileIrBytecodeFragmentEntry" + ) + else() + message(FATAL_ERROR "Unknown cuTile register kind '${register}'") + endif() + + _cutile_fragment_tag_header_files(fragment_tag_header_files ${_CUTILE_FRAGMENT_TAG_HEADER_FILES}) + + string(CONFIGURE "${artifact_basename}" _artifact_basename @ONLY) + set(_artifact_stem "${_CUTILE_KERNEL_BASENAME}_${_artifact_basename}") + set(_artifact_file "${_CUTILE_OUTPUT_DIRECTORY}/${_artifact_stem}.${artifact_ext}") + set(_embedded_header "${_CUTILE_OUTPUT_DIRECTORY}/${_artifact_stem}_${register}.h") + set(_fragment_cpp "${_CUTILE_OUTPUT_DIRECTORY}/${_artifact_stem}_${register}.cpp") + set(embedded_header_file "${_artifact_stem}_${register}.h") + + _cutile_make_python_args(_python_args) + + set(_export_python_executable "${Python3_EXECUTABLE}") + if(DEFINED python_executable AND NOT "${python_executable}" STREQUAL "") + string(CONFIGURE "${python_executable}" _export_python_executable @ONLY) + endif() + + if(DEFINED prebuilt_artifact AND NOT "${prebuilt_artifact}" STREQUAL "") + string(CONFIGURE "${prebuilt_artifact}" _prebuilt_artifact @ONLY) + if(NOT IS_ABSOLUTE "${_prebuilt_artifact}") + set(_prebuilt_artifact "${_CUTILE_KERNEL_DIR}/${_prebuilt_artifact}") + endif() + add_custom_command( + OUTPUT "${_artifact_file}" + COMMAND "${CMAKE_COMMAND}" -E copy_if_different "${_prebuilt_artifact}" "${_artifact_file}" + DEPENDS "${_prebuilt_artifact}" + COMMENT "Copying prebuilt cuTile ${_CUTILE_KERNEL_BASENAME} ${output_format} ${data_type}" + VERBATIM + ) + else() + add_custom_command( + OUTPUT "${_artifact_file}" + COMMAND "${_export_python_executable}" "${_CUTILE_KERNEL_DIR}/${_CUTILE_EXPORT_SCRIPT}" + "${_artifact_file}" ${_python_args} + WORKING_DIRECTORY "${_CUTILE_KERNEL_DIR}" + DEPENDS "${_CUTILE_KERNEL_DIR}/${_CUTILE_EXPORT_SCRIPT}" + "${_CUTILE_KERNEL_DIR}/${_CUTILE_KERNEL_PYTHON}" + COMMENT "Exporting cuTile ${_CUTILE_KERNEL_BASENAME} ${output_format} ${data_type}" + VERBATIM + ) + endif() + + add_custom_command( + OUTPUT "${_embedded_header}" + COMMAND "${CUTILE_BIN2C}" --const --name ${bin2c_symbol} --static "${_artifact_file}" > + "${_embedded_header}" + DEPENDS "${_artifact_file}" + VERBATIM + ) + + configure_file( + "${CMAKE_CURRENT_FUNCTION_LIST_DIR}/register_cutile_fragment.cpp.in" "${_fragment_cpp}" @ONLY + ) + list(APPEND ${source_list_var} "${_embedded_header}" "${_fragment_cpp}") + set(${source_list_var} + "${${source_list_var}}" + PARENT_SCOPE + ) +endfunction() + +function(generate_cutile_kernels source_list_var) + set(options) + set(one_value KERNEL_DIR KERNEL_BASENAME KERNEL_PYTHON EXPORT_SCRIPT OUTPUT_DIRECTORY + MATRIX_JSON_FILE FRAGMENT_TAG_FORMAT_CUBIN FRAGMENT_TAG_FORMAT_TILEIR + ) + set(multi_value FRAGMENT_TAG_HEADER_FILES) + cmake_parse_arguments(_CUTILE "${options}" "${one_value}" "${multi_value}" ${ARGN}) + + if(NOT _CUTILE_KERNEL_BASENAME) + message(FATAL_ERROR "generate_cutile_kernels: KERNEL_BASENAME is required") + endif() + if(NOT _CUTILE_KERNEL_PYTHON) + message(FATAL_ERROR "generate_cutile_kernels: KERNEL_PYTHON is required") + endif() + + _cutile_kernels_setup( + MATRIX_JSON_FILE "${_CUTILE_MATRIX_JSON_FILE}" OUTPUT_DIRECTORY "${_CUTILE_OUTPUT_DIRECTORY}" + ) + if(NOT _CUTILE_SETUP_OK) + # This function's parent is cpp/CMakeLists.txt. Propagate the disabled feature state there so + # the compile definition cannot retain a stale value from a previous generator invocation. + set(CUVS_CUTILE_ENABLED + 0 + PARENT_SCOPE + ) + set(${source_list_var} + "" + PARENT_SCOPE + ) + return() + endif() + + compute_matrix_product(matrix_product MATRIX_JSON_FILE "${_CUTILE_MATRIX_JSON_FILE}") + + set(_matrix_tiles_header "${_CUTILE_OUTPUT_DIRECTORY}/fused_1nn_cutile_tiles.hpp") + _cutile_generate_matrix_tiles_header("${_matrix_tiles_header}" "${_CUTILE_MATRIX_JSON_FILE}") + + string(JSON len LENGTH "${matrix_product}") + math(EXPR last "${len} - 1") + + # cmake-lint: disable=C0103,E1120 + foreach(i RANGE "${last}") + string(JSON matrix_json_entry GET "${matrix_product}" "${i}") + process_cutile_matrix_entry( + "${source_list_var}" + KERNEL_DIR "${_CUTILE_KERNEL_DIR}" + KERNEL_BASENAME "${_CUTILE_KERNEL_BASENAME}" + KERNEL_PYTHON "${_CUTILE_KERNEL_PYTHON}" + EXPORT_SCRIPT "${_CUTILE_EXPORT_SCRIPT}" + OUTPUT_DIRECTORY "${_CUTILE_OUTPUT_DIRECTORY}" + FRAGMENT_TAG_FORMAT_CUBIN "${_CUTILE_FRAGMENT_TAG_FORMAT_CUBIN}" + FRAGMENT_TAG_FORMAT_TILEIR "${_CUTILE_FRAGMENT_TAG_FORMAT_TILEIR}" + FRAGMENT_TAG_HEADER_FILES ${_CUTILE_FRAGMENT_TAG_HEADER_FILES} + MATRIX_JSON_ENTRY "${matrix_json_entry}" + ) + endforeach() + + set(CUVS_CUTILE_ENABLED + 1 + PARENT_SCOPE + ) + set(${source_list_var} + "${${source_list_var}}" + PARENT_SCOPE + ) +endfunction() 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..4415c6a7a8 --- /dev/null +++ b/cpp/cmake/modules/generate_cutile_tile_metadata.py @@ -0,0 +1,54 @@ +#!/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 +from pathlib import Path + + +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 = {} + for entry in json.loads(args.matrix.read_text()): + default_tile = entry.get("_tile", [{}])[0] + for data in entry["_data"]: + for abi in entry["_abi"]: + tile = tuple( + abi.get(k, default_tile.get(k)) + for k in ("tile_m", "tile_n", "tile_k") + ) + if any(value is None for value in tile): + raise ValueError("missing cuTile tile geometry") + for exported in entry["_export"]: + suffix = f"{data['data_abbrev']}_{exported.get('arch_tag', 'tileir')}_{abi['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/cmake/modules/register_cutile_fragment.cpp.in b/cpp/cmake/modules/register_cutile_fragment.cpp.in new file mode 100644 index 0000000000..7206e88e57 --- /dev/null +++ b/cpp/cmake/modules/register_cutile_fragment.cpp.in @@ -0,0 +1,31 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include "@embedded_header_file@" +#include + +@fragment_tag_header_files@ + + namespace +{ + using fragment_tag = @fragment_tag@; + using fragment_entry = @fragment_entry_type@; + +} // namespace + +template <> +const uint8_t* const fragment_entry::data = @bin2c_symbol@; + +template <> +const size_t fragment_entry::length = sizeof(@bin2c_symbol@); + +template <> +const int fragment_entry::tile_m = @tile_m@; + +template <> +const int fragment_entry::tile_n = @tile_n@; + +template <> +const int fragment_entry::tile_k = @tile_k@; diff --git a/cpp/include/cuvs/cluster/kmeans.hpp b/cpp/include/cuvs/cluster/kmeans.hpp index ceb42d1e71..251ff40f69 100644 --- a/cpp/include/cuvs/cluster/kmeans.hpp +++ b/cpp/include/cuvs/cluster/kmeans.hpp @@ -787,6 +787,15 @@ void predict(raft::resources const& handle, bool normalize_weight, raft::host_scalar_view inertia); +void predict(raft::resources const& handle, + const kmeans::params& params, + raft::device_matrix_view X, + std::optional> sample_weight, + raft::device_matrix_view centroids, + raft::device_vector_view labels, + bool normalize_weight, + raft::host_scalar_view inertia); + /** * @brief Predict the closest cluster each sample in X belongs to. * @@ -838,10 +847,10 @@ void predict(raft::resources const& handle, */ void predict(raft::resources const& handle, const kmeans::params& params, - raft::device_matrix_view X, - std::optional> sample_weight, - raft::device_matrix_view centroids, - raft::device_vector_view labels, + raft::device_matrix_view X, + std::optional> sample_weight, + raft::device_matrix_view centroids, + raft::device_vector_view labels, bool normalize_weight, raft::host_scalar_view inertia); @@ -952,6 +961,24 @@ void predict(raft::resources const& handle, * @param[out] inertia Sum of squared distances of samples to * their closest cluster center. */ +void predict(raft::resources const& handle, + const kmeans::params& params, + raft::device_matrix_view X, + std::optional> sample_weight, + raft::device_matrix_view centroids, + raft::device_vector_view labels, + bool normalize_weight, + raft::host_scalar_view inertia); + +void predict(raft::resources const& handle, + const kmeans::params& params, + raft::device_matrix_view X, + std::optional> sample_weight, + raft::device_matrix_view centroids, + raft::device_vector_view labels, + bool normalize_weight, + raft::host_scalar_view inertia); + void predict(raft::resources const& handle, const kmeans::params& params, raft::device_matrix_view X, diff --git a/cpp/include/cuvs/detail/jit_lto/CutileFragmentEntry.hpp b/cpp/include/cuvs/detail/jit_lto/CutileFragmentEntry.hpp new file mode 100644 index 0000000000..724662c9dd --- /dev/null +++ b/cpp/include/cuvs/detail/jit_lto/CutileFragmentEntry.hpp @@ -0,0 +1,119 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +#include +#include +#include + +namespace cuvs::detail::jit_lto { + +/** cuTile GEMM-style block geometry embedded in generated static fragment specializations. */ +struct CutileTileConfig { + int tile_m; + int tile_n; + int tile_k; +}; + +/** Embedded CUDA binary module (cubin), loaded directly via cudaLibraryLoadData. */ +struct CubinFragmentEntry { + virtual ~CubinFragmentEntry() = default; + + virtual const uint8_t* get_data() const = 0; + + virtual size_t get_length() const = 0; + + virtual const char* get_key() const = 0; + + virtual int get_cc_major() const = 0; + + virtual int get_cc_minor() const = 0; + + virtual int get_tile_m() const { return 0; } + + virtual int get_tile_n() const { return 0; } + + virtual int get_tile_k() const { return 0; } +}; + +template +struct StaticCubinFragmentEntry final : CubinFragmentEntry { + const uint8_t* get_data() const override { return StaticCubinFragmentEntry::data; } + + size_t get_length() const override { return StaticCubinFragmentEntry::length; } + + const char* get_key() const override + { + return typeid(StaticCubinFragmentEntry).name(); + } + + int get_cc_major() const override { return FragmentTag::cc_major; } + + int get_cc_minor() const override { return FragmentTag::cc_minor; } + + int get_tile_m() const override { return tile_m; } + + int get_tile_n() const override { return tile_n; } + + int get_tile_k() const override { return tile_k; } + + static const int tile_m; + static const int tile_n; + static const int tile_k; + + static const uint8_t* const data; + static const size_t length; +}; + +/** Embedded TileIR bytecode, JIT-compiled by the driver when no matching cubin exists. */ +struct TileIrBytecodeFragmentEntry { + virtual ~TileIrBytecodeFragmentEntry() = default; + + virtual const uint8_t* get_data() const = 0; + + virtual size_t get_length() const = 0; + + virtual const char* get_key() const = 0; + + virtual int get_tile_m() const { return 0; } + + virtual int get_tile_n() const { return 0; } + + virtual int get_tile_k() const { return 0; } +}; + +template +struct StaticTileIrBytecodeFragmentEntry final : TileIrBytecodeFragmentEntry { + const uint8_t* get_data() const override + { + return StaticTileIrBytecodeFragmentEntry::data; + } + + size_t get_length() const override + { + return StaticTileIrBytecodeFragmentEntry::length; + } + + const char* get_key() const override + { + return typeid(StaticTileIrBytecodeFragmentEntry).name(); + } + + int get_tile_m() const override { return tile_m; } + + int get_tile_n() const override { return tile_n; } + + int get_tile_k() const override { return tile_k; } + + static const int tile_m; + static const int tile_n; + static const int tile_k; + + static const uint8_t* const data; + static const size_t length; +}; + +} // namespace cuvs::detail::jit_lto diff --git a/cpp/include/cuvs/detail/jit_lto/TileAlgorithmPlanner.hpp b/cpp/include/cuvs/detail/jit_lto/TileAlgorithmPlanner.hpp new file mode 100644 index 0000000000..fb6025fd64 --- /dev/null +++ b/cpp/include/cuvs/detail/jit_lto/TileAlgorithmPlanner.hpp @@ -0,0 +1,74 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +#include "CutileFragmentEntry.hpp" + +#include +#include +#include +#include +#include +#include +#include + +#include + +namespace cuvs::detail::jit_lto { + +struct CutileRuntimeCapabilities; + +struct TileLauncherCache { + std::shared_mutex mutex; + std::unordered_map> launchers; + // Cache expected compatibility misses so an unsupported module is not loaded on every call. + // Unexpected CUDA errors are raised rather than inserted here. + std::unordered_set unavailable_launchers; +}; + +/** Loads prebuilt cubins or TileIR bytecode directly through the CUDA library API. */ +struct TileAlgorithmPlanner { + TileAlgorithmPlanner(std::string entrypoint, TileLauncherCache& launcher_cache) + : entrypoint_(std::move(entrypoint)), launcher_cache_(launcher_cache) + { + } + + virtual ~TileAlgorithmPlanner() = default; + + std::shared_ptr get_launcher(); + + /** Returns nullptr when no module can be loaded for the current device (does not RAFT_FAIL). */ + std::shared_ptr try_get_launcher(); + + template + void add_static_fragment() + { + cubin_fragments_.push_back(std::make_unique>()); + } + + template + void add_static_tileir_fragment() + { + tileir_fragment_ = std::make_unique>(); + } + + /** Tile geometry from the cubin or TileIR fragment that would load on this device. */ + CutileTileConfig tile_config() const; + + protected: + std::vector> cubin_fragments_; + std::unique_ptr tileir_fragment_; + + private: + std::string get_planner_key(const CutileRuntimeCapabilities* capabilities) const; + + std::shared_ptr build(const CutileRuntimeCapabilities* capabilities); + + std::string entrypoint_; + TileLauncherCache& launcher_cache_; +}; + +} // namespace cuvs::detail::jit_lto diff --git a/cpp/include/cuvs/detail/jit_lto/cutile_arch_tags.hpp b/cpp/include/cuvs/detail/jit_lto/cutile_arch_tags.hpp new file mode 100644 index 0000000000..2b378dac78 --- /dev/null +++ b/cpp/include/cuvs/detail/jit_lto/cutile_arch_tags.hpp @@ -0,0 +1,54 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +#ifndef CUVS_CUTILE_ENABLED +#define CUVS_CUTILE_ENABLED 0 +#endif + +namespace cuvs::detail::jit_lto { + +#if CUVS_CUTILE_ENABLED + +/** Must stay in sync with cuTile matrix _arch entries and planner add_static_fragment calls. */ +struct cutile_arch_8_0 { + static constexpr int cc_major = 8; + static constexpr int cc_minor = 0; +}; + +struct cutile_arch_8_6 { + static constexpr int cc_major = 8; + static constexpr int cc_minor = 6; +}; + +struct cutile_arch_9_0 { + static constexpr int cc_major = 9; + static constexpr int cc_minor = 0; +}; + +struct cutile_arch_10_0 { + static constexpr int cc_major = 10; + static constexpr int cc_minor = 0; +}; + +struct cutile_arch_12_0 { + static constexpr int cc_major = 12; + static constexpr int cc_minor = 0; +}; + +inline bool is_embedded_cubin_arch(int cc_major, int cc_minor) +{ + if (cc_minor < 0) { return false; } + return cc_major == 8 || cc_major == 9 || cc_major == 10 || cc_major == 12; +} + +#else + +inline bool is_embedded_cubin_arch(int, int) { return false; } + +#endif + +} // namespace cuvs::detail::jit_lto diff --git a/cpp/include/cuvs/detail/jit_lto/cutile_module.hpp b/cpp/include/cuvs/detail/jit_lto/cutile_module.hpp new file mode 100644 index 0000000000..bf26e1c9c5 --- /dev/null +++ b/cpp/include/cuvs/detail/jit_lto/cutile_module.hpp @@ -0,0 +1,109 @@ +/* + * 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 +#include + +#include +#include + +namespace cuvs::detail::jit_lto { + +struct CutileModuleImage { + const uint8_t* data; + size_t size; +}; + +/** + * Selects the newest compatible cubin in the device's compute-capability major family. + * + * CUDA cubins are forward compatible across minor revisions within a major family, so an SM 8.9 + * device can load SM 8.6 SASS and an SM 12.1 device can load SM 12.0 SASS. + */ +inline const CubinFragmentEntry* find_compatible_cubin_fragment( + int cc_major, + int cc_minor, + const std::vector>& cubin_fragments) +{ + const CubinFragmentEntry* best = nullptr; + for (const auto& fragment : cubin_fragments) { + if (fragment->get_cc_major() != cc_major || fragment->get_cc_minor() > cc_minor) { continue; } + if (best == nullptr || fragment->get_cc_minor() > best->get_cc_minor()) { + best = fragment.get(); + } + } + return best; +} + +/** Selects compatible prebuilt SASS for the device, or TileIR when the driver can JIT it. */ +inline std::optional resolve_cutile_module_image( + const CutileRuntimeCapabilities& capabilities, + const std::vector>& cubin_fragments, + const TileIrBytecodeFragmentEntry* tileir_fragment) +{ + if (const auto* fragment = find_compatible_cubin_fragment( + capabilities.cc_major, capabilities.cc_minor, cubin_fragments)) { + return CutileModuleImage{fragment->get_data(), fragment->get_length()}; + } + if (tileir_fragment != nullptr && tileir_fallback_available(capabilities.driver_version)) { + return CutileModuleImage{tileir_fragment->get_data(), tileir_fragment->get_length()}; + } + return std::nullopt; +} + +inline bool is_expected_cutile_unavailable(cudaError_t status) +{ + switch (status) { + case cudaErrorInvalidDeviceFunction: + case cudaErrorInvalidPtx: + case cudaErrorNoKernelImageForDevice: + case cudaErrorSymbolNotFound: + case cudaErrorUnsupportedPtxVersion: + case cudaErrorCallRequiresNewerDriver: + case cudaErrorSharedObjectSymbolNotFound: + case cudaErrorSharedObjectInitFailed: + case cudaErrorJitCompilerNotFound: return true; + default: return false; + } +} + +/** + * Loads a cuTile launcher, returning null for an expected module/JIT compatibility rejection. + * Unexpected CUDA failures retain the normal RAFT exception behavior. + */ +inline std::shared_ptr try_load_cutile_launcher( + const CutileModuleImage& image, const std::string& kernel_symbol) +{ + cudaLibrary_t library{}; + auto load_status = + cudaLibraryLoadData(&library, image.data, nullptr, nullptr, 0, nullptr, nullptr, 0); + if (load_status != cudaSuccess) { + if (is_expected_cutile_unavailable(load_status)) { return nullptr; } + RAFT_CUDA_TRY(load_status); + } + + cudaKernel_t kernel{}; + load_status = cudaLibraryGetKernel(&kernel, library, kernel_symbol.c_str()); + if (load_status != cudaSuccess) { + RAFT_CUDA_TRY(cudaLibraryUnload(library)); + if (is_expected_cutile_unavailable(load_status)) { return nullptr; } + RAFT_CUDA_TRY(load_status); + } + + return std::make_shared(kernel, library); +} + +} // namespace cuvs::detail::jit_lto diff --git a/cpp/include/cuvs/detail/jit_lto/cutile_smoke_fragments.hpp b/cpp/include/cuvs/detail/jit_lto/cutile_smoke_fragments.hpp new file mode 100644 index 0000000000..3b52f3daf8 --- /dev/null +++ b/cpp/include/cuvs/detail/jit_lto/cutile_smoke_fragments.hpp @@ -0,0 +1,15 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +#pragma once + +namespace cuvs::detail::jit_lto { + +template +struct fragment_tag_cutile_smoke_add_cubin { + static constexpr int cc_major = ArchTag::cc_major; + static constexpr int cc_minor = ArchTag::cc_minor; +}; + +} // namespace cuvs::detail::jit_lto 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..807dc50e24 --- /dev/null +++ b/cpp/include/cuvs/detail/jit_lto/fused_distance_nn/fused_1nn_fragments.hpp @@ -0,0 +1,66 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +#include + +#include + +#include +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 fused_1nn_data_tag; + +template <> +struct fused_1nn_data_tag { + using type = cuvs::neighbors::detail::tag_f; +}; + +template <> +struct fused_1nn_data_tag { + using type = cuvs::neighbors::detail::tag_h; +}; + +template +using fused_1nn_data_tag_t = typename fused_1nn_data_tag::type; + +template +struct fused_1nn_index_tag; + +template <> +struct fused_1nn_index_tag { + using type = cuvs::neighbors::detail::tag_index_i32; +}; + +template <> +struct fused_1nn_index_tag { + using type = cuvs::neighbors::detail::tag_index_i64; +}; + +template +using fused_1nn_index_tag_t = typename fused_1nn_index_tag::type; + +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/include/cuvs/detail/jit_lto/tileir_compat.hpp b/cpp/include/cuvs/detail/jit_lto/tileir_compat.hpp new file mode 100644 index 0000000000..8e5e599069 --- /dev/null +++ b/cpp/include/cuvs/detail/jit_lto/tileir_compat.hpp @@ -0,0 +1,115 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +#ifndef CUVS_CUTILE_ENABLED +#define CUVS_CUTILE_ENABLED 0 +#endif + +#include +#include + +#include + +namespace cuvs::detail::jit_lto { + +/** Runtime/device properties that determine cuTile image selection and launch eligibility. */ +struct CutileRuntimeCapabilities { + int device; + int cc_major; + int cc_minor; + int driver_version; +}; + +inline bool query_current_cutile_runtime_capabilities(CutileRuntimeCapabilities& capabilities) +{ + if (cudaGetDevice(&capabilities.device) != cudaSuccess) { return false; } + if (cudaDeviceGetAttribute(&capabilities.cc_major, + cudaDevAttrComputeCapabilityMajor, + capabilities.device) != cudaSuccess) { + return false; + } + if (cudaDeviceGetAttribute(&capabilities.cc_minor, + cudaDevAttrComputeCapabilityMinor, + capabilities.device) != cudaSuccess) { + return false; + } + return cudaDriverGetVersion(&capabilities.driver_version) == cudaSuccess; +} + +/** Minimum CUDA driver version (from cudaDriverGetVersion) for TileIR JIT of embedded bytecode. */ +inline constexpr int kMinTileIrJitDriverVersion = 13010; // CUDA 13.1 / driver >= 590.44 + +/** Minimum CUDA runtime version (from cudaRuntimeGetVersion) for cuTile integration. */ +inline constexpr int kMinCutileRuntimeVersion = 13000; + +inline constexpr bool library_built_with_cutile() +{ +#if CUVS_CUTILE_ENABLED + return true; +#else + return false; +#endif +} + +inline bool runtime_cuda13_or_newer() +{ + int runtime_version = 0; + if (cudaRuntimeGetVersion(&runtime_version) != cudaSuccess) { return false; } + return runtime_version >= kMinCutileRuntimeVersion; +} + +/** True when this build embeds cuTile artifacts and the runtime is CUDA 13+. */ +inline bool cutile_integration_enabled() +{ + return library_built_with_cutile() && runtime_cuda13_or_newer(); +} + +/** True when this build embeds compatible SASS in the device's compute-capability major family. */ +inline bool has_embedded_cubin_for_arch(int cc_major, int cc_minor) +{ + return is_embedded_cubin_arch(cc_major, cc_minor); +} + +/** True when the driver can JIT-compile embedded TileIR bytecode at load time. */ +inline bool tileir_fallback_available(int driver_version) +{ + return driver_version >= kMinTileIrJitDriverVersion; +} + +/** + * True when a cuTile launch may be attempted for the given device: cuTile is enabled, the runtime + * is CUDA 13+, and either compatible same-family SASS exists (no driver JIT required) or the + * driver can JIT the embedded TileIR bytecode fallback. + */ +#if CUVS_CUTILE_ENABLED +inline bool cutile_launch_available_for_arch(int cc_major, int cc_minor, int driver_version) +{ + if (!runtime_cuda13_or_newer()) { return false; } + // The exported fused-1NN kernels require Ampere-or-newer tensor-core semantics, and the current + // integration is validated only through the SM12 family. + if (cc_major < 8 || cc_major > 12) { return false; } + if (has_embedded_cubin_for_arch(cc_major, cc_minor)) { return true; } + return tileir_fallback_available(driver_version); +} +#else +inline constexpr bool cutile_launch_available_for_arch(int, int, int) { return false; } +#endif + +#if CUVS_CUTILE_ENABLED +inline bool cutile_launch_available_on_current_device() +{ + CutileRuntimeCapabilities capabilities{}; + if (!query_current_cutile_runtime_capabilities(capabilities)) { return false; } + return cutile_launch_available_for_arch( + capabilities.cc_major, capabilities.cc_minor, capabilities.driver_version); +} +#else +/** Compile-time false when cuTile is not built; use in if constexpr to skip cuTile-only paths. */ +inline constexpr bool cutile_launch_available_on_current_device() { return false; } +#endif + +} // namespace cuvs::detail::jit_lto diff --git a/cpp/src/cluster/detail/kmeans.cuh b/cpp/src/cluster/detail/kmeans.cuh index e3ffb4a439..208784762d 100644 --- a/cpp/src/cluster/detail/kmeans.cuh +++ b/cpp/src/cluster/detail/kmeans.cuh @@ -51,6 +51,7 @@ #include #include #include +#include namespace cuvs::cluster::kmeans::detail { @@ -684,8 +685,7 @@ void kmeans_fit( DataT* cur_centroids_ptr = cur_centroids_buf.data(); DataT* new_centroids_ptr = new_centroids_buf.data(); - auto minClusterAndDistance = raft::make_device_vector, IndexT>( - handle, device_buffer_samples); + rmm::device_uvector assignment_storage(0, stream); auto L2NormBatch = raft::make_device_vector(handle, device_buffer_samples); auto batch_weights_buf = raft::make_device_vector(handle, device_buffer_samples); rmm::device_uvector L2NormBuf_OR_DistBuf(0, stream); @@ -855,9 +855,6 @@ void kmeans_fit( auto batch_weights_view = cur_batch_weights(static_cast(data_batch.offset()), wt_data, cur_batch_size); - auto minCAD_view = raft::make_device_vector_view, IndexT>( - minClusterAndDistance.data_handle(), cur_batch_size); - if constexpr (!data_on_device) { if (need_compute_norms) { if (!norms_cached) { @@ -885,7 +882,7 @@ void kmeans_fit( metric, iter_params.batch_samples, iter_params.batch_centroids, - minCAD_view, + assignment_storage, l2_const_view, L2NormBuf_OR_DistBuf, ws, @@ -1019,13 +1016,13 @@ void fit(raft::resources const& handle, kmeans_fit(handle, params, X, sample_weight, centroids, inertia, n_iter); } -template +template void kmeans_predict(raft::resources const& handle, const cuvs::cluster::kmeans::params& pams, raft::device_matrix_view X, std::optional> sample_weight, raft::device_matrix_view centroids, - raft::device_vector_view labels, + raft::device_vector_view labels, bool normalize_weight, raft::host_scalar_view inertia) { @@ -1073,8 +1070,6 @@ void kmeans_predict(raft::resources const& handle, raft::make_const_mdspan(weight.view())); } - auto minClusterAndDistance = - raft::make_device_vector, IndexT>(handle, n_samples); rmm::device_uvector L2NormBuf_OR_DistBuf(0, stream); // L2 norm of X: ||x||^2 @@ -1084,49 +1079,86 @@ void kmeans_predict(raft::resources const& handle, raft::linalg::norm(handle, X, L2NormX.view()); } - // computes minClusterAndDistance[0:n_samples) where minClusterAndDistance[i] - // is a pair where - // 'key' is index to a sample in 'centroids' (index of the nearest - // centroid) and 'value' is the distance between the sample 'X[i]' and the - // 'centroid[key]' auto l2normx_view = raft::make_device_vector_view(L2NormX.data_handle(), n_samples); - cuvs::cluster::kmeans::detail::minClusterAndDistanceCompute( - handle, - X, - centroids, - minClusterAndDistance.view(), - l2normx_view, - L2NormBuf_OR_DistBuf, - pams.metric, - pams.batch_samples, - pams.batch_centroids, - workspace); - - // calculate cluster cost phi_x(C) rmm::device_scalar clusterCostD(stream); - raft::linalg::map( - handle, - minClusterAndDistance.view(), - [=] __device__(const raft::KeyValuePair kvp, DataT wt) { - raft::KeyValuePair res; - res.value = kvp.value * wt; - res.key = kvp.key; - return res; - }, - raft::make_const_mdspan(minClusterAndDistance.view()), - raft::make_const_mdspan(weight.view())); - - cuvs::cluster::kmeans::detail::computeClusterCost( - handle, - minClusterAndDistance.view(), - workspace, - raft::make_device_scalar_view(clusterCostD.data()), - raft::value_op{}, - raft::add_op{}); - - raft::linalg::map( - handle, labels, raft::key_op{}, raft::make_const_mdspan(minClusterAndDistance.view())); + const auto requirements = get_fused_1nn_requirements( + handle, X, centroids, metric, pams.batch_samples, pams.batch_centroids); + + if (requirements.output_layout == Fused1nnOutputLayout::Soa) { + auto nearest_dist = raft::make_device_vector(handle, n_samples); + if constexpr (std::is_same_v) { + minClusterAndDistanceCompute(handle, + X, + centroids, + labels, + nearest_dist.view(), + l2normx_view, + L2NormBuf_OR_DistBuf, + metric, + pams.batch_samples, + pams.batch_centroids, + workspace, + requirements); + } else { + auto index_labels = raft::make_device_vector(handle, n_samples); + minClusterAndDistanceCompute(handle, + X, + centroids, + index_labels.view(), + nearest_dist.view(), + l2normx_view, + L2NormBuf_OR_DistBuf, + metric, + pams.batch_samples, + pams.batch_centroids, + workspace, + requirements); + raft::linalg::map( + handle, labels, raft::cast_op{}, raft::make_const_mdspan(index_labels.view())); + } + raft::linalg::map(handle, + nearest_dist.view(), + raft::mul_op{}, + raft::make_const_mdspan(nearest_dist.view()), + raft::make_const_mdspan(weight.view())); + computeClusterCost(handle, + nearest_dist.view(), + workspace, + raft::make_device_scalar_view(clusterCostD.data()), + raft::identity_op{}, + raft::add_op{}); + } else { + using KvpT = raft::KeyValuePair; + auto nearest = raft::make_device_vector(handle, n_samples); + minClusterAndDistanceComputeKvp(handle, + X, + centroids, + nearest.view(), + l2normx_view, + L2NormBuf_OR_DistBuf, + metric, + pams.batch_samples, + pams.batch_centroids, + workspace, + requirements); + auto* nearest_ptr = nearest.data_handle(); + raft::linalg::map_offset(handle, labels, [nearest_ptr] __device__(IndexT i) { + return static_cast(nearest_ptr[i].key); + }); + auto* weights = weight.data_handle(); + cuda::counting_iterator indices(IndexT{0}); + cuda::transform_iterator weighted_dist(indices, + [nearest_ptr, weights] __device__(IndexT i) -> DataT { + return nearest_ptr[i].value * weights[i]; + }); + computeClusterCostFromIterator(handle, + weighted_dist, + n_samples, + workspace, + raft::make_device_scalar_view(clusterCostD.data()), + raft::add_op{}); + } inertia[0] = clusterCostD.value(stream); } diff --git a/cpp/src/cluster/detail/kmeans_balanced.cuh b/cpp/src/cluster/detail/kmeans_balanced.cuh index 272d45f13d..1e84f6e7e1 100644 --- a/cpp/src/cluster/detail/kmeans_balanced.cuh +++ b/cpp/src/cluster/detail/kmeans_balanced.cuh @@ -5,6 +5,7 @@ #pragma once +#include "../kmeans.cuh" #include "kmeans_common.cuh" #include @@ -43,6 +44,7 @@ #include #include +#include #include #include #include @@ -52,6 +54,131 @@ namespace cuvs::cluster::kmeans::detail { +template +bool predict_core_min_cluster(const raft::resources& handle, + raft::device_matrix_view X, + raft::device_matrix_view centroids, + raft::device_vector_view X_norm, + cuvs::distance::DistanceType metric, + LabelT* labels, + rmm::device_uvector& L2NormBuf_OR_DistBuf, + rmm::device_uvector& workspace, + rmm::device_async_resource_ref mr, + const MathT* cutile_x_norm) +{ + auto n_rows = X.extent(0); + const auto requirements = get_fused_1nn_requirements(handle, X, centroids, metric); + + if (requirements.output_layout == Fused1nnOutputLayout::Soa) { + // cuTile indices are int32. For 32-bit labels, request that native representation directly; + // cluster labels are never negative, so no conversion or copy is required. + if constexpr (std::is_same_v || std::is_same_v) { + if (n_rows <= static_cast(std::numeric_limits::max())) { + using CutileIdxT = int; + const auto cutile_rows = static_cast(n_rows); + const auto cutile_cols = static_cast(X.extent(1)); + const auto cutile_clusters = static_cast(centroids.extent(0)); + auto cutile_X = raft::make_device_matrix_view( + X.data_handle(), cutile_rows, cutile_cols); + auto cutile_centroids = raft::make_device_matrix_view( + centroids.data_handle(), cutile_clusters, cutile_cols); + const auto cutile_requirements = + get_fused_1nn_requirements(handle, cutile_X, cutile_centroids, metric); + RAFT_EXPECTS(cutile_requirements.output_layout == Fused1nnOutputLayout::Soa, + "resolved cuTile plan changed while adapting 32-bit labels"); + + auto nearest_dist = raft::make_device_mdarray( + handle, mr, raft::make_extents(cutile_rows)); + auto* cutile_labels = reinterpret_cast(labels); + auto labels_view = + raft::make_device_vector_view(cutile_labels, cutile_rows); + auto cutile_norm = + raft::make_device_vector_view(X_norm.data_handle(), cutile_rows); + minClusterAndDistanceCompute(handle, + cutile_X, + cutile_centroids, + labels_view, + nearest_dist.view(), + cutile_norm, + L2NormBuf_OR_DistBuf, + metric, + 0, + 0, + workspace, + cutile_requirements, + cutile_x_norm); + return true; + } + } + + constexpr bool label_storage_is_cutile_index = std::is_same_v; + auto nearest_dist = + raft::make_device_mdarray(handle, mr, raft::make_extents(n_rows)); + + if constexpr (label_storage_is_cutile_index) { + auto* cutile_labels = reinterpret_cast(labels); + auto labels_view = raft::make_device_vector_view(cutile_labels, n_rows); + minClusterAndDistanceCompute(handle, + X, + centroids, + labels_view, + nearest_dist.view(), + X_norm, + L2NormBuf_OR_DistBuf, + metric, + 0, + 0, + workspace, + requirements, + cutile_x_norm); + } else { + auto nearest_idx = + raft::make_device_mdarray(handle, mr, raft::make_extents(n_rows)); + minClusterAndDistanceCompute(handle, + X, + centroids, + nearest_idx.view(), + nearest_dist.view(), + X_norm, + L2NormBuf_OR_DistBuf, + metric, + 0, + 0, + workspace, + requirements, + cutile_x_norm); + raft::copy( + handle, raft::make_device_vector_view(labels, n_rows), nearest_idx.view()); + } + return true; + } + + // The generic pairwise-distance assignment path does not implement InnerProduct. Preserve the + // historical GEMM-plus-argmin fallback in predict_core after any final cuTile rejection. + if (metric == cuvs::distance::DistanceType::InnerProduct) { return false; } + + using KvpT = raft::KeyValuePair; + auto nearest = + raft::make_device_mdarray(handle, mr, raft::make_extents(n_rows)); + minClusterAndDistanceComputeKvp(handle, + X, + centroids, + nearest.view(), + X_norm, + L2NormBuf_OR_DistBuf, + metric, + 0, + 0, + workspace, + requirements); + auto* nearest_ptr = nearest.data_handle(); + raft::linalg::map_offset( + handle, + raft::make_device_vector_view(labels, n_rows), + [nearest_ptr] __device__(IdxT i) -> LabelT { return static_cast(nearest_ptr[i].key); }); + return true; +} + /** * @brief Predict labels for the dataset; floating-point types only. * @@ -82,6 +209,7 @@ inline std::enable_if_t> predict_core( IdxT dim, const MathT* dataset, const MathT* dataset_norm, + const MathT* dataset_cutile_norm, IdxT n_rows, LabelT* labels, rmm::device_async_resource_ref mr) @@ -99,55 +227,65 @@ inline std::enable_if_t> predict_core( raft::make_device_matrix_view(centers, n_clusters, dim); auto X_norm_view = raft::make_device_vector_view(dataset_norm, n_rows); - auto minClusterAndDistance = raft::make_device_mdarray, IdxT>( - handle, mr, raft::make_extents(n_rows)); - - cuvs::cluster::kmeans::detail::minClusterAndDistanceCompute( - handle, - X_view, - centroids_view, - minClusterAndDistance.view(), - X_norm_view, - L2NormBuf_OR_DistBuf, - params.metric, - 0, // batch_samples (unused for fused reduction) - 0, // batch_centroids (unused for fused reduction) - workspace); - - // Copy keys to output labels - raft::linalg::map(handle, - raft::make_const_mdspan(minClusterAndDistance.view()), - raft::make_device_vector_view(labels, n_rows), - raft::compose_op, raft::key_op>()); + predict_core_min_cluster(handle, + X_view, + centroids_view, + X_norm_view, + params.metric, + labels, + L2NormBuf_OR_DistBuf, + workspace, + mr, + dataset_cutile_norm); break; } case cuvs::distance::DistanceType::InnerProduct: { - // TODO: pass buffer - rmm::device_uvector distances(n_rows * n_clusters, stream, mr); + rmm::device_uvector L2NormBuf_OR_DistBuf(0, stream, mr); + rmm::device_uvector workspace(0, stream, mr); - MathT alpha = -1.0; - MathT beta = 0.0; + auto X_view = raft::make_device_matrix_view(dataset, n_rows, dim); + auto centroids_view = + raft::make_device_matrix_view(centers, n_clusters, dim); + auto X_norm_view = raft::make_device_vector_view(dataset_norm, n_rows); - raft::linalg::gemm(handle, - true, - false, - n_clusters, - n_rows, - dim, - &alpha, - centers, - dim, - dataset, - dim, - &beta, - distances.data(), - n_clusters, - stream); + if (!predict_core_min_cluster(handle, + X_view, + centroids_view, + X_norm_view, + params.metric, + labels, + L2NormBuf_OR_DistBuf, + workspace, + mr, + dataset_cutile_norm)) { + rmm::device_uvector distances( + static_cast(n_rows) * static_cast(n_clusters), stream, mr); + + MathT alpha = -1.0; + MathT beta = 0.0; + + raft::linalg::gemm(handle, + true, + false, + n_clusters, + n_rows, + dim, + &alpha, + centers, + dim, + dataset, + dim, + &beta, + distances.data(), + n_clusters, + stream); - auto distances_const_view = raft::make_device_matrix_view( - distances.data(), n_rows, n_clusters); - auto labels_view = raft::make_device_vector_view(labels, n_rows); - raft::matrix::argmin(handle, distances_const_view, labels_view); + auto distances_const_view = + raft::make_device_matrix_view( + distances.data(), n_rows, n_clusters); + auto labels_view = raft::make_device_vector_view(labels, n_rows); + raft::matrix::argmin(handle, distances_const_view, labels_view); + } break; } default: { @@ -164,57 +302,125 @@ inline std::enable_if_t> predict_core( * * @tparam MathT type of the centroids and mapped data * @tparam IdxT index type + * @tparam LabelT label type * * @param[in] n_clusters number of clusters in kmeans clustering * @param[in] n_rows Number of samples in the dataset * @param[in] dim Number of features in the dataset * @param[in] metric Distance metric - * @param[in] needs_conversion Whether the data needs to be converted to MathT + * @param[in] data_is_math_type Whether the input data already uses MathT * @return A suggested minibatch size and the expected memory cost per-row (in bytes) */ -template +template auto calc_minibatch_size(const raft::resources& handle, IdxT n_clusters, IdxT n_rows, IdxT dim, cuvs::distance::DistanceType metric, - bool needs_conversion) -> std::tuple + bool data_is_math_type) -> std::tuple { n_clusters = std::max(1, n_clusters); // Estimate memory needs per row (i.e element of the batch). - size_t mem_per_row = 0; + auto saturating_add = [](size_t a, size_t b) { + return b > std::numeric_limits::max() - a ? std::numeric_limits::max() : a + b; + }; + auto saturating_multiply = [](size_t a, size_t b) { + return b != 0 && a > std::numeric_limits::max() / b ? std::numeric_limits::max() + : a * b; + }; + size_t common_mem_per_row = 0; + size_t path_mem_per_row = 0; + size_t fixed_bytes = 0; switch (metric) { case distance::DistanceType::L2Expanded: - case distance::DistanceType::L2SqrtExpanded: { - if (use_fused(handle, n_rows, n_clusters, dim)) { - // fusedL2NN needs a mutex and a key-value pair for each row. - mem_per_row += sizeof(int); - mem_per_row += sizeof(raft::KeyValuePair); - } else { - // unfused path needs a full GEMM output (distance matrix row). - mem_per_row += sizeof(MathT) * n_clusters; + case distance::DistanceType::L2SqrtExpanded: + case distance::DistanceType::CosineExpanded: + case distance::DistanceType::InnerProduct: { + const auto fused_path = use_fused(handle, n_rows, n_clusters, dim, metric); + + if (metric != distance::DistanceType::InnerProduct) { + // predict may need a minibatch-sized input-norm buffer before entering predict_core. + common_mem_per_row += sizeof(MathT); + fixed_bytes = saturating_multiply(sizeof(MathT), static_cast(n_clusters)); + } + + auto path_bytes_per_row = [&](FusedDistancePath path) { + size_t bytes = 0; + switch (path) { + case FusedDistancePath::Cutile: + // cuTile writes separate distance and index arrays. + bytes += sizeof(MathT); + if constexpr (!std::is_same_v) { bytes += sizeof(IdxT); } + if constexpr (std::is_same_v) { + // cuTile converts chunk-local int32 indices to int64 output indices. + bytes += sizeof(int); + } + if constexpr (std::is_same_v) { + if (metric != distance::DistanceType::InnerProduct) { + // TF32-compatible row norms are materialized for cuTile L2/cosine. + bytes += sizeof(MathT); + } + } + break; + case FusedDistancePath::Cutlass: + // CUTLASS writes native KVP assignments and uses one mutex per row. + bytes += sizeof(int); + bytes += sizeof(raft::KeyValuePair); + break; + case FusedDistancePath::Unfused: + // Unfused assignment needs a full distance matrix row. + bytes = saturating_add( + bytes, saturating_multiply(sizeof(MathT), static_cast(n_clusters))); + if (metric != distance::DistanceType::InnerProduct) { + bytes += sizeof(raft::KeyValuePair); + } + break; + } + return bytes; + }; + + path_mem_per_row = path_bytes_per_row(fused_path); + if (fused_path == FusedDistancePath::Cutile) { + const auto prop = raft::resource::get_device_properties(handle); + const auto fallback = use_legacy_fused(prop.major, n_rows, n_clusters, metric); + path_mem_per_row = std::max(path_mem_per_row, path_bytes_per_row(fallback)); + + // Hopper changes from unfused to CUTLASS at 4096 rows. A final minibatch can fall below + // that threshold even when full batches use CUTLASS, so reserve for both real outcomes. + if (prop.major == 9 && n_clusters < IdxT{4096} && n_rows >= IdxT{4096}) { + path_mem_per_row = + std::max(path_mem_per_row, path_bytes_per_row(FusedDistancePath::Unfused)); + } } } break; // Other metrics require storing a distance matrix. default: { - mem_per_row += sizeof(MathT) * n_clusters; + path_mem_per_row = saturating_multiply(sizeof(MathT), static_cast(n_clusters)); } } // If we need to convert to MathT, space required for the converted batch. - if (!needs_conversion) { mem_per_row += sizeof(MathT) * dim; } + if (!data_is_math_type) { + common_mem_per_row = saturating_add( + common_mem_per_row, saturating_multiply(sizeof(MathT), static_cast(dim))); + } + const size_t mem_per_row = saturating_add(common_mem_per_row, path_mem_per_row); // Heuristic: calculate the minibatch size in order to use at most 80% or 512MB workspace memory. // We go below 1GB here as the allocation is mostly done in a single chunk which // is problematic if e.g. a pool allocator manages its own chunks <= 1GB. const auto free_ws_size = raft::resource::get_workspace_free_bytes(handle); - const auto available_ws_size = + const auto workspace_limit = std::min((free_ws_size * size_t{8}) / size_t{10}, size_t{1} << 29); + const auto available_ws_size = + workspace_limit > fixed_bytes ? workspace_limit - fixed_bytes : size_t{1}; IdxT minibatch_size = std::max(IdxT{1}, static_cast(available_ws_size / mem_per_row)); - minibatch_size = raft::round_down_safe(minibatch_size, IdxT{64}); + if (minibatch_size >= IdxT{64}) { + minibatch_size = raft::round_down_safe(minibatch_size, IdxT{64}); + } minibatch_size = std::min(minibatch_size, n_rows); return std::make_tuple(minibatch_size, mem_per_row); } @@ -354,6 +560,34 @@ void compute_norm(const raft::resources& handle, norm_fin_op); } +/** Computes TF32-compatible norms for cuTile, converting to float when necessary. */ +template +void compute_cutile_norm(const raft::resources& handle, + float* dataset_norm, + const T* dataset, + IdxT dim, + IdxT n_rows, + MappingOpT mapping_op, + bool take_sqrt, + rmm::device_async_resource_ref mr) +{ + auto stream = raft::resource::get_cuda_stream(handle); + rmm::device_uvector mapped_dataset(0, stream, mr); + const float* dataset_ptr = nullptr; + if constexpr (std::is_same_v) { + dataset_ptr = dataset; + } else { + mapped_dataset.resize(static_cast(n_rows) * static_cast(dim), stream); + raft::linalg::map( + handle, + raft::make_device_vector_view(dataset, n_rows * dim), + raft::make_device_vector_view(mapped_dataset.data(), n_rows * dim), + mapping_op); + dataset_ptr = mapped_dataset.data(); + } + computeCutileRowNorms(handle, dataset_ptr, dataset_norm, n_rows, dim, take_sqrt); +} + /** * @brief Predict labels for the dataset. * @@ -386,13 +620,14 @@ void predict(const raft::resources& handle, LabelT* labels, MappingOpT mapping_op, std::optional mr = std::nullopt, - const MathT* dataset_norm = nullptr) + const MathT* dataset_norm = nullptr, + const MathT* dataset_cutile_norm = nullptr) { auto stream = raft::resource::get_cuda_stream(handle); raft::common::nvtx::range fun_scope( "predict(%zu, %u)", static_cast(n_rows), n_clusters); auto mem_res = mr.value_or(raft::resource::get_workspace_resource_ref(handle)); - auto [max_minibatch_size, _mem_per_row] = calc_minibatch_size( + auto [max_minibatch_size, _mem_per_row] = calc_minibatch_size( handle, n_clusters, n_rows, dim, params.metric, std::is_same_v); rmm::device_uvector cur_dataset( std::is_same_v ? 0 : max_minibatch_size * dim, stream, mem_res); @@ -449,6 +684,7 @@ void predict(const raft::resources& handle, dim, cur_dataset_ptr, dataset_norm_ptr, + dataset_cutile_norm == nullptr ? nullptr : dataset_cutile_norm + offset, minibatch_size, labels + offset, mem_res); @@ -778,6 +1014,7 @@ void balancing_em_iters(const raft::resources& handle, IdxT dim, const T* dataset, const MathT* dataset_norm, + const MathT* dataset_cutile_norm, IdxT n_rows, IdxT n_clusters, MathT* cluster_centers, @@ -846,7 +1083,8 @@ void balancing_em_iters(const raft::resources& handle, cluster_labels, mapping_op, device_memory, - dataset_norm); + dataset_norm, + dataset_cutile_norm); // M: Maximization step - calculate optimal cluster centers calc_centers_and_sizes(handle, cluster_centers, @@ -880,7 +1118,8 @@ void build_clusters(const raft::resources& handle, CounterT* cluster_sizes, MappingOpT mapping_op, rmm::device_async_resource_ref device_memory, - const MathT* dataset_norm = nullptr) + const MathT* dataset_norm = nullptr, + const MathT* dataset_cutile_norm = nullptr) { auto stream = raft::resource::get_cuda_stream(handle); // "randomly" initialize labels @@ -910,6 +1149,7 @@ void build_clusters(const raft::resources& handle, dim, dataset, dataset_norm, + dataset_cutile_norm, n_rows, n_clusters, cluster_centers, @@ -1011,6 +1251,7 @@ auto build_fine_clusters(const raft::resources& handle, IdxT dim, const T* dataset_mptr, const MathT* dataset_norm_mptr, + const MathT* dataset_cutile_norm_mptr, const LabelT* labels_mptr, IdxT n_rows, const IdxT* fine_clusters_nums, @@ -1031,9 +1272,12 @@ auto build_fine_clusters(const raft::resources& handle, auto large_ws = raft::resource::get_large_workspace_resource_ref(handle); rmm::device_uvector mc_trainset_buf(mesocluster_size_max * dim, stream, large_ws); rmm::device_uvector mc_trainset_norm_buf(mesocluster_size_max, stream, device_memory); - auto mc_trainset_ids = mc_trainset_ids_buf.data(); - auto mc_trainset = mc_trainset_buf.data(); - auto mc_trainset_norm = mc_trainset_norm_buf.data(); + rmm::device_uvector mc_trainset_cutile_norm_buf( + dataset_cutile_norm_mptr == nullptr ? 0 : mesocluster_size_max, stream, device_memory); + auto mc_trainset_ids = mc_trainset_ids_buf.data(); + auto mc_trainset = mc_trainset_buf.data(); + auto mc_trainset_norm = mc_trainset_norm_buf.data(); + auto mc_trainset_cutile_norm = mc_trainset_cutile_norm_buf.data(); // label (cluster ID) of each vector rmm::device_uvector mc_trainset_labels(mesocluster_size_max, stream, device_memory); @@ -1077,6 +1321,13 @@ auto build_fine_clusters(const raft::resources& handle, mc_trainset_ids + k, dataset_norm_mptr, mc_trainset_norm); + if (dataset_cutile_norm_mptr != nullptr) { + thrust::gather(raft::resource::get_thrust_policy(handle), + mc_trainset_ids, + mc_trainset_ids + k, + dataset_cutile_norm_mptr, + mc_trainset_cutile_norm); + } } build_clusters(handle, @@ -1090,7 +1341,8 @@ auto build_fine_clusters(const raft::resources& handle, mc_trainset_csizes_tmp.data(), mapping_op, device_memory, - mc_trainset_norm); + mc_trainset_norm, + dataset_cutile_norm_mptr == nullptr ? nullptr : mc_trainset_cutile_norm); raft::copy(handle, raft::make_device_vector_view(cluster_centers + (dim * fine_clusters_csum[i]), @@ -1146,12 +1398,14 @@ void build_hierarchical(const raft::resources& handle, // TODO: Remove the explicit managed memory- we shouldn't be creating this on the user's behalf. rmm::mr::managed_memory_resource managed_memory; rmm::device_async_resource_ref device_memory = raft::resource::get_workspace_resource_ref(handle); - auto [max_minibatch_size, mem_per_row] = calc_minibatch_size( + auto [max_minibatch_size, mem_per_row] = calc_minibatch_size( handle, n_clusters, n_rows, dim, params.metric, std::is_same_v); // Precompute the L2 norm of the dataset if relevant and not yet computed. rmm::device_uvector dataset_norm_buf(0, stream, device_memory); - const MathT* dataset_norm = nullptr; + rmm::device_uvector dataset_cutile_norm_buf(0, stream, device_memory); + const MathT* dataset_norm = nullptr; + const MathT* dataset_cutile_norm = nullptr; if ((params.metric == cuvs::distance::DistanceType::L2Expanded || params.metric == cuvs::distance::DistanceType::L2SqrtExpanded || params.metric == cuvs::distance::DistanceType::CosineExpanded)) { @@ -1178,6 +1432,26 @@ void build_hierarchical(const raft::resources& handle, device_memory); } dataset_norm = (const MathT*)dataset_norm_buf.data(); + + if constexpr (std::is_same_v) { + if (use_fused(handle, n_rows, n_clusters, dim, params.metric) == + FusedDistancePath::Cutile) { + dataset_cutile_norm_buf.resize(n_rows, stream); + const bool take_sqrt = params.metric == cuvs::distance::DistanceType::CosineExpanded; + for (IdxT offset = 0; offset < n_rows; offset += max_minibatch_size) { + const IdxT minibatch_size = std::min(max_minibatch_size, n_rows - offset); + compute_cutile_norm(handle, + dataset_cutile_norm_buf.data() + offset, + dataset + dim * offset, + dim, + minibatch_size, + mapping_op, + take_sqrt, + device_memory); + } + dataset_cutile_norm = dataset_cutile_norm_buf.data(); + } + } } /* Temporary workaround to cub::DeviceHistogram not supporting any type that isn't natively @@ -1201,7 +1475,8 @@ void build_hierarchical(const raft::resources& handle, mesocluster_sizes_buf.data(), mapping_op, device_memory, - dataset_norm); + dataset_norm, + dataset_cutile_norm); } auto mesocluster_sizes = mesocluster_sizes_buf.data(); @@ -1233,6 +1508,7 @@ void build_hierarchical(const raft::resources& handle, dim, dataset, dataset_norm, + dataset_cutile_norm, mesocluster_labels, n_rows, fine_clusters_nums.data(), @@ -1271,6 +1547,7 @@ void build_hierarchical(const raft::resources& handle, dim, dataset, dataset_norm, + dataset_cutile_norm, n_rows, n_clusters, cluster_centers, diff --git a/cpp/src/cluster/detail/kmeans_common.cuh b/cpp/src/cluster/detail/kmeans_common.cuh index ab3ef0a05a..b7c286edee 100644 --- a/cpp/src/cluster/detail/kmeans_common.cuh +++ b/cpp/src/cluster/detail/kmeans_common.cuh @@ -5,8 +5,10 @@ #pragma once #include "../../distance/distance.cuh" +#include "../../distance/fused_distance_nn.cuh" #include #include +#include #include #include @@ -57,29 +59,101 @@ namespace cuvs::cluster::kmeans::detail { +template +inline constexpr bool is_cutile_fused_data_type_v = + std::is_same_v || std::is_same_v; + +using FusedDistancePath = cuvs::distance::detail::Fused1nnBackend; + +/** Native result representation selected by fused 1-NN. */ +enum class Fused1nnOutputLayout : std::uint8_t { + /** Separate index and distance arrays (cuTile). */ + Soa, + /** Native key/value-pair array (CUTLASS/SIMT and the unfused reducer). */ + Kvp, +}; + +/** Norm representation required by the resolved fused-1NN implementation. */ +enum class Fused1nnNormPolicy : std::uint8_t { + Default, + Tf32, +}; + /** - * @brief Returns true if the fused distance NN implementation should be used. + * Resolved fused-1NN storage and execution requirements. * - * On Ampere (SM <= 8.x) always use fused. - * On Hopper (SM 9.x) use fused when m or n >= 4096. - * On Blackwell (SM >= 10.x) use unfused. + * KMeans uses the result layout and byte counts to reuse its existing raw buffers. The fused + * implementation owns execution of the resolved path. + */ +template +struct Fused1nnRequirements { + FusedDistancePath path{}; + Fused1nnOutputLayout output_layout{}; + Fused1nnNormPolicy norm_policy{}; + size_t result_bytes{}; + size_t result_alignment{}; + size_t distance_offset{}; + size_t workspace_bytes{}; + size_t workspace_alignment{}; + IndexT sample_tile{}; + IndexT centroid_tile{}; +}; + +inline constexpr bool uses_fused_distance_nn(FusedDistancePath path) +{ + return path != FusedDistancePath::Unfused; +} + +/** + * @brief Select the pre-cuTile assignment path. + * + * This is also the fallback after a pointer-aware cuTile launch probe fails. + */ +template +constexpr FusedDistancePath use_legacy_fused(int cc_major, + IdxT m, + IdxT n, + cuvs::distance::DistanceType metric) +{ + return cuvs::distance::detail::fused_1nn_legacy_backend(cc_major, m, n, metric); +} + +template +FusedDistancePath use_legacy_fused(const raft::resources& handle, + IdxT m, + IdxT n, + cuvs::distance::DistanceType metric) +{ + const auto prop = raft::resource::get_device_properties(handle); + return cuvs::distance::detail::fused_1nn_legacy_backend(prop.major, m, n, metric); +} + +/** + * @brief Selects the fused-distance assignment path for KMeans. + * + * With CUDA 13, float/half use cuTile whenever the build and device support it. CUDA 12 and a + * failed CUDA 13 cuTile probe use the historical CUTLASS/unfused heuristic. */ template -bool use_fused(const raft::resources& handle, IdxT m, IdxT n, IdxT k) +FusedDistancePath use_fused( + const raft::resources& handle, IdxT m, IdxT n, IdxT k, cuvs::distance::DistanceType metric) { - cudaDeviceProp prop; - prop = raft::resource::get_device_properties(handle); - if (prop.major <= 8) { - // Use fused for Ampere or before - return true; - } else if (prop.major == 9 && (m >= 4096 || n >= 4096)) { - // On Hopper if m, n are bigger than 4096, use fused - return true; - } else if (prop.major >= 10) { - // On Blackwell onwards, use unfused - return false; + (void)k; + +#if CUDART_VERSION >= 13000 + if constexpr (is_cutile_fused_data_type_v) { + if constexpr (cuvs::detail::jit_lto::library_built_with_cutile()) { + const bool dimensions_fit_i32 = n <= static_cast(std::numeric_limits::max()) && + k <= static_cast(std::numeric_limits::max()); + if (dimensions_fit_i32 && + cuvs::detail::jit_lto::cutile_launch_available_on_current_device()) { + return FusedDistancePath::Cutile; + } + } } - return false; +#endif + + return use_legacy_fused(handle, m, n, metric); } template @@ -207,28 +281,22 @@ IndexT getCentroidsBatchSize(int batch_centroids, IndexT n_local_clusters) return (minVal == 0) ? n_local_clusters : minVal; } -template -void computeClusterCost(raft::resources const& handle, - raft::device_vector_view minClusterDistance, - rmm::device_uvector& workspace, - raft::device_scalar_view clusterCost, - MainOpT main_op, - ReductionOpT reduction_op) +template +void computeClusterCostFromIterator(raft::resources const& handle, + InputIteratorT input, + IndexT n, + rmm::device_uvector& workspace, + raft::device_scalar_view clusterCost, + ReductionOpT reduction_op) { cudaStream_t stream = raft::resource::get_cuda_stream(handle); - cuda::transform_iterator itr(minClusterDistance.data_handle(), main_op); - size_t temp_storage_bytes = 0; RAFT_CUDA_TRY(cub::DeviceReduce::Reduce(nullptr, temp_storage_bytes, - itr, + input, clusterCost.data_handle(), - minClusterDistance.size(), + n, reduction_op, OutputT(), stream)); @@ -237,14 +305,31 @@ void computeClusterCost(raft::resources const& handle, RAFT_CUDA_TRY(cub::DeviceReduce::Reduce(workspace.data(), temp_storage_bytes, - itr, + input, clusterCost.data_handle(), - minClusterDistance.size(), + n, reduction_op, OutputT(), stream)); } +template +void computeClusterCost(raft::resources const& handle, + raft::device_vector_view minClusterDistance, + rmm::device_uvector& workspace, + raft::device_scalar_view clusterCost, + MainOpT main_op, + ReductionOpT reduction_op) +{ + cuda::transform_iterator input(minClusterDistance.data_handle(), main_op); + computeClusterCostFromIterator( + handle, input, minClusterDistance.size(), workspace, clusterCost, reduction_op); +} + template void sampleCentroids(raft::resources const& handle, raft::device_matrix_view X, @@ -339,8 +424,19 @@ void pairwise_distance_kmeans(raft::resources const& handle, DataT, raft::layout_c_contiguous, IndexT>(handle, X, centroids, pairwiseDistance); + } else if (metric == cuvs::distance::DistanceType::L2Unexpanded) { + if constexpr (std::is_same_v) { + cuvs::distance::distance(handle, X, centroids, pairwiseDistance); + } else { + RAFT_FAIL("L2Unexpanded KMeans distance requires int32-indexed batches"); + } } else { - RAFT_FAIL("kmeans requires L2Expanded or L2SqrtExpanded distance, have %i", + RAFT_FAIL("kmeans requires L2Expanded, L2SqrtExpanded, or L2Unexpanded distance, have %i", static_cast(metric)); } } @@ -378,34 +474,60 @@ void shuffleAndGather(raft::resources const& handle, stream); } -// Calculates a pair for every sample in input 'X' where key is an -// index to an sample in 'centroids' (index of the nearest centroid) and 'value' -// is the distance between the sample and the 'centroid[key]' +// Calculates nearest centroid index and distance for every sample in input 'X'. +template +Fused1nnRequirements get_fused_1nn_requirements( + raft::resources const& handle, + raft::device_matrix_view X, + raft::device_matrix_view centroids, + cuvs::distance::DistanceType metric, + int batch_samples = 0, + int batch_centroids = 0); + template -void minClusterAndDistanceCompute( +void minClusterAndDistanceCompute(raft::resources const& handle, + raft::device_matrix_view X, + raft::device_matrix_view centroids, + raft::device_vector_view nearest_idx, + raft::device_vector_view nearest_dist, + raft::device_vector_view L2NormX, + rmm::device_uvector& L2NormBuf_OR_DistBuf, + cuvs::distance::DistanceType metric, + int batch_samples, + int batch_centroids, + rmm::device_uvector& workspace, + const Fused1nnRequirements& requirements, + const DataT* cutile_x_norm = nullptr); + +template +void minClusterAndDistanceComputeKvp( raft::resources const& handle, raft::device_matrix_view X, raft::device_matrix_view centroids, - raft::device_vector_view, IndexT> minClusterAndDistance, + raft::device_vector_view, IndexT> nearest, raft::device_vector_view L2NormX, rmm::device_uvector& L2NormBuf_OR_DistBuf, cuvs::distance::DistanceType metric, int batch_samples, int batch_centroids, - rmm::device_uvector& workspace); - -#define EXTERN_TEMPLATE_MIN_CLUSTER_AND_DISTANCE(DataT, IndexT) \ - extern template void minClusterAndDistanceCompute( \ - raft::resources const& handle, \ - raft::device_matrix_view X, \ - raft::device_matrix_view centroids, \ - raft::device_vector_view, IndexT> minClusterAndDistance, \ - raft::device_vector_view L2NormX, \ - rmm::device_uvector& L2NormBuf_OR_DistBuf, \ - cuvs::distance::DistanceType metric, \ - int batch_samples, \ - int batch_centroids, \ - rmm::device_uvector& workspace); + rmm::device_uvector& workspace, + const Fused1nnRequirements& requirements); + +#define EXTERN_TEMPLATE_MIN_CLUSTER_AND_DISTANCE(DataT, IndexT) \ + extern template void minClusterAndDistanceCompute( \ + raft::resources const& handle, \ + raft::device_matrix_view X, \ + raft::device_matrix_view centroids, \ + raft::device_vector_view nearest_idx, \ + raft::device_vector_view nearest_dist, \ + raft::device_vector_view L2NormX, \ + rmm::device_uvector& L2NormBuf_OR_DistBuf, \ + cuvs::distance::DistanceType metric, \ + int batch_samples, \ + int batch_centroids, \ + rmm::device_uvector& workspace, \ + const Fused1nnRequirements& requirements, \ + const DataT* cutile_x_norm); EXTERN_TEMPLATE_MIN_CLUSTER_AND_DISTANCE(float, int64_t) EXTERN_TEMPLATE_MIN_CLUSTER_AND_DISTANCE(float, int) @@ -414,6 +536,14 @@ EXTERN_TEMPLATE_MIN_CLUSTER_AND_DISTANCE(double, int) #undef EXTERN_TEMPLATE_MIN_CLUSTER_AND_DISTANCE +template +void computeCutileRowNorms(raft::resources const& handle, + const float* matrix, + float* norms, + IndexT n_rows, + IndexT n_cols, + bool take_sqrt); + template void minClusterDistanceCompute(raft::resources const& handle, raft::device_matrix_view X, @@ -457,45 +587,56 @@ void countSamplesInCluster(raft::resources const& handle, { cudaStream_t stream = raft::resource::get_cuda_stream(handle); auto n_samples = X.extent(0); - auto n_features = X.extent(1); auto n_clusters = centroids.extent(0); - // stores (key, value) pair corresponding to each sample where - // - key is the index of nearest cluster - // - value is the distance to the nearest cluster - auto minClusterAndDistance = - raft::make_device_vector, IndexT>(handle, n_samples); - - // temporary buffer to store distance matrix, destructor releases the resource rmm::device_uvector L2NormBuf_OR_DistBuf(0, stream); - - // computes minClusterAndDistance[0:n_samples) where minClusterAndDistance[i] - // is a pair where - // 'key' is index to an sample in 'centroids' (index of the nearest - // centroid) and 'value' is the distance between the sample 'X[i]' and the - // 'centroid[key]' - cuvs::cluster::kmeans::detail::minClusterAndDistanceCompute( - handle, - X, - (raft::device_matrix_view)centroids, - minClusterAndDistance.view(), - L2NormX, - L2NormBuf_OR_DistBuf, - params.metric, - params.batch_samples, - params.batch_centroids, - workspace); - - cuda::transform_iterator itr(minClusterAndDistance.data_handle(), - cuvs::cluster::kmeans::detail::KeyValueIndexOp{}); - - // count # of samples in each cluster - countLabels(handle, - itr, - sampleCountInCluster.data_handle(), - (IndexT)n_samples, - (IndexT)n_clusters, - workspace); + auto centroids_const = raft::make_const_mdspan(centroids); + const auto requirements = get_fused_1nn_requirements( + handle, X, centroids_const, params.metric, params.batch_samples, params.batch_centroids); + + auto count_labels = [&](auto labels) { + countLabels(handle, + labels, + sampleCountInCluster.data_handle(), + static_cast(n_samples), + static_cast(n_clusters), + workspace); + }; + + if (requirements.output_layout == Fused1nnOutputLayout::Soa) { + auto nearest_idx = raft::make_device_vector(handle, n_samples); + auto nearest_dist = raft::make_device_vector(handle, n_samples); + minClusterAndDistanceCompute(handle, + X, + centroids_const, + nearest_idx.view(), + nearest_dist.view(), + L2NormX, + L2NormBuf_OR_DistBuf, + params.metric, + params.batch_samples, + params.batch_centroids, + workspace, + requirements); + count_labels(nearest_idx.data_handle()); + } else { + using KvpT = raft::KeyValuePair; + auto nearest = raft::make_device_vector(handle, n_samples); + minClusterAndDistanceComputeKvp(handle, + X, + centroids_const, + nearest.view(), + L2NormX, + L2NormBuf_OR_DistBuf, + params.metric, + params.batch_samples, + params.batch_centroids, + workspace, + requirements); + auto labels = + thrust::make_transform_iterator(nearest.data_handle(), KeyValueIndexOp{}); + count_labels(labels); + } } /** @@ -676,7 +817,7 @@ __device__ void check_convergence(raft::device_scalar_view clusteri * @param[in] batch_samples_param Batch-samples param forwarded to minClusterAndDistanceCompute * @param[in] batch_centroids_param Batch-centroids param forwarded to * minClusterAndDistanceCompute - * @param[inout] minClusterAndDistance Work buffer [batch_size] + * @param[inout] assignment_storage Native SoA or KVP assignment storage for one batch * @param[in] L2NormBatch Precomputed data norms [batch_size] * @param[inout] L2NormBuf_OR_DistBuf Resizable scratch * @param[inout] workspace Resizable scratch @@ -685,66 +826,103 @@ __device__ void check_convergence(raft::device_scalar_view clusteri * @param[inout] clustering_cost Running cost scalar (device) (added into) */ template -void process_batch( - raft::resources const& handle, - raft::device_matrix_view batch_data, - raft::device_vector_view batch_weights, - raft::device_matrix_view centroids, - cuvs::distance::DistanceType metric, - int batch_samples_param, - int batch_centroids_param, - raft::device_vector_view, IndexT> minClusterAndDistance, - raft::device_vector_view L2NormBatch, - rmm::device_uvector& L2NormBuf_OR_DistBuf, - rmm::device_uvector& workspace, - raft::device_matrix_view centroid_sums, - raft::device_vector_view weight_per_cluster, - raft::device_scalar_view clustering_cost, - rmm::device_uvector& batch_workspace) +void process_batch(raft::resources const& handle, + raft::device_matrix_view batch_data, + raft::device_vector_view batch_weights, + raft::device_matrix_view centroids, + cuvs::distance::DistanceType metric, + int batch_samples_param, + int batch_centroids_param, + rmm::device_uvector& assignment_storage, + raft::device_vector_view L2NormBatch, + rmm::device_uvector& L2NormBuf_OR_DistBuf, + rmm::device_uvector& workspace, + raft::device_matrix_view centroid_sums, + raft::device_vector_view weight_per_cluster, + raft::device_scalar_view clustering_cost, + rmm::device_uvector& batch_workspace) { - cudaStream_t stream = raft::resource::get_cuda_stream(handle); + cudaStream_t stream = raft::resource::get_cuda_stream(handle); + const auto n_samples = batch_data.extent(0); + const auto requirements = get_fused_1nn_requirements( + handle, batch_data, centroids, metric, batch_samples_param, batch_centroids_param); + auto batch_cost = raft::make_device_scalar(handle, DataT{0}); - minClusterAndDistanceCompute(handle, - batch_data, - centroids, - minClusterAndDistance, - L2NormBatch, - L2NormBuf_OR_DistBuf, - metric, - batch_samples_param, - batch_centroids_param, - workspace); - - KeyValueIndexOp conversion_op; - thrust::transform_iterator, - const raft::KeyValuePair*> - labels_itr(minClusterAndDistance.data_handle(), conversion_op); - - compute_centroid_adjustments(handle, - batch_data, - batch_weights, - labels_itr, - static_cast(centroid_sums.extent(0)), - centroid_sums, - weight_per_cluster, - batch_workspace, - /*reset_sums=*/false); - - raft::linalg::map( - handle, - minClusterAndDistance, - [=] __device__(const raft::KeyValuePair kvp, DataT wt) { - raft::KeyValuePair res; - res.value = kvp.value * wt; - res.key = kvp.key; - return res; - }, - raft::make_const_mdspan(minClusterAndDistance), - batch_weights); + if (requirements.output_layout == Fused1nnOutputLayout::Soa) { + const auto dist_offset = requirements.distance_offset; + if (assignment_storage.size() < requirements.result_bytes) { + assignment_storage.resize(requirements.result_bytes, stream); + } + auto* nearest_idx = reinterpret_cast(assignment_storage.data()); + auto* nearest_dist = reinterpret_cast(assignment_storage.data() + dist_offset); + auto nearest_idx_view = raft::make_device_vector_view(nearest_idx, n_samples); + auto nearest_dist_view = raft::make_device_vector_view(nearest_dist, n_samples); + minClusterAndDistanceCompute(handle, + batch_data, + centroids, + nearest_idx_view, + nearest_dist_view, + L2NormBatch, + L2NormBuf_OR_DistBuf, + metric, + batch_samples_param, + batch_centroids_param, + workspace, + requirements); + compute_centroid_adjustments(handle, + batch_data, + batch_weights, + nearest_idx, + static_cast(centroid_sums.extent(0)), + centroid_sums, + weight_per_cluster, + batch_workspace, + /*reset_sums=*/false); + auto* weights = batch_weights.data_handle(); + cuda::counting_iterator indices(IndexT{0}); + cuda::transform_iterator weighted_dist(indices, + [nearest_dist, weights] __device__(IndexT i) -> DataT { + return nearest_dist[i] * weights[i]; + }); + computeClusterCostFromIterator( + handle, weighted_dist, n_samples, workspace, batch_cost.view(), raft::add_op{}); + } else { + using KvpT = raft::KeyValuePair; + if (assignment_storage.size() < requirements.result_bytes) { + assignment_storage.resize(requirements.result_bytes, stream); + } + auto* nearest = reinterpret_cast(assignment_storage.data()); + auto nearest_view = raft::make_device_vector_view(nearest, n_samples); + minClusterAndDistanceComputeKvp(handle, + batch_data, + centroids, + nearest_view, + L2NormBatch, + L2NormBuf_OR_DistBuf, + metric, + batch_samples_param, + batch_centroids_param, + workspace, + requirements); + auto labels = thrust::make_transform_iterator(nearest, KeyValueIndexOp{}); + compute_centroid_adjustments(handle, + batch_data, + batch_weights, + labels, + static_cast(centroid_sums.extent(0)), + centroid_sums, + weight_per_cluster, + batch_workspace, + /*reset_sums=*/false); + auto* weights = batch_weights.data_handle(); + cuda::counting_iterator indices(IndexT{0}); + cuda::transform_iterator weighted_dist( + indices, + [nearest, weights] __device__(IndexT i) -> DataT { return nearest[i].value * weights[i]; }); + computeClusterCostFromIterator( + handle, weighted_dist, n_samples, workspace, batch_cost.view(), raft::add_op{}); + } - auto batch_cost = raft::make_device_scalar(handle, DataT{0}); - computeClusterCost( - handle, minClusterAndDistance, workspace, batch_cost.view(), raft::value_op{}, raft::add_op{}); raft::linalg::add(clustering_cost.data_handle(), clustering_cost.data_handle(), batch_cost.data_handle(), diff --git a/cpp/src/cluster/detail/kmeans_mg.cuh b/cpp/src/cluster/detail/kmeans_mg.cuh index dbe2c23039..1de8a401a7 100644 --- a/cpp/src/cluster/detail/kmeans_mg.cuh +++ b/cpp/src/cluster/detail/kmeans_mg.cuh @@ -214,8 +214,7 @@ void mnmg_fit( auto sqrd_norm_error_dev = raft::make_device_scalar(dev_res, DataT{0}); IndexT alloc_batch_size = device_buffer_samples; auto batch_weights = raft::make_device_vector(dev_res, alloc_batch_size); - auto minClusterAndDistance = - raft::make_device_vector, IndexT>(dev_res, alloc_batch_size); + rmm::device_uvector assignment_storage(0, stream); auto L2NormBatch = raft::make_device_vector(dev_res, data_on_device ? IndexT{0} : alloc_batch_size); rmm::device_uvector L2NormBuf_OR_DistBuf(0, stream); @@ -448,10 +447,6 @@ void mnmg_fit( L2NormBatch_const = raft::make_const_mdspan(norm_slice); } - auto minClusterAndDistance_view = - raft::make_device_vector_view, IndexT>( - minClusterAndDistance.data_handle(), current_batch_size); - cuvs::cluster::kmeans::detail::process_batch( dev_res, batch_data_view, @@ -460,7 +455,7 @@ void mnmg_fit( metric, params.batch_samples, params.batch_centroids, - minClusterAndDistance_view, + assignment_storage, L2NormBatch_const, L2NormBuf_OR_DistBuf, workspace, @@ -555,7 +550,8 @@ void mnmg_fit( batch_data_view, rank_centroids_const, raft::make_device_scalar_view(batch_clustering_cost.data_handle()), - batch_sw); + batch_sw, + cuvs::distance::DistanceType::L2Unexpanded); raft::linalg::add(dev_res, raft::make_const_mdspan(clustering_cost.view()), diff --git a/cpp/src/cluster/detail/minClusterDistanceCompute.cu b/cpp/src/cluster/detail/minClusterDistanceCompute.cu index ee3cc3cdfd..4d9403a6d3 100644 --- a/cpp/src/cluster/detail/minClusterDistanceCompute.cu +++ b/cpp/src/cluster/detail/minClusterDistanceCompute.cu @@ -7,190 +7,305 @@ #include "../../distance/unfused_distance_nn.cuh" #include "kmeans_common.cuh" +#include #include +#include +#include + namespace cuvs::cluster::kmeans::detail { -// Calculates a pair for every sample in input 'X' where key is an -// index to an sample in 'centroids' (index of the nearest centroid) and 'value' -// is the distance between the sample and the 'centroids[key]'. +namespace { + +__device__ __forceinline__ float round_to_tf32(float value) +{ +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800 + return nvcuda::wmma::__float_to_tf32(value); +#else + return value; +#endif +} + +struct tf32_square_op { + template + __device__ float operator()(float value, IndexT) const + { + const float rounded = round_to_tf32(value); + return rounded * rounded; + } +}; + +template +void compute_tf32_row_norms(raft::resources const& handle, + const float* matrix, + float* norms, + IndexT n_rows, + IndexT n_cols, + bool take_sqrt) +{ + if (n_rows == 0) { return; } + auto matrix_view = raft::make_device_matrix_view(matrix, n_rows, n_cols); + auto norms_view = raft::make_device_vector_view(norms, n_rows); + if (take_sqrt) { + raft::linalg::coalesced_reduction(handle, + matrix_view, + norms_view, + 0.0f, + false, + tf32_square_op{}, + raft::add_op{}, + raft::sqrt_op{}); + } else { + raft::linalg::coalesced_reduction(handle, + matrix_view, + norms_view, + 0.0f, + false, + tf32_square_op{}, + raft::add_op{}, + raft::identity_op{}); + } +} + +} // namespace + +template +void computeCutileRowNorms(raft::resources const& handle, + const float* matrix, + float* norms, + IndexT n_rows, + IndexT n_cols, + bool take_sqrt) +{ + compute_tf32_row_norms(handle, matrix, norms, n_rows, n_cols, take_sqrt); +} + template -void minClusterAndDistanceCompute( +Fused1nnRequirements get_fused_1nn_requirements( raft::resources const& handle, raft::device_matrix_view X, raft::device_matrix_view centroids, - raft::device_vector_view, IndexT> minClusterAndDistance, - raft::device_vector_view L2NormX, - rmm::device_uvector& L2NormBuf_OR_DistBuf, cuvs::distance::DistanceType metric, int batch_samples, - int batch_centroids, - rmm::device_uvector& workspace) + int batch_centroids) { - cudaStream_t stream = raft::resource::get_cuda_stream(handle); - auto n_samples = X.extent(0); - auto n_features = X.extent(1); - auto n_clusters = centroids.extent(0); - const bool can_use_fused_path = metric == cuvs::distance::DistanceType::L2Expanded || - metric == cuvs::distance::DistanceType::L2SqrtExpanded || - metric == cuvs::distance::DistanceType::CosineExpanded; - - if (can_use_fused_path) { - L2NormBuf_OR_DistBuf.resize(n_clusters, stream); - auto centroidsNorm = - raft::make_device_vector_view(L2NormBuf_OR_DistBuf.data(), n_clusters); - - if (metric == cuvs::distance::DistanceType::CosineExpanded) { - raft::linalg::norm( - handle, centroids, centroidsNorm, raft::sqrt_op{}); - } else { - raft::linalg::norm( - handle, centroids, centroidsNorm); + const auto path = cuvs::distance::detail::resolve_fused_1nn_backend(handle, + X.data_handle(), + centroids.data_handle(), + X.extent(0), + centroids.extent(0), + X.extent(1), + metric); + + Fused1nnRequirements requirements{}; + requirements.path = path; + const cuvs::distance::detail::Top1nnTuning default_tuning{}; + requirements.sample_tile = std::min( + getDataBatchSize(batch_samples, X.extent(0)), + static_cast(std::min(default_tuning.unfused.row_tile, + static_cast(std::numeric_limits::max())))); + requirements.centroid_tile = std::min( + getCentroidsBatchSize(batch_centroids, centroids.extent(0)), + static_cast(std::min(default_tuning.unfused.candidate_tile, + static_cast(std::numeric_limits::max())))); + requirements.workspace_alignment = alignof(int); + + if (path == FusedDistancePath::Cutile) { + requirements.output_layout = Fused1nnOutputLayout::Soa; + requirements.norm_policy = + std::is_same_v ? Fused1nnNormPolicy::Tf32 : Fused1nnNormPolicy::Default; + requirements.result_alignment = 16; + requirements.distance_offset = + raft::alignTo(sizeof(IndexT) * static_cast(X.extent(0)), size_t{16}); + requirements.result_bytes = + requirements.distance_offset + sizeof(DataT) * static_cast(X.extent(0)); + if constexpr (std::is_same_v) { + requirements.workspace_bytes = + sizeof(int) * + cuvs::distance::detail::fused_1nn_cutile_index_workspace_rows(X.extent(0)); } + } else { + requirements.output_layout = Fused1nnOutputLayout::Kvp; + requirements.norm_policy = Fused1nnNormPolicy::Default; + requirements.result_alignment = alignof(raft::KeyValuePair); + requirements.result_bytes = + sizeof(raft::KeyValuePair) * static_cast(X.extent(0)); + if (path == FusedDistancePath::Cutlass) { + requirements.workspace_bytes = sizeof(int) * static_cast(X.extent(0)); + } else if (path == FusedDistancePath::Unfused && + (metric == cuvs::distance::DistanceType::L2Expanded || + metric == cuvs::distance::DistanceType::L2SqrtExpanded || + metric == cuvs::distance::DistanceType::CosineExpanded)) { + auto sample_tile = requirements.sample_tile; + const auto centroid_tile = requirements.centroid_tile; + sample_tile = std::min(sample_tile, std::numeric_limits::max() / centroid_tile); + const size_t distance_bytes = + sizeof(DataT) * static_cast(sample_tile) * static_cast(centroid_tile); + requirements.workspace_alignment = + std::max(alignof(DataT), alignof(raft::KeyValuePair)); + requirements.workspace_bytes = + raft::alignTo(distance_bytes, alignof(raft::KeyValuePair)); + if (centroid_tile < centroids.extent(0)) { + requirements.workspace_bytes += + sizeof(raft::KeyValuePair) * static_cast(sample_tile); + } + } + } + return requirements; +} - raft::KeyValuePair initial_value(0, std::numeric_limits::max()); - raft::matrix::fill(handle, minClusterAndDistance, initial_value); - - const bool use_fused_path = - use_fused(handle, n_samples, n_clusters, n_features); - - if (use_fused_path) { - workspace.resize((sizeof(int)) * n_samples, stream); - - cuvs::distance::fusedDistanceNNMinReduce, IndexT>( - minClusterAndDistance.data_handle(), - X.data_handle(), - centroids.data_handle(), - L2NormX.data_handle(), - centroidsNorm.data_handle(), - n_samples, - n_clusters, - n_features, - (void*)workspace.data(), - metric != cuvs::distance::DistanceType::L2Expanded, - false, - true, - metric, - 0.0f, - stream); - } else { - auto dataBatchSize = getDataBatchSize(batch_samples, n_samples); - auto centroidsBatchSize = getCentroidsBatchSize(batch_centroids, n_clusters); - - // The unfused reduction indexes its distance matrix with IndexT. - dataBatchSize = - std::min(dataBatchSize, std::numeric_limits::max() / centroidsBatchSize); - - workspace.resize(sizeof(DataT) * dataBatchSize * centroidsBatchSize, stream); - - using KeyValueT = raft::KeyValuePair; - const bool tileCentroids = centroidsBatchSize < n_clusters; - rmm::device_uvector batchMinClusterAndDistance(tileCentroids ? dataBatchSize : 0, - stream); - - for (IndexT dIdx = 0; dIdx < n_samples;) { - auto ns = std::min(dataBatchSize, n_samples - dIdx); - auto minClusterAndDistanceView = raft::make_device_vector_view( - minClusterAndDistance.data_handle() + dIdx, ns); - - for (IndexT cIdx = 0; cIdx < n_clusters;) { - auto nc = std::min(centroidsBatchSize, n_clusters - cIdx); - auto batchMin = tileCentroids ? batchMinClusterAndDistance.data() - : minClusterAndDistanceView.data_handle(); - - cuvs::distance::unfusedDistanceNNMinReduce( - handle, - batchMin, - X.data_handle() + dIdx * n_features, - centroids.data_handle() + cIdx * n_features, - L2NormX.data_handle() + dIdx, - centroidsNorm.data_handle() + cIdx, - ns, - nc, - n_features, - (void*)workspace.data(), - metric != cuvs::distance::DistanceType::L2Expanded, - tileCentroids, - true, - metric, - 0.0f, - stream); - - if (tileCentroids) { - // Convert tile-local centroid indices and merge the tile minima. - auto batchMinView = - raft::make_device_vector_view(batchMin, ns); - raft::linalg::map( - handle, - minClusterAndDistanceView, - [cIdx] __device__(KeyValueT current, KeyValueT batch) { - batch.key += cIdx; - return batch.value < current.value ? batch : current; - }, - raft::make_const_mdspan(minClusterAndDistanceView), - batchMinView); +template +void min_cluster_and_distance_compute_impl(raft::resources const& handle, + raft::device_matrix_view X, + raft::device_matrix_view centroids, + IndexT* nearest_idx, + DataT* nearest_dist, + raft::KeyValuePair* native_kvp, + raft::device_vector_view L2NormX, + rmm::device_uvector& L2NormBuf_OR_DistBuf, + cuvs::distance::DistanceType metric, + int batch_samples, + int batch_centroids, + rmm::device_uvector& workspace, + const Fused1nnRequirements& requirements, + const DataT* cutile_x_norm) +{ + cudaStream_t stream = raft::resource::get_cuda_stream(handle); + auto n_samples = X.extent(0); + auto n_features = X.extent(1); + auto n_clusters = centroids.extent(0); + const bool is_l2_cos = metric == cuvs::distance::DistanceType::L2Expanded || + metric == cuvs::distance::DistanceType::L2SqrtExpanded || + metric == cuvs::distance::DistanceType::CosineExpanded; + const auto fused_path = requirements.path; + const bool cutile_ready = fused_path == FusedDistancePath::Cutile; + if (workspace.size() < requirements.workspace_bytes) { + workspace.resize(requirements.workspace_bytes, stream); + } + if (cutile_ready) { + RAFT_EXPECTS(native_kvp == nullptr && nearest_idx != nullptr && nearest_dist != nullptr, + "cuTile fused 1-NN requires native separate index and distance outputs"); + } else { + RAFT_EXPECTS(native_kvp != nullptr && nearest_idx == nullptr && nearest_dist == nullptr, + "CUTLASS and unfused 1-NN require their native KVP output"); + } + + if (is_l2_cos || cutile_ready) { + const DataT* x_norm_ptr = nullptr; + const DataT* centroids_norm_ptr = nullptr; + if (is_l2_cos) { + x_norm_ptr = L2NormX.data_handle(); + if constexpr (std::is_same_v) { + if (cutile_ready) { + constexpr size_t norm_alignment = 16 / sizeof(float); + const bool take_sqrt = metric == cuvs::distance::DistanceType::CosineExpanded; + size_t centroid_offset = 0; + if (cutile_x_norm == nullptr) { + centroid_offset = raft::alignTo(static_cast(n_samples), norm_alignment); + } + L2NormBuf_OR_DistBuf.resize(centroid_offset + static_cast(n_clusters), stream); + auto* tf32_x_norms = cutile_x_norm == nullptr ? L2NormBuf_OR_DistBuf.data() + : const_cast(cutile_x_norm); + auto* tf32_centroid_norms = L2NormBuf_OR_DistBuf.data() + centroid_offset; + if (cutile_x_norm == nullptr) { + compute_tf32_row_norms( + handle, X.data_handle(), tf32_x_norms, n_samples, n_features, take_sqrt); } - cIdx += nc; + { + compute_tf32_row_norms(handle, + centroids.data_handle(), + tf32_centroid_norms, + n_clusters, + n_features, + take_sqrt); + } + x_norm_ptr = tf32_x_norms; + centroids_norm_ptr = tf32_centroid_norms; + } else { + L2NormBuf_OR_DistBuf.resize(n_clusters, stream); + centroids_norm_ptr = L2NormBuf_OR_DistBuf.data(); + } + } else { + L2NormBuf_OR_DistBuf.resize(n_clusters, stream); + centroids_norm_ptr = L2NormBuf_OR_DistBuf.data(); + } + + if (!cutile_ready) { + auto centroids_norm = + raft::make_device_vector_view(L2NormBuf_OR_DistBuf.data(), n_clusters); + if (metric == cuvs::distance::DistanceType::CosineExpanded) { + raft::linalg::norm( + handle, centroids, centroids_norm, raft::sqrt_op{}); + } else { + raft::linalg::norm( + handle, centroids, centroids_norm); } - dIdx += ns; } } + + const bool needs_index_workspace = cutile_ready && std::is_same_v; + cuvs::distance::detail::Top1nnTuning tuning{}; + tuning.unfused.row_tile = static_cast(requirements.sample_tile); + tuning.unfused.candidate_tile = static_cast(requirements.centroid_tile); + cuvs::distance::top_1_nn( + handle, + cutile_ready ? nearest_idx : nullptr, + cutile_ready ? nearest_dist : nullptr, + X.data_handle(), + centroids.data_handle(), + x_norm_ptr, + centroids_norm_ptr, + n_samples, + n_clusters, + n_features, + tuning, + !cutile_ready || needs_index_workspace ? (void*)workspace.data() : nullptr, + workspace.size(), + metric != cuvs::distance::DistanceType::L2Expanded, + false, + true, + metric, + 0.0f, + fused_path, + native_kvp, + stream); } else { auto dataBatchSize = getDataBatchSize(batch_samples, n_samples); auto centroidsBatchSize = getCentroidsBatchSize(batch_centroids, n_clusters); - // TODO: Unless pool allocator is used, passing in a workspace for this - // isn't really increasing performance because this needs to do a re-allocation - // anyways. ref https://github.com/rapidsai/raft/issues/930 L2NormBuf_OR_DistBuf.resize(dataBatchSize * centroidsBatchSize, stream); - // pairwiseDistance[ns x nc] - tensor wrapper around the distance buffer auto pairwiseDistance = raft::make_device_matrix_view( L2NormBuf_OR_DistBuf.data(), dataBatchSize, centroidsBatchSize); - raft::KeyValuePair initial_value(0, std::numeric_limits::max()); - raft::matrix::fill(handle, minClusterAndDistance, initial_value); + using KeyValueT = raft::KeyValuePair; + auto* kvp_output = native_kvp; + auto kvp_output_view = raft::make_device_vector_view(kvp_output, n_samples); + KeyValueT initial_value(0, std::numeric_limits::max()); + raft::matrix::fill(handle, kvp_output_view, initial_value); - // tile over the input dataset for (IndexT dIdx = 0; dIdx < n_samples; dIdx += dataBatchSize) { - // # of samples for the current batch auto ns = std::min((IndexT)dataBatchSize, n_samples - dIdx); - // 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); - // minClusterAndDistanceView [ns x n_clusters] - auto minClusterAndDistanceView = - raft::make_device_vector_view, IndexT>( - minClusterAndDistance.data_handle() + dIdx, ns); + auto temp_kvp_view = raft::make_device_vector_view(kvp_output + dIdx, ns); - // tile over the centroids for (IndexT cIdx = 0; cIdx < n_clusters; cIdx += centroidsBatchSize) { - // # of centroids for the current batch auto nc = std::min((IndexT)centroidsBatchSize, n_clusters - cIdx); - // 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); - // pairwiseDistanceView [ns x nc] - view representing the pairwise - // distance for current batch auto pairwiseDistanceView = raft::make_device_matrix_view(pairwiseDistance.data_handle(), ns, nc); - // calculate pairwise distance between current tile of cluster centroids - // and input dataset pairwise_distance_kmeans( handle, datasetView, centroidsView, pairwiseDistanceView, metric); - // argmin reduction returning pair - // calculates the closest centroid and the distance to the closest - // centroid raft::linalg::coalescedReduction( - minClusterAndDistanceView.data_handle(), + temp_kvp_view.data_handle(), pairwiseDistanceView.data_handle(), pairwiseDistanceView.extent(1), pairwiseDistanceView.extent(0), @@ -198,7 +313,7 @@ void minClusterAndDistanceCompute( stream, true, [=] __device__(const DataT val, const IndexT i) { - raft::KeyValuePair pair; + KeyValueT pair; pair.key = cIdx + i; pair.value = val; return pair; @@ -210,18 +325,100 @@ void minClusterAndDistanceCompute( } } -#define INSTANTIATE_MIN_CLUSTER_AND_DISTANCE(DataT, IndexT) \ - template void minClusterAndDistanceCompute( \ - raft::resources const& handle, \ - raft::device_matrix_view X, \ - raft::device_matrix_view centroids, \ - raft::device_vector_view, IndexT> minClusterAndDistance, \ - raft::device_vector_view L2NormX, \ - rmm::device_uvector& L2NormBuf_OR_DistBuf, \ - cuvs::distance::DistanceType metric, \ - int batch_samples, \ - int batch_centroids, \ - rmm::device_uvector& workspace); +template +void minClusterAndDistanceCompute(raft::resources const& handle, + raft::device_matrix_view X, + raft::device_matrix_view centroids, + raft::device_vector_view nearest_idx, + raft::device_vector_view nearest_dist, + raft::device_vector_view L2NormX, + rmm::device_uvector& L2NormBuf_OR_DistBuf, + cuvs::distance::DistanceType metric, + int batch_samples, + int batch_centroids, + rmm::device_uvector& workspace, + const Fused1nnRequirements& requirements, + const DataT* cutile_x_norm) +{ + RAFT_EXPECTS(requirements.output_layout == Fused1nnOutputLayout::Soa, + "resolved fused 1-NN plan requires separate output arrays"); + if constexpr (is_cutile_fused_data_type_v) { + if (requirements.path == FusedDistancePath::Cutile) { + RAFT_EXPECTS(cuvs::distance::detail::can_launch_fused_1nn_tile(nearest_idx.data_handle(), + nearest_dist.data_handle(), + X.data_handle(), + centroids.data_handle(), + X.extent(0), + centroids.extent(0), + X.extent(1), + metric), + "resolved cuTile plan has incompatible output storage"); + } + } + + min_cluster_and_distance_compute_impl(handle, + X, + centroids, + nearest_idx.data_handle(), + nearest_dist.data_handle(), + nullptr, + L2NormX, + L2NormBuf_OR_DistBuf, + metric, + batch_samples, + batch_centroids, + workspace, + requirements, + cutile_x_norm); +} + +template +void minClusterAndDistanceComputeKvp( + raft::resources const& handle, + raft::device_matrix_view X, + raft::device_matrix_view centroids, + raft::device_vector_view, IndexT> nearest, + raft::device_vector_view L2NormX, + rmm::device_uvector& L2NormBuf_OR_DistBuf, + cuvs::distance::DistanceType metric, + int batch_samples, + int batch_centroids, + rmm::device_uvector& workspace, + const Fused1nnRequirements& requirements) +{ + RAFT_EXPECTS(requirements.output_layout == Fused1nnOutputLayout::Kvp, + "resolved fused 1-NN plan requires KVP output"); + min_cluster_and_distance_compute_impl(handle, + X, + centroids, + nullptr, + nullptr, + nearest.data_handle(), + L2NormX, + L2NormBuf_OR_DistBuf, + metric, + batch_samples, + batch_centroids, + workspace, + requirements, + nullptr); +} + +#define INSTANTIATE_MIN_CLUSTER_AND_DISTANCE(DataT, IndexT) \ + template void minClusterAndDistanceCompute( \ + raft::resources const& handle, \ + raft::device_matrix_view X, \ + raft::device_matrix_view centroids, \ + raft::device_vector_view nearest_idx, \ + raft::device_vector_view nearest_dist, \ + raft::device_vector_view L2NormX, \ + rmm::device_uvector& L2NormBuf_OR_DistBuf, \ + cuvs::distance::DistanceType metric, \ + int batch_samples, \ + int batch_centroids, \ + rmm::device_uvector& workspace, \ + const Fused1nnRequirements& requirements, \ + const DataT* cutile_x_norm); INSTANTIATE_MIN_CLUSTER_AND_DISTANCE(float, int64_t) INSTANTIATE_MIN_CLUSTER_AND_DISTANCE(double, int64_t) @@ -230,6 +427,48 @@ INSTANTIATE_MIN_CLUSTER_AND_DISTANCE(double, int) #undef INSTANTIATE_MIN_CLUSTER_AND_DISTANCE +template void computeCutileRowNorms( + raft::resources const&, const float*, float*, int, int, bool); +template void computeCutileRowNorms( + raft::resources const&, const float*, float*, int64_t, int64_t, bool); + +#define INSTANTIATE_FUSED_1NN_REQUIREMENTS(DataT, IndexT) \ + template Fused1nnRequirements get_fused_1nn_requirements( \ + raft::resources const&, \ + raft::device_matrix_view, \ + raft::device_matrix_view, \ + cuvs::distance::DistanceType, \ + int, \ + int); + +INSTANTIATE_FUSED_1NN_REQUIREMENTS(float, int64_t) +INSTANTIATE_FUSED_1NN_REQUIREMENTS(double, int64_t) +INSTANTIATE_FUSED_1NN_REQUIREMENTS(float, int) +INSTANTIATE_FUSED_1NN_REQUIREMENTS(double, int) + +#undef INSTANTIATE_FUSED_1NN_REQUIREMENTS + +#define INSTANTIATE_MIN_CLUSTER_AND_DISTANCE_KVP(DataT, IndexT) \ + template void minClusterAndDistanceComputeKvp( \ + raft::resources const&, \ + raft::device_matrix_view, \ + raft::device_matrix_view, \ + raft::device_vector_view, IndexT>, \ + raft::device_vector_view, \ + rmm::device_uvector&, \ + cuvs::distance::DistanceType, \ + int, \ + int, \ + rmm::device_uvector&, \ + const Fused1nnRequirements&); + +INSTANTIATE_MIN_CLUSTER_AND_DISTANCE_KVP(float, int64_t) +INSTANTIATE_MIN_CLUSTER_AND_DISTANCE_KVP(double, int64_t) +INSTANTIATE_MIN_CLUSTER_AND_DISTANCE_KVP(float, int) +INSTANTIATE_MIN_CLUSTER_AND_DISTANCE_KVP(double, int) + +#undef INSTANTIATE_MIN_CLUSTER_AND_DISTANCE_KVP + template void minClusterDistanceCompute(raft::resources const& handle, raft::device_matrix_view X, @@ -247,51 +486,123 @@ void minClusterDistanceCompute(raft::resources const& handle, auto n_features = X.extent(1); auto n_clusters = centroids.extent(0); - bool is_fused = metric == cuvs::distance::DistanceType::L2Expanded || - metric == cuvs::distance::DistanceType::L2SqrtExpanded || - metric == cuvs::distance::DistanceType::CosineExpanded; - - raft::matrix::fill(handle, minClusterDistance, std::numeric_limits::max()); - - if (is_fused) { - L2NormBuf_OR_DistBuf.resize(n_clusters, stream); - auto centroidsNorm = - raft::make_device_vector_view(L2NormBuf_OR_DistBuf.data(), n_clusters); + const bool is_l2_cos = metric == cuvs::distance::DistanceType::L2Expanded || + metric == cuvs::distance::DistanceType::L2SqrtExpanded || + metric == cuvs::distance::DistanceType::CosineExpanded; + + FusedDistancePath fused_path = + is_l2_cos ? use_fused(handle, n_samples, n_clusters, n_features, metric) + : FusedDistancePath::Unfused; + bool cutile_ready = false; + if constexpr (is_cutile_fused_data_type_v) { + if (fused_path == FusedDistancePath::Cutile) { + cutile_ready = + cuvs::distance::detail::can_launch_fused_1nn_tile(static_cast(nullptr), + minClusterDistance.data_handle(), + X.data_handle(), + centroids.data_handle(), + n_samples, + n_clusters, + n_features, + metric); + if (!cutile_ready) { fused_path = use_legacy_fused(handle, n_samples, n_clusters, metric); } + } + } - if (metric == cuvs::distance::DistanceType::CosineExpanded) { - raft::linalg::norm( - handle, - raft::make_device_matrix_view( - centroids.data_handle(), centroids.extent(0), centroids.extent(1)), - centroidsNorm, - raft::sqrt_op{}); + if (uses_fused_distance_nn(fused_path)) { + const DataT* x_norm_ptr = L2NormX.data_handle(); + const DataT* centroids_norm_ptr; + if constexpr (std::is_same_v) { + if (cutile_ready) { + constexpr size_t norm_alignment = 16 / sizeof(float); + const size_t x_norm_storage = raft::alignTo(static_cast(n_samples), norm_alignment); + L2NormBuf_OR_DistBuf.resize(x_norm_storage + static_cast(n_clusters), stream); + auto* tf32_x_norms = L2NormBuf_OR_DistBuf.data(); + auto* tf32_centroid_norms = tf32_x_norms + x_norm_storage; + const bool take_sqrt = metric == cuvs::distance::DistanceType::CosineExpanded; + compute_tf32_row_norms( + handle, X.data_handle(), tf32_x_norms, n_samples, n_features, take_sqrt); + compute_tf32_row_norms( + handle, centroids.data_handle(), tf32_centroid_norms, n_clusters, n_features, take_sqrt); + x_norm_ptr = tf32_x_norms; + centroids_norm_ptr = tf32_centroid_norms; + } else { + L2NormBuf_OR_DistBuf.resize(n_clusters, stream); + centroids_norm_ptr = L2NormBuf_OR_DistBuf.data(); + } } else { - raft::linalg::norm( - handle, - raft::make_device_matrix_view( - centroids.data_handle(), centroids.extent(0), centroids.extent(1)), - centroidsNorm); + L2NormBuf_OR_DistBuf.resize(n_clusters, stream); + centroids_norm_ptr = L2NormBuf_OR_DistBuf.data(); } - workspace.resize(sizeof(int) * n_samples, stream); + if (!cutile_ready) { + auto centroids_norm = + raft::make_device_vector_view(L2NormBuf_OR_DistBuf.data(), n_clusters); + if (metric == cuvs::distance::DistanceType::CosineExpanded) { + raft::linalg::norm( + handle, + raft::make_device_matrix_view( + centroids.data_handle(), centroids.extent(0), centroids.extent(1)), + centroids_norm, + raft::sqrt_op{}); + } else { + raft::linalg::norm( + handle, + raft::make_device_matrix_view( + centroids.data_handle(), centroids.extent(0), centroids.extent(1)), + centroids_norm); + } + } - cuvs::distance::fusedDistanceNNMinReduce( - minClusterDistance.data_handle(), - X.data_handle(), - centroids.data_handle(), - L2NormX.data_handle(), - centroidsNorm.data_handle(), - n_samples, - n_clusters, - n_features, - (void*)workspace.data(), - metric != cuvs::distance::DistanceType::L2Expanded, - false, - true, - metric, - 0.0f, - stream); + if (!cutile_ready) { workspace.resize(sizeof(int) * static_cast(n_samples), stream); } + + cuvs::distance::detail::Top1nnTuning tuning{}; + if (cutile_ready) { + cuvs::distance::top_1_nn(handle, + nullptr, + minClusterDistance.data_handle(), + X.data_handle(), + centroids.data_handle(), + x_norm_ptr, + centroids_norm_ptr, + n_samples, + n_clusters, + n_features, + tuning, + nullptr, + 0, + metric != cuvs::distance::DistanceType::L2Expanded, + false, + true, + metric, + 0.0f, + FusedDistancePath::Cutile, + nullptr, + stream); + } else { + cuvs::distance::top_1_nn(nullptr, + minClusterDistance.data_handle(), + X.data_handle(), + centroids.data_handle(), + x_norm_ptr, + centroids_norm_ptr, + n_samples, + n_clusters, + n_features, + tuning, + (void*)workspace.data(), + workspace.size(), + metric != cuvs::distance::DistanceType::L2Expanded, + false, + true, + metric, + 0.0f, + FusedDistancePath::Cutlass, + nullptr, + stream); + } } else { + raft::matrix::fill(handle, minClusterDistance, std::numeric_limits::max()); auto dataBatchSize = getDataBatchSize(batch_samples, n_samples); auto centroidsBatchSize = getCentroidsBatchSize(batch_centroids, n_clusters); @@ -300,8 +611,6 @@ void minClusterDistanceCompute(raft::resources const& handle, auto pairwiseDistance = raft::make_device_matrix_view( L2NormBuf_OR_DistBuf.data(), dataBatchSize, centroidsBatchSize); - // tile over the input data and calculate distance matrix [n_samples x - // n_clusters] for (IndexT dIdx = 0; dIdx < n_samples; dIdx += dataBatchSize) { auto ns = std::min((IndexT)dataBatchSize, n_samples - dIdx); @@ -311,7 +620,6 @@ void minClusterDistanceCompute(raft::resources const& handle, auto minClusterDistanceView = raft::make_device_vector_view(minClusterDistance.data_handle() + dIdx, ns); - // tile over the centroids for (IndexT cIdx = 0; cIdx < n_clusters; cIdx += centroidsBatchSize) { auto nc = std::min((IndexT)centroidsBatchSize, n_clusters - cIdx); diff --git a/cpp/src/cluster/kmeans.cuh b/cpp/src/cluster/kmeans.cuh index f6e2c7d819..572104f8eb 100644 --- a/cpp/src/cluster/kmeans.cuh +++ b/cpp/src/cluster/kmeans.cuh @@ -14,9 +14,13 @@ #include #include #include +#include #include +#include +#include #include +#include namespace cuvs::cluster::kmeans { @@ -61,6 +65,7 @@ using KeyValueIndexOp = cuvs::cluster::kmeans::detail::KeyValueIndexOp +template void predict(raft::resources const& handle, const kmeans::params& params, raft::device_matrix_view X, std::optional> sample_weight, raft::device_matrix_view centroids, - raft::device_vector_view labels, + raft::device_vector_view labels, bool normalize_weight, raft::host_scalar_view inertia); -#define EXTERN_TEMPLATE_PREDICT(DataT, IndexT) \ - extern template void predict( \ +#define EXTERN_TEMPLATE_PREDICT(DataT, IndexT, LabelT) \ + extern template void predict( \ raft::resources const& handle, \ const kmeans::params& params, \ raft::device_matrix_view X, \ std::optional> sample_weight, \ raft::device_matrix_view centroids, \ - raft::device_vector_view labels, \ + raft::device_vector_view labels, \ bool normalize_weight, \ raft::host_scalar_view inertia); -EXTERN_TEMPLATE_PREDICT(double, int) -EXTERN_TEMPLATE_PREDICT(double, int64_t) -EXTERN_TEMPLATE_PREDICT(float, int) -EXTERN_TEMPLATE_PREDICT(float, int64_t) +EXTERN_TEMPLATE_PREDICT(double, int, int) +EXTERN_TEMPLATE_PREDICT(double, int, int64_t) +EXTERN_TEMPLATE_PREDICT(double, int64_t, int) +EXTERN_TEMPLATE_PREDICT(double, int64_t, int64_t) +EXTERN_TEMPLATE_PREDICT(float, int, int) +EXTERN_TEMPLATE_PREDICT(float, int, int64_t) +EXTERN_TEMPLATE_PREDICT(float, int64_t, int) +EXTERN_TEMPLATE_PREDICT(float, int64_t, int64_t) #undef EXTERN_TEMPLATE_PREDICT @@ -331,6 +340,7 @@ void min_cluster_distance(raft::resources const& handle, * @param[in] centroids Cluster centroids [n_clusters x n_features] * @param[out] cost Sum of squared distances to nearest centroid (device) * @param[in] sample_weight Optional per-sample weights [n_samples] + * @param[in] metric Squared-L2 implementation used to evaluate inertia */ template void cluster_cost( @@ -338,23 +348,63 @@ void cluster_cost( raft::device_matrix_view X, raft::device_matrix_view centroids, raft::device_scalar_view cost, - std::optional> sample_weight = std::nullopt) + std::optional> sample_weight = std::nullopt, + cuvs::distance::DistanceType metric = cuvs::distance::DistanceType::L2Unexpanded) { auto stream = raft::resource::get_cuda_stream(handle); auto n_clusters = centroids.extent(0); auto n_samples = X.extent(0); auto n_features = X.extent(1); + RAFT_EXPECTS(metric == cuvs::distance::DistanceType::L2Expanded || + metric == cuvs::distance::DistanceType::L2Unexpanded, + "cluster_cost requires a squared-L2 distance metric"); + + if constexpr (std::is_same_v) { + if (metric == cuvs::distance::DistanceType::L2Unexpanded) { + constexpr IndexT max_i32 = std::numeric_limits::max(); + RAFT_EXPECTS(n_clusters > 0 && n_clusters <= max_i32 && n_features <= max_i32, + "stable cluster_cost requires n_clusters and n_features to fit in int32"); + + raft::matrix::fill(handle, cost, DataT{0}); + auto batch_cost = raft::make_device_scalar(handle, DataT{0}); + auto centroids_i32 = raft::make_device_matrix_view( + centroids.data_handle(), static_cast(n_clusters), static_cast(n_features)); + // The i32 path indexes both X[batch_rows, n_features] and its distance workspace + // [batch_rows, n_clusters]. + const IndexT max_batch_rows = max_i32 / std::max(n_clusters, n_features); + + for (IndexT offset = 0; offset < n_samples; offset += max_batch_rows) { + const int batch_rows = static_cast(std::min(max_batch_rows, n_samples - offset)); + auto X_i32 = raft::make_device_matrix_view( + X.data_handle() + offset * n_features, batch_rows, static_cast(n_features)); + + std::optional> batch_weights = std::nullopt; + if (sample_weight.has_value()) { + batch_weights = raft::make_device_vector_view( + sample_weight->data_handle() + offset, batch_rows); + } + + raft::matrix::fill(handle, batch_cost.view(), DataT{0}); + cluster_cost( + handle, X_i32, centroids_i32, batch_cost.view(), batch_weights, metric); + raft::linalg::add( + cost.data_handle(), cost.data_handle(), batch_cost.data_handle(), 1, stream); + } + return; + } + } + rmm::device_uvector workspace(n_samples * sizeof(IndexT), stream); auto x_norms = raft::make_device_vector(handle, n_samples); - raft::linalg::norm(handle, X, x_norms.view()); + if (metric == cuvs::distance::DistanceType::L2Expanded) { + raft::linalg::norm(handle, X, x_norms.view()); + } auto min_cluster_distance = raft::make_device_vector(handle, n_samples); rmm::device_uvector l2_norm_or_distance_buffer(0, stream); - auto metric = cuvs::distance::DistanceType::L2Expanded; - cuvs::cluster::kmeans::min_cluster_distance( handle, X, @@ -435,22 +485,23 @@ void cluster_cost( * */ template -void min_cluster_and_distance( - raft::resources const& handle, - raft::device_matrix_view X, - raft::device_matrix_view centroids, - raft::device_vector_view, IndexT> minClusterAndDistance, - raft::device_vector_view L2NormX, - rmm::device_uvector& L2NormBuf_OR_DistBuf, - cuvs::distance::DistanceType metric, - int batch_samples, - int batch_centroids, - rmm::device_uvector& workspace) +void min_cluster_and_distance(raft::resources const& handle, + raft::device_matrix_view X, + raft::device_matrix_view centroids, + raft::device_vector_view nearest_idx, + raft::device_vector_view nearest_dist, + raft::device_vector_view L2NormX, + rmm::device_uvector& L2NormBuf_OR_DistBuf, + cuvs::distance::DistanceType metric, + int batch_samples, + int batch_centroids, + rmm::device_uvector& workspace) { cuvs::cluster::kmeans::detail::minClusterAndDistanceCompute(handle, X, centroids, - minClusterAndDistance, + nearest_idx, + nearest_dist, L2NormX, L2NormBuf_OR_DistBuf, metric, diff --git a/cpp/src/cluster/kmeans_impl.cuh b/cpp/src/cluster/kmeans_impl.cuh index 49eafac5f2..df3b5eedab 100644 --- a/cpp/src/cluster/kmeans_impl.cuh +++ b/cpp/src/cluster/kmeans_impl.cuh @@ -24,17 +24,17 @@ void fit(raft::resources const& handle, handle, params, X, sample_weight, centroids, inertia, n_iter); } -template +template void predict(raft::resources const& handle, const kmeans::params& params, raft::device_matrix_view X, std::optional> sample_weight, raft::device_matrix_view centroids, - raft::device_vector_view labels, + raft::device_vector_view labels, bool normalize_weight, raft::host_scalar_view inertia) { - cuvs::cluster::kmeans::detail::kmeans_predict( + cuvs::cluster::kmeans::detail::kmeans_predict( handle, params, X, sample_weight, centroids, labels, normalize_weight, inertia); } diff --git a/cpp/src/cluster/kmeans_predict_double.cu b/cpp/src/cluster/kmeans_predict_double.cu index 52d120a232..31531ceee5 100644 --- a/cpp/src/cluster/kmeans_predict_double.cu +++ b/cpp/src/cluster/kmeans_predict_double.cu @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2024-2025, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ @@ -9,19 +9,21 @@ namespace cuvs::cluster::kmeans { -#define INSTANTIATE_PREDICT(DataT, IndexT) \ - template void predict( \ +#define INSTANTIATE_PREDICT(DataT, IndexT, LabelT) \ + template void predict( \ raft::resources const& handle, \ const kmeans::params& params, \ raft::device_matrix_view X, \ std::optional> sample_weight, \ raft::device_matrix_view centroids, \ - raft::device_vector_view labels, \ + raft::device_vector_view labels, \ bool normalize_weight, \ raft::host_scalar_view inertia); -INSTANTIATE_PREDICT(double, int) -INSTANTIATE_PREDICT(double, int64_t) +INSTANTIATE_PREDICT(double, int, int) +INSTANTIATE_PREDICT(double, int, int64_t) +INSTANTIATE_PREDICT(double, int64_t, int) +INSTANTIATE_PREDICT(double, int64_t, int64_t) #undef INSTANTIATE_PREDICT @@ -35,7 +37,33 @@ void predict(raft::resources const& handle, raft::host_scalar_view inertia) { - cuvs::cluster::kmeans::predict( + cuvs::cluster::kmeans::predict( + handle, params, X, sample_weight, centroids, labels, normalize_weight, inertia); +} + +void predict(raft::resources const& handle, + const kmeans::params& params, + raft::device_matrix_view X, + std::optional> sample_weight, + raft::device_matrix_view centroids, + raft::device_vector_view labels, + bool normalize_weight, + raft::host_scalar_view inertia) +{ + cuvs::cluster::kmeans::predict( + handle, params, X, sample_weight, centroids, labels, normalize_weight, inertia); +} + +void predict(raft::resources const& handle, + const kmeans::params& params, + raft::device_matrix_view X, + std::optional> sample_weight, + raft::device_matrix_view centroids, + raft::device_vector_view labels, + bool normalize_weight, + raft::host_scalar_view inertia) +{ + cuvs::cluster::kmeans::predict( handle, params, X, sample_weight, centroids, labels, normalize_weight, inertia); } @@ -49,7 +77,7 @@ void predict(raft::resources const& handle, raft::host_scalar_view inertia) { - cuvs::cluster::kmeans::predict( + cuvs::cluster::kmeans::predict( handle, params, X, sample_weight, centroids, labels, normalize_weight, inertia); } } // namespace cuvs::cluster::kmeans diff --git a/cpp/src/cluster/kmeans_predict_float.cu b/cpp/src/cluster/kmeans_predict_float.cu index 30812aa141..27d33d2e0e 100644 --- a/cpp/src/cluster/kmeans_predict_float.cu +++ b/cpp/src/cluster/kmeans_predict_float.cu @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2022-2025, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ @@ -9,19 +9,21 @@ namespace cuvs::cluster::kmeans { -#define INSTANTIATE_PREDICT(DataT, IndexT) \ - template void predict( \ +#define INSTANTIATE_PREDICT(DataT, IndexT, LabelT) \ + template void predict( \ raft::resources const& handle, \ const kmeans::params& params, \ raft::device_matrix_view X, \ std::optional> sample_weight, \ raft::device_matrix_view centroids, \ - raft::device_vector_view labels, \ + raft::device_vector_view labels, \ bool normalize_weight, \ raft::host_scalar_view inertia); -INSTANTIATE_PREDICT(float, int) -INSTANTIATE_PREDICT(float, int64_t) +INSTANTIATE_PREDICT(float, int, int) +INSTANTIATE_PREDICT(float, int, int64_t) +INSTANTIATE_PREDICT(float, int64_t, int) +INSTANTIATE_PREDICT(float, int64_t, int64_t) #undef INSTANTIATE_PREDICT @@ -35,7 +37,33 @@ void predict(raft::resources const& handle, raft::host_scalar_view inertia) { - cuvs::cluster::kmeans::predict( + cuvs::cluster::kmeans::predict( + handle, params, X, sample_weight, centroids, labels, normalize_weight, inertia); +} + +void predict(raft::resources const& handle, + const kmeans::params& params, + raft::device_matrix_view X, + std::optional> sample_weight, + raft::device_matrix_view centroids, + raft::device_vector_view labels, + bool normalize_weight, + raft::host_scalar_view inertia) +{ + cuvs::cluster::kmeans::predict( + handle, params, X, sample_weight, centroids, labels, normalize_weight, inertia); +} + +void predict(raft::resources const& handle, + const kmeans::params& params, + raft::device_matrix_view X, + std::optional> sample_weight, + raft::device_matrix_view centroids, + raft::device_vector_view labels, + bool normalize_weight, + raft::host_scalar_view inertia) +{ + cuvs::cluster::kmeans::predict( handle, params, X, sample_weight, centroids, labels, normalize_weight, inertia); } @@ -49,7 +77,7 @@ void predict(raft::resources const& handle, raft::host_scalar_view inertia) { - cuvs::cluster::kmeans::predict( + cuvs::cluster::kmeans::predict( handle, params, X, sample_weight, centroids, labels, normalize_weight, inertia); } } // namespace cuvs::cluster::kmeans diff --git a/cpp/src/detail/jit_lto/TileAlgorithmPlanner.cpp b/cpp/src/detail/jit_lto/TileAlgorithmPlanner.cpp new file mode 100644 index 0000000000..65363d6fe3 --- /dev/null +++ b/cpp/src/detail/jit_lto/TileAlgorithmPlanner.cpp @@ -0,0 +1,137 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include +#include +#include +#include + +#include +#include + +#include +#include + +namespace cuvs::detail::jit_lto { + +namespace { + +template +CutileTileConfig tile_config_from_fragment(const FragmentT* fragment, const std::string& entrypoint) +{ + if (fragment == nullptr) { + RAFT_FAIL("cuTile planner '%s' has no registered fragments", entrypoint.c_str()); + } + const int tile_m = fragment->get_tile_m(); + const int tile_n = fragment->get_tile_n(); + const int tile_k = fragment->get_tile_k(); + if (tile_m <= 0 || tile_n <= 0 || tile_k <= 0) { + RAFT_FAIL( + "cuTile planner '%s' is missing tile geometry in its static fragment (check " + "register_cutile_fragment.cpp generation)", + entrypoint.c_str()); + } + return CutileTileConfig{tile_m, tile_n, tile_k}; +} + +} // namespace + +std::shared_ptr TileAlgorithmPlanner::try_get_launcher() +{ + CutileRuntimeCapabilities capabilities{}; + const auto* current_capabilities = + query_current_cutile_runtime_capabilities(capabilities) ? &capabilities : nullptr; + auto launch_key = this->get_planner_key(current_capabilities); + + { + std::shared_lock read_lock(launcher_cache_.mutex); + if (launcher_cache_.unavailable_launchers.count(launch_key)) { return nullptr; } + if (auto it = launcher_cache_.launchers.find(launch_key); + it != launcher_cache_.launchers.end()) { + return it->second; + } + } + + std::unique_lock write_lock(launcher_cache_.mutex); + if (launcher_cache_.unavailable_launchers.count(launch_key)) { return nullptr; } + if (auto it = launcher_cache_.launchers.find(launch_key); it != launcher_cache_.launchers.end()) { + return it->second; + } + + auto launcher = this->build(current_capabilities); + if (!launcher) { + launcher_cache_.unavailable_launchers.insert(launch_key); + return nullptr; + } + launcher_cache_.launchers[launch_key] = launcher; + return launcher; +} + +std::shared_ptr TileAlgorithmPlanner::get_launcher() +{ + auto launcher = try_get_launcher(); + if (!launcher) { + RAFT_FAIL("Failed to build launcher for kernel entrypoint: %s", entrypoint_.c_str()); + } + return launcher; +} + +std::string TileAlgorithmPlanner::get_planner_key( + const CutileRuntimeCapabilities* capabilities) const +{ + std::string key = entrypoint_; + for (const auto& fragment : cubin_fragments_) { + key += fragment->get_key(); + } + if (tileir_fragment_) { key += tileir_fragment_->get_key(); } + + if (capabilities != nullptr) { + key += ":device=" + std::to_string(capabilities->device); + key += ":cc=" + std::to_string(capabilities->cc_major) + "." + + std::to_string(capabilities->cc_minor); + if (const auto* fragment = cuvs::detail::jit_lto::find_compatible_cubin_fragment( + capabilities->cc_major, capabilities->cc_minor, cubin_fragments_)) { + key += ":cubin=" + std::to_string(fragment->get_cc_major()) + "." + + std::to_string(fragment->get_cc_minor()); + } else { + key += ":tileir"; + } + key += ":driver=" + std::to_string(capabilities->driver_version); + } + return key; +} + +CutileTileConfig TileAlgorithmPlanner::tile_config() const +{ + CutileRuntimeCapabilities capabilities{}; + if (query_current_cutile_runtime_capabilities(capabilities)) { + if (const auto* fragment = cuvs::detail::jit_lto::find_compatible_cubin_fragment( + capabilities.cc_major, capabilities.cc_minor, cubin_fragments_)) { + return tile_config_from_fragment(fragment, entrypoint_); + } + } + + if (tileir_fragment_) { return tile_config_from_fragment(tileir_fragment_.get(), entrypoint_); } + + if (!cubin_fragments_.empty()) { + return tile_config_from_fragment(cubin_fragments_.front().get(), entrypoint_); + } + + RAFT_FAIL("cuTile planner '%s' has no registered fragments", entrypoint_.c_str()); +} + +std::shared_ptr TileAlgorithmPlanner::build( + const CutileRuntimeCapabilities* capabilities) +{ + if (capabilities == nullptr) { return nullptr; } + + auto image = cuvs::detail::jit_lto::resolve_cutile_module_image( + *capabilities, cubin_fragments_, tileir_fragment_.get()); + if (!image) { return nullptr; } + + return cuvs::detail::jit_lto::try_load_cutile_launcher(*image, entrypoint_); +} + +} // namespace cuvs::detail::jit_lto diff --git a/cpp/src/detail/jit_lto/cutile_smoke/cutile_smoke_matrix.json b/cpp/src/detail/jit_lto/cutile_smoke/cutile_smoke_matrix.json new file mode 100644 index 0000000000..3047d111d2 --- /dev/null +++ b/cpp/src/detail/jit_lto/cutile_smoke/cutile_smoke_matrix.json @@ -0,0 +1,16 @@ +[ + { + "_data": [{"data_type": "float", "data_abbrev": "f"}], + "_metric": [{"metric": "add", "metric_abbrev": "add"}], + "_index": [{"index_type": "int32", "index_abbrev": "i32"}], + "_abi": [{"abi": "contiguous", "abi_abbrev": "contiguous"}], + "_tile": [{"tile_m": 256, "tile_n": 1, "tile_k": 1}], + "_export": [ + {"output_format": "cubin", "register": "cubin", "gpu_code": "sm_80", "arch_tag": "cutile_arch_8_0", "artifact_basename": "@gpu_code@", "artifact_ext": "cubin"}, + {"output_format": "cubin", "register": "cubin", "gpu_code": "sm_86", "arch_tag": "cutile_arch_8_6", "artifact_basename": "@gpu_code@", "artifact_ext": "cubin"}, + {"output_format": "cubin", "register": "cubin", "gpu_code": "sm_90", "arch_tag": "cutile_arch_9_0", "artifact_basename": "@gpu_code@", "artifact_ext": "cubin"}, + {"output_format": "cubin", "register": "cubin", "gpu_code": "sm_100", "arch_tag": "cutile_arch_10_0", "artifact_basename": "@gpu_code@", "artifact_ext": "cubin"}, + {"output_format": "cubin", "register": "cubin", "gpu_code": "sm_120", "arch_tag": "cutile_arch_12_0", "artifact_basename": "@gpu_code@", "artifact_ext": "cubin"} + ] + } +] diff --git a/cpp/src/detail/jit_lto/cutile_smoke/export_smoke.py b/cpp/src/detail/jit_lto/cutile_smoke/export_smoke.py new file mode 100644 index 0000000000..edce04dfde --- /dev/null +++ b/cpp/src/detail/jit_lto/cutile_smoke/export_smoke.py @@ -0,0 +1,76 @@ +# ============================================================================= +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# ============================================================================= +"""Export the standalone cuTile embedding smoke kernel.""" + +from __future__ import annotations + +import argparse +from pathlib import Path +import sys + +import cuda.tile as ct +from cuda.tile.compilation import ( + ArrayConstraint, + CallingConvention, + KernelSignature, + export_kernel, +) + +sys.path.insert(0, str(Path(__file__).resolve().parent)) + +from smoke_kernel import TILE_SIZE, cutile_smoke_add + + +def _array_constraint() -> ArrayConstraint: + return ArrayConstraint( + ct.float32, + ndim=1, + index_dtype=ct.int32, + stride_lower_bound_incl=(None,), + alias_groups=(), + may_alias_internally=False, + stride_constant=(1,), + stride_divisible_by=(1,), + shape_divisible_by=(TILE_SIZE,), + base_addr_divisible_by=16, + ) + + +def _signature() -> KernelSignature: + array = _array_constraint() + return KernelSignature( + parameters=[array, array, array], + calling_convention=CallingConvention.cutile_python_v1(), + ).with_symbol("cutile_smoke_add") + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("output_file", type=Path) + parser.add_argument("--format", choices=("cubin",), required=True) + parser.add_argument("--data-type", choices=("float",), required=True) + parser.add_argument("--metric", choices=("add",), required=True) + parser.add_argument("--index-type", choices=("int32",), 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", required=True) + args = parser.parse_args() + + if (args.tile_m, args.tile_n, args.tile_k) != (TILE_SIZE, 1, 1): + raise ValueError("cutile smoke kernel requires a 256x1x1 tile") + + export_kernel( + kernel=cutile_smoke_add, + signatures=[_signature()], + output_file=str(args.output_file), + gpu_code=args.gpu_code, + output_format=args.format, + ) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/cpp/src/detail/jit_lto/cutile_smoke/smoke_kernel.py b/cpp/src/detail/jit_lto/cutile_smoke/smoke_kernel.py new file mode 100644 index 0000000000..9b4d6f056a --- /dev/null +++ b/cpp/src/detail/jit_lto/cutile_smoke/smoke_kernel.py @@ -0,0 +1,17 @@ +# ============================================================================= +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# ============================================================================= + +import cuda.tile as ct + + +TILE_SIZE = 256 + + +@ct.kernel +def cutile_smoke_add(lhs, rhs, output): + block = ct.bid(0) + lhs_tile = ct.load(lhs, block, TILE_SIZE) + rhs_tile = ct.load(rhs, block, TILE_SIZE) + ct.store(output, block, lhs_tile + rhs_tile) diff --git a/cpp/src/distance/detail/fused_distance_nn.cuh b/cpp/src/distance/detail/fused_distance_nn.cuh index f9dbd968ec..d8cfde6e0c 100644 --- a/cpp/src/distance/detail/fused_distance_nn.cuh +++ b/cpp/src/distance/detail/fused_distance_nn.cuh @@ -1,11 +1,12 @@ /* - * 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 +#include "fused_distance_nn/cutile/fused_1nn_tile.hpp" #include "fused_distance_nn/cutlass_base.cuh" #include "fused_distance_nn/fused_cosine_nn.cuh" #include "fused_distance_nn/fused_l2_nn.cuh" @@ -13,20 +14,103 @@ #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 // 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 fused 1-NN primitive. */ +enum class Fused1nnBackend : std::uint8_t { + Cutile, + Cutlass, + Unfused, + Auto, +}; + +/** 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{}; +}; + +/** Select the pre-cuTile implementation used by algorithm-level AUTO dispatch. */ +template +constexpr Fused1nnBackend fused_1nn_legacy_backend(int cc_major, + IdxT m, + IdxT n, + cuvs::distance::DistanceType metric) +{ + const bool legacy_fused_metric = metric == cuvs::distance::DistanceType::L2Expanded || + metric == cuvs::distance::DistanceType::L2SqrtExpanded || + metric == cuvs::distance::DistanceType::CosineExpanded; + if (!legacy_fused_metric) { return Fused1nnBackend::Unfused; } + if (cc_major <= 8 || (cc_major == 9 && (m >= 4096 || n >= 4096))) { + return Fused1nnBackend::Cutlass; + } + return Fused1nnBackend::Unfused; +} + +/** Resolve AUTO before allocating backend-native output and workspace storage. */ +template +Fused1nnBackend resolve_fused_1nn_backend(const raft::resources& handle, + const DataT* x, + const DataT* y, + IdxT m, + IdxT n, + IdxT k, + cuvs::distance::DistanceType metric) +{ +#if CUDART_VERSION >= 13000 + if constexpr (is_fused_1nn_cutile_data_v) { + if (can_launch_fused_1nn_tile(x, y, m, n, k, metric)) { return Fused1nnBackend::Cutile; } + } +#endif + const auto prop = raft::resource::get_device_properties(handle); + return fused_1nn_legacy_backend(prop.major, m, n, metric); +} + +/** + * Output-independent backend probe. Call this before allocating backend-native result storage. + * cuTile delegates to its launcher/ABI probe; CUTLASS is available for the legacy L2/cosine + * fused primitive only. + */ +template +bool can_launch_fused_1nn_backend(Fused1nnBackend backend, + const DataT* x, + const DataT* y, + IdxT m, + IdxT n, + IdxT k, + cuvs::distance::DistanceType metric) +{ + if (backend == Fused1nnBackend::Cutile) { + if constexpr (is_fused_1nn_cutile_data_v) { + return can_launch_fused_1nn_tile(x, y, m, n, k, metric); + } + return false; + } + const bool supported_metric = metric == cuvs::distance::DistanceType::L2Expanded || + metric == cuvs::distance::DistanceType::L2SqrtExpanded || + metric == cuvs::distance::DistanceType::CosineExpanded; + return (backend == Fused1nnBackend::Cutlass || backend == Fused1nnBackend::Unfused) && + supported_metric && 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..1c3eacc6bf --- /dev/null +++ b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_cutile_matrix.json @@ -0,0 +1,515 @@ +[ + { + "_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 + } + ], + "_data": [ + { + "data_type": "float", + "data_abbrev": "f" + } + ], + "_metric": [ + { + "metric": "runtime" + } + ], + "_index": [ + { + "index_type": "int32", + "index_abbrev": "i32" + } + ], + "_export": [ + { + "output_format": "cubin", + "artifact_ext": "cubin", + "artifact_basename": "@data_type@_@index_abbrev@_@abi_abbrev@_@gpu_code@", + "register": "cubin", + "gpu_code": "sm_80", + "cc_major": 8, + "cc_minor": 0, + "arch_tag": "cutile_arch_8_0" + }, + { + "output_format": "cubin", + "artifact_ext": "cubin", + "artifact_basename": "@data_type@_@index_abbrev@_@abi_abbrev@_@gpu_code@", + "register": "cubin", + "gpu_code": "sm_86", + "cc_major": 8, + "cc_minor": 6, + "arch_tag": "cutile_arch_8_6" + } + ] + }, + { + "_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 + } + ], + "_data": [ + { + "data_type": "half", + "data_abbrev": "h" + } + ], + "_metric": [ + { + "metric": "runtime" + } + ], + "_index": [ + { + "index_type": "int32", + "index_abbrev": "i32" + } + ], + "_export": [ + { + "output_format": "cubin", + "artifact_ext": "cubin", + "artifact_basename": "@data_type@_@index_abbrev@_@abi_abbrev@_@gpu_code@", + "register": "cubin", + "gpu_code": "sm_80", + "cc_major": 8, + "cc_minor": 0, + "arch_tag": "cutile_arch_8_0" + } + ] + }, + { + "_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 + } + ], + "_data": [ + { + "data_type": "half", + "data_abbrev": "h" + } + ], + "_metric": [ + { + "metric": "runtime" + } + ], + "_index": [ + { + "index_type": "int32", + "index_abbrev": "i32" + } + ], + "_export": [ + { + "output_format": "cubin", + "artifact_ext": "cubin", + "artifact_basename": "@data_type@_@index_abbrev@_@abi_abbrev@_@gpu_code@", + "register": "cubin", + "gpu_code": "sm_86", + "cc_major": 8, + "cc_minor": 6, + "arch_tag": "cutile_arch_8_6" + } + ] + }, + { + "_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 + } + ], + "_data": [ + { + "data_type": "float", + "data_abbrev": "f" + } + ], + "_metric": [ + { + "metric": "runtime" + } + ], + "_index": [ + { + "index_type": "int32", + "index_abbrev": "i32" + } + ], + "_export": [ + { + "output_format": "cubin", + "artifact_ext": "cubin", + "artifact_basename": "@data_type@_@index_abbrev@_@abi_abbrev@_@gpu_code@", + "register": "cubin", + "gpu_code": "sm_90", + "cc_major": 9, + "cc_minor": 0, + "arch_tag": "cutile_arch_9_0" + } + ] + }, + { + "_abi": [ + { + "matrix_layout": "relaxed", + "abi_abbrev": "relaxed", + "abi_tag": "cutile_abi_relaxed" + }, + { + "matrix_layout": "strict", + "abi_abbrev": "strict", + "abi_tag": "cutile_abi_strict" + } + ], + "_data": [ + { + "data_type": "half", + "data_abbrev": "h" + } + ], + "_metric": [ + { + "metric": "runtime" + } + ], + "_index": [ + { + "index_type": "int32", + "index_abbrev": "i32" + } + ], + "_tile": [ + { + "tile_m": 128, + "tile_n": 128, + "tile_k": 128 + } + ], + "_export": [ + { + "output_format": "cubin", + "artifact_ext": "cubin", + "artifact_basename": "@data_type@_@index_abbrev@_@abi_abbrev@_@gpu_code@", + "register": "cubin", + "gpu_code": "sm_90", + "cc_major": 9, + "cc_minor": 0, + "arch_tag": "cutile_arch_9_0" + } + ] + }, + { + "_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 + } + ], + "_data": [ + { + "data_type": "float", + "data_abbrev": "f" + } + ], + "_metric": [ + { + "metric": "runtime" + } + ], + "_index": [ + { + "index_type": "int32", + "index_abbrev": "i32" + } + ], + "_export": [ + { + "output_format": "cubin", + "artifact_ext": "cubin", + "artifact_basename": "@data_type@_@index_abbrev@_@abi_abbrev@_@gpu_code@", + "register": "cubin", + "gpu_code": "sm_100", + "cc_major": 10, + "cc_minor": 0, + "arch_tag": "cutile_arch_10_0" + } + ] + }, + { + "_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 + } + ], + "_data": [ + { + "data_type": "half", + "data_abbrev": "h" + } + ], + "_metric": [ + { + "metric": "runtime" + } + ], + "_index": [ + { + "index_type": "int32", + "index_abbrev": "i32" + } + ], + "_export": [ + { + "output_format": "cubin", + "artifact_ext": "cubin", + "artifact_basename": "@data_type@_@index_abbrev@_@abi_abbrev@_@gpu_code@", + "register": "cubin", + "gpu_code": "sm_100", + "cc_major": 10, + "cc_minor": 0, + "arch_tag": "cutile_arch_10_0" + } + ] + }, + { + "_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 + } + ], + "_data": [ + { + "data_type": "float", + "data_abbrev": "f" + } + ], + "_metric": [ + { + "metric": "runtime" + } + ], + "_index": [ + { + "index_type": "int32", + "index_abbrev": "i32" + } + ], + "_export": [ + { + "output_format": "cubin", + "artifact_ext": "cubin", + "artifact_basename": "@data_type@_@index_abbrev@_@abi_abbrev@_@gpu_code@", + "register": "cubin", + "gpu_code": "sm_120", + "cc_major": 12, + "cc_minor": 0, + "arch_tag": "cutile_arch_12_0" + } + ] + }, + { + "_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 + } + ], + "_data": [ + { + "data_type": "half", + "data_abbrev": "h" + } + ], + "_metric": [ + { + "metric": "runtime" + } + ], + "_index": [ + { + "index_type": "int32", + "index_abbrev": "i32" + } + ], + "_export": [ + { + "output_format": "cubin", + "artifact_ext": "cubin", + "artifact_basename": "@data_type@_@index_abbrev@_@abi_abbrev@_@gpu_code@", + "register": "cubin", + "gpu_code": "sm_120", + "cc_major": 12, + "cc_minor": 0, + "arch_tag": "cutile_arch_12_0" + } + ] + }, + { + "_abi": [ + { + "matrix_layout": "relaxed", + "abi_abbrev": "relaxed", + "abi_tag": "cutile_abi_relaxed" + }, + { + "matrix_layout": "strict", + "abi_abbrev": "strict", + "abi_tag": "cutile_abi_strict" + } + ], + "_data": [ + { + "data_type": "half", + "data_abbrev": "h" + }, + { + "data_type": "float", + "data_abbrev": "f" + } + ], + "_metric": [ + { + "metric": "runtime" + } + ], + "_index": [ + { + "index_type": "int32", + "index_abbrev": "i32" + } + ], + "_tile": [ + { + "tile_m": 128, + "tile_n": 128, + "tile_k": 32 + } + ], + "_export": [ + { + "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" + } + ] + } +] 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..0755788121 --- /dev/null +++ b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_planner.hpp @@ -0,0 +1,130 @@ +/* + * 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 = fused_1nn_data_tag_t; + 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..90f241ff01 --- /dev/null +++ b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_tile.cu @@ -0,0 +1,470 @@ +/* + * 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 +bool 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, + cudaStream_t stream) +{ + if constexpr (!std::is_same_v && !std::is_same_v) { return false; } + + if (nearest_dist == nullptr) { return false; } + + Fused1nnTilePlanner planner; + planner.add_entrypoint(); + planner.add_tileir_fallback(); + auto launcher = planner.try_get_launcher(); + if (!launcher) { return false; } + 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: return false; + } + + 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()); + return true; +} + +template +bool try_fused_1nn_tile_dispatch(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) +{ + return launch_fused_1nn_tile( + nearest_idx, nearest_dist, x, y, xn, yn, m, n, k, metric, is_sqrt, stream); +} + +} // namespace + +template + requires is_fused_1nn_cutile_data_v +bool can_launch_fused_1nn_tile( + 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 +bool can_launch_fused_1nn_tile(IdxT* nearest_idx, + DataT* nearest_dist, + const DataT* x, + const DataT* y, + IdxT m, + IdxT n, + IdxT k, + cuvs::distance::DistanceType metric) +{ + if (!can_launch_fused_1nn_tile(x, y, m, n, k, metric)) { return false; } + if (nearest_dist == nullptr || !is_16_byte_aligned(nearest_dist)) { return false; } + if constexpr (std::is_same_v) { + if (!is_16_byte_aligned(nearest_idx)) { return false; } + } + 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)); + if (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)) { + return false; + } + return true; +} + +template + requires is_fused_1nn_cutile_data_v +bool can_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) +{ + if (!can_launch_fused_1nn_tile(nearest_idx, nearest_dist, x, y, m, n, k, metric)) { + return false; + } + if (metric == cuvs::distance::DistanceType::InnerProduct) { return true; } + if (xn == nullptr || yn == nullptr) { return false; } + if (!is_16_byte_aligned(xn) || !is_16_byte_aligned(yn)) { return false; } + const auto xn_bytes = checked_tensor_bytes(m, IdxT{1}, sizeof(*xn)); + const auto yn_bytes = checked_tensor_bytes(n, IdxT{1}, sizeof(*yn)); + const auto dist_bytes = checked_tensor_bytes(m, IdxT{1}, sizeof(DataT)); + const auto idx_bytes = checked_tensor_bytes(m, IdxT{1}, sizeof(IdxT)); + return !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); +} + +template + requires is_fused_1nn_cutile_data_v +bool try_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) +{ + if (!can_launch_fused_1nn_tile(nearest_idx, nearest_dist, x, y, xn, yn, m, n, k, metric)) { + return false; + } + + 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) { + return try_fused_1nn_tile_dispatch( + nearest_idx, nearest_dist, x, y, xn, yn, m, n, k, metric, is_sqrt, stream); + } + return try_fused_1nn_tile_dispatch( + nearest_idx, nearest_dist, x, y, xn, yn, m, n, k, metric, is_sqrt, stream); + } else { + if (nearest_idx != nullptr && index_workspace == nullptr) { return false; } + if (!is_16_byte_aligned(index_workspace)) { return false; } + const auto workspace_bytes = checked_tensor_bytes(m, IdxT{1}, sizeof(int)); + const auto x_bytes = checked_tensor_bytes(m, k, sizeof(DataT)); + const auto y_bytes = checked_tensor_bytes(n, k, sizeof(DataT)); + const auto norm_x_bytes = checked_tensor_bytes(m, IdxT{1}, sizeof(*xn)); + const auto norm_y_bytes = checked_tensor_bytes(n, IdxT{1}, sizeof(*yn)); + const auto dist_bytes = checked_tensor_bytes(m, IdxT{1}, sizeof(DataT)); + const auto idx_bytes = checked_tensor_bytes(m, IdxT{1}, sizeof(IdxT)); + if (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, norm_x_bytes) || + byte_ranges_overlap(index_workspace, workspace_bytes, yn, norm_y_bytes) || + byte_ranges_overlap(index_workspace, workspace_bytes, nearest_dist, dist_bytes) || + byte_ranges_overlap(index_workspace, workspace_bytes, nearest_idx, idx_bytes)) { + return false; + } + + // Keep every chunk offset 16-byte aligned for x, xn, and nearest_dist. + 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 == nullptr ? nullptr : nearest_dist + offset; + + const bool launched = + use_strict_abi + ? try_fused_1nn_tile_dispatch(tmp_idx, + batch_dist, + batch_x, + y, + batch_xn, + yn, + batch_m, + static_cast(n), + static_cast(k), + metric, + is_sqrt, + stream) + : try_fused_1nn_tile_dispatch(tmp_idx, + batch_dist, + batch_x, + y, + batch_xn, + yn, + batch_m, + static_cast(n), + static_cast(k), + metric, + is_sqrt, + stream); + if (!launched) { return false; } + + if (nearest_idx != nullptr) { + raft::linalg::unaryOp( + nearest_idx + offset, tmp_idx, batch_m, raft::cast_op{}, stream); + } + offset += batch_m64; + } + return true; + } +} + +#define CUVS_INST_CAN_LAUNCH_FUSED_1NN_TILE_INPUTS(DataT, IdxT) \ + template CUVS_EXPORT bool can_launch_fused_1nn_tile( \ + const DataT*, const DataT*, IdxT, IdxT, IdxT, cuvs::distance::DistanceType) + +CUVS_INST_CAN_LAUNCH_FUSED_1NN_TILE_INPUTS(float, int); +CUVS_INST_CAN_LAUNCH_FUSED_1NN_TILE_INPUTS(float, int64_t); +CUVS_INST_CAN_LAUNCH_FUSED_1NN_TILE_INPUTS(half, int); +CUVS_INST_CAN_LAUNCH_FUSED_1NN_TILE_INPUTS(half, int64_t); + +#undef CUVS_INST_CAN_LAUNCH_FUSED_1NN_TILE_INPUTS + +#define CUVS_INST_CAN_LAUNCH_FUSED_1NN_TILE_PREFLIGHT(DataT, IdxT) \ + template CUVS_EXPORT bool can_launch_fused_1nn_tile( \ + IdxT*, DataT*, const DataT*, const DataT*, IdxT, IdxT, IdxT, cuvs::distance::DistanceType) + +CUVS_INST_CAN_LAUNCH_FUSED_1NN_TILE_PREFLIGHT(float, int); +CUVS_INST_CAN_LAUNCH_FUSED_1NN_TILE_PREFLIGHT(float, int64_t); +CUVS_INST_CAN_LAUNCH_FUSED_1NN_TILE_PREFLIGHT(half, int); +CUVS_INST_CAN_LAUNCH_FUSED_1NN_TILE_PREFLIGHT(half, int64_t); + +#undef CUVS_INST_CAN_LAUNCH_FUSED_1NN_TILE_PREFLIGHT + +#define CUVS_INST_CAN_LAUNCH_FUSED_1NN_TILE(DataT, IdxT) \ + template CUVS_EXPORT bool can_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) + +CUVS_INST_CAN_LAUNCH_FUSED_1NN_TILE(float, int); +CUVS_INST_CAN_LAUNCH_FUSED_1NN_TILE(float, int64_t); +CUVS_INST_CAN_LAUNCH_FUSED_1NN_TILE(half, int); +CUVS_INST_CAN_LAUNCH_FUSED_1NN_TILE(half, int64_t); + +#undef CUVS_INST_CAN_LAUNCH_FUSED_1NN_TILE + +#define CUVS_INST_TRY_FUSED_1NN_TILE(DataT, IdxT) \ + template CUVS_EXPORT bool try_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_TRY_FUSED_1NN_TILE(float, int); +CUVS_INST_TRY_FUSED_1NN_TILE(float, int64_t); +CUVS_INST_TRY_FUSED_1NN_TILE(half, int); +CUVS_INST_TRY_FUSED_1NN_TILE(half, int64_t); + +#undef CUVS_INST_TRY_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..add5f7a94e --- /dev/null +++ b/cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_tile.hpp @@ -0,0 +1,164 @@ +/* + * 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 + +#ifndef CUVS_CUTILE_ENABLED +#define CUVS_CUTILE_ENABLED 0 +#endif + +namespace cuvs { +namespace distance { +namespace detail { + +template +inline constexpr bool is_fused_1nn_cutile_data_v = + std::is_same_v || std::is_same_v; + +// Tensor-core products accumulate in FP32; FP16 norms must remain FP32 through the epilogue. +template +using fused_1nn_cutile_norm_t = std::conditional_t, float, DataT>; + +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); +} + +#if CUVS_CUTILE_ENABLED +/** + * 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 can_launch_fused_1nn_tile( + const DataT* x, const DataT* y, IdxT m, IdxT n, IdxT k, cuvs::distance::DistanceType metric); + +/** + * Return whether the supplied problem can use cuTile without fallback scratch. + * + * The result includes runtime/device support, exported ABI constraints, and launcher construction. + * A successful probe populates the shared launcher cache used by try_fused_1nn_tile. + * An int64 output index still requires an int32 workspace sized to the largest launch chunk. + */ +template + requires is_fused_1nn_cutile_data_v +bool can_launch_fused_1nn_tile(IdxT* nearest_idx, + DataT* nearest_dist, + const DataT* x, + const DataT* y, + IdxT m, + IdxT n, + IdxT k, + cuvs::distance::DistanceType metric); + +/** + * Return whether the supplied problem and existing norm buffers can use cuTile. + * + * The overload without norm pointers is a preflight probe for callers that allocate aligned norm + * buffers only after the remaining launch requirements have been validated. + */ +template + requires is_fused_1nn_cutile_data_v +bool can_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); + +template + requires is_fused_1nn_cutile_data_v +bool try_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); +#else +template +bool can_launch_fused_1nn_tile( + const DataT*, const DataT*, IdxT, IdxT, IdxT, cuvs::distance::DistanceType) +{ + return false; +} + +template +bool can_launch_fused_1nn_tile( + IdxT*, DataT*, const DataT*, const DataT*, IdxT, IdxT, IdxT, cuvs::distance::DistanceType) +{ + return false; +} + +template +bool can_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) +{ + return false; +} + +template +bool try_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) +{ + return false; +} +#endif + +} // namespace detail +} // namespace distance +} // namespace cuvs diff --git a/cpp/src/distance/detail/fused_distance_nn/fused_cosine_nn.cuh b/cpp/src/distance/detail/fused_distance_nn/fused_cosine_nn.cuh index 12f4f17cac..692d6e3471 100644 --- a/cpp/src/distance/detail/fused_distance_nn/fused_cosine_nn.cuh +++ b/cpp/src/distance/detail/fused_distance_nn/fused_cosine_nn.cuh @@ -1,11 +1,11 @@ /* - * 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/cosine.cuh" // ops::l2_exp_distance_op +#include "../distance_ops/cosine.cuh" // ops::cosine_distance_op #include "../pairwise_distance_base.cuh" // PairwiseDistances #include "cutlass_base.cuh" #include "helper_structs.cuh" diff --git a/cpp/src/distance/detail/fused_distance_nn/fused_l2_nn.cuh b/cpp/src/distance/detail/fused_distance_nn/fused_l2_nn.cuh index f1aad72110..142bc57909 100644 --- a/cpp/src/distance/detail/fused_distance_nn/fused_l2_nn.cuh +++ b/cpp/src/distance/detail/fused_distance_nn/fused_l2_nn.cuh @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2021-2025, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2021-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/src/distance/detail/fused_distance_nn/helper_structs.cuh b/cpp/src/distance/detail/fused_distance_nn/helper_structs.cuh index 762c720568..19a358a27a 100644 --- a/cpp/src/distance/detail/fused_distance_nn/helper_structs.cuh +++ b/cpp/src/distance/detail/fused_distance_nn/helper_structs.cuh @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2021-2024, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2021-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/src/distance/detail/fused_distance_nn/predicated_tile_iterator_reduced_vec.h b/cpp/src/distance/detail/fused_distance_nn/predicated_tile_iterator_reduced_vec.h index caa6a36d53..6b089493e5 100644 --- a/cpp/src/distance/detail/fused_distance_nn/predicated_tile_iterator_reduced_vec.h +++ b/cpp/src/distance/detail/fused_distance_nn/predicated_tile_iterator_reduced_vec.h @@ -1,7 +1,7 @@ // clang-format off /* * SPDX-FileCopyrightText: Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 AND BSD-3-Clause */ // clang-format on diff --git a/cpp/src/distance/detail/pairwise_matrix/dispatch-ext.cuh b/cpp/src/distance/detail/pairwise_matrix/dispatch-ext.cuh index c93a2f3f2b..f7b47e132c 100644 --- a/cpp/src/distance/detail/pairwise_matrix/dispatch-ext.cuh +++ b/cpp/src/distance/detail/pairwise_matrix/dispatch-ext.cuh @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2023, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2023, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ #pragma once diff --git a/cpp/src/distance/distance-ext.cuh b/cpp/src/distance/distance-ext.cuh index e3841d2caa..1b9637420d 100644 --- a/cpp/src/distance/distance-ext.cuh +++ b/cpp/src/distance/distance-ext.cuh @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2018-2024, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2018-2024, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ #pragma once diff --git a/cpp/src/distance/distance.cu b/cpp/src/distance/distance.cu index 964f569ede..565b655142 100644 --- a/cpp/src/distance/distance.cu +++ b/cpp/src/distance/distance.cu @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2018-2024, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2018-2024, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/src/distance/fused_distance_nn-inl.cuh b/cpp/src/distance/fused_distance_nn-inl.cuh index 3fa80a9b60..f2a50a22fc 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 @@ -289,6 +293,54 @@ void fusedDistanceNNMinReduce(OutT* min, float metric_arg, cudaStream_t stream) { + if constexpr (std::is_same_v>) { + detail::Top1nnTuning tuning{}; + top_1_nn(nullptr, + nullptr, + x, + y, + xn, + yn, + m, + n, + k, + tuning, + workspace, + 0, + sqrt, + initOutBuffer, + isRowMajor, + metric, + metric_arg, + detail::Fused1nnBackend::Cutlass, + min, + stream); + return; + } else if constexpr (std::is_same_v) { + detail::Top1nnTuning tuning{}; + top_1_nn(nullptr, + min, + x, + y, + xn, + yn, + m, + n, + k, + tuning, + workspace, + 0, + sqrt, + initOutBuffer, + isRowMajor, + metric, + metric_arg, + detail::Fused1nnBackend::Cutlass, + nullptr, + stream); + return; + } + MinAndDistanceReduceOp redOp; KVPMinReduce pairRedOp; @@ -311,6 +363,234 @@ void fusedDistanceNNMinReduce(OutT* min, stream); } +template +void top_1_nn(raft::resources const& handle, + IdxT* nearest_idx, + DataT* nearest_dist, + 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::Fused1nnBackend backend, + raft::KeyValuePair* cutlass_kvp_output, + cudaStream_t stream) +{ + RAFT_EXPECTS(is_row_major, "fusedDistanceNN only supports row-major inputs"); + RAFT_EXPECTS(backend != detail::Fused1nnBackend::Auto, + "top_1_nn requires AUTO to be resolved before dispatch"); + if (backend == detail::Fused1nnBackend::Cutile) { + RAFT_EXPECTS(cutlass_kvp_output == nullptr, + "cuTile top_1_nn requires its native separate output arrays"); + if constexpr (detail::is_fused_1nn_cutile_data_v && + std::is_same_v>) { + if constexpr (std::is_same_v) { + const auto required_workspace_bytes = + sizeof(int) * detail::fused_1nn_cutile_index_workspace_rows(m); + RAFT_EXPECTS(workspace != nullptr && workspace_bytes >= required_workspace_bytes, + "cuTile top_1_nn workspace is too small for int64 output batching"); + } + const bool launched = detail::try_fused_1nn_tile( + nearest_idx, nearest_dist, x, y, xn, yn, m, n, k, metric, sqrt, workspace, stream); + RAFT_EXPECTS(launched, + "Requested cuTile fused 1-NN backend is unavailable for this input/device"); + return; + } else { + RAFT_FAIL("Requested cuTile fused 1-NN backend does not support these data/norm types"); + } + } + RAFT_EXPECTS(detail::can_launch_fused_1nn_backend(backend, x, y, m, n, k, metric), + "Requested fused 1-NN backend is unavailable for this input"); + RAFT_EXPECTS(metric != cuvs::distance::DistanceType::InnerProduct, + "Only cuTile top_1_nn supports InnerProduct (as a maximum reduction)"); + RAFT_EXPECTS(nearest_idx == nullptr && nearest_dist == nullptr, + "CUTLASS and unfused top_1_nn require their native KVP output buffer"); + constexpr bool matching_norm_type = std::is_same_v; + RAFT_EXPECTS(matching_norm_type, "CUTLASS and unfused top_1_nn require matching norm types"); + if constexpr (matching_norm_type) { + if (backend == detail::Fused1nnBackend::Unfused) { + RAFT_EXPECTS(cutlass_kvp_output != nullptr, + "Unfused top_1_nn requires its native KVP output buffer"); + RAFT_EXPECTS(tuning.unfused.row_tile > 0 && tuning.unfused.candidate_tile > 0, + "Unfused top_1_nn tile dimensions must be positive"); + + const auto max_row_tile = static_cast(m); + const auto max_candidate_tile = static_cast(n); + const auto row_tile = static_cast(std::min(tuning.unfused.row_tile, max_row_tile)); + const auto candidate_tile = + static_cast(std::min(tuning.unfused.candidate_tile, max_candidate_tile)); + const auto distance_workspace_bytes = static_cast(row_tile) * + static_cast(candidate_tile) * + sizeof(DataT); + const auto candidate_min_offset = + raft::alignTo(distance_workspace_bytes, alignof(raft::KeyValuePair)); + const auto candidate_min_bytes = + candidate_tile < n + ? static_cast(row_tile) * sizeof(raft::KeyValuePair) + : std::size_t{0}; + const auto required_workspace_bytes = candidate_min_offset + candidate_min_bytes; + RAFT_EXPECTS(workspace != nullptr && workspace_bytes >= required_workspace_bytes, + "Unfused top_1_nn workspace is smaller than its configured tile"); + + using KeyValueT = raft::KeyValuePair; + auto* candidate_min = + candidate_tile < n + ? reinterpret_cast(static_cast(workspace) + candidate_min_offset) + : nullptr; + 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 output = + raft::make_device_vector_view(cutlass_kvp_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 ? 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, + output, + [candidate_offset] __device__(KeyValueT current, KeyValueT candidate) { + candidate.key += candidate_offset; + return candidate.value < current.value ? candidate : current; + }, + raft::make_const_mdspan(output), + candidate_output); + } + } + } + return; + } + RAFT_EXPECTS(backend == detail::Fused1nnBackend::Cutlass, "Unknown fused 1-NN backend"); + RAFT_EXPECTS(cutlass_kvp_output != nullptr, + "CUTLASS fused 1-NN requires its native KVP output buffer"); + RAFT_EXPECTS( + workspace != nullptr && workspace_bytes >= sizeof(int) * static_cast(m), + "CUTLASS top_1_nn workspace is too small"); + MinAndDistanceReduceOp red_op; + KVPMinReduce pair_red_op; + fusedDistanceNN, IdxT>(cutlass_kvp_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); + } +} + +template +void top_1_nn(IdxT* nearest_idx, + DataT* nearest_dist, + const DataT* x, + const DataT* y, + const NormT* xn, + const NormT* yn, + IdxT m, + IdxT n, + IdxT k, + const detail::Top1nnTuning&, + void* workspace, + std::size_t, + bool sqrt, + bool init_out_buffer, + bool is_row_major, + cuvs::distance::DistanceType metric, + float metric_arg, + detail::Fused1nnBackend backend, + raft::KeyValuePair* cutlass_kvp_output, + cudaStream_t stream) +{ + RAFT_EXPECTS(backend == detail::Fused1nnBackend::Cutlass, + "The no-handle top_1_nn compatibility path supports CUTLASS only"); + RAFT_EXPECTS(is_row_major, "fusedDistanceNN only supports row-major inputs"); + RAFT_EXPECTS(detail::can_launch_fused_1nn_backend(backend, x, y, m, n, k, metric), + "Requested CUTLASS fused 1-NN backend is unavailable for this input"); + constexpr bool matching_norm_type = std::is_same_v; + RAFT_EXPECTS(matching_norm_type, "CUTLASS top_1_nn requires matching norm types"); + + if constexpr (std::is_same_v) { + MinAndDistanceReduceOp red_op; + KVPMinReduce pair_red_op; + if (cutlass_kvp_output != nullptr) { + fusedDistanceNN, IdxT>(cutlass_kvp_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); + return; + } + RAFT_EXPECTS(nearest_idx == nullptr && nearest_dist != nullptr, + "CUTLASS scalar top_1_nn requires a distance output and no index output"); + fusedDistanceNN(nearest_dist, + x, + y, + xn, + yn, + m, + n, + k, + workspace, + red_op, + pair_red_op, + sqrt, + init_out_buffer, + is_row_major, + metric, + metric_arg, + stream); + } +} + /** @} */ } // namespace distance diff --git a/cpp/src/distance/top_1_nn.cu b/cpp/src/distance/top_1_nn.cu new file mode 100644 index 0000000000..5e4daef89f --- /dev/null +++ b/cpp/src/distance/top_1_nn.cu @@ -0,0 +1,62 @@ +/* + * 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(DataT, IdxT, NormT) \ + template CUVS_EXPORT void top_1_nn(raft::resources const&, \ + IdxT*, \ + DataT*, \ + const DataT*, \ + const DataT*, \ + const NormT*, \ + const NormT*, \ + IdxT, \ + IdxT, \ + IdxT, \ + const detail::Top1nnTuning&, \ + void*, \ + std::size_t, \ + bool, \ + bool, \ + bool, \ + DistanceType, \ + float, \ + detail::Fused1nnBackend, \ + raft::KeyValuePair*, \ + cudaStream_t); \ + template CUVS_EXPORT void top_1_nn(IdxT*, \ + DataT*, \ + const DataT*, \ + const DataT*, \ + const NormT*, \ + const NormT*, \ + IdxT, \ + IdxT, \ + IdxT, \ + const detail::Top1nnTuning&, \ + void*, \ + std::size_t, \ + bool, \ + bool, \ + bool, \ + DistanceType, \ + float, \ + detail::Fused1nnBackend, \ + raft::KeyValuePair*, \ + cudaStream_t) + +CUVS_INSTANTIATE_TOP_1_NN(float, int, float); +CUVS_INSTANTIATE_TOP_1_NN(float, int64_t, float); +CUVS_INSTANTIATE_TOP_1_NN(double, int, double); +CUVS_INSTANTIATE_TOP_1_NN(double, int64_t, double); +CUVS_INSTANTIATE_TOP_1_NN(half, int, float); +CUVS_INSTANTIATE_TOP_1_NN(half, int64_t, float); + +#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..3c7bdbe3a3 --- /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 + +#include + +namespace cuvs::distance { + +/** Dispatch 1-NN to a selected backend using backend-native output storage. */ +template +CUVS_EXPORT void top_1_nn(raft::resources const& handle, + IdxT* nearest_idx, + DataT* nearest_dist, + 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::Fused1nnBackend backend, + raft::KeyValuePair* cutlass_kvp_output, + cudaStream_t stream); + +/** CUTLASS-only overload used by the no-handle legacy compatibility wrapper. */ +template +CUVS_EXPORT void top_1_nn(IdxT* nearest_idx, + DataT* nearest_dist, + 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::Fused1nnBackend backend, + raft::KeyValuePair* cutlass_kvp_output, + cudaStream_t stream); + +#define CUVS_EXTERN_TOP_1_NN(DataT, IdxT, NormT) \ + extern template void top_1_nn(raft::resources const&, \ + IdxT*, \ + DataT*, \ + const DataT*, \ + const DataT*, \ + const NormT*, \ + const NormT*, \ + IdxT, \ + IdxT, \ + IdxT, \ + const detail::Top1nnTuning&, \ + void*, \ + std::size_t, \ + bool, \ + bool, \ + bool, \ + DistanceType, \ + float, \ + detail::Fused1nnBackend, \ + raft::KeyValuePair*, \ + cudaStream_t); \ + extern template void top_1_nn(IdxT*, \ + DataT*, \ + const DataT*, \ + const DataT*, \ + const NormT*, \ + const NormT*, \ + IdxT, \ + IdxT, \ + IdxT, \ + const detail::Top1nnTuning&, \ + void*, \ + std::size_t, \ + bool, \ + bool, \ + bool, \ + DistanceType, \ + float, \ + detail::Fused1nnBackend, \ + raft::KeyValuePair*, \ + cudaStream_t) + +CUVS_EXTERN_TOP_1_NN(float, int, float); +CUVS_EXTERN_TOP_1_NN(float, int64_t, float); +CUVS_EXTERN_TOP_1_NN(double, int, double); +CUVS_EXTERN_TOP_1_NN(double, int64_t, double); +CUVS_EXTERN_TOP_1_NN(half, int, float); +CUVS_EXTERN_TOP_1_NN(half, int64_t, float); + +#undef CUVS_EXTERN_TOP_1_NN + +} // namespace cuvs::distance diff --git a/cpp/tests/CMakeLists.txt b/cpp/tests/CMakeLists.txt index e5715bea71..7d3720be08 100644 --- a/cpp/tests/CMakeLists.txt +++ b/cpp/tests/CMakeLists.txt @@ -136,6 +136,23 @@ ConfigureTest( PERCENT 100 ) +ConfigureTest( + NAME CUTILE_SMOKE_TEST + PATH detail/jit_lto/cutile_smoke.cu + GPUS 1 + PERCENT 100 +) +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 + # exporting the cuTile internals from libcuvs. + target_sources( + CUTILE_SMOKE_TEST PRIVATE ${cutile_smoke_files} + "${CUVS_SOURCE_DIR}/src/detail/jit_lto/TileAlgorithmPlanner.cpp" + ) + target_include_directories(CUTILE_SMOKE_TEST PRIVATE "${cutile_smoke_generated_dir}") +endif() + ConfigureTest( NAME NEIGHBORS_ANN_IVF_FLAT_UDF_TEST PATH neighbors/ann_ivf_flat/test_udf.cu diff --git a/cpp/tests/cluster/kmeans_predict_batching.cu b/cpp/tests/cluster/kmeans_predict_batching.cu index e22cbc9dcd..b740c3f998 100644 --- a/cpp/tests/cluster/kmeans_predict_batching.cu +++ b/cpp/tests/cluster/kmeans_predict_batching.cu @@ -136,8 +136,13 @@ TEST(KMeansPredict, BatchParametersPreserveResultsAndReduceUnfusedAllocations) // predict selects fused or unfused 1-NN according to the architecture heuristic. The batching // parameters only affect the unfused path, so every GPU checks the results while allocation // reductions are required only when this problem shape dispatches to unfused 1-NN. - const bool uses_unfused_path = - !detail::use_fused(handle, test_n_samples, test_n_clusters, test_n_features); + const auto fused_path = + detail::use_fused(handle, + test_n_samples, + test_n_clusters, + test_n_features, + cuvs::distance::DistanceType::L2Expanded); + const bool uses_unfused_path = !detail::uses_fused_distance_nn(fused_path); for (std::size_t i = 1; i < batch_configs.size(); ++i) { auto config = batch_configs[i]; diff --git a/cpp/tests/detail/jit_lto/cutile_smoke.cu b/cpp/tests/detail/jit_lto/cutile_smoke.cu new file mode 100644 index 0000000000..843835e22e --- /dev/null +++ b/cpp/tests/detail/jit_lto/cutile_smoke.cu @@ -0,0 +1,137 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include +#include +#include +#include + +#include + +#include + +#include +#include +#include + +namespace cuvs::detail::jit_lto { + +#if !CUVS_CUTILE_ENABLED + +TEST(CutileSmoke, DisabledBuild) +{ + GTEST_SKIP() << "cuTile embedded kernels are disabled in this build"; +} + +#else + +namespace { + +template +using smoke_fragment = StaticCubinFragmentEntry>; + +std::vector> make_smoke_fragments() +{ + std::vector> fragments; + fragments.emplace_back(std::make_unique>()); + fragments.emplace_back(std::make_unique>()); + fragments.emplace_back(std::make_unique>()); + fragments.emplace_back(std::make_unique>()); + fragments.emplace_back(std::make_unique>()); + return fragments; +} + +void add_smoke_fragments(TileAlgorithmPlanner& planner) +{ + planner.add_static_fragment>(); + planner.add_static_fragment>(); + planner.add_static_fragment>(); + planner.add_static_fragment>(); + planner.add_static_fragment>(); +} + +} // namespace + +TEST(CutileSmoke, ResolvesEveryEmbeddedArchitecture) +{ + auto fragments = make_smoke_fragments(); + + EXPECT_EQ(find_compatible_cubin_fragment(8, 0, fragments), fragments[0].get()); + EXPECT_EQ(find_compatible_cubin_fragment(8, 9, fragments), fragments[1].get()); + EXPECT_EQ(find_compatible_cubin_fragment(9, 0, fragments), fragments[2].get()); + EXPECT_EQ(find_compatible_cubin_fragment(10, 0, fragments), fragments[3].get()); + EXPECT_EQ(find_compatible_cubin_fragment(12, 1, fragments), fragments[4].get()); + EXPECT_EQ(find_compatible_cubin_fragment(7, 5, fragments), nullptr); +} + +TEST(CutileSmoke, LaunchesCompatibleCubin) +{ + CutileRuntimeCapabilities capabilities{}; + if (!query_current_cutile_runtime_capabilities(capabilities)) { + GTEST_SKIP() << "No CUDA device is available"; + } + + auto fragments = make_smoke_fragments(); + if (find_compatible_cubin_fragment(capabilities.cc_major, capabilities.cc_minor, fragments) == + nullptr) { + GTEST_SKIP() << "No embedded smoke cubin is compatible with this device"; + } + + TileLauncherCache cache; + TileAlgorithmPlanner planner{"cutile_smoke_add", cache}; + add_smoke_fragments(planner); + auto launcher = planner.try_get_launcher(); + ASSERT_NE(launcher, nullptr); + + cudaStream_t stream = nullptr; + constexpr int count = 256; + std::array host_lhs{}; + std::array host_rhs{}; + std::array host_output{}; + for (int i = 0; i < count; ++i) { + host_lhs[i] = static_cast(i); + host_rhs[i] = static_cast(count - i); + } + + float* lhs = nullptr; + float* rhs = nullptr; + float* output = nullptr; + ASSERT_EQ(cudaMalloc(&lhs, sizeof(host_lhs)), cudaSuccess); + ASSERT_EQ(cudaMalloc(&rhs, sizeof(host_rhs)), cudaSuccess); + ASSERT_EQ(cudaMalloc(&output, sizeof(host_output)), cudaSuccess); + ASSERT_EQ(cudaMemcpy(lhs, host_lhs.data(), sizeof(host_lhs), cudaMemcpyHostToDevice), + cudaSuccess); + ASSERT_EQ(cudaMemcpy(rhs, host_rhs.data(), sizeof(host_rhs), cudaMemcpyHostToDevice), + cudaSuccess); + + using smoke_kernel_t = void(void*, int, int, void*, int, int, void*, int, int); + launcher->template dispatch(stream, + dim3{1, 1, 1}, + dim3{1, 1, 1}, + 0, + static_cast(lhs), + count, + 1, + static_cast(rhs), + count, + 1, + static_cast(output), + count, + 1); + ASSERT_EQ(cudaGetLastError(), cudaSuccess); + ASSERT_EQ(cudaMemcpy(host_output.data(), output, sizeof(host_output), cudaMemcpyDeviceToHost), + cudaSuccess); + ASSERT_EQ(cudaFree(lhs), cudaSuccess); + ASSERT_EQ(cudaFree(rhs), cudaSuccess); + ASSERT_EQ(cudaFree(output), cudaSuccess); + + for (const auto value : host_output) { + EXPECT_FLOAT_EQ(value, static_cast(count)); + } +} + +#endif + +} // namespace cuvs::detail::jit_lto diff --git a/cpp/tests/neighbors/distance_nn.cu b/cpp/tests/neighbors/distance_nn.cu index f31f3ebacf..42d59b6af4 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 */ @@ -18,6 +18,14 @@ namespace cuvs::neighbors { enum class ImplType { fused, unfused }; +template +void vector_compare_soa(raft::resources const& handle, + const raft::KeyValuePair* ref, + const IdxT* indices, + const AccT* distances, + IdxT n, + ComparisonSummary& summary); + template struct NNInputs { IdxT m; @@ -27,6 +35,9 @@ struct NNInputs { bool sqrt; uint64_t rng_seed; double tol; + cuvs::distance::detail::Fused1nnBackend backend = + cuvs::distance::detail::Fused1nnBackend::Cutlass; + cuvs::distance::detail::Top1nnTuning tuning{}; }; __global__ void fill_int8(int8_t* buff, int len, int seed_offset) @@ -50,13 +61,17 @@ 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)}, + cutile_idx{raft::make_device_vector(handle, m)}, + cutile_dist{raft::make_device_vector(handle, m)} { } @@ -87,7 +102,16 @@ class NNTest : public ::testing::TestWithParam> { } if constexpr (impl == ImplType::fused) { - workspace_size = m * sizeof(IdxT); + workspace_size = m * sizeof(int); + if (backend == cuvs::distance::detail::Fused1nnBackend::Unfused) { + using KvpT = raft::KeyValuePair; + const auto row_tile = std::min(m, tuning.unfused.row_tile); + const auto candidate_tile = std::min(n, tuning.unfused.candidate_tile); + const auto distance_bytes = row_tile * candidate_tile * sizeof(AccT); + workspace_size = + raft::alignTo(distance_bytes, alignof(KvpT)) + + (candidate_tile < static_cast(n) ? row_tile * sizeof(KvpT) : std::size_t{0}); + } } else if constexpr (impl == ImplType::unfused) { workspace_size = m * n * sizeof(AccT); } @@ -114,21 +138,35 @@ 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::Fused1nnBackend::Cutile && + !cuvs::distance::detail::can_launch_fused_1nn_backend( + backend, x.data_handle(), y.data_handle(), m, n, k, metric)) { + GTEST_SKIP() << "cuTile is not available for this device/input"; + } + cuvs::distance::top_1_nn( + handle, + backend == cuvs::distance::detail::Fused1nnBackend::Cutile ? cutile_idx.data_handle() + : nullptr, + backend == cuvs::distance::detail::Fused1nnBackend::Cutile ? cutile_dist.data_handle() + : nullptr, + 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, + backend == cuvs::distance::detail::Fused1nnBackend::Cutile ? nullptr : out.data_handle(), + stream); } else { static_assert(sizeof(DataT) == 0, "fusedDistanceNNMinReduce is not implemented for datatype other than float"); @@ -156,7 +194,20 @@ 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::Fused1nnBackend::Cutile) { + vector_compare_soa(handle, + ref_out.data_handle(), + cutile_idx.data_handle(), + cutile_dist.data_handle(), + m, + summary); + } else { + 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 +221,16 @@ class NNTest : public ::testing::TestWithParam> { IdxT k; DistanceType metric; bool sqrt; + cuvs::distance::detail::Fused1nnBackend 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 cutile_idx; + raft::device_vector cutile_dist; size_t workspace_size; }; @@ -195,6 +250,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::Fused1nnBackend::Unfused; + inputs.push_back(input); + } +#if CUVS_CUTILE_ENABLED + for (auto input : input_fp32) { + input.backend = cuvs::distance::detail::Fused1nnBackend::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 +274,7 @@ 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)); // 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..81ac21c1f7 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 */ @@ -73,6 +73,28 @@ RAFT_KERNEL ref_nn_kernel( IdxT tid = threadIdx.x + blockIdx.x * IdxT(blockDim.x); for (IdxT m = tid; m < M; m += (blockDim.x * gridDim.x)) { + if (metric == DistanceType::InnerProduct) { + IdxT max_index = N + 1; + AccT max_score = min_val(); + for (IdxT n = 0; n < N; n++) { + AccT score = AccT(0.0); + for (IdxT k = 0; k < K; k++) { + score += AccT(A[m * K + k]) * AccT(B[n * K + k]); + } + if (score > max_score) { + max_score = score; + max_index = n; + } + } + if constexpr (std::is_fundamental::value) { + out[m] = max_score; + } else { + out[m].key = max_index; + out[m].value = max_score; + } + continue; + } + IdxT min_index = N + 1; AccT min_dist = max_val(); @@ -207,4 +229,28 @@ void vector_compare( } } +template +void vector_compare_soa(raft::resources const& handle, + const raft::KeyValuePair* ref, + const IdxT* indices, + const AccT* distances, + IdxT n, + ComparisonSummary& summary) +{ + auto ref_h = raft::make_host_vector, IdxT>(n); + auto idx_h = raft::make_host_vector(n); + auto dist_h = raft::make_host_vector(n); + auto stream = raft::resource::get_cuda_stream(handle); + raft::copy(ref_h.data_handle(), ref, n, stream); + raft::copy(idx_h.data_handle(), indices, n, stream); + raft::copy(dist_h.data_handle(), distances, n, stream); + raft::resource::sync_stream(handle, stream); + summary.init(); + for (IdxT i = 0; i < n; ++i) { + const auto a = static_cast(ref_h(i).value); + const auto b = static_cast(dist_h(i)); + summary.update(std::abs(a - b), i, a, b, ref_h(i).key != idx_h(i)); + } +} + } // namespace cuvs::neighbors diff --git a/dependencies.yaml b/dependencies.yaml index 86ccd33d06..08b1209411 100644 --- a/dependencies.yaml +++ b/dependencies.yaml @@ -13,6 +13,7 @@ files: - checks - clang - cuda + - cutile_python - cuda_version - depends_on_cuda_python - depends_on_cupy @@ -41,6 +42,7 @@ files: - build_py_cuvs - clang - cuda + - cutile_python - cuda_version - depends_on_cuda_python - depends_on_cupy @@ -78,6 +80,7 @@ files: includes: - clang - cuda + - cutile_python - cuda_version - depends_on_cupy - docs @@ -139,6 +142,7 @@ files: table: tool.rapids-build-backend key: requires includes: + - cutile_python - depends_on_libraft - depends_on_librmm - depends_on_libkvikio @@ -425,6 +429,48 @@ dependencies: - libcusolver-dev - libcusparse-dev - libnvjitlink-dev + cutile_python: + specific: + - output_types: conda + matrices: + - matrix: + cuda: "12.*" + packages: + - matrix: + cuda: "13.3" + packages: + - cutile-python + - cuda-tileiras + - matrix: + cuda: "13.*" + packages: + - cutile-python + - cuda-tileiras + - matrix: + packages: + - cutile-python + - cuda-tileiras + - output_types: [requirements, pyproject] + matrices: + - matrix: + cuda: "12.*" + packages: + - matrix: + cuda: "13.3" + packages: + - cuda-tile + - cuda-toolkit[tileiras]==13.3.* + - matrix: + cuda: "13.*" + packages: + - &cutile_python_cu13 cuda-tile + - &cutile_toolkit_cu13 cuda-toolkit[tileiras]==13.* + # if no matching matrix selectors passed, list the CUDA 13 requirement + # (as a source of documentation in the generated pyproject.toml) + - matrix: + packages: + - *cutile_python_cu13 + - *cutile_toolkit_cu13 cuda_wheels: specific: # cuVS needs 'nvJitLink>={whatever-cuvs-was-built-against}' at runtime, and mixing diff --git a/python/libcuvs/pyproject.toml b/python/libcuvs/pyproject.toml index 9e050be9bc..42b5e8e213 100644 --- a/python/libcuvs/pyproject.toml +++ b/python/libcuvs/pyproject.toml @@ -83,6 +83,8 @@ regex = "(?P.*)" build-backend = "scikit_build_core.build" requires = [ "cmake>=4.0", + "cuda-tile", + "cuda-toolkit[tileiras]==13.*", "libkvikio==26.10.*,>=0.0.0a0", "libraft==26.10.*,>=0.0.0a0", "librmm==26.10.*,>=0.0.0a0",