[Opt] Support Sm70 on V100

NVIDIA-Merlin · Sep 12, 2023 · 36d8bd0 · 36d8bd0
1 parent 1152d4f
commit 36d8bd0
Show file tree

Hide file tree

Showing 6 changed files with 32 additions and 20 deletions.
diff --git a/benchmark/merlin_hashtable_benchmark.cc.cu b/benchmark/merlin_hashtable_benchmark.cc.cu
@@ -480,7 +480,7 @@ void test_main(std::vector<API_Select>& apis, const size_t dim,
   options.dim = dim;
   options.max_hbm_for_vectors = nv::merlin::GB(hbm4values);
   options.io_by_cpu = io_by_cpu;
-  using Table = nv::merlin::HashTable<K, V, S, EvictStrategy::kLru>;
+  using Table = nv::merlin::HashTable<K, V, S, EvictStrategy::kLru, Sm80>;
 
   std::shared_ptr<Table> table = std::make_shared<Table>();
   table->init(options);

diff --git a/include/merlin/core_kernels/lookup.cuh b/include/merlin/core_kernels/lookup.cuh
@@ -667,6 +667,12 @@ struct LookupValueBufConfig<Sm80> {
   static constexpr uint32_t size_pipeline_v2 = 128 * sizeof(float);
 };
 
+template <>
+struct LookupValueBufConfig<Sm70> {
+  static constexpr uint32_t size_pipeline_v1 = 112 * sizeof(float);
+  static constexpr uint32_t size_pipeline_v2 = 64 * sizeof(float);
+};
+
 template <typename K, typename V, typename S = uint64_t,
           typename ArchTag = Sm80>
 struct SelectPipelineLookupKernelWithIO {

diff --git a/include/merlin/core_kernels/update.cuh b/include/merlin/core_kernels/update.cuh
@@ -567,12 +567,17 @@ struct ValueConfig_Update<Sm80> {
   static constexpr uint32_t size_pipeline = 128 * sizeof(byte4);
 };
 
-template <typename K, typename V, typename S, int Strategy, typename ArchTag>
-struct KernelSelector_Update;
+template <>
+struct ValueConfig_Update<Sm70> {
+  // Value size greater than it will bring poor performance for TLP.
+  static constexpr uint32_t size_tlp = 8 * sizeof(byte4);
+  // Value size greater than it will reduce the occupancy for Pipeline.
+  // When the value is very high, the kernel will fail to launch.
+  static constexpr uint32_t size_pipeline = 64 * sizeof(byte4);
+};
 
-template <typename K, typename V, typename S, int Strategy>
-struct KernelSelector_Update<K, V, S, Strategy, Sm80> {
-  using ArchTag = Sm80;
+template <typename K, typename V, typename S, int Strategy, typename ArchTag>
+struct KernelSelector_Update {
   using ValueConfig = ValueConfig_Update<ArchTag>;
   using Params = Params_Update<K, V, S>;
 

diff --git a/include/merlin/core_kernels/update_score.cuh b/include/merlin/core_kernels/update_score.cuh
@@ -502,13 +502,8 @@ struct Launch_Pipeline_UpdateScore {
   }
 };
 
-/// TODO: support more arch.
-template <typename K, typename V, typename S, int Strategy, typename ArchTag>
-struct KernelSelector_UpdateScore;
-
 template <typename K, typename V, typename S, int Strategy>
-struct KernelSelector_UpdateScore<K, V, S, Strategy, Sm80> {
-  using ArchTag = Sm80;
+struct KernelSelector_UpdateScore {
   using Params = Params_UpdateScore<K, V, S>;
 
   static bool callable(bool unique_key, uint32_t bucket_size) {

diff --git a/include/merlin/core_kernels/upsert_and_evict.cuh b/include/merlin/core_kernels/upsert_and_evict.cuh
@@ -1234,12 +1234,19 @@ struct ValueConfig_UpsertAndEvict<Sm80> {
   static constexpr uint32_t size_pipeline = 128 * sizeof(byte4);
 };
 
-template <typename K, typename V, typename S, int Strategy, typename ArchTag>
-struct KernelSelector_UpsertAndEvict;
+template <>
+struct ValueConfig_UpsertAndEvict<Sm70> {
+  // Value size greater than it will bring poor performance for TLPv1.
+  static constexpr uint32_t size_tlp_v1 = 16 * sizeof(byte4);
+  // Value size greater than it will bring wrong result for TLPv2.
+  static constexpr uint32_t size_tlp_v2 = 32 * sizeof(byte4);
+  // Value size greater than it will reduce the occupancy for Pipeline.
+  // When the value is very high, the kernel will fail to launch.
+  static constexpr uint32_t size_pipeline = 64 * sizeof(byte4);
+};
 
-template <typename K, typename V, typename S, int Strategy>
-struct KernelSelector_UpsertAndEvict<K, V, S, Strategy, Sm80> {
-  using ArchTag = Sm80;
+template <typename K, typename V, typename S, int Strategy, typename ArchTag>
+struct KernelSelector_UpsertAndEvict {
   using ValueConfig = ValueConfig_UpsertAndEvict<ArchTag>;
   using Params = Params_UpsertAndEvict<K, V, S>;
 

diff --git a/include/merlin_hashtable.cuh b/include/merlin_hashtable.cuh
@@ -1072,9 +1072,8 @@ class HashTable {
       if (((step_counter++) % kernel_select_interval_) == 0) {
         load_factor = fast_load_factor(0, stream, false);
       }
-      using Selector =
-          KernelSelector_UpdateScore<key_type, value_type, score_type,
-                                     evict_strategy, ArchTag>;
+      using Selector = KernelSelector_UpdateScore<key_type, value_type,
+                                                  score_type, evict_strategy>;
       if (Selector::callable(unique_key,
                              static_cast<uint32_t>(options_.max_bucket_size))) {
         typename Selector::Params kernelParams(