[FEA] JIT LTO Pairwise Distances

cpp/CMakeLists.txt

@@ -848,6 +848,93 @@ if(NOT BUILD_CPU_ONLY)
     OUTPUT_DIRECTORY "${CMAKE_CURRENT_BINARY_DIR}/generated_kernels/cagra/filter"
     KERNEL_LINK_LIBRARIES jit_lto_kernel_usage_requirements
   )
+  set(distance_ns "cuvs::distance::detail")
+  set(pairwise_matrix_jit_dir
+      "${CMAKE_CURRENT_SOURCE_DIR}/src/distance/detail/pairwise_matrix/jit_lto_kernels"
+  )
+  generate_jit_lto_kernels(
+    jit_lto_files
+    NAME_FORMAT
+      "pairwise_matrix_@distance_abbrev@_data_@data_abbrev@_acc_@acc_abbrev@_out_@out_abbrev@_index_@index_abbrev@_layout_@layout_abbrev@_veclen_@veclen@"
+    MATRIX_JSON_FILE "${pairwise_matrix_jit_dir}/compute_distance_epilog_matrix.json"
+    KERNEL_INPUT_FILE "${pairwise_matrix_jit_dir}/pairwise_matrix_kernel.cu.in"
+    FRAGMENT_TAG_FORMAT
+      "${distance_ns}::fragment_tag_pairwise_matrix<${distance_ns}::tag_distance_@distance_abbrev@, ${neighbors_ns}::tag_@data_abbrev@, ${neighbors_ns}::tag_@acc_abbrev@, ${neighbors_ns}::tag_@out_abbrev@, ${neighbors_ns}::tag_index_@index_abbrev@, ${distance_ns}::tag_fin_op_identity, ${distance_ns}::tag_layout_@layout_abbrev@, @veclen@>"
+    FRAGMENT_TAG_HEADER_FILES "<cuvs/detail/jit_lto/pairwise_matrix/pairwise_matrix_fragments.hpp>"
+                              "<cuvs/detail/jit_lto/common_fragments.hpp>"
+    OUTPUT_DIRECTORY "${CMAKE_CURRENT_BINARY_DIR}/generated_kernels/distance/pairwise_matrix/kernel"
+    KERNEL_LINK_LIBRARIES jit_lto_kernel_usage_requirements
+  )
+  generate_jit_lto_kernels(
+    jit_lto_files
+    NAME_FORMAT
+      "pairwise_matrix_@distance_abbrev@_data_@data_abbrev@_acc_@acc_abbrev@_out_@out_abbrev@_index_@index_abbrev@_fin_op_rbf_layout_@layout_abbrev@_veclen_@veclen@"
+    MATRIX_JSON_FILE "${pairwise_matrix_jit_dir}/compute_distance_epilog_rbf_matrix.json"
+    KERNEL_INPUT_FILE "${pairwise_matrix_jit_dir}/pairwise_matrix_rbf_kernel.cu.in"
+    FRAGMENT_TAG_FORMAT
+      "${distance_ns}::fragment_tag_pairwise_matrix<${distance_ns}::tag_distance_@distance_abbrev@, ${neighbors_ns}::tag_@data_abbrev@, ${neighbors_ns}::tag_@acc_abbrev@, ${neighbors_ns}::tag_@out_abbrev@, ${neighbors_ns}::tag_index_@index_abbrev@, ${distance_ns}::tag_fin_op_rbf, ${distance_ns}::tag_layout_@layout_abbrev@, @veclen@>"
+    FRAGMENT_TAG_HEADER_FILES "<cuvs/detail/jit_lto/pairwise_matrix/pairwise_matrix_fragments.hpp>"
+                              "<cuvs/detail/jit_lto/common_fragments.hpp>"
+    OUTPUT_DIRECTORY
+      "${CMAKE_CURRENT_BINARY_DIR}/generated_kernels/distance/pairwise_matrix/kernel_rbf"
+    KERNEL_LINK_LIBRARIES jit_lto_kernel_usage_requirements
+  )
+  generate_jit_lto_kernels(
+    jit_lto_files
+    NAME_FORMAT
+      "pairwise_matrix_compute_distance_@distance_abbrev@_data_@data_abbrev@_acc_@acc_abbrev@_index_@index_abbrev@"
+    MATRIX_JSON_FILE "${pairwise_matrix_jit_dir}/compute_distance_matrix.json"
+    KERNEL_INPUT_FILE "${pairwise_matrix_jit_dir}/compute_distance_kernel.cu.in"
+    FRAGMENT_TAG_FORMAT
+      "${distance_ns}::fragment_tag_compute_distance<${distance_ns}::tag_distance_@distance_abbrev@, ${neighbors_ns}::tag_@data_abbrev@, ${neighbors_ns}::tag_@acc_abbrev@, ${neighbors_ns}::tag_index_@index_abbrev@>"
+    FRAGMENT_TAG_HEADER_FILES "<cuvs/detail/jit_lto/pairwise_matrix/pairwise_matrix_fragments.hpp>"
+                              "<cuvs/detail/jit_lto/common_fragments.hpp>"
+    OUTPUT_DIRECTORY
+      "${CMAKE_CURRENT_BINARY_DIR}/generated_kernels/distance/pairwise_matrix/compute_distance"
+    KERNEL_LINK_LIBRARIES jit_lto_kernel_usage_requirements
+  )
+  generate_jit_lto_kernels(
+    jit_lto_files
+    NAME_FORMAT
+      "pairwise_matrix_compute_distance_@distance_abbrev@_data_@data_abbrev@_acc_@acc_abbrev@_index_@index_abbrev@"
+    MATRIX_JSON_FILE "${pairwise_matrix_jit_dir}/compute_distance_rbf_matrix.json"
+    KERNEL_INPUT_FILE "${pairwise_matrix_jit_dir}/compute_distance_kernel.cu.in"
+    FRAGMENT_TAG_FORMAT
+      "${distance_ns}::fragment_tag_compute_distance<${distance_ns}::tag_distance_@distance_abbrev@, ${neighbors_ns}::tag_@data_abbrev@, ${neighbors_ns}::tag_@acc_abbrev@, ${neighbors_ns}::tag_index_@index_abbrev@>"
+    FRAGMENT_TAG_HEADER_FILES "<cuvs/detail/jit_lto/pairwise_matrix/pairwise_matrix_fragments.hpp>"
+                              "<cuvs/detail/jit_lto/common_fragments.hpp>"
+    OUTPUT_DIRECTORY
+      "${CMAKE_CURRENT_BINARY_DIR}/generated_kernels/distance/pairwise_matrix/compute_distance_rbf"
+    KERNEL_LINK_LIBRARIES jit_lto_kernel_usage_requirements
+  )
+  generate_jit_lto_kernels(
+    jit_lto_files
+    NAME_FORMAT
+      "pairwise_matrix_compute_distance_epilog_@distance_abbrev@_data_@data_abbrev@_acc_@acc_abbrev@_index_@index_abbrev@_layout_@layout_abbrev@_veclen_@veclen@"
+    MATRIX_JSON_FILE "${pairwise_matrix_jit_dir}/compute_distance_epilog_matrix.json"
+    KERNEL_INPUT_FILE "${pairwise_matrix_jit_dir}/compute_distance_epilog_kernel.cu.in"
+    FRAGMENT_TAG_FORMAT
+      "${distance_ns}::fragment_tag_compute_distance_epilog<${distance_ns}::tag_distance_@distance_abbrev@, ${neighbors_ns}::tag_@data_abbrev@, ${neighbors_ns}::tag_@acc_abbrev@, ${neighbors_ns}::tag_index_@index_abbrev@, ${distance_ns}::tag_layout_@layout_abbrev@, @veclen@>"
+    FRAGMENT_TAG_HEADER_FILES "<cuvs/detail/jit_lto/pairwise_matrix/pairwise_matrix_fragments.hpp>"
+                              "<cuvs/detail/jit_lto/common_fragments.hpp>"
+    OUTPUT_DIRECTORY
+      "${CMAKE_CURRENT_BINARY_DIR}/generated_kernels/distance/pairwise_matrix/compute_distance_epilog"
+    KERNEL_LINK_LIBRARIES jit_lto_kernel_usage_requirements
+  )
+  generate_jit_lto_kernels(
+    jit_lto_files
+    NAME_FORMAT
+      "pairwise_matrix_compute_distance_epilog_@distance_abbrev@_data_@data_abbrev@_acc_@acc_abbrev@_index_@index_abbrev@_layout_@layout_abbrev@_veclen_@veclen@"
+    MATRIX_JSON_FILE "${pairwise_matrix_jit_dir}/compute_distance_epilog_rbf_matrix.json"
+    KERNEL_INPUT_FILE "${pairwise_matrix_jit_dir}/compute_distance_epilog_kernel.cu.in"
+    FRAGMENT_TAG_FORMAT
+      "${distance_ns}::fragment_tag_compute_distance_epilog<${distance_ns}::tag_distance_@distance_abbrev@, ${neighbors_ns}::tag_@data_abbrev@, ${neighbors_ns}::tag_@acc_abbrev@, ${neighbors_ns}::tag_index_@index_abbrev@, ${distance_ns}::tag_layout_@layout_abbrev@, @veclen@>"
+    FRAGMENT_TAG_HEADER_FILES "<cuvs/detail/jit_lto/pairwise_matrix/pairwise_matrix_fragments.hpp>"
+                              "<cuvs/detail/jit_lto/common_fragments.hpp>"
+    OUTPUT_DIRECTORY
+      "${CMAKE_CURRENT_BINARY_DIR}/generated_kernels/distance/pairwise_matrix/compute_distance_epilog_rbf"
+    KERNEL_LINK_LIBRARIES jit_lto_kernel_usage_requirements
+  )
   endblock()
 
   # Note that this matrix contains an `arch_includes` placeholder, since we don't currently have a

cpp/include/cuvs/detail/jit_lto/common_fragments.hpp

@@ -9,13 +9,15 @@ namespace cuvs::neighbors::detail {
 
 struct tag_f {};
 struct tag_h {};
+struct tag_d {};
 struct tag_i8 {};
 struct tag_u8 {};
 struct tag_filter_none {};
 struct tag_filter_bitset {};
 
 struct tag_bitset_u32 {};
 
+struct tag_index_i32 {};
 struct tag_index_u32 {};
 struct tag_index_i64 {};
 

cpp/include/cuvs/detail/jit_lto/pairwise_matrix/pairwise_matrix_fragments.hpp

@@ -0,0 +1,51 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+#pragma once
+
+namespace cuvs::distance::detail {
+
+struct tag_layout_row {};
+struct tag_layout_col {};
+
+struct tag_fin_op_identity {};
+struct tag_fin_op_rbf {};
+
+struct tag_distance_canberra {};
+struct tag_distance_correlation {};
+struct tag_distance_cosine {};
+struct tag_distance_hamming_unexpanded {};
+struct tag_distance_hellinger_expanded {};
+struct tag_distance_jensen_shannon {};
+struct tag_distance_kl_divergence {};
+struct tag_distance_l1 {};
+struct tag_distance_l2_expanded {};
+struct tag_distance_l2_unexpanded {};
+struct tag_distance_l_inf {};
+struct tag_distance_lp_unexpanded {};
+struct tag_distance_russel_rao {};
+
+template <typename DistanceTag,
+          typename DataTag,
+          typename AccTag,
+          typename OutTag,
+          typename IndexTag,
+          typename FinOpTag,
+          typename LayoutTag,
+          int Veclen>
+struct fragment_tag_pairwise_matrix {};
+
+template <typename DistanceTag, typename DataTag, typename AccTag, typename IndexTag>
+struct fragment_tag_compute_distance {};
+
+template <typename DistanceTag,
+          typename DataTag,
+          typename AccTag,
+          typename IndexTag,
+          typename LayoutTag,
+          int Veclen>
+struct fragment_tag_compute_distance_epilog {};
+
+}  // namespace cuvs::distance::detail

cpp/src/distance/detail/distance.cuh

@@ -7,8 +7,6 @@
 
 #include "distance_ops/all_ops.cuh"
 #include "pairwise_matrix/dispatch.cuh"
-#include "pairwise_matrix/dispatch_sm60.cuh"
-#include "pairwise_matrix/dispatch_sm80.cuh"
 #include <cuvs/distance/distance.hpp>
 #include <raft/core/operators.hpp>
 #include <raft/core/resource/cuda_stream.hpp>

cpp/src/distance/detail/distance_ops/correlation.cuh

@@ -1,10 +1,11 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2023, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION.
  * SPDX-License-Identifier: Apache-2.0
  */
 
 #pragma once
 
+#include <raft/core/operators.hpp>            // raft::sqrt
 #include <raft/util/cuda_dev_essentials.cuh>  // DI
 
 #include <cuda_fp16.h>

cpp/src/distance/detail/distance_ops/hellinger.cuh

@@ -1,9 +1,10 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2023, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION.
  * SPDX-License-Identifier: Apache-2.0
  */
 
 #pragma once
+#include <raft/core/operators.hpp>            // raft::sqrt
 #include <raft/util/cuda_dev_essentials.cuh>  // DI
 
 namespace cuvs::distance::detail::ops {

cpp/src/distance/detail/distance_ops/l1.cuh

@@ -1,9 +1,10 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2023, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION.
  * SPDX-License-Identifier: Apache-2.0
  */
 
 #pragma once
+#include <raft/core/operators.hpp>            // raft::abs
 #include <raft/util/cuda_dev_essentials.cuh>  // DI
 
 namespace cuvs::distance::detail::ops {

cpp/src/distance/detail/distance_ops/l2_unexp.cuh

@@ -1,10 +1,11 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2023, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION.
  * SPDX-License-Identifier: Apache-2.0
  */
 
 #pragma once
 
+#include <raft/core/operators.hpp>            // raft::sqrt
 #include <raft/util/cuda_dev_essentials.cuh>  // DI
 
 #include <cuda_fp16.h>

cpp/src/distance/detail/distance_ops/l_inf.cuh

@@ -1,10 +1,11 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2023, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION.
  * SPDX-License-Identifier: Apache-2.0
  */
 
 #pragma once
 
+#include <raft/core/operators.hpp>            // raft::abs, raft::max
 #include <raft/util/cuda_dev_essentials.cuh>  // DI
 
 namespace cuvs::distance::detail::ops {

cpp/src/distance/detail/pairwise_distance_base.cuh

@@ -1,8 +1,11 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2022-2023, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2022-2026, NVIDIA CORPORATION.
  * SPDX-License-Identifier: Apache-2.0
  */
 #pragma once
+#ifdef CUVS_DISTANCE_PAIRWISE_USE_JIT
+#include "pairwise_matrix/jit_lto_kernels/device_functions.cuh"
+#endif
 #include <raft/linalg/contractions.cuh>       // raft::linalg::Contractions_NT
 #include <raft/util/cuda_dev_essentials.cuh>  // ceildiv
 #include <raft/util/cuda_rt_essentials.hpp>   // RAFT_CUDA_TRY
@@ -150,7 +153,12 @@ struct PairwiseDistances : public BaseClass {
           // Calculate distance_op epilog.
           // Use .template to disambiguate (See:
           // https://en.cppreference.com/w/cpp/language/dependent_name)
+#ifdef CUVS_DISTANCE_PAIRWISE_USE_JIT
+          compute_distance_epilog<Policy, OpT, AccT, IdxT>(
+            distance_op, acc, regxn, regyn, tile_idx_n, tile_idx_m);
+#else
           distance_op.template epilog<Policy>(acc, regxn, regyn, tile_idx_n, tile_idx_m);
+#endif
           // And any possible additional epilogs
           epilog_op(acc, regxn, regyn, tile_idx_n, tile_idx_m);
         } else {
@@ -159,7 +167,12 @@ struct PairwiseDistances : public BaseClass {
           // Calculate distance_op epilog.
           // Use .template to disambiguate (See:
           // https://en.cppreference.com/w/cpp/language/dependent_name)
+#ifdef CUVS_DISTANCE_PAIRWISE_USE_JIT
+          compute_distance_epilog<Policy, OpT, AccT, IdxT>(
+            distance_op, acc, nullptr, nullptr, tile_idx_n, tile_idx_m);
+#else
           distance_op.template epilog<Policy>(acc, nullptr, nullptr, tile_idx_n, tile_idx_m);
+#endif
           // And any possible additional epilogs
           epilog_op(acc, nullptr, nullptr, tile_idx_n, tile_idx_m);
         }
@@ -203,7 +216,12 @@ struct PairwiseDistances : public BaseClass {
       for (int i = 0; i < P::AccRowsPerTh; ++i) {
 #pragma unroll
         for (int j = 0; j < P::AccColsPerTh; ++j) {
+#ifdef CUVS_DISTANCE_PAIRWISE_USE_JIT
+          compute_distance<OpT, DataT, AccT, IdxT>(
+            distance_op, acc[i][j], reg_x[i][v], reg_y[j][v]);
+#else
           distance_op.core(acc[i][j], reg_x[i][v], reg_y[j][v]);
+#endif
         }
       }
     }