Repository: rapidsai/cudf (language: python)

Issue:
[BUG] Integer overflow in compiled binops for large tables.
Consider:

```c++
#include <cudf/binaryop.hpp>
#include <cudf/column/column_factories.hpp>
#include <cudf/scalar/scalar.hpp>
#include <cudf/scalar/scalar_factories.hpp>
#include <cudf/types.hpp>
#include <vector>

int main(void) {

  cudf::size_type num_rows = 1799989091;
  auto s = cudf::make_default_constructed_scalar(
      cudf::data_type(cudf::type_id::TIMESTAMP_DAYS));
  auto a = cudf::make_column_from_scalar(*s, num_rows);
  auto b = cudf::make_column_from_scalar(*s, num_rows);
  auto mask = cudf::binary_operation(a->view(), b->view(),
                                     cudf::binary_operator::GREATER,
                                     cudf::data_type(cudf::type_id::BOOL8));
}
```

Run under compute-sanitizer, this produces:
```
========= Invalid __global__ read of size 4 bytes
=========     at void cudf::binops::compiled::for_each_kernel<cudf::binops::compiled::binary_op_device_dispatcher<cudf::binops::compiled::ops::Greater>>(int, T1)+0x220d0
=========     by thread (739,0,0) in block (452470,0,0)
=========     Address 0x77685800338c is out of bounds
=========     and is 6475992180 bytes before the nearest allocation at 0x7769da000000 of size 1799989104 bytes
=========     Saved host backtrace up to driver entry point at kernel launch time
=========     Host Frame: [0x2f26ef]
=========                in /usr/lib/x86_64-linux-gnu/libcuda.so.1
=========     Host Frame: [0x15a13]
=========                in /home/coder/.conda/envs/rapids/lib/libcudart.so.12
=========     Host Frame:cudaLaunchKernel [0x75750]
=========                in /home/coder/.conda/envs/rapids/lib/libcudart.so.12
=========     Host Frame:void cudf::binops::compiled::for_each<cudf::binops::compiled::binary_op_device_dispatcher<cudf::binops::compiled::ops::Greater> >(rmm::cuda_stream_view, int, cudf::binops::compiled::binary_op_device_dispatcher<cudf::binops::compiled::ops::Greater>) [0x4385b9]
=========                in /home/coder/cudf/cpp/build/conda/cuda-12.5/release/libcudf.so
=========     Host Frame:void cudf::binops::compiled::apply_binary_op<cudf::binops::compiled::ops::Greater>(cudf::mutable_column_view&, cudf::column_view const&, cudf::column_view const&, bool, bool, rmm::cuda_stream_view) [0x43946b]
=========                in /home/coder/cudf/cpp/build/conda/cuda-12.5/release/libcudf.so
```

The relevant part of the binops code is:

```c++
template <typename Functor>
CUDF_KERNEL void for_each_kernel(cudf::size_type size, Functor f)
{
  int tid    = threadIdx.x;
  int blkid  = blockIdx.x;
  int blksz  = blockDim.x;
  int gridsz = gridDim.x;

  int start = tid + blkid * blksz;
  int step  = blksz * gridsz;

#pragma unroll
  for (cudf::size_type i = start; i < size; i += step) {
    f(i);
  }
}

/**
 * @brief Launches Simplified for_each kernel with maximum occupancy grid dimensions.
 *
 * @tparam Functor
 * @param stream CUDA stream used for device memory operations and kernel launches.
 * @param size number of elements to process.
 * @param f Functor object to call for each element.
 */
template <typename Functor>
void for_each(rmm::cuda_stream_view stream, cudf::size_type size, Functor f)
{
  int block_size;
  int min_grid_size;
  CUDF_CUDA_TRY(
    cudaOccupancyMaxPotentialBlockSize(&min_grid_size, &block_size, for_each_kernel<decltype(f)>));
  // 2 elements per thread.
  int const grid_size = util::div_rounding_up_safe(size, 2 * block_size);
  for_each_kernel<<<grid_size, block_size, 0, stream.value()>>>(size, std::forward<Functor&&>(f));
}
```

On my system, assuming the for each kernel specializes to `f := out[i] = left[i] > right[i];`, then I get:

```
block_size = 768;
grid_size = 1171868;
```
and
```
// This is used as the step in the grid-stride loop
step = block_size * grid_size = 899994624;
```

But, `num_rows + step > std::numeric_limits<size_type>::max()`, so the grid-stride loop overflows in the increment step.



Discussion / hints:
This is a bug of the same flavor as many we've fixed in https://github.com/rapidsai/cudf/issues/10368. This one just hadn't been addressed yet.

Fix patch:
diff --git a/cpp/src/binaryop/compiled/binary_ops.cuh b/cpp/src/binaryop/compiled/binary_ops.cuh
index c6af0c3c58a..06987139188 100644
--- a/cpp/src/binaryop/compiled/binary_ops.cuh
+++ b/cpp/src/binaryop/compiled/binary_ops.cuh
@@ -21,7 +21,7 @@
 
 #include <cudf/column/column_device_view.cuh>
 #include <cudf/column/column_view.hpp>
-#include <cudf/detail/utilities/integer_utils.hpp>
+#include <cudf/detail/utilities/cuda.cuh>
 #include <cudf/unary.hpp>
 
 #include <rmm/cuda_stream_view.hpp>
@@ -253,16 +253,11 @@ struct binary_op_double_device_dispatcher {
 template <typename Functor>
 CUDF_KERNEL void for_each_kernel(cudf::size_type size, Functor f)
 {
-  int tid    = threadIdx.x;
-  int blkid  = blockIdx.x;
-  int blksz  = blockDim.x;
-  int gridsz = gridDim.x;
-
-  int start = tid + blkid * blksz;
-  int step  = blksz * gridsz;
+  auto start        = cudf::detail::grid_1d::global_thread_id();
+  auto const stride = cudf::detail::grid_1d::grid_stride();
 
 #pragma unroll
-  for (cudf::size_type i = start; i < size; i += step) {
+  for (auto i = start; i < size; i += stride) {
     f(i);
   }
 }
@@ -282,9 +277,9 @@ void for_each(rmm::cuda_stream_view stream, cudf::size_type size, Functor f)
   int min_grid_size;
   CUDF_CUDA_TRY(
     cudaOccupancyMaxPotentialBlockSize(&min_grid_size, &block_size, for_each_kernel<decltype(f)>));
-  // 2 elements per thread.
-  int const grid_size = util::div_rounding_up_safe(size, 2 * block_size);
-  for_each_kernel<<<grid_size, block_size, 0, stream.value()>>>(size, std::forward<Functor&&>(f));
+  auto grid = cudf::detail::grid_1d(size, block_size, 2 /* elements_per_thread */);
+  for_each_kernel<<<grid.num_blocks, grid.num_threads_per_block, 0, stream.value()>>>(
+    size, std::forward<Functor&&>(f));
 }
 
 template <class BinaryOperator>


Test patch:
diff --git a/cpp/tests/binaryop/binop-compiled-test.cpp b/cpp/tests/binaryop/binop-compiled-test.cpp
index 3bd67001c16..7cdce1ff735 100644
--- a/cpp/tests/binaryop/binop-compiled-test.cpp
+++ b/cpp/tests/binaryop/binop-compiled-test.cpp
@@ -23,9 +23,11 @@
 #include <cudf_test/testing_main.hpp>
 #include <cudf_test/type_lists.hpp>
 
+#include <cudf/aggregation.hpp>
 #include <cudf/binaryop.hpp>
 #include <cudf/detail/iterator.cuh>
 #include <cudf/fixed_point/fixed_point.hpp>
+#include <cudf/reduction.hpp>
 #include <cudf/types.hpp>
 
 #include <thrust/iterator/counting_iterator.h>
@@ -820,4 +822,24 @@ TEST_F(BinaryOperationCompiledTest_NullOpsString, NullMin_Vector_Vector)
   CUDF_TEST_EXPECT_COLUMNS_EQUIVALENT(expected, result->view());
 }
 
+TEST(BinaryOperationCompiledTest, LargeColumnNoOverflow)
+{
+  cudf::size_type num_rows{1'799'989'091};
+  auto big = cudf::make_column_from_scalar(
+    cudf::numeric_scalar<cudf::id_to_type<cudf::type_id::INT8>>{10, true}, num_rows);
+  auto small = cudf::make_column_from_scalar(
+    cudf::numeric_scalar<cudf::id_to_type<cudf::type_id::INT8>>{1, true}, num_rows);
+
+  auto mask = cudf::binary_operation(big->view(),
+                                     small->view(),
+                                     cudf::binary_operator::GREATER,
+                                     cudf::data_type{cudf::type_id::BOOL8});
+
+  auto agg = cudf::make_sum_aggregation<cudf::reduce_aggregation>();
+  auto result =
+    cudf::reduce(mask->view(), *agg, cudf::data_type{cudf::type_to_id<cudf::size_type>()});
+  auto got = static_cast<cudf::numeric_scalar<cudf::size_type>*>(result.get())->value();
+  EXPECT_EQ(num_rows, got);
+}
+
 CUDF_TEST_PROGRAM_MAIN()