Skip to content

Replace hand-rolled output iterators with cuda::tabulate_output_iterator #23856

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
rapids-bot merged 1 commit into from
Aug 28, 2026
+29 −96
Merged

Replace hand-rolled output iterators with cuda::tabulate_output_iterator #23856

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
76 changes: 20 additions & 56 deletions cpp/src/copying/contiguous_split.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -895,57 +895,19 @@ struct split_key_functor {
};

/**
* @brief Output iterator for writing values to the dst_offset field of the
* dst_buf_info struct
* @brief Writes values to the dst_offset field of the dst_buf_info struct
*/
struct dst_offset_output_iterator {
struct set_dst_offset_fn {
dst_buf_info* c;
using value_type = std::size_t;
using difference_type = std::size_t;
using pointer = std::size_t*;
using reference = std::size_t&;
using iterator_category = thrust::output_device_iterator_tag;

dst_offset_output_iterator operator+ __host__ __device__(int i) { return {c + i}; }

dst_offset_output_iterator& operator++ __host__ __device__()
{
c++;
return *this;
}

reference operator[] __device__(int i) { return dereference(c + i); }
reference operator* __device__() { return dereference(c); }

private:
reference __device__ dereference(dst_buf_info* c) { return c->dst_offset; }
__device__ void operator()(size_type i, std::size_t value) const { c[i].dst_offset = value; }
};

/**
* @brief Output iterator for writing values to the valid_count field of the
* dst_buf_info struct
* @brief Writes values to the valid_count field of the dst_buf_info struct
*/
struct dst_valid_count_output_iterator {
struct set_valid_count_fn {
dst_buf_info* c;
using value_type = size_type;
using difference_type = size_type;
using pointer = size_type*;
using reference = size_type&;
using iterator_category = thrust::output_device_iterator_tag;

dst_valid_count_output_iterator operator+ __host__ __device__(int i) { return {c + i}; }

dst_valid_count_output_iterator& operator++ __host__ __device__()
{
c++;
return *this;
}

reference operator[] __device__(int i) { return dereference(c + i); }
reference operator* __device__() { return dereference(c); }

private:
reference __device__ dereference(dst_buf_info* c) { return c->valid_count; }
__device__ void operator()(size_type i, size_type value) const { c[i].valid_count = value; }
};

/**
Expand Down Expand Up @@ -1357,12 +1319,13 @@ std::unique_ptr<packed_partition_buf_size_and_dst_buf_info> compute_splits(
auto values =
cudf::detail::make_counting_transform_iterator(0, buf_size_functor{d_dst_buf_info});

thrust::exclusive_scan_by_key(rmm::exec_policy_nosync(stream, temp_mr),
keys,
keys + num_bufs,
values,
dst_offset_output_iterator{d_dst_buf_info},
std::size_t{0});
thrust::exclusive_scan_by_key(
rmm::exec_policy_nosync(stream, temp_mr),
keys,
keys + num_bufs,
values,
cuda::make_tabulate_output_iterator(set_dst_offset_fn{d_dst_buf_info}),
std::size_t{0});
}

partition_buf_size_and_dst_buf_info->copy_to_host();
Expand Down Expand Up @@ -1882,12 +1845,13 @@ struct contiguous_split_state {
cuda::proclaim_return_type<size_type>(
[] __device__(dst_buf_info const& info) { return info.valid_count; }));

thrust::reduce_by_key(rmm::exec_policy_nosync(stream, temp_mr),
keys,
keys + num_batches_total,
values,
cuda::make_discard_iterator(),
dst_valid_count_output_iterator{d_orig_dst_buf_info.data()});
thrust::reduce_by_key(
rmm::exec_policy_nosync(stream, temp_mr),
keys,
keys + num_batches_total,
values,
cuda::make_discard_iterator(),
cuda::make_tabulate_output_iterator(set_valid_count_fn{d_orig_dst_buf_info.data()}));

detail::cuda_memcpy<dst_buf_info>(h_orig_dst_buf_info, d_orig_dst_buf_info, stream);

Expand Down
6 changes: 3 additions & 3 deletions cpp/src/io/parquet/reader_impl_preprocess.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -696,7 +696,7 @@ void reader_impl::generate_list_column_row_counts(is_estimate_row_counts is_esti
key_input,
key_input + pass.pages.size(),
page_input,
chunk_row_output_iter{pass.pages.device_ptr()});
cuda::make_tabulate_output_iterator(set_chunk_row_fn{pass.pages.device_ptr()}));

// To compensate for the list row size estimates, force the row count on the last page for each
// column chunk (each rowgroup) such that it ends on the real known row count. this is so that
Expand Down Expand Up @@ -810,7 +810,7 @@ void reader_impl::preprocess_subpass_pages(read_mode mode, size_t chunk_read_lim
key_input,
key_input + pass.pages.size(),
page_input,
chunk_row_output_iter{pass.pages.device_ptr()});
cuda::make_tabulate_output_iterator(set_chunk_row_fn{pass.pages.device_ptr()}));

// copy chunk_row into the subpass pages
// only need to do this if we are not processing the whole pass in one subpass
Expand Down Expand Up @@ -1149,7 +1149,7 @@ cudf::detail::host_vector<size_t> reader_impl::calculate_page_string_offsets()
page_keys,
page_keys + subpass.pages.size(),
val_iter,
page_offset_output_iter{subpass.pages.device_ptr()});
cuda::make_tabulate_output_iterator(set_str_offset_fn{subpass.pages.device_ptr()}));

// now sum up page sizes
rmm::device_uvector<int> reduce_keys(d_col_sizes.size(), _stream);
Expand Down
43 changes: 6 additions & 37 deletions cpp/src/io/parquet/reader_impl_preprocess_utils.cuh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -387,24 +387,12 @@ struct get_reduction_key {
/**
* @brief Writes to the chunk_row field of the PageInfo struct
*/
struct chunk_row_output_iter {
struct set_chunk_row_fn {
PageInfo* p;
using value_type = size_type;
using difference_type = size_type;
using pointer = size_type*;
using reference = size_type&;
using iterator_category = thrust::output_device_iterator_tag;

CUDF_HOST_DEVICE constexpr inline chunk_row_output_iter operator+(int i) const { return {p + i}; }

CUDF_HOST_DEVICE constexpr inline chunk_row_output_iter& operator++()
__device__ constexpr void operator()(size_type i, size_type value) const
{
p++;
return *this;
p[i].chunk_row = value;
}

__device__ constexpr inline reference operator[](int i) { return p[i].chunk_row; }
__device__ constexpr inline reference operator*() { return p->chunk_row; }
};

/**
Expand Down Expand Up @@ -479,30 +467,11 @@ struct page_to_string_size {
};

/**
* @brief Functor to access and update the str_offset field of the PageInfo struct
* @brief Writes to the str_offset field of the PageInfo struct
*/
struct page_offset_output_iter {
struct set_str_offset_fn {
PageInfo* p;

using value_type = size_t;
using difference_type = size_t;
using pointer = size_t*;
using reference = size_t&;
using iterator_category = thrust::output_device_iterator_tag;

CUDF_HOST_DEVICE constexpr inline page_offset_output_iter operator+(int i) const
{
return {p + i};
}

CUDF_HOST_DEVICE constexpr inline page_offset_output_iter& operator++()
{
p++;
return *this;
}

__device__ constexpr inline reference operator[](int i) { return p[i].str_offset; }
__device__ constexpr inline reference operator*() { return p->str_offset; }
__device__ constexpr void operator()(size_type i, size_t value) const { p[i].str_offset = value; }
};

/**
Expand Down
Loading