Pushing changes to GitHub Pages.

NVIDIA-Merlin · Aug 21, 2024 · 2a5e350 · 2a5e350
1 parent f267e0a
commit 2a5e350
Show file tree

Hide file tree

Showing 3 changed files with 83 additions and 11 deletions.
diff --git a/review/pr-200/_sources/api/program_listing_file_merlin_hashtable.cuh.rst.txt b/review/pr-200/_sources/api/program_listing_file_merlin_hashtable.cuh.rst.txt
@@ -1631,11 +1631,47 @@ Program Listing for File merlin_hashtable.cuh
          int grid_size = std::min(
              sm_cnt_ * max_threads_per_block_ / options_.block_size,
              static_cast<int>(SAFE_GET_GRID_SIZE(n, options_.block_size)));
-         const int TILE_SIZE = 32;
-         dump_kernel_v2<key_type, value_type, score_type, PredFunctor, TILE_SIZE>
-             <<<grid_size, options_.block_size, 0, stream>>>(
-                 d_table_, table_->buckets, pattern, threshold, keys, values,
-                 scores, offset, n, d_counter);
+         if (sizeof(V) == sizeof(float) && dim() >= 32 && dim() % 4 == 0) {
+           if (dim() >= 128) {
+             const int TILE_SIZE = 32;
+             dump_kernel_v2_vectorized<key_type, value_type, score_type, PredFunctor, TILE_SIZE>
+                 <<<grid_size, options_.block_size, 0, stream>>>(
+                     d_table_, table_->buckets, pattern, threshold, keys, values,
+                     scores, offset, n, d_counter);
+           } else if (dim() >= 64) {
+             const int TILE_SIZE = 16;
+             dump_kernel_v2_vectorized<key_type, value_type, score_type, PredFunctor, TILE_SIZE>
+                 <<<grid_size, options_.block_size, 0, stream>>>(
+                     d_table_, table_->buckets, pattern, threshold, keys, values,
+                     scores, offset, n, d_counter);
+           } else {
+             const int TILE_SIZE = 8;
+             dump_kernel_v2_vectorized<key_type, value_type, score_type, PredFunctor, TILE_SIZE>
+                 <<<grid_size, options_.block_size, 0, stream>>>(
+                     d_table_, table_->buckets, pattern, threshold, keys, values,
+                     scores, offset, n, d_counter);
+           }
+         } else {
+           if (dim() >= 32) {
+             const int TILE_SIZE = 32;
+             dump_kernel_v2<key_type, value_type, score_type, PredFunctor, TILE_SIZE>
+                 <<<grid_size, options_.block_size, 0, stream>>>(
+                     d_table_, table_->buckets, pattern, threshold, keys, values,
+                     scores, offset, n, d_counter);
+           } else if (dim() >= 16) {
+             const int TILE_SIZE = 16;
+             dump_kernel_v2<key_type, value_type, score_type, PredFunctor, TILE_SIZE>
+                 <<<grid_size, options_.block_size, 0, stream>>>(
+                     d_table_, table_->buckets, pattern, threshold, keys, values,
+                     scores, offset, n, d_counter);
+           } else {
+             const int TILE_SIZE = 8;
+             dump_kernel_v2<key_type, value_type, score_type, PredFunctor, TILE_SIZE>
+                 <<<grid_size, options_.block_size, 0, stream>>>(
+                     d_table_, table_->buckets, pattern, threshold, keys, values,
+                     scores, offset, n, d_counter);
+           }
+         }
        } else {
          const size_t score_size = scores ? sizeof(score_type) : 0;
          const size_t kvm_size =