From dccb586111651b5e0fa1ad7845d73f593670fa07 Mon Sep 17 00:00:00 2001
From: Nghia Truong <nghiatruong.vn@gmail.com>
Date: Tue, 9 Aug 2022 12:44:17 -0700
Subject: [PATCH] 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 https://github.com/NVIDIA/spark-rapids/issues/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 https://github.com/rapidsai/cudf/issues/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)
---
 .../cudf/detail/labeling/label_segments.cuh   | 36 +++++++++++--------
 1 file changed, 22 insertions(+), 14 deletions(-)

diff --git a/cpp/include/cudf/detail/labeling/label_segments.cuh b/cpp/include/cudf/detail/labeling/label_segments.cuh
index e30f5b3ee91..5a901cc4e3f 100644
--- a/cpp/include/cudf/detail/labeling/label_segments.cuh
+++ b/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);
 }