Merge pull request #11502 from rapidsai/branch-22.08

[gpuCI] Forward-merge branch-22.08 to branch-22.10 [skip gpuci]

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

@@ -69,10 +69,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;
@@ -85,19 +87,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);
 }