Fix out-of-bound access in cudf::detail::label_segments (#11497)

In `cudf::detail::label_segments`, when the input lists column has empty/nulls lists at the end of the column, its `offsets` column will contain out-of-bound indices. This leads to invalid memory access bug. Such bug is elusive and doesn't show up consistently. Test failures reported in NVIDIA/spark-rapids#6249 are due to this.

The existing unit tests already cover such corner case. Unfortunately, the bug didn't show up until being tested on some systems. Even that, it was very difficult to reproduce it.

Closes #11495.

Authors:
   - Nghia Truong (https://github.com/ttnghia)

Approvers:
   - Tobias Ribizel (https://github.com/upsj)
   - Bradley Dice (https://github.com/bdice)
   - Jim Brennan (https://github.com/jbrennan333)
   - Alessandro Bellina (https://github.com/abellina)
   - Karthikeyan (https://github.com/karthikeyann)

cpp/include/cudf/detail/labeling/label_segments.cuh

@@ -23,6 +23,7 @@
 #include <thrust/distance.h>
 #include <thrust/for_each.h>
 #include <thrust/iterator/constant_iterator.h>
+#include <thrust/iterator/iterator_traits.h>
 #include <thrust/reduce.h>
 #include <thrust/scan.h>
 #include <thrust/uninitialized_fill.h>
@@ -67,10 +68,12 @@ void label_segments(InputIterator offsets_begin,
                     OutputIterator label_end,
                     rmm::cuda_stream_view stream)
 {
+  auto const num_labels = thrust::distance(label_begin, label_end);
+
   // If the output array is empty, that means we have all empty segments.
   // In such cases, we must terminate immediately. Otherwise, the `for_each` loop below may try to
   // access memory of the output array, resulting in "illegal memory access" error.
-  if (thrust::distance(label_begin, label_end) == 0) { return; }
+  if (num_labels == 0) { return; }
 
   // When the output array is not empty, always fill it with `0` value first.
   using OutputType = typename thrust::iterator_value<OutputIterator>::type;
@@ -83,19 +86,24 @@ void label_segments(InputIterator offsets_begin,
   // very large segment.
   if (thrust::distance(offsets_begin, offsets_end) <= 2) { return; }
 
-  thrust::for_each(rmm::exec_policy(stream),
-                   offsets_begin + 1,  // exclude the first offset value
-                   offsets_end - 1,    // exclude the last offset value
-                   [offsets = offsets_begin, output = label_begin] __device__(auto const idx) {
-                     // Zero-normalized offsets.
-                     auto const dst_idx = idx - (*offsets);
-
-                     // Scatter value `1` to the index at (idx - offsets[0]).
-                     // In case we have repeated offsets (i.e., we have empty segments), this
-                     // `atomicAdd` call will make sure the label values corresponding to these
-                     // empty segments will be skipped in the output.
-                     atomicAdd(&output[dst_idx], OutputType{1});
-                   });
+  thrust::for_each(
+    rmm::exec_policy(stream),
+    offsets_begin + 1,  // exclude the first offset value
+    offsets_end - 1,    // exclude the last offset value
+    [num_labels = static_cast<typename thrust::iterator_value<InputIterator>::type>(num_labels),
+     offsets    = offsets_begin,
+     output     = label_begin] __device__(auto const idx) {
+      // Zero-normalized offsets.
+      auto const dst_idx = idx - (*offsets);
+
+      // Scatter value `1` to the index at (idx - offsets[0]).
+      // Note that we need to check for out of bound, since the offset values may be invalid due to
+      // empty segments at the end.
+      // In case we have repeated offsets (i.e., we have empty segments), this `atomicAdd` call will
+      // make sure the label values corresponding to these empty segments will be skipped in the
+      // output.
+      if (dst_idx < num_labels) { atomicAdd(&output[dst_idx], OutputType{1}); }
+    });
   thrust::inclusive_scan(rmm::exec_policy(stream), label_begin, label_end, label_begin);
 }