Restrict active histogram channels to channel count (#1796)

Fixes: #1792
NVIDIA · Jun 4, 2024 · 4d215e5 · 4d215e5
1 parent 0958e8b
commit 4d215e5
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 1 deletion.
diff --git a/cub/cub/device/device_histogram.cuh b/cub/cub/device/device_histogram.cuh
@@ -422,6 +422,7 @@ struct DeviceHistogram
   //! - The input is a sequence of *pixel* structures, where each pixel comprises
   //!   a record of ``NUM_CHANNELS`` consecutive data samples
   //!   (e.g., an *RGBA* pixel).
+  //! - ``NUM_CHANNELS`` can be up to 4.
   //! - Of the ``NUM_CHANNELS`` specified, the function will only compute
   //!   histograms for the first ``NUM_ACTIVE_CHANNELS``
   //!   (e.g., only *RGB* histograms from *RGBA* pixel samples).
@@ -621,6 +622,7 @@ struct DeviceHistogram
   //!
   //! - The input is a sequence of *pixel* structures, where each pixel
   //!   comprises a record of ``NUM_CHANNELS`` consecutive data samples (e.g., an *RGBA* pixel).
+  //! - ``NUM_CHANNELS`` can be up to 4.
   //! - Of the ``NUM_CHANNELS`` specified, the function will only compute
   //!   histograms for the first ``NUM_ACTIVE_CHANNELS`` (e.g., only *RGB*
   //!   histograms from *RGBA* pixel samples).
@@ -1180,6 +1182,7 @@ struct DeviceHistogram
   //!
   //! - The input is a sequence of *pixel* structures, where each pixel
   //!   comprises a record of ``NUM_CHANNELS`` consecutive data samples (e.g., an *RGBA* pixel).
+  //! - ``NUM_CHANNELS`` can be up to 4.
   //! - Of the ``NUM_CHANNELS`` specified, the function will only compute
   //!   histograms for the first ``NUM_ACTIVE_CHANNELS`` (e.g., *RGB* histograms from *RGBA* pixel samples).
   //! - The number of histogram bins for channel\ :sub:`i` is ``num_levels[i] - 1``.
@@ -1360,6 +1363,7 @@ struct DeviceHistogram
   //!
   //! - The input is a sequence of *pixel* structures, where each pixel comprises
   //!   a record of ``NUM_CHANNELS`` consecutive data samples (e.g., an *RGBA* pixel).
+  //! - ``NUM_CHANNELS`` can be up to 4.
   //! - Of the ``NUM_CHANNELS`` specified, the function will only compute
   //!   histograms for the first ``NUM_ACTIVE_CHANNELS`` (e.g., *RGB* histograms from *RGBA* pixel samples).
   //! - A two-dimensional *region of interest* within ``d_samples`` can be

diff --git a/cub/cub/device/dispatch/dispatch_histogram.cuh b/cub/cub/device/dispatch/dispatch_histogram.cuh
@@ -563,7 +563,8 @@ template <int NUM_CHANNELS,
 struct DispatchHistogram : SelectedPolicy
 {
   static_assert(NUM_CHANNELS <= 4, "Histograms only support up to 4 channels");
-  static_assert(NUM_ACTIVE_CHANNELS <= 4, "Histograms only support up to 4 active channels");
+  static_assert(NUM_ACTIVE_CHANNELS <= NUM_CHANNELS,
+                "Active channels must be at most the number of total channels of the input samples");
 
 public:
   //---------------------------------------------------------------------