Use cuda::std::array in histogram APIs

Fixes: #1765
NVIDIA · Mar 1, 2025 · 6059dd3 · 6059dd3
1 parent 4a3a219
commit 6059dd3
Show file tree

Hide file tree

Showing 3 changed files with 125 additions and 160 deletions.
diff --git a/cub/cub/device/device_histogram.cuh b/cub/cub/device/device_histogram.cuh
@@ -45,6 +45,8 @@
 #include <cub/detail/nvtx.cuh>
 #include <cub/device/dispatch/dispatch_histogram.cuh>
 
+#include <cuda/std/array>
+
 #include <iterator>
 #include <limits>
 
@@ -195,10 +197,10 @@ struct DeviceHistogram
       d_temp_storage,
       temp_storage_bytes,
       d_samples,
-      &d_histogram,
-      &num_levels,
-      &lower_level,
-      &upper_level,
+      ::cuda::std::array<CounterT*, 1>{{d_histogram}},
+      ::cuda::std::array<int, 1>{{num_levels}},
+      ::cuda::std::array<LevelT, 1>{{lower_level}},
+      ::cuda::std::array<LevelT, 1>{{upper_level}},
       num_samples,
       static_cast<OffsetT>(1),
       sizeof(SampleT) * num_samples,
@@ -346,10 +348,10 @@ struct DeviceHistogram
       d_temp_storage,
       temp_storage_bytes,
       d_samples,
-      &d_histogram,
-      &num_levels,
-      &lower_level,
-      &upper_level,
+      ::cuda::std::array<CounterT*, 1>{{d_histogram}},
+      ::cuda::std::array<int, 1>{{num_levels}},
+      ::cuda::std::array<LevelT, 1>{{lower_level}},
+      ::cuda::std::array<LevelT, 1>{{upper_level}},
       num_row_samples,
       num_rows,
       row_stride_bytes,
@@ -501,10 +503,10 @@ struct DeviceHistogram
     void* d_temp_storage,
     size_t& temp_storage_bytes,
     SampleIteratorT d_samples,
-    CounterT* d_histogram[NUM_ACTIVE_CHANNELS],
-    const int num_levels[NUM_ACTIVE_CHANNELS],
-    const LevelT lower_level[NUM_ACTIVE_CHANNELS],
-    const LevelT upper_level[NUM_ACTIVE_CHANNELS],
+    ::cuda::std::array<CounterT*, NUM_ACTIVE_CHANNELS> d_histogram,
+    ::cuda::std::array<int, NUM_ACTIVE_CHANNELS> num_levels,
+    ::cuda::std::array<LevelT, NUM_ACTIVE_CHANNELS> lower_level,
+    ::cuda::std::array<LevelT, NUM_ACTIVE_CHANNELS> upper_level,
     OffsetT num_pixels,
     cudaStream_t stream = 0)
   {
@@ -688,10 +690,10 @@ struct DeviceHistogram
     void* d_temp_storage,
     size_t& temp_storage_bytes,
     SampleIteratorT d_samples,
-    CounterT* d_histogram[NUM_ACTIVE_CHANNELS],
-    const int num_levels[NUM_ACTIVE_CHANNELS],
-    const LevelT lower_level[NUM_ACTIVE_CHANNELS],
-    const LevelT upper_level[NUM_ACTIVE_CHANNELS],
+    ::cuda::std::array<CounterT*, NUM_ACTIVE_CHANNELS> d_histogram,
+    ::cuda::std::array<int, NUM_ACTIVE_CHANNELS> num_levels,
+    ::cuda::std::array<LevelT, NUM_ACTIVE_CHANNELS> lower_level,
+    ::cuda::std::array<LevelT, NUM_ACTIVE_CHANNELS> upper_level,
     OffsetT num_row_pixels,
     OffsetT num_rows,
     size_t row_stride_bytes,
@@ -855,9 +857,9 @@ struct DeviceHistogram
       d_temp_storage,
       temp_storage_bytes,
       d_samples,
-      &d_histogram,
-      &num_levels,
-      &d_levels,
+      ::cuda::std::array<CounterT*, 1>{{d_histogram}},
+      ::cuda::std::array<int, 1>{{num_levels}},
+      ::cuda::std::array<const LevelT*, 1>{{d_levels}},
       num_samples,
       (OffsetT) 1,
       (size_t) (sizeof(SampleT) * num_samples),
@@ -993,9 +995,9 @@ struct DeviceHistogram
       d_temp_storage,
       temp_storage_bytes,
       d_samples,
-      &d_histogram,
-      &num_levels,
-      &d_levels,
+      ::cuda::std::array<CounterT*, 1>{{d_histogram}},
+      ::cuda::std::array<int, 1>{{num_levels}},
+      ::cuda::std::array<const LevelT*, 1>{{d_levels}},
       num_row_samples,
       num_rows,
       row_stride_bytes,
@@ -1138,9 +1140,9 @@ struct DeviceHistogram
     void* d_temp_storage,
     size_t& temp_storage_bytes,
     SampleIteratorT d_samples,
-    CounterT* d_histogram[NUM_ACTIVE_CHANNELS],
-    const int num_levels[NUM_ACTIVE_CHANNELS],
-    const LevelT* const d_levels[NUM_ACTIVE_CHANNELS],
+    ::cuda::std::array<CounterT*, NUM_ACTIVE_CHANNELS> d_histogram,
+    ::cuda::std::array<int, NUM_ACTIVE_CHANNELS> num_levels,
+    ::cuda::std::array<const LevelT*, NUM_ACTIVE_CHANNELS> d_levels,
     OffsetT num_pixels,
     cudaStream_t stream = 0)
   {
@@ -1315,9 +1317,9 @@ struct DeviceHistogram
     void* d_temp_storage,
     size_t& temp_storage_bytes,
     SampleIteratorT d_samples,
-    CounterT* d_histogram[NUM_ACTIVE_CHANNELS],
-    const int num_levels[NUM_ACTIVE_CHANNELS],
-    const LevelT* const d_levels[NUM_ACTIVE_CHANNELS],
+    ::cuda::std::array<CounterT*, NUM_ACTIVE_CHANNELS> d_histogram,
+    ::cuda::std::array<int, NUM_ACTIVE_CHANNELS> num_levels,
+    ::cuda::std::array<const LevelT*, NUM_ACTIVE_CHANNELS> d_levels,
     OffsetT num_row_pixels,
     OffsetT num_rows,
     size_t row_stride_bytes,