Omg

gevtushenko · May 18, 2023 · cc746ab · cc746ab
1 parent 30d4e0c
commit cc746ab
Showing 1 changed file with 36 additions and 4 deletions.
diff --git a/benchmarks/bench/segmented_sort/keys.cu b/benchmarks/bench/segmented_sort/keys.cu
@@ -11,6 +11,10 @@
 // %RANGE% TUNE_RANK_ALGORITHM ra 0:4:1
 // %RANGE% TUNE_LOAD ld 0:2:1
 // %RANGE% TUNE_TRANSPOSE trp 0:1:1
+// %RANGE% TUNE_S_LOAD sld 0:2:1
+// %RANGE% TUNE_S_TRANSPOSE strp 0:1:1
+// %RANGE% TUNE_M_LOAD mld 0:2:1
+// %RANGE% TUNE_M_TRANSPOSE mtrp 0:1:1
 
 #include <cub/device/device_segmented_sort.cuh>
 
@@ -39,12 +43,40 @@
 #define TUNE_LOAD_MODIFIER cub::LOAD_CA
 #endif // TUNE_LOAD 
 
+#if TUNE_S_LOAD == 0
+#define TUNE_S_LOAD_MODIFIER cub::LOAD_DEFAULT
+#elif TUNE_S_LOAD == 1
+#define TUNE_S_LOAD_MODIFIER cub::LOAD_LDG
+#else // TUNE_S_LOAD == 2
+#define TUNE_S_LOAD_MODIFIER cub::LOAD_CA
+#endif // TUNE_S_LOAD 
+
+#if TUNE_M_LOAD == 0
+#define TUNE_M_LOAD_MODIFIER cub::LOAD_DEFAULT
+#elif TUNE_M_LOAD == 1
+#define TUNE_M_LOAD_MODIFIER cub::LOAD_LDG
+#else // TUNE_M_LOAD == 2
+#define TUNE_M_LOAD_MODIFIER cub::LOAD_CA
+#endif // TUNE_M_LOAD 
+
 #if TUNE_TRANSPOSE == 0
 #define TUNE_LOAD_ALGORITHM cub::BLOCK_LOAD_DIRECT
 #else // TUNE_TRANSPOSE == 1
 #define TUNE_LOAD_ALGORITHM cub::BLOCK_LOAD_WARP_TRANSPOSE
 #endif // TUNE_TRANSPOSE
 
+#if TUNE_S_TRANSPOSE == 0
+#define TUNE_S_LOAD_ALGORITHM cub::WarpLoadAlgorithm::WARP_LOAD_DIRECT
+#else // TUNE_S_TRANSPOSE == 1
+#define TUNE_S_LOAD_ALGORITHM cub::WarpLoadAlgorithm::WARP_LOAD_TRANSPOSE
+#endif // TUNE_S_TRANSPOSE
+
+#if TUNE_M_TRANSPOSE == 0
+#define TUNE_M_LOAD_ALGORITHM cub::WarpLoadAlgorithm::WARP_LOAD_DIRECT
+#else // TUNE_M_TRANSPOSE == 1
+#define TUNE_M_LOAD_ALGORITHM cub::WarpLoadAlgorithm::WARP_LOAD_TRANSPOSE
+#endif // TUNE_M_TRANSPOSE
+
 template <class KeyT>
 struct device_seg_sort_policy_hub
 {
@@ -76,14 +108,14 @@ struct device_seg_sort_policy_hub
       // Small policy
       cub::AgentSubWarpMergeSortPolicy<TUNE_SW_THREADS,
                                        ITEMS_PER_SMALL_THREAD,
-                                       cub::WarpLoadAlgorithm::WARP_LOAD_DIRECT,
-                                       cub::CacheLoadModifier::LOAD_DEFAULT>,
+                                       TUNE_S_LOAD_ALGORITHM,
+                                       TUNE_S_LOAD_MODIFIER>,
 
       // Medium policy
       cub::AgentSubWarpMergeSortPolicy<TUNE_MW_THREADS,
                                        ITEMS_PER_MEDIUM_THREAD,
-                                       cub::WarpLoadAlgorithm::WARP_LOAD_DIRECT,
-                                       cub::CacheLoadModifier::LOAD_DEFAULT>>;
+                                       TUNE_M_LOAD_ALGORITHM,
+                                       TUNE_M_LOAD_MODIFIER>>;
   };
 
   using MaxPolicy = Policy350;