Prevent grid stride loop overflow in libcudf kernels

[nvdbaranec]

(updated Aug 2023), (updated Nov 2024)

Background

We found a kernel indexing overflow issue, first discovered in the fused_concatenate kernels (#10333) and this issue is present in a number of our CUDA kernels that take the following form:

size_type output_index = threadIdx.x + blockIdx.x * blockDim.x;  
while (output_index < output_size) {
  output_index += blockDim.x * gridDim.x;
}

If we have an output_size of say 1.2 billion and a grid size that's the same, the following happens: Some late thread id, say 1.19 billion attempts to add 1.2 billion (blockDim.x * gridDim.x) and overflows the size_type (signed 32 bits).

We made a round of fixes in #10448, and then later found another instance of this error in #13838. Our first pass of investigation was not adequate to contain the issue, so we need to take another close look.

Part 1 - First pass fix kernels with this issue

Source file	Kernels	Status
copying/concatenate.cu	fused_concatenate_kernel	#10448
valid_if.cuh	valid_if_kernel	#10448
scatter.cu	marking_bitmask_kernel	#10448
replace/nulls.cu	replace_nulls_strings	#10448
replace/nulls.cu	replace_nulls	#10448
rolling/rolling_detail.cuh	gpu_rolling	#10448
rolling/jit/kernel.cu	gpu_rolling_new	#10448
transform/compute_column.cu	compute_column_kernel	#10448
copying/concatenate.cu	fused_concatenate_string_offset_kernel	#13838
replace/replace.cu	replace_strings_first_pass replace_strings_second_pass replace_kernel	#13905
copying/concatenate.cu	concatenate_masks_kernel fused_concatenate_string_offset_kernel fused_concatenate_string_chars_kernel fused_concatenate_kernel (int64)	#13906
hash/helper_functions.cuh	init_hashtbl	#13895
null_mask.cu	set_null_mask_kernel copy_offset_bitmask count_set_bits_kernel	#13895
transform/row_bit_count.cu	compute_row_sizes	#13895
multibyte_split.cu	multibyte_split_init_kernel multibyte_split_seed_kernel (auto??) multibyte_split_kernel	#13910
IO modules: parquet, orc, json		#13910
io/utilities/parsing_utils.cu	count_and_set_positions (uint64_t)	#13910
conditional_join_kernels.cuh	compute_conditional_join_output_size conditional_join	#13971
merge.cu	materialize_merged_bitmask_kernel	#13972
partitioning.cu	compute_row_partition_numbers compute_row_output_locations copy_block_partitions	#13973
json_path.cu	get_json_object_kernel	#13962
tdigest	compute_percentiles_kernel (int)	#13962
strings/attributes.cu	count_characters_parallel_fn	#13968
strings/convert/convert_urls.cu	url_decode_char_counter (int) url_decode_char_replacer (int)	#13968
text/subword/data_normalizer.cu	kernel_data_normalizer (uint32_t)	#13915
text/subword/subword_tokenize.cu	kernel_compute_tensor_metadata (uint32_t)	#13915
text/subword/wordpiece_tokenizer.cu	init_data_and_mark_word_start_and_ends (uint32_t) mark_string_start_and_ends (uint32_t) kernel_wordpiece_tokenizer (uint32_t)	#13915

Part 2 - Take another pass over more challenging kernels

Source file	Kernels	Status
null_mask.cuh	subtract_set_bits_range_boundaries_kernel
valid_if.cuh	valid_if_n_kernel
copy_if_else.cuh	copy_if_else_kernel
gather.cuh	gather_chars_fn_string_parallel
more?	search gridDim.x or blockDim.x to find more examples

Part 2b - results of searching

Source file: line content	Overflow risk (yes/no)	Status
cpp/src/binaryop/compiled/binary_ops.cuh	yes	#17354
cpp/src/binaryop/jit/kernel.cu	yes	#17420
include/cudf/detail/copy_if.cuh: int tid = threadIdx.x + per_thread * block_size * blockIdx.x; include/cudf/detail/copy_if.cuh: int tid = threadIdx.x + per_thread * block_size * blockIdx.x; include/cudf/detail/copy_if.cuh: int const wid = threadIdx.x / cudf::detail::warp_size; include/cudf/detail/copy_if.cuh: int const lane = threadIdx.x % cudf::detail::warp_size;		17520
include/cudf/detail/copy_if_else.cuh: int const lane_id = threadIdx.x % cudf::detail::warp_size;	no	17473
include/cudf/detail/copy_range.cuh: int const lane_id = threadIdx.x % warp_size; include/cudf/detail/copy_range.cuh: cudf::size_type const tid = threadIdx.x + blockIdx.x * blockDim.x; include/cudf/detail/copy_range.cuh: source_idx += blockDim.x * gridDim.x;	yes	#17409
include/cudf/detail/null_mask.cuh: auto const tid = threadIdx.x + blockIdx.x * blockDim.x; include/cudf/detail/null_mask.cuh: destination_word_index += blockDim.x * gridDim.x) { include/cudf/detail/null_mask.cuh: size_type const tid = threadIdx.x + blockIdx.x * blockDim.x; include/cudf/detail/null_mask.cuh: range_id += blockDim.x * gridDim.x;	no
include/cudf/detail/valid_if.cuh: auto block_offset = blockIdx.x * blockDim.x; include/cudf/detail/valid_if.cuh: auto const thread_idx = block_offset + threadIdx.x; include/cudf/detail/valid_if.cuh: block_offset += blockDim.x * gridDim.x;	no
include/cudf/strings/detail/gather.cuh: int global_thread_id = blockIdx.x * blockDim.x + threadIdx.x; include/cudf/strings/detail/gather.cuh: int nwarps = gridDim.x * blockDim.x / cudf::detail::warp_size; include/cudf/strings/detail/gather.cuh: size_type begin_out_string_idx = blockIdx.x * strings_per_threadblock;		#17404
src/join/mixed_join_kernel.cuh: cudf::size_type outer_row_index = threadIdx.x + blockIdx.x * block_size; src/join/mixed_join_size_kernel.cuh: cudf::size_type const start_idx = threadIdx.x + blockIdx.x * block_size;		#17633
src/partitioning/partitioning.cu: size_type partition_number = threadIdx.x; src/partitioning/partitioning.cu: size_type const write_location = partition_number * gridDim.x + blockIdx.x; src/partitioning/partitioning.cu: block_partition_offsets[partition_number * gridDim.x + blockIdx.x]; src/partitioning/partitioning.cu: if (ELEMENTS_PER_THREAD * threadIdx.x + i < num_partitions) { src/partitioning/partitioning.cu: block_partition_sizes[blockIdx.x + (ELEMENTS_PER_THREAD * threadIdx.x + i) * gridDim.x]; src/partitioning/partitioning.cu: if (ELEMENTS_PER_THREAD * threadIdx.x + i < num_partitions) { src/partitioning/partitioning.cu: scanned_block_partition_sizes[ipartition * gridDim.x + blockIdx.x];	yes	#17473
src/strings/regex/utilities.cuh: auto const thread_idx = threadIdx.x + blockIdx.x * blockDim.x; src/strings/regex/utilities.cuh: auto const thread_idx = threadIdx.x + blockIdx.x * blockDim.x;		#17404
src/strings/search/find.cu: size_type const idx = static_cast<size_type>(threadIdx.x + blockIdx.x * blockDim.x); src/strings/search/find.cu: size_type const idx = static_cast<size_type>(threadIdx.x + blockIdx.x * blockDim.x);		#17404
src/text/subword/data_normalizer.cu: uint32_t* block_base = code_points + blockIdx.x * blockDim.x * MAX_NEW_CHARS;		#17404
src/text/vocabulary_tokenize.cu: auto const idx = static_cast<std::size_t>(threadIdx.x + blockIdx.x * blockDim.x);		#17404
src/transform/compute_column.cu: &intermediate_storage[threadIdx.x * device_expression_data.num_intermediates];	no
src/transform/jit/kernel.cu: thread_index_type const start = threadIdx.x + blockIdx.x * blockDim.x;<br>src/transform/jit/kernel.cu: thread_index_type const stride = blockDim.x * gridDim.x;		#17473
src/transform/row_bit_count.cu: int const tid = threadIdx.x + blockIdx.x * blockDim.x; src/transform/row_bit_count.cu: row_span* my_branch_stack = thread_branch_stacks + (threadIdx.x * max_branch_depth);		#17473

Additional information

There are also a number of kernels that have this pattern but probably don't ever overflow because they are indexing by bitmask words. (Example)
Additional, In this kernel, source_idx probably overflows, but harmlessly.

A snippet of code to see this in action:

size_type const size = 1200000000;
auto big = cudf::make_fixed_width_column(data_type{type_id::INT32}, size, mask_state::UNALLOCATED);  
auto x = cudf::rolling_window(*big, 1, 1, 1, cudf::detail::sum_aggregation{}); 

Note: rmm may mask out of bounds accesses in some cases, so it's helpful to run with the plain cuda allocator.

[nvdbaranec]

The fix in basically all of these cases is quite simple: just make the index a size_t

[jrhemstad]

I'd love to just add an algorithm to do this.

https://godbolt.org/z/hK95z7zff

[harrism]

Or even a simple range helper for for loops: https://github.com/harrism/hemi#simple-grid-stride-loops

[harrism]

The fix in basically all of these cases is quite simple: just make the index a size_t

I think the general approach should be:

1. Use an algorithm (thrust:: or std::) if possible before ever writing a custom kernel -- this way you write a per-element device functor instead, and indexing is handled for you.
2. If a custom kernel must be written, it should use device-side algorithms instead of raw loops.
3. If a raw grid-stride loop is required and an existing algorithm won't work, we should provide utilities to abstract the iteration and or the range and use auto for the type to avoid these mistakes.

[github-actions]

This issue has been labeled inactive-30d due to no recent activity in the past 30 days. Please close this issue if no further response or action is needed. Otherwise, please respond with a comment indicating any updates or changes to the original issue and/or confirm this issue still needs to be addressed. This issue will be labeled inactive-90d if there is no activity in the next 60 days.

[nvdbaranec]

Still relevant.

[harrism]

Created https://github.com/harrism/ranger as a solution to this. Needs to be moved into libcudf.

[github-actions]

This issue has been labeled inactive-30d due to no recent activity in the past 30 days. Please close this issue if no further response or action is needed. Otherwise, please respond with a comment indicating any updates or changes to the original issue and/or confirm this issue still needs to be addressed. This issue will be labeled inactive-90d if there is no activity in the next 60 days.

[nvdbaranec]

Still relevant.

[github-actions]

This issue has been labeled inactive-30d due to no recent activity in the past 30 days. Please close this issue if no further response or action is needed. Otherwise, please respond with a comment indicating any updates or changes to the original issue and/or confirm this issue still needs to be addressed. This issue will be labeled inactive-90d if there is no activity in the next 60 days.

[nvdbaranec]

Still relevant.

[github-actions]

This issue has been labeled inactive-90d due to no recent activity in the past 90 days. Please close this issue if no further response or action is needed. Otherwise, please respond with a comment indicating any updates or changes to the original issue and/or confirm this issue still needs to be addressed.

[GregoryKimball]

Thanks @harrism for creating the ranger repo! Do you think we are ready to kick off an integration with libcudf, or does ranger need more development first?

[wence-]

wrt attempting to find locations where this might be happening. In host code, clang and gcc will warn if you add -Wsign-conversion (not covered by -Wall -Wextra) under some circumstances. Unfortunately there is no such option for nvcc.

#include <cstdint>
int what(int upper)
{
  int i = 0; // no warning if this is a std::int64_t
  unsigned int stride = 10;
  while (i < upper) {
    i = i + stride; // clang warns for this, so does gcc
  }
  i = 0;
  while (i < upper) {
    i += stride; // gcc warns for this, clang does not.
  }
  return i;
}