workaround only for sm120 and 121

Author: divyegala

cpp/include/raft/linalg/detail/cublaslt_wrappers.hpp

@@ -10,6 +10,7 @@
 #include <raft/core/resource/cublaslt_handle.hpp>
 #include <raft/core/resource/cuda_stream.hpp>
 #include <raft/core/resource/custom_resource.hpp>
+#include <raft/core/resource/device_properties.hpp>
 #include <raft/core/resources.hpp>
 #include <raft/util/cache.hpp>
 #include <raft/util/cuda_data_type.hpp>
@@ -104,12 +105,16 @@ struct matmul_key_hash {
  * cuBLASLt 13.6 may select algorithm 68 once A's physical span reaches 2^31 elements. That
  * algorithm fails during execution for FP32, so select the next ranked heuristic instead.
  */
-inline auto needs_cublaslt_13_6_workaround(const matmul_key_t& args, std::size_t version) noexcept
-  -> bool
+inline auto needs_cublaslt_13_6_workaround(const matmul_key_t& args,
+                                           std::size_t version,
+                                           int device_major,
+                                           int device_minor) noexcept -> bool
 {
   constexpr uint64_t max_safe_span = (uint64_t{1} << 31) - 1;
   const auto a_columns             = args.trans_a ? args.m : args.k;
-  return version == 130600 && args.lda != 0 && a_columns > max_safe_span / args.lda;
+  const bool is_sm120_or_sm121     = device_major == 12 && (device_minor == 0 || device_minor == 1);
+  return version == 130600 && is_sm120_or_sm121 && args.lda != 0 &&
+         a_columns > max_safe_span / args.lda;
 }
 
 inline auto get_cublaslt_algorithm_id(const cublasLtMatmulHeuristicResult_t& heuristic) -> int
@@ -212,7 +217,9 @@ struct matmul_desc {
     bool use_cublaslt_13_6_workaround = false;
     if constexpr (std::is_same_v<S, float> && std::is_same_v<A, float> &&
                   std::is_same_v<B, float> && std::is_same_v<C, float>) {
-      use_cublaslt_13_6_workaround = needs_cublaslt_13_6_workaround(args, cublasLtGetVersion());
+      const auto& device_properties = resource::get_device_properties(res);
+      use_cublaslt_13_6_workaround  = needs_cublaslt_13_6_workaround(
+        args, cublasLtGetVersion(), device_properties.major, device_properties.minor);
     }
 
     constexpr int workaround_heuristic_results = 2;

cpp/tests/linalg/gemm_basic.cpp

@@ -166,31 +166,53 @@ TEST(Raft, GemmPointerModeDeviceDefaults) { test_gemm_pointer_mode_device(false,
 TEST(Raft, GemmCublasLt136WorkaroundPredicate)
 {
   constexpr std::size_t affected_version = 130600;
+  constexpr int affected_device_major    = 12;
+  constexpr int affected_device_minor    = 1;
   const detail::matmul_key_t below_boundary{134217727, 1, 2, 16, 1, 134217727, true, true};
   const detail::matmul_key_t at_boundary{134217728, 1, 2, 16, 1, 134217728, true, true};
   const detail::matmul_key_t above_boundary{134217729, 1, 2, 16, 1, 134217729, true, true};
 
-  EXPECT_FALSE(detail::needs_cublaslt_13_6_workaround(below_boundary, affected_version));
-  EXPECT_TRUE(detail::needs_cublaslt_13_6_workaround(at_boundary, affected_version));
-  EXPECT_TRUE(detail::needs_cublaslt_13_6_workaround(above_boundary, affected_version));
-  EXPECT_FALSE(detail::needs_cublaslt_13_6_workaround(at_boundary, 130599));
-  EXPECT_FALSE(detail::needs_cublaslt_13_6_workaround(at_boundary, 130601));
+  const auto needs_workaround = [&](const auto& args) {
+    return detail::needs_cublaslt_13_6_workaround(
+      args, affected_version, affected_device_major, affected_device_minor);
+  };
+
+  EXPECT_FALSE(needs_workaround(below_boundary));
+  EXPECT_TRUE(needs_workaround(at_boundary));
+  EXPECT_TRUE(needs_workaround(above_boundary));
+  EXPECT_FALSE(detail::needs_cublaslt_13_6_workaround(
+    at_boundary, 130599, affected_device_major, affected_device_minor));
+  EXPECT_FALSE(detail::needs_cublaslt_13_6_workaround(
+    at_boundary, 130601, affected_device_major, affected_device_minor));
 
   auto different_output    = at_boundary;
   different_output.trans_b = false;
   different_output.n       = 2;
   different_output.ldb     = 7;
   different_output.ldc     = 11;
-  EXPECT_TRUE(detail::needs_cublaslt_13_6_workaround(different_output, affected_version));
+  EXPECT_TRUE(needs_workaround(different_output));
 
   const detail::matmul_key_t non_transposed_below{2, 1, 134217727, 16, 134217727, 2, false, false};
   const detail::matmul_key_t non_transposed_at{2, 1, 134217728, 16, 134217728, 2, false, false};
-  EXPECT_FALSE(detail::needs_cublaslt_13_6_workaround(non_transposed_below, affected_version));
-  EXPECT_TRUE(detail::needs_cublaslt_13_6_workaround(non_transposed_at, affected_version));
+  EXPECT_FALSE(needs_workaround(non_transposed_below));
+  EXPECT_TRUE(needs_workaround(non_transposed_at));
 
   auto invalid_lda = at_boundary;
   invalid_lda.lda  = 0;
-  EXPECT_FALSE(detail::needs_cublaslt_13_6_workaround(invalid_lda, affected_version));
+  EXPECT_FALSE(needs_workaround(invalid_lda));
+}
+
+TEST(Raft, GemmCublasLt136WorkaroundArchitectures)
+{
+  constexpr std::size_t affected_version = 130600;
+  const detail::matmul_key_t at_boundary{134217728, 1, 2, 16, 1, 134217728, true, true};
+
+  EXPECT_TRUE(detail::needs_cublaslt_13_6_workaround(at_boundary, affected_version, 12, 0));
+  EXPECT_TRUE(detail::needs_cublaslt_13_6_workaround(at_boundary, affected_version, 12, 1));
+  EXPECT_FALSE(detail::needs_cublaslt_13_6_workaround(at_boundary, affected_version, 7, 5));
+  EXPECT_FALSE(detail::needs_cublaslt_13_6_workaround(at_boundary, affected_version, 8, 0));
+  EXPECT_FALSE(detail::needs_cublaslt_13_6_workaround(at_boundary, affected_version, 8, 9));
+  EXPECT_FALSE(detail::needs_cublaslt_13_6_workaround(at_boundary, affected_version, 9, 0));
 }
 
 }  // namespace raft::linalg