Skip to content

Commit

Permalink
Fix some libcudf calls to cudf::detail::gather (#11963)
Browse files Browse the repository at this point in the history
Fixes a couple source files that were calling gather by type-dispatching directly to the internal `column_gatherer` functor instead of using the `cudf::detail::gather` function(s). This simplifies the code and improves maintenance. For example, extra code to resolve the null-mask is eliminated since the appropriate `cudf::detail::gather` call does this automatically.
No function has changed, just code cleanup.

Authors:
  - David Wendt (https://github.com/davidwendt)

Approvers:
  - Mark Harris (https://github.com/harrism)
  - Nghia Truong (https://github.com/ttnghia)
  - Karthikeyan (https://github.com/karthikeyann)

URL: #11963
  • Loading branch information
davidwendt authored Oct 26, 2022
1 parent c146d21 commit fac35b4
Show file tree
Hide file tree
Showing 2 changed files with 19 additions and 39 deletions.
41 changes: 11 additions & 30 deletions cpp/src/lists/copying/gather.cu
Original file line number Diff line number Diff line change
Expand Up @@ -100,36 +100,17 @@ std::unique_ptr<column> gather_list_leaf(column_view const& column,
size_type gather_map_size = gd.gather_map_size;

// call the normal gather
auto leaf_column = cudf::type_dispatcher<dispatch_storage_type>(
column.type(),
cudf::detail::column_gatherer{},
column,
gather_map_begin,
gather_map_begin + gather_map_size,
// note : we don't need to bother checking for out-of-bounds here since
// our inputs at this stage aren't coming from the user.
false,
stream,
mr);

// the column_gatherer doesn't create the null mask because it expects
// that will be done in the gather_bitmask() step. however, gather_bitmask()
// only happens at the root level, and by definition this column is a
// leaf. so we have to generate the bitmask ourselves.
// TODO : it might make sense to expose a gather() function that takes a column_view and
// returns a column that does this work correctly.
size_type null_count = column.null_count();
if (null_count > 0) {
auto list_cdv = column_device_view::create(column, stream);
auto validity = cudf::detail::valid_if(
gather_map_begin,
gather_map_begin + gd.gather_map_size,
[cdv = *list_cdv] __device__(int index) { return cdv.is_valid(index) ? true : false; },
stream,
mr);

leaf_column->set_null_mask(std::move(validity.first), validity.second);
}
// note : we don't need to bother checking for out-of-bounds here since
// our inputs at this stage aren't coming from the user.
auto gather_table = cudf::detail::gather(cudf::table_view({column}),
gather_map_begin,
gather_map_begin + gather_map_size,
out_of_bounds_policy::DONT_CHECK,
stream,
mr);
auto leaf_column = std::move(gather_table->release().front());

if (column.null_count() == 0) { leaf_column->set_null_mask(rmm::device_buffer{}, 0); }

return leaf_column;
}
Expand Down
17 changes: 8 additions & 9 deletions cpp/src/partitioning/partitioning.cu
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,7 @@
#include <cub/cub.cuh>
#include <cudf/column/column_factories.hpp>
#include <cudf/copying.hpp>
#include <cudf/detail/gather.hpp>
#include <cudf/detail/nvtx/ranges.hpp>
#include <cudf/detail/scatter.cuh>
#include <cudf/detail/utilities/cuda.cuh>
Expand Down Expand Up @@ -436,15 +437,13 @@ struct copy_block_partitions_dispatcher {
grid_size,
stream);

// Use gather instead for non-fixed width types
return type_dispatcher(input.type(),
detail::column_gatherer{},
input,
gather_map.begin(),
gather_map.end(),
false,
stream,
mr);
auto gather_table = cudf::detail::gather(cudf::table_view({input}),
gather_map,
out_of_bounds_policy::DONT_CHECK,
cudf::detail::negative_index_policy::NOT_ALLOWED,
stream,
mr);
return std::move(gather_table->release().front());
}
};

Expand Down

0 comments on commit fac35b4

Please sign in to comment.