fix get_cutlass install path, make changes as per review comments

rapidsai · rapids-bot · Nov 16, 2022 · Oct 20, 2022 · Oct 21, 2022 · Oct 21, 2022
commit 7b512f99bbfffee488ad329062787b017ba2f087
@@ -25,11 +25,12 @@ function(find_and_configure_cutlass)
       set(CUTLASS_ENABLE_CUBLAS OFF CACHE BOOL "Disable CUTLASS to build with cuBLAS library.")
 
       rapids_cpm_find(NvidiaCutlass ${PKG_VERSION}
-          GLOBAL_TARGETS     nvidia::cutlass::cutlass
+          GLOBAL_TARGETS  nvidia::cutlass::cutlass
           CPM_ARGS
-            GIT_REPOSITORY   ${PKG_REPOSITORY}
-            GIT_TAG          ${PKG_PINNED_TAG}
-            OPTIONS
+          GIT_REPOSITORY  ${PKG_REPOSITORY}
+          GIT_TAG         ${PKG_PINNED_TAG}
+          GIT_SHALLOW     TRUE
+          OPTIONS
               "CUDAToolkit_ROOT ${CUDAToolkit_LIBRARY_DIR}"
       )
 
@@ -51,7 +52,8 @@ function(find_and_configure_cutlass)
 
     # Tell cmake where it can find the generated NvidiaCutlass-config.cmake we wrote.
     include("${rapids-cmake-dir}/export/find_package_root.cmake")
-    rapids_export_find_package_root(INSTALL NvidiaCutlass [=[${CMAKE_CURRENT_LIST_DIR}]=] raft-distance-exports)
+    rapids_export_find_package_root(INSTALL NvidiaCutlass [=[${CMAKE_CURRENT_LIST_DIR}/../]=] raft-distance-exports)
+    rapids_export_find_package_root(BUILD NvidiaCutlass [=[${CMAKE_CURRENT_LIST_DIR}]=] raft-distance-exports)
 endfunction()
 
 if(NOT RAFT_CUTLASS_GIT_TAG)

@@ -82,7 +82,7 @@ void cosineImpl(const DataT* x,
                 FinalLambda fin_op,
                 cudaStream_t stream)
 {
-  const auto deviceVersion = getMajorMinorVersion();
+  const auto deviceVersion = getComputeCapability();
   if (deviceVersion.first >= 8) {
     using CosineOp_ = CosineOp<DataT, AccT>;
     CosineOp_ cosine_dist_op;
@@ -228,6 +228,10 @@ void cosineAlgo1(Index_ m,
 {
   auto norm_op = [] __device__(AccType in) { return raft::mySqrt(in); };
 
+  // raft distance support inputs as float/double and output as uint8_t/float/double.
+  static_assert(!((sizeof(OutType) > 1) && (sizeof(AccType) != sizeof(OutType))),
+                "OutType can be uint8_t, float, double,"
+                "if sizeof(OutType) > 1 then sizeof(AccType) == sizeof(OutType).");
   typedef typename std::conditional<sizeof(OutType) == 1, OutType, AccType>::type CosOutType;
   CosOutType* pDcast = reinterpret_cast<CosOutType*>(pD);
 

@@ -648,6 +648,10 @@ void distance(const InType* x,
   using final_op_type = default_fin_op<AccType, OutType, Index_>;
   final_op_type fin_op;
 
+  // raft distance support inputs as float/double and output as uint8_t/float/double.
+  static_assert(! ((sizeof(OutType) > 1) && (sizeof(AccType) != sizeof(OutType))),
+                "OutType can be uint8_t, float, double,"
+                "if sizeof(OutType) > 1 then sizeof(AccType) == sizeof(OutType).");
   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());

@@ -88,7 +88,7 @@ void euclideanExpImpl(const DataT* x,
                       FinalLambda fin_op,
                       cudaStream_t stream)
 {
-  const auto deviceVersion = getMajorMinorVersion();
+  const auto deviceVersion = getComputeCapability();
   if (deviceVersion.first >= 8) {
     using L2Op = L2ExpandedOp<DataT, AccT>;
     L2Op L2_dist_op(sqrt);
@@ -245,6 +245,10 @@ void euclideanAlgo1(Index_ m,
 {
   auto norm_op = [] __device__(InType in) { return in; };
 
+  // raft distance support inputs as float/double and output as uint8_t/float/double.
+  static_assert(! ((sizeof(OutType) > 1) && (sizeof(AccType) != sizeof(OutType))),
+                "OutType can be uint8_t, float, double,"
+                "if sizeof(OutType) > 1 then sizeof(AccType) == sizeof(OutType).");
   typedef typename std::conditional<sizeof(OutType) == 1, OutType, AccType>::type ExpOutType;
   ExpOutType* pDcast = reinterpret_cast<ExpOutType*>(pD);
 

@@ -19,6 +19,13 @@
 #pragma GCC diagnostic push
 #pragma GCC diagnostic ignored "-Wstrict-aliasing"
 
+// We define CUTLASS_NAMESPACE in case
+// RAFT cmake is not used
+#ifndef CUTLASS_NAMESPACE
+#define cutlass raft_cutlass 
+#endif
+
+
 #include <rmm/device_uvector.hpp>
 
 #include <cutlass/cutlass.h>

@@ -17,9 +17,14 @@
 /*! \file
   \brief Epilogue for threadblock scoped GEMMs using Tensor Ops.
 
-  The epilogue rearranges the result of a matrix product through shared memory to match canonical
-  tensor layouts in global memory. Epilogues support conversion and reduction operations.
-
+This is adapted from DefaultEpilogueWithBroadcastTensorOp from CUTLASS 2.9.0
+(https://github.com/NVIDIA/cutlass/blob/master/include/cutlass/epilogue/threadblock/default_epilogue_with_broadcast.h#L75)
+
+This epilogue allows us to load norm buffers using PredicatedTileIteratorNormVec 
+and EpilogueWithBroadcast used for distances L2/cosine as well as applies user-define elementwise operation.
+-- A norm load is provided PredicatedTileIteratorNormVec
+-- B norm load is provided by EpilogueWithBroadcast
+-- elementwise operation is provided by OutputOp
 */
 
 #pragma once

@@ -18,6 +18,9 @@
 /*! \file
   \brief Functor performing distance operations used by epilogues of pairwise distance
   * kernels.
+* This is adapted from LinearCombinationBiasElementwise from CUTLASS 2.9.0
+* customized for applying elementwise distance formula on accumulated GEMM value
+* and applying user-defined final custom operation on the distance value.
 */
 
 #pragma once

@@ -80,7 +80,7 @@ struct PairwiseDistanceGemm {
 
   // This code section describes how threadblocks are scheduled on GPU
   /// Threadblock-level swizzling operator
-  using ThreadblockSwizzle = cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<>;  // <- ??
+  using ThreadblockSwizzle = cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<>;
 
   /// data layout for final output matrix.
   // we keep this same layout even for column major inputs
@@ -179,7 +179,7 @@ struct PairwiseDistanceGemm<double,
 
   // This code section describes how threadblocks are scheduled on GPU
   /// Threadblock-level swizzling operator
-  using ThreadblockSwizzle = cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<>;  // <- ??
+  using ThreadblockSwizzle = cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<>;
 
   /// data layout for final output matrix.
   // we keep this same layout even for column major inputs

@@ -17,8 +17,13 @@
 /*! \file
   \brief Epilogue for threadblock scoped GEMMs using Tensor Ops.
 
-  The epilogue rearranges the result of a matrix product through shared memory to match canonical
-  tensor layouts in global memory. Epilogues support conversion and reduction operations.
+This file contains a customized version of PredicatedTileIterator from CUTLASS 2.9.0
+(https://github.com/NVIDIA/cutlass/blob/v2.9.0/include/cutlass/epilogue/threadblock/predicated_tile_iterator.h#L75)
+
+Changes:
+- added `Layout_` template param
+- Only the row index is used to load the data in load_with_byte_offset().
+  This way the same normalization data is used across all columns in a row.
 
 */
 

@@ -356,7 +356,7 @@ inline int getMultiProcessorCount()
 }
 
 /** helper method to get major minor compute capability version */
-inline std::pair<int, int> getMajorMinorVersion()
+inline std::pair<int, int> getComputeCapability()
 {
   int devId;
   RAFT_CUDA_TRY(cudaGetDevice(&devId));

@@ -150,6 +150,10 @@ class DistanceAdjTest : public ::testing::TestWithParam<DistanceAdjInputs<DataTy
 
  protected:
   DistanceAdjInputs<DataType> params;
+  // We use uint8_t even if the output in this test is a bool because
+  // cutlass doesn't support bool as output buffer yet. In cuda
+  // sizeof(bool) is 1 byte hence it doesn't increase
+  // memory consumption if we use uint8_t instead of bool.
   rmm::device_uvector<uint8_t> dist_ref;
   rmm::device_uvector<uint8_t> dist;
   raft::handle_t handle;