Add dispatch based on compute architecture

cpp/include/raft/distance/detail/distance.cuh

@@ -27,24 +27,12 @@
 
 #include <raft/core/operators.hpp>
 
-#include <raft/distance/detail/distance_ops/canberra.cuh>
-#include <raft/distance/detail/distance_ops/correlation.cuh>
-#include <raft/distance/detail/distance_ops/cosine.cuh>
-#include <raft/distance/detail/distance_ops/hamming.cuh>
-#include <raft/distance/detail/distance_ops/hellinger.cuh>
-#include <raft/distance/detail/distance_ops/jensen_shannon.cuh>
-#include <raft/distance/detail/distance_ops/kl_divergence.cuh>
-#include <raft/distance/detail/distance_ops/l1.cuh>
-#include <raft/distance/detail/distance_ops/l2_exp.cuh>
-#include <raft/distance/detail/distance_ops/l2_unexp.cuh>
-#include <raft/distance/detail/distance_ops/l_inf.cuh>
-#include <raft/distance/detail/distance_ops/lp_unexp.cuh>
-#include <raft/distance/detail/distance_ops/russel_rao.cuh>
-
+#include <raft/distance/detail/distance_ops/all_ops.cuh>
 #include <raft/distance/detail/pairwise_matrix/dispatch.cuh>
 
 #include <raft/distance/distance_types.hpp>
 #include <raft/linalg/gemm.cuh>
+#include <raft/util/arch.cuh>
 #include <raft/util/cuda_utils.cuh>
 #include <rmm/device_uvector.hpp>
 
@@ -126,7 +114,7 @@ void distance_impl(raft::resources const& handle,
   const DataT* y_norm = nullptr;
 
   cudaStream_t stream = raft::resource::get_cuda_stream(handle);
-  distance_matrix_dispatch<decltype(distance_op), DataT, AccT, OutT, FinOpT, IdxT>(
+  pairwise_matrix_dispatch<decltype(distance_op), DataT, AccT, OutT, FinOpT, IdxT>(
     distance_op, m, n, k, x, y, x_norm, y_norm, out, fin_op, stream, is_row_major);
 }
 
@@ -205,7 +193,7 @@ void distance_impl(raft::resources const& handle,
 
   using OpT = ops::correlation_distance_op<DataT, AccT, IdxT>;
   OpT corr_op(is_row_major, sq_norm_col_vec, sq_norm_row_vec, m, n, k);
-  distance_matrix_dispatch<decltype(corr_op), DataT, AccT, OutT, FinOpT, IdxT>(
+  pairwise_matrix_dispatch<decltype(corr_op), DataT, AccT, OutT, FinOpT, IdxT>(
     corr_op, m, n, k, x, y, norm_col_vec, norm_row_vec, out, fin_op, stream, is_row_major);
 }
 
@@ -248,34 +236,9 @@ void distance_impl(raft::resources const& handle,
       norm_A, x, k, m, raft::linalg::L2Norm, is_row_major, stream, raft::sqrt_op{});
   }
 
-  // On CUDA 12:
-  // - always execute normal kernel
-  //
-  // On CUDA 11 and below:
-  // - execute CUTLASS-based kernel on SM_80 and above
-  // - execute normal kernel otherwise.
-
-  if constexpr (__CUDACC_VER_MAJOR__ == 12) {
-    // Always execute legacy kernels on CUDA 12
-    ops::cosine_distance_op<DataT, AccT, IdxT> distance_op{};
-    distance_matrix_dispatch<decltype(distance_op), DataT, AccT, OutT, FinOpT, IdxT>(
-      distance_op, m, n, k, x, y, norm_A, norm_B, out, fin_op, stream, is_row_major);
-  } else {
-    const auto deviceVersion = getComputeCapability();
-    if (deviceVersion.first >= 8) {
-      // If device is SM_80 or later, use CUTLASS-based kernel.
-      using Op = ops::cosine_cutlass_op<DataT, AccT>;
-      Op distance_op{};
-
-      distance_matrix_cutlass_dispatch<decltype(distance_op), DataT, AccT, OutT, FinOpT, IdxT>(
-        distance_op, m, n, k, x, y, norm_A, norm_B, out, fin_op, stream, is_row_major);
-    } else {
-      // Else use "legacy" cosine kernel
-      ops::cosine_distance_op<DataT, AccT, IdxT> distance_op{};
-      distance_matrix_dispatch<decltype(distance_op), DataT, AccT, OutT, FinOpT, IdxT>(
-        distance_op, m, n, k, x, y, norm_A, norm_B, out, fin_op, stream, is_row_major);
-    }
-  }
+  ops::cosine_distance_op<DataT, AccT, IdxT> distance_op{};
+  pairwise_matrix_dispatch<decltype(distance_op), DataT, AccT, OutT, FinOpT, IdxT>(
+    distance_op, m, n, k, x, y, norm_A, norm_B, out, fin_op, stream, is_row_major);
 }
 
 template <typename DataT, typename AccT, typename OutT, typename FinOpT, typename IdxT = int>
@@ -300,7 +263,7 @@ void distance_impl(raft::resources const& handle,
 
   cudaStream_t stream = raft::resource::get_cuda_stream(handle);
 
-  distance_matrix_dispatch<decltype(distance_op), DataT, AccT, OutT, FinOpT, IdxT>(
+  pairwise_matrix_dispatch<decltype(distance_op), DataT, AccT, OutT, FinOpT, IdxT>(
     distance_op, m, n, k, x, y, x_norm, y_norm, out, fin_op, stream, is_row_major);
 }
 
@@ -362,7 +325,7 @@ void distance_impl(raft::resources const& handle,
   const DataT* x_norm = nullptr;
   const DataT* y_norm = nullptr;
 
-  distance_matrix_dispatch<decltype(distance_op), DataT, AccT, OutT, FinOpT, IdxT>(
+  pairwise_matrix_dispatch<decltype(distance_op), DataT, AccT, OutT, FinOpT, IdxT>(
     distance_op, m, n, k, x, y, x_norm, y_norm, out, fin_op, stream, is_row_major);
 
   // Finally revert sqrt of x and y
@@ -394,7 +357,7 @@ void distance_impl(raft::resources const& handle,
 
   cudaStream_t stream = raft::resource::get_cuda_stream(handle);
 
-  distance_matrix_dispatch<decltype(distance_op), DataT, AccT, OutT, FinOpT, IdxT>(
+  pairwise_matrix_dispatch<decltype(distance_op), DataT, AccT, OutT, FinOpT, IdxT>(
     distance_op, m, n, k, x, y, x_norm, y_norm, out, fin_op, stream, is_row_major);
 }
 
@@ -438,7 +401,7 @@ void distance_impl(raft::resources const& handle,
   const DataT* x_norm = nullptr;
   const DataT* y_norm = nullptr;
 
-  distance_matrix_dispatch<decltype(kl_divergence), DataT, AccT, OutT, FinOpT, IdxT>(
+  pairwise_matrix_dispatch<decltype(kl_divergence), DataT, AccT, OutT, FinOpT, IdxT>(
     kl_divergence, m, n, k, x, y, x_norm, y_norm, out, fin_op, stream, is_row_major);
 
   if (x != y) {
@@ -469,7 +432,7 @@ void distance_impl(raft::resources const& handle,
   const DataT* y_norm = nullptr;
 
   cudaStream_t stream = raft::resource::get_cuda_stream(handle);
-  distance_matrix_dispatch<decltype(distance_op), DataT, AccT, OutT, FinOpT, IdxT>(
+  pairwise_matrix_dispatch<decltype(distance_op), DataT, AccT, OutT, FinOpT, IdxT>(
     distance_op, m, n, k, x, y, x_norm, y_norm, out, fin_op, stream, is_row_major);
 }
 
@@ -514,34 +477,9 @@ void distance_impl_l2_expanded(  // NOTE: different name
       norm_A, x, k, m, raft::linalg::L2Norm, is_row_major, stream, raft::identity_op{});
   }
 
-  // On CUDA 12:
-  // - always execute normal kernel
-  //
-  // On CUDA 11 and below:
-  // - execute CUTLASS-based kernel on SM_80 and above
-  // - execute normal kernel otherwise.
-
-  if constexpr (__CUDACC_VER_MAJOR__ == 12) {
-    // Always execute legacy kernels on CUDA 12
-    ops::l2_exp_distance_op<DataT, AccT, IdxT> l2_op(perform_sqrt);
-    distance_matrix_dispatch<decltype(l2_op), DataT, AccT, OutT, FinOpT, IdxT>(
-      l2_op, m, n, k, x, y, norm_A, norm_B, out, fin_op, stream, is_row_major);
-  } else {
-    const auto deviceVersion = getComputeCapability();
-    if (deviceVersion.first >= 8) {
-      // If device is SM_80 or later, use CUTLASS-based kernel.
-      using L2Op = ops::l2_exp_cutlass_op<DataT, AccT>;
-      L2Op l2_op(perform_sqrt);
-
-      distance_matrix_cutlass_dispatch<decltype(l2_op), DataT, AccT, OutT, FinOpT, IdxT>(
-        l2_op, m, n, k, x, y, norm_A, norm_B, out, fin_op, stream, is_row_major);
-    } else {
-      // Else use "legacy" L2
-      ops::l2_exp_distance_op<DataT, AccT, IdxT> l2_op(perform_sqrt);
-      distance_matrix_dispatch<decltype(l2_op), DataT, AccT, OutT, FinOpT, IdxT>(
-        l2_op, m, n, k, x, y, norm_A, norm_B, out, fin_op, stream, is_row_major);
-    }
-  }
+  ops::l2_exp_distance_op<DataT, AccT, IdxT> distance_op{perform_sqrt};
+  pairwise_matrix_dispatch<decltype(distance_op), DataT, AccT, OutT, FinOpT, IdxT>(
+    distance_op, m, n, k, x, y, norm_A, norm_B, out, fin_op, stream, is_row_major);
 }
 
 template <typename DataT, typename AccT, typename OutT, typename FinOpT, typename IdxT = int>
@@ -610,7 +548,7 @@ void distance_impl(raft::resources const& handle,
 
   cudaStream_t stream = raft::resource::get_cuda_stream(handle);
 
-  distance_matrix_dispatch<decltype(l2_op), DataT, AccT, OutT, FinOpT, IdxT>(
+  pairwise_matrix_dispatch<decltype(l2_op), DataT, AccT, OutT, FinOpT, IdxT>(
     l2_op, m, n, k, x, y, norm_A, norm_B, out, fin_op, stream, is_row_major);
 }
 
@@ -638,7 +576,7 @@ void distance_impl(raft::resources const& handle,
 
   cudaStream_t stream = raft::resource::get_cuda_stream(handle);
 
-  distance_matrix_dispatch<decltype(l2_op), DataT, AccT, OutT, FinOpT, IdxT>(
+  pairwise_matrix_dispatch<decltype(l2_op), DataT, AccT, OutT, FinOpT, IdxT>(
     l2_op, m, n, k, x, y, norm_A, norm_B, out, fin_op, stream, is_row_major);
 }
 
@@ -664,7 +602,7 @@ void distance_impl(raft::resources const& handle,
 
   cudaStream_t stream = raft::resource::get_cuda_stream(handle);
 
-  distance_matrix_dispatch<decltype(distance_op), DataT, AccT, OutT, FinOpT, IdxT>(
+  pairwise_matrix_dispatch<decltype(distance_op), DataT, AccT, OutT, FinOpT, IdxT>(
     distance_op, m, n, k, x, y, x_norm, y_norm, out, fin_op, stream, is_row_major);
 }
 
@@ -690,7 +628,7 @@ void distance_impl(raft::resources const& handle,
 
   cudaStream_t stream = raft::resource::get_cuda_stream(handle);
 
-  distance_matrix_dispatch<decltype(distance_op), DataT, AccT, OutT, FinOpT, IdxT>(
+  pairwise_matrix_dispatch<decltype(distance_op), DataT, AccT, OutT, FinOpT, IdxT>(
     distance_op, m, n, k, x, y, x_norm, y_norm, out, fin_op, stream, is_row_major);
 }
 
@@ -716,7 +654,7 @@ void distance_impl(raft::resources const& handle,
 
   cudaStream_t stream = raft::resource::get_cuda_stream(handle);
 
-  distance_matrix_dispatch<decltype(distance_op), DataT, AccT, OutT, FinOpT, IdxT>(
+  pairwise_matrix_dispatch<decltype(distance_op), DataT, AccT, OutT, FinOpT, IdxT>(
     distance_op, m, n, k, x, y, x_norm, y_norm, out, fin_op, stream, is_row_major);
 }
 

cpp/include/raft/distance/detail/distance_ops/all_ops.cuh

@@ -0,0 +1,35 @@
+/*
+ * Copyright (c) 2023, NVIDIA CORPORATION.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#pragma once
+
+// Defines a named requirement "has_cutlass_op"
+#include <raft/distance/detail/distance_ops/cutlass.cuh>
+
+// The distance operations:
+#include <raft/distance/detail/distance_ops/canberra.cuh>
+#include <raft/distance/detail/distance_ops/correlation.cuh>
+#include <raft/distance/detail/distance_ops/cosine.cuh>
+#include <raft/distance/detail/distance_ops/hamming.cuh>
+#include <raft/distance/detail/distance_ops/hellinger.cuh>
+#include <raft/distance/detail/distance_ops/jensen_shannon.cuh>
+#include <raft/distance/detail/distance_ops/kl_divergence.cuh>
+#include <raft/distance/detail/distance_ops/l1.cuh>
+#include <raft/distance/detail/distance_ops/l2_exp.cuh>
+#include <raft/distance/detail/distance_ops/l2_unexp.cuh>
+#include <raft/distance/detail/distance_ops/l_inf.cuh>
+#include <raft/distance/detail/distance_ops/lp_unexp.cuh>
+#include <raft/distance/detail/distance_ops/russel_rao.cuh>

cpp/include/raft/distance/detail/distance_ops/canberra.cuh

@@ -27,8 +27,12 @@ namespace raft::distance::detail::ops {
  *
  *  c_ij = sum_k |x_ik - y_kj| / ( |x_ik| + |y_kj| )
  */
-template <typename DataT, typename AccT, typename IdxT>
+template <typename DataType, typename AccType, typename IdxType>
 struct canberra_distance_op {
+  using DataT = DataType;
+  using AccT  = AccType;
+  using IdxT  = IdxType;
+
   // Load norms of input data
   static constexpr bool use_norms = false;
   // Whether the core function requires so many instructions that it makes sense

cpp/include/raft/distance/detail/distance_ops/correlation.cuh

@@ -28,8 +28,12 @@ namespace raft::distance::detail::ops {
  *           /
  *           (|| x - mean(x) ||_2 || y - mean(y) ||_2)
  */
-template <typename DataT, typename AccT, typename IdxT>
+template <typename DataType, typename AccType, typename IdxType>
 struct correlation_distance_op {
+  using DataT = DataType;
+  using AccT  = AccType;
+  using IdxT  = IdxType;
+
   const DataT* x2n;
   const DataT* y2n;
   IdxT m;

cpp/include/raft/distance/detail/distance_ops/cosine.cuh

@@ -20,15 +20,30 @@
 
 namespace raft::distance::detail::ops {
 
+// Epilogue operator for CUTLASS based kernel
+template <typename DataT, typename AccT>
+struct cosine_cutlass_op {
+  __device__ cosine_cutlass_op() noexcept {}
+  __device__ AccT operator()(DataT& aNorm, const DataT& bNorm, DataT& accVal) const noexcept
+  {
+    return static_cast<AccT>(1.0) - (AccT)(accVal / (aNorm * bNorm));
+  }
+  __device__ AccT operator()(DataT aData) const noexcept { return aData; }
+};
+
 /**
  * @brief the expanded cosine distance matrix calculation
  *
  * It computes the following equation:
  *
  * d(x, y) = 1 - (x ⋅ y) / ( ||x||_2 ||y||_2)
  */
-template <typename DataT, typename AccT, typename IdxT>
+template <typename DataType, typename AccType, typename IdxType>
 struct cosine_distance_op {
+  using DataT = DataType;
+  using AccT  = AccType;
+  using IdxT  = IdxType;
+
   // Load norms of input data
   static constexpr bool use_norms = true;
   // Whether the core function requires so many instructions that it makes sense
@@ -60,16 +75,8 @@ struct cosine_distance_op {
       }
     }
   }
-};
 
-template <typename DataT, typename AccT>
-struct cosine_cutlass_op {
-  __device__ cosine_cutlass_op() noexcept {}
-  __device__ AccT operator()(DataT& aNorm, const DataT& bNorm, DataT& accVal) const noexcept
-  {
-    return static_cast<AccT>(1.0) - (AccT)(accVal / (aNorm * bNorm));
-  }
-  __device__ AccT operator()(DataT aData) const noexcept { return aData; }
+  cosine_cutlass_op<DataT, AccT> get_cutlass_op() { return cosine_cutlass_op<DataT, AccT>(); }
 };
 
 }  // namespace raft::distance::detail::ops

cpp/include/raft/distance/detail/distance_ops/cutlass.cuh

@@ -0,0 +1,40 @@
+/*
+ * Copyright (c) 2023, NVIDIA CORPORATION.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#pragma once
+
+#include <type_traits>
+
+namespace raft::distance::detail::ops {
+
+// This file defines the named requirement "has_cutlass_op" that can be used to
+// determine if a distance operation has a CUTLASS op that can be used to pass
+// to CUTLASS. Examples of distance operations that satisfy this requirement are
+// cosine_distance_op and l2_exp_distance_op.
+
+// Primary template handles types that do not support CUTLASS.
+// This pattern is described in:
+// https://en.cppreference.com/w/cpp/types/void_t
+template <typename, typename = void>
+struct has_cutlass_op : std::false_type {
+};
+
+// Specialization recognizes types that do support CUTLASS
+template <typename T>
+struct has_cutlass_op<T, std::void_t<decltype(&T::get_cutlass_op)>> : std::true_type {
+};
+
+}  // namespace raft::distance::detail::ops

cpp/include/raft/distance/detail/distance_ops/hamming.cuh

@@ -26,8 +26,12 @@ namespace raft::distance::detail::ops {
  *
  *    c_ij = sum_k (x_ik != y_kj) / k
  */
-template <typename DataT, typename AccT, typename IdxT>
+template <typename DataType, typename AccType, typename IdxType>
 struct hamming_distance_op {
+  using DataT = DataType;
+  using AccT  = AccType;
+  using IdxT  = IdxType;
+
   IdxT k;
 
   hamming_distance_op(IdxT k_) noexcept : k(k_) {}

cpp/include/raft/distance/detail/distance_ops/hellinger.cuh

@@ -27,8 +27,12 @@ namespace raft::distance::detail::ops {
  *  c_ij = sqrt(1 - sum_k sqrt(x_ik * y_kj))
  *
  */
-template <typename DataT, typename AccT, typename IdxT>
+template <typename DataType, typename AccType, typename IdxType>
 struct hellinger_distance_op {
+  using DataT = DataType;
+  using AccT  = AccType;
+  using IdxT  = IdxType;
+
   // Load norms of input data
   static constexpr bool use_norms = false;
   // Whether the core function requires so many instructions that it makes sense