cutlass based euclidean expanded, cosine kernels

rapidsai · rapids-bot · Nov 16, 2022 · Oct 20, 2022 · Oct 21, 2022 · Oct 21, 2022
commit 20648c599ec8a95eca23820ce2a5ac05f7b86c55
@@ -162,7 +162,8 @@ add_library(raft::raft ALIAS raft)
 
 target_include_directories(raft INTERFACE
         "$<BUILD_INTERFACE:${RAFT_SOURCE_DIR}/include>"
-        "$<INSTALL_INTERFACE:include>")
+        "$<INSTALL_INTERFACE:include>"
+        "${CUTLASS_DIR}/include")
 
 # Keep RAFT as lightweight as possible.
 # Only CUDA libs and rmm should
@@ -225,21 +226,21 @@ set_target_properties(raft_distance PROPERTIES EXPORT_NAME distance)
 if(RAFT_COMPILE_DIST_LIBRARY)
   add_library(raft_distance_lib
     src/distance/pairwise_distance.cu
-    src/distance/specializations/detail/canberra.cu
-    src/distance/specializations/detail/chebyshev.cu
-    src/distance/specializations/detail/correlation.cu
-    src/distance/specializations/detail/cosine.cu
-    src/distance/specializations/detail/hamming_unexpanded.cu
-    src/distance/specializations/detail/hellinger_expanded.cu
-    src/distance/specializations/detail/jensen_shannon_float_float_float_int.cu
-    src/distance/specializations/detail/jensen_shannon_float_float_float_uint32.cu
-    src/distance/specializations/detail/jensen_shannon_double_double_double_int.cu
-    src/distance/specializations/detail/kl_divergence_float_float_float_int.cu
-    src/distance/specializations/detail/kl_divergence_float_float_float_uint32.cu
-    src/distance/specializations/detail/kl_divergence_double_double_double_int.cu
-    src/distance/specializations/detail/l1_float_float_float_int.cu
-    src/distance/specializations/detail/l1_float_float_float_uint32.cu
-    src/distance/specializations/detail/l1_double_double_double_int.cu
+#    src/distance/specializations/detail/canberra.cu
+#    src/distance/specializations/detail/chebyshev.cu
+#    src/distance/specializations/detail/correlation.cu
+#    src/distance/specializations/detail/cosine.cu
+#    src/distance/specializations/detail/hamming_unexpanded.cu
+#    src/distance/specializations/detail/hellinger_expanded.cu
+#    src/distance/specializations/detail/jensen_shannon_float_float_float_int.cu
+#    src/distance/specializations/detail/jensen_shannon_float_float_float_uint32.cu
+#    src/distance/specializations/detail/jensen_shannon_double_double_double_int.cu
+#    src/distance/specializations/detail/kl_divergence_float_float_float_int.cu
+#    src/distance/specializations/detail/kl_divergence_float_float_float_uint32.cu
+#    src/distance/specializations/detail/kl_divergence_double_double_double_int.cu
+#    src/distance/specializations/detail/l1_float_float_float_int.cu
+#    src/distance/specializations/detail/l1_float_float_float_uint32.cu
+#    src/distance/specializations/detail/l1_double_double_double_int.cu
     src/distance/specializations/detail/l2_expanded_float_float_float_int.cu
     src/distance/specializations/detail/l2_expanded_float_float_float_uint32.cu
     src/distance/specializations/detail/l2_expanded_double_double_double_int.cu
@@ -249,12 +250,12 @@ if(RAFT_COMPILE_DIST_LIBRARY)
     src/distance/specializations/detail/l2_sqrt_unexpanded_float_float_float_int.cu
     src/distance/specializations/detail/l2_sqrt_unexpanded_float_float_float_uint32.cu
     src/distance/specializations/detail/l2_sqrt_unexpanded_double_double_double_int.cu
-    src/distance/specializations/detail/l2_unexpanded_double_double_double_int.cu
-    src/distance/specializations/detail/l2_unexpanded_float_float_float_uint32.cu
-    src/distance/specializations/detail/l2_unexpanded_float_float_float_int.cu
-    src/distance/specializations/detail/lp_unexpanded_double_double_double_int.cu
-    src/distance/specializations/detail/lp_unexpanded_float_float_float_uint32.cu
-    src/distance/specializations/detail/lp_unexpanded_float_float_float_int.cu
+#    src/distance/specializations/detail/l2_unexpanded_double_double_double_int.cu
+#    src/distance/specializations/detail/l2_unexpanded_float_float_float_uint32.cu
+#    src/distance/specializations/detail/l2_unexpanded_float_float_float_int.cu
+#    src/distance/specializations/detail/lp_unexpanded_double_double_double_int.cu
+#    src/distance/specializations/detail/lp_unexpanded_float_float_float_uint32.cu
+#    src/distance/specializations/detail/lp_unexpanded_float_float_float_int.cu
   )
   set_target_properties(
     raft_distance_lib
@@ -305,16 +306,16 @@ set_target_properties(raft_nn PROPERTIES EXPORT_NAME nn)
 
 if(RAFT_COMPILE_NN_LIBRARY)
   add_library(raft_nn_lib
-          src/nn/specializations/ball_cover.cu
-          src/nn/specializations/detail/ball_cover_lowdim_pass_one_2d.cu
-          src/nn/specializations/detail/ball_cover_lowdim_pass_two_2d.cu
-          src/nn/specializations/detail/ball_cover_lowdim_pass_one_3d.cu
-          src/nn/specializations/detail/ball_cover_lowdim_pass_two_3d.cu
+#          src/nn/specializations/ball_cover.cu
+#          src/nn/specializations/detail/ball_cover_lowdim_pass_one_2d.cu
+#          src/nn/specializations/detail/ball_cover_lowdim_pass_two_2d.cu
+#          src/nn/specializations/detail/ball_cover_lowdim_pass_one_3d.cu
+#          src/nn/specializations/detail/ball_cover_lowdim_pass_two_3d.cu
           src/nn/specializations/fused_l2_knn_long_float_true.cu
           src/nn/specializations/fused_l2_knn_long_float_false.cu
           src/nn/specializations/fused_l2_knn_int_float_true.cu
           src/nn/specializations/fused_l2_knn_int_float_false.cu
-          src/nn/specializations/knn.cu
+#          src/nn/specializations/knn.cu
           )
   set_target_properties(
     raft_nn_lib

diff --git a/cpp/include/raft/core/cudart_utils.hpp b/cpp/include/raft/core/cudart_utils.hpp
@@ -354,6 +354,19 @@ inline int getMultiProcessorCount()
   return mpCount;
 }
 
+/** helper method to get max usable shared mem per block parameter */
+inline std::pair <int,int> getMajorMinorVersion()
+{
+  int devId;
+  RAFT_CUDA_TRY(cudaGetDevice(&devId));
+  int majorVer, minorVer;
+  RAFT_CUDA_TRY(cudaDeviceGetAttribute(&majorVer, cudaDevAttrComputeCapabilityMajor, devId));
+  RAFT_CUDA_TRY(cudaDeviceGetAttribute(&minorVer, cudaDevAttrComputeCapabilityMinor, devId));
+
+  return std::make_pair(majorVer, minorVer);
+}
+
+
 /** helper method to convert an array on device to a string on host */
 template <typename T>
 std::string arr2Str(const T* arr, int size, std::string name, cudaStream_t stream, int width = 4)

@@ -18,11 +18,24 @@
 
 #include <raft/distance/detail/pairwise_distance_base.cuh>
 #include <raft/linalg/norm.cuh>
+#include <raft/distance/detail/pairwise_distance_cutlass_base.cuh>
 
 namespace raft {
 namespace distance {
 namespace detail {
 
+template <typename DataT, typename AccT>
+struct CosineOp {
+    __device__ __host__ CosineOp() { }
+    __device__ __host__ AccT operator() (DataT &aNorm, const DataT &bNorm, DataT &accVal) const {
+        return static_cast<AccT>(1.0) - (AccT) (accVal / (aNorm * bNorm));
+    }
+    __device__ __host__ AccT operator() (DataT aData) const {
+        return aData;
+    }
+};
+
+
 /**
  * @brief the cosine distance matrix calculation implementer
  *  It computes the following equation:
@@ -71,61 +84,72 @@ void cosineImpl(const DataT* x,
                 FinalLambda fin_op,
                 cudaStream_t stream)
 {
-  typedef typename raft::linalg::Policy4x4<DataT, VecLen>::Policy RowPolicy;
-  typedef typename raft::linalg::Policy4x4<DataT, VecLen>::ColPolicy ColPolicy;
-
-  typedef typename std::conditional<isRowMajor, RowPolicy, ColPolicy>::type KPolicy;
-
-  dim3 blk(KPolicy::Nthreads);
-
-  // Accumulation operation lambda
-  auto core_lambda = [] __device__(AccT & acc, DataT & x, DataT & y) { acc += x * y; };
-
-  // epilogue operation lambda for final value calculation
-  auto epilog_lambda = [] __device__(AccT acc[KPolicy::AccRowsPerTh][KPolicy::AccColsPerTh],
-                                     DataT * regxn,
-                                     DataT * regyn,
-                                     IdxT gridStrideX,
-                                     IdxT gridStrideY) {
-#pragma unroll
-    for (int i = 0; i < KPolicy::AccRowsPerTh; ++i) {
-#pragma unroll
-      for (int j = 0; j < KPolicy::AccColsPerTh; ++j) {
-        acc[i][j] = acc[i][j] / (regxn[i] * regyn[j]);
-      }
-    }
-  };
+  const auto deviceVersion  = getMajorMinorVersion();
+  if (deviceVersion.first >= 8) {
+    using CosineOp_ = CosineOp<DataT, AccT>;
+    CosineOp_ cosine_dist_op;
+
+    cutlassDistanceKernel<DataT, AccT, OutT, IdxT, VecLen, FinalLambda, CosineOp_, isRowMajor>(
+                    x, y, xn, yn, m, n, k, lda, ldb, ldd, dOutput, fin_op, cosine_dist_op, stream);
 
-  constexpr size_t shmemSize =
-    KPolicy::SmemSize + ((KPolicy::Mblk + KPolicy::Nblk) * sizeof(DataT));
-  if (isRowMajor) {
-    auto cosineRowMajor = pairwiseDistanceMatKernel<true,
-                                                    DataT,
-                                                    AccT,
-                                                    OutT,
-                                                    IdxT,
-                                                    KPolicy,
-                                                    decltype(core_lambda),
-                                                    decltype(epilog_lambda),
-                                                    FinalLambda,
-                                                    true>;
-    dim3 grid           = launchConfigGenerator<KPolicy>(m, n, shmemSize, cosineRowMajor);
-    cosineRowMajor<<<grid, blk, shmemSize, stream>>>(
-      x, y, xn, yn, m, n, k, lda, ldb, ldd, dOutput, core_lambda, epilog_lambda, fin_op);
   } else {
-    auto cosineColMajor = pairwiseDistanceMatKernel<true,
-                                                    DataT,
-                                                    AccT,
-                                                    OutT,
-                                                    IdxT,
-                                                    KPolicy,
-                                                    decltype(core_lambda),
-                                                    decltype(epilog_lambda),
-                                                    FinalLambda,
-                                                    false>;
-    dim3 grid           = launchConfigGenerator<KPolicy>(m, n, shmemSize, cosineColMajor);
-    cosineColMajor<<<grid, blk, shmemSize, stream>>>(
-      x, y, xn, yn, m, n, k, lda, ldb, ldd, dOutput, core_lambda, epilog_lambda, fin_op);
+
+    typedef typename raft::linalg::Policy4x4<DataT, VecLen>::Policy RowPolicy;
+    typedef typename raft::linalg::Policy4x4<DataT, VecLen>::ColPolicy ColPolicy;
+
+    typedef typename std::conditional<isRowMajor, RowPolicy, ColPolicy>::type KPolicy;
+
+    dim3 blk(KPolicy::Nthreads);
+
+    // Accumulation operation lambda
+    auto core_lambda = [] __device__(AccT & acc, DataT & x, DataT & y) { acc += x * y; };
+
+    // epilogue operation lambda for final value calculation
+    auto epilog_lambda = [] __device__(AccT acc[KPolicy::AccRowsPerTh][KPolicy::AccColsPerTh],
+                                      DataT * regxn,
+                                      DataT * regyn,
+                                      IdxT gridStrideX,
+                                      IdxT gridStrideY) {
+  #pragma unroll
+      for (int i = 0; i < KPolicy::AccRowsPerTh; ++i) {
+  #pragma unroll
+        for (int j = 0; j < KPolicy::AccColsPerTh; ++j) {
+          acc[i][j] = 1.0 - (acc[i][j] / (regxn[i] * regyn[j]) );
+        }
+      }
+    };
+
+    constexpr size_t shmemSize =
+      KPolicy::SmemSize + ((KPolicy::Mblk + KPolicy::Nblk) * sizeof(DataT));
+    if (isRowMajor) {
+      auto cosineRowMajor = pairwiseDistanceMatKernel<true,
+                                                      DataT,
+                                                      AccT,
+                                                      OutT,
+                                                      IdxT,
+                                                      KPolicy,
+                                                      decltype(core_lambda),
+                                                      decltype(epilog_lambda),
+                                                      FinalLambda,
+                                                      true>;
+      dim3 grid           = launchConfigGenerator<KPolicy>(m, n, shmemSize, cosineRowMajor);
+      cosineRowMajor<<<grid, blk, shmemSize, stream>>>(
+        x, y, xn, yn, m, n, k, lda, ldb, ldd, dOutput, core_lambda, epilog_lambda, fin_op);
+    } else {
+      auto cosineColMajor = pairwiseDistanceMatKernel<true,
+                                                      DataT,
+                                                      AccT,
+                                                      OutT,
+                                                      IdxT,
+                                                      KPolicy,
+                                                      decltype(core_lambda),
+                                                      decltype(epilog_lambda),
+                                                      FinalLambda,
+                                                      false>;
+      dim3 grid           = launchConfigGenerator<KPolicy>(m, n, shmemSize, cosineColMajor);
+      cosineColMajor<<<grid, blk, shmemSize, stream>>>(
+        x, y, xn, yn, m, n, k, lda, ldb, ldd, dOutput, core_lambda, epilog_lambda, fin_op);
+    }
   }
 
   RAFT_CUDA_TRY(cudaGetLastError());
@@ -207,13 +231,7 @@ void cosineAlgo1(Index_ m,
 {
   auto norm_op = [] __device__(AccType in) { return raft::mySqrt(in); };
 
-  // Wrap fin_op to allow computing 1 - pA before calling fin_op
-  auto wrapped_fin_op = [fin_op] __device__(AccType d_val, Index_ g_d_idx) {
-    return fin_op(static_cast<AccType>(1.0) - d_val, g_d_idx);
-  };
-
-  typedef std::is_same<OutType, bool> is_bool;
-  typedef typename std::conditional<is_bool::value, OutType, AccType>::type CosOutType;
+  typedef typename std::conditional<sizeof(OutType) == 1, OutType, AccType>::type CosOutType;
   CosOutType* pDcast = reinterpret_cast<CosOutType*>(pD);
 
   ASSERT(
@@ -234,12 +252,12 @@ void cosineAlgo1(Index_ m,
 
   if (isRowMajor) {
     lda = k, ldb = k, ldd = n;
-    cosine<InType, AccType, CosOutType, Index_, decltype(wrapped_fin_op), true>(
-      m, n, k, lda, ldb, ldd, pA, pB, col_vec, row_vec, pDcast, wrapped_fin_op, stream);
+    cosine<InType, AccType, CosOutType, Index_, FinalLambda, true>(
+      m, n, k, lda, ldb, ldd, pA, pB, col_vec, row_vec, pDcast, fin_op, stream);
   } else {
     lda = n, ldb = m, ldd = m;
-    cosine<InType, AccType, CosOutType, Index_, decltype(wrapped_fin_op), false>(
-      n, m, k, lda, ldb, ldd, pB, pA, row_vec, col_vec, pDcast, wrapped_fin_op, stream);
+    cosine<InType, AccType, CosOutType, Index_, FinalLambda, false>(
+      n, m, k, lda, ldb, ldd, pB, pA, row_vec, col_vec, pDcast, fin_op, stream);
   }
 }
 

@@ -615,6 +615,16 @@ void distance(const InType* x,
  * @note if workspace is passed as nullptr, this will return in
  *  worksize, the number of bytes of workspace required
  */
+
+template <typename AccType, typename OutType, typename Index>
+struct default_fin_op {
+    __host__ __device__  default_fin_op() { };
+    // functor signature.
+    __host__ __device__ OutType operator()(AccType d_val, Index g_d_idx) const {
+      return d_val;
+    }
+};
+
 template <raft::distance::DistanceType distanceType,
           typename InType,
           typename AccType,
@@ -632,9 +642,11 @@ void distance(const InType* x,
               bool isRowMajor   = true,
               InType metric_arg = 2.0f)
 {
-  auto default_fin_op = [] __device__(AccType d_val, Index_ g_d_idx) { return d_val; };
-  distance<distanceType, InType, AccType, OutType, decltype(default_fin_op), Index_>(
-    x, y, dist, m, n, k, workspace, worksize, default_fin_op, stream, isRowMajor, metric_arg);
+  using final_op_type = default_fin_op<AccType, OutType, Index_>;
+  final_op_type fin_op;
+
+  distance<distanceType, InType, AccType, OutType, final_op_type, Index_>(
+    x, y, dist, m, n, k, workspace, worksize, fin_op, stream, isRowMajor, metric_arg);
   RAFT_CUDA_TRY(cudaPeekAtLastError());
 }