Adds new virtual shared memory facility to DeviceMergeSort (#1117)

* Adds new vsmem facility to DeviceMergeSort * fixes narrowing type conversion * addresses review comments * moves vsmem tests to own tu * disable benches for int128 value types * adds discard_memory hints * reverts unrelated formatting changes
NVIDIA · Dec 6, 2023 · d855743 · d855743
1 parent 9eca399
commit d855743
Show file tree

Hide file tree

Showing 8 changed files with 819 additions and 852 deletions.
diff --git a/cub/benchmarks/bench/merge_sort/pairs.cu b/cub/benchmarks/bench/merge_sort/pairs.cu
@@ -156,15 +156,17 @@ using key_types = all_types;
 #ifdef TUNE_ValueT
 using value_types = nvbench::type_list<TUNE_ValueT>;
 #else // !defined(TUNE_ValueT)
-using value_types = nvbench::type_list<int8_t,
-                                       int16_t,
-                                       int32_t,
-                                       int64_t
+using value_types =
+  nvbench::type_list<int8_t,
+                     int16_t,
+                     int32_t,
+                     int64_t
 #if NVBENCH_HELPER_HAS_I128
-                                       ,
-                                       int128_t
+// nvcc currently hangs for __int128 value type with the fallback policy of {CTA: 64, IPT: 1}. NVBug 4384075
+//  ,
+//  int128_t
 #endif
-                                       >;
+                     >;
 #endif // TUNE_ValueT
 
 NVBENCH_BENCH_TYPES(pairs, NVBENCH_TYPE_AXES(key_types, value_types, offset_types))

diff --git a/cub/cub/agent/agent_merge_sort.cuh b/cub/cub/agent/agent_merge_sort.cuh
@@ -114,8 +114,6 @@ struct AgentBlockSort
   static constexpr int BLOCK_THREADS = Policy::BLOCK_THREADS;
   static constexpr int ITEMS_PER_THREAD = Policy::ITEMS_PER_THREAD;
   static constexpr int ITEMS_PER_TILE = Policy::ITEMS_PER_TILE;
-  static constexpr int SHARED_MEMORY_SIZE =
-    static_cast<int>(sizeof(TempStorage));
 
   //---------------------------------------------------------------------
   // Per thread data
@@ -428,8 +426,6 @@ struct AgentMerge
   static constexpr int BLOCK_THREADS = Policy::BLOCK_THREADS;
   static constexpr int ITEMS_PER_THREAD = Policy::ITEMS_PER_THREAD;
   static constexpr int ITEMS_PER_TILE = Policy::ITEMS_PER_TILE;
-  static constexpr int SHARED_MEMORY_SIZE =
-    static_cast<int>(sizeof(TempStorage));
 
   //---------------------------------------------------------------------
   // Per thread data