Update sort groupby to use non-atomic operation #9035
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -21,20 +21,133 @@ | |
| #include <cudf/column/column_view.hpp> | ||
| #include <cudf/detail/aggregation/aggregation.cuh> | ||
| #include <cudf/detail/iterator.cuh> | ||
| #include <cudf/detail/valid_if.cuh> | ||
| #include <cudf/types.hpp> | ||
| #include <cudf/utilities/span.hpp> | ||
|
|
||
| #include <rmm/cuda_stream_view.hpp> | ||
| #include <rmm/exec_policy.hpp> | ||
|
|
||
| #include <thrust/functional.h> | ||
| #include <thrust/iterator/counting_iterator.h> | ||
| #include <thrust/iterator/discard_iterator.h> | ||
| #include <thrust/iterator/transform_output_iterator.h> | ||
| #include <thrust/iterator/zip_iterator.h> | ||
| #include <thrust/reduce.h> | ||
|
|
||
| namespace cudf { | ||
| namespace groupby { | ||
| namespace detail { | ||
|
|
||
| // ArgMin binary operator with tuple of (value, index) | ||
| template <typename T> | ||
| struct ArgMin { | ||
| CUDA_HOST_DEVICE_CALLABLE auto operator()(thrust::tuple<T, size_type> const& lhs, | ||
| thrust::tuple<T, size_type> const& rhs) const | ||
| { | ||
| if (thrust::get<1>(lhs) == cudf::detail::ARGMIN_SENTINEL) | ||
| return rhs; | ||
| else if (thrust::get<1>(rhs) == cudf::detail::ARGMIN_SENTINEL) | ||
| return lhs; | ||
| else | ||
| return thrust::get<0>(lhs) < thrust::get<0>(rhs) ? lhs : rhs; | ||
| } | ||
| }; | ||
|
|
||
| // ArgMax binary operator with tuple of (value, index) | ||
| template <typename T> | ||
| struct ArgMax { | ||
| CUDA_HOST_DEVICE_CALLABLE auto operator()(thrust::tuple<T, size_type> const& lhs, | ||
| thrust::tuple<T, size_type> const& rhs) const | ||
| { | ||
| if (thrust::get<1>(lhs) == cudf::detail::ARGMIN_SENTINEL) | ||
| return rhs; | ||
| else if (thrust::get<1>(rhs) == cudf::detail::ARGMIN_SENTINEL) | ||
| return lhs; | ||
| else | ||
| return thrust::get<0>(lhs) > thrust::get<0>(rhs) ? lhs : rhs; | ||
| } | ||
| }; | ||
|
|
||
| struct get_tuple_second_element { | ||
| template <typename T> | ||
| __device__ size_type operator()(thrust::tuple<T, size_type> const& rhs) const | ||
| { | ||
| return thrust::get<1>(rhs); | ||
| } | ||
| }; | ||
|
|
||
| /** | ||
| * @brief Functor to store the boolean value to null mask. | ||
| */ | ||
| struct bool_to_nullmask { | ||
| mutable_column_device_view d_result; | ||
| __device__ void operator()(size_type i, bool rhs) | ||
| { | ||
| if (rhs) { | ||
| d_result.set_valid(i); | ||
| } else { | ||
| d_result.set_null(i); | ||
| } | ||
| } | ||
| }; | ||
|
|
||
| /** | ||
| * @brief Returns index for non-null element, and SENTINEL for null element in a column. | ||
| * | ||
| */ | ||
| struct null_as_sentinel { | ||
| column_device_view const col; | ||
| size_type const SENTINEL; | ||
| __device__ size_type operator()(size_type i) const { return col.is_null(i) ? SENTINEL : i; } | ||
| }; | ||
|
|
||
| /** | ||
| * @brief Value accessor for column which supports dictionary column too. | ||
| * | ||
| * @tparam T Type of the underlying column. For dictionary column, type of the key column. | ||
| */ | ||
| template <typename T> | ||
| struct value_accessor { | ||
| column_device_view const col; | ||
| bool const is_dict; | ||
| value_accessor(column_device_view const& col) : col(col), is_dict(cudf::is_dictionary(col.type())) | ||
| { | ||
| } | ||
|
|
||
| __device__ T value(size_type i) const | ||
| { | ||
| if (is_dict) { | ||
| auto keys = col.child(dictionary_column_view::keys_column_index); | ||
| return keys.element<T>(static_cast<size_type>(col.element<dictionary32>(i))); | ||
| } else { | ||
| return col.element<T>(i); | ||
| } | ||
| } | ||
| __device__ auto operator()(size_type i) const { return value(i); } | ||
| }; | ||
|
|
||
| /** | ||
| * @brief Null replaced value accessor for column which supports dictionary column too. | ||
| * For null value, returns null `init` value | ||
| * | ||
| * @tparam T Type of the underlying column. For dictionary column, type of the key column. | ||
| */ | ||
| template <typename T> | ||
| struct null_replaced_value_accessor : value_accessor<T> { | ||
|
There was a problem hiding this comment. Same question here. Does the existing There was a problem hiding this comment.
There was a problem hiding this comment. Can you use this? cudf/cpp/include/cudf/dictionary/detail/iterator.cuh Lines 110 to 112 in 8b02ca3 There was a problem hiding this comment. It would be better to use the indices for any cudf operations where possible for both run-time and compile-time performance. For example, sorting in general only needs the indices. to get the indices column_view decorated with the offset, size, and validity-mask appropriately. There was a problem hiding this comment. hash groupby produces base type column as output. There was a problem hiding this comment. This sounds correct to me. Aggregates like min/max return values that already exist in the column so the output would have the same keys as the input. Whereas, sum/prod create totally new values. There was a problem hiding this comment. Also, here is an example using the the dictionary-pair-iterator along with a null-replacement transformer. cudf/cpp/src/reductions/simple.cuh Lines 142 to 146 in f0fa255 I'm inclined to prefer your approach here instead since it simplifies the caller to one value-accessor. The only thing that makes me nervous is that col.element<dictionary32>(i) would be included/inlined for every type and that function contains it's own type-dispatcher call in it. But technically every type is potentially a dictionary key type so I think the same amount of code is generated either way. Anyway, it may be worth looking into using this null-replacement accessor in the reductions code too.
There was a problem hiding this comment.
There was a problem hiding this comment. Dictionary index types can technically be any unsigned integer type. The |
||
| using super_t = value_accessor<T>; | ||
| bool const has_nulls; | ||
| T const init; | ||
| null_replaced_value_accessor(column_device_view const& col, T const& init, bool const has_nulls) | ||
| : super_t(col), init(init), has_nulls(has_nulls) | ||
| { | ||
| } | ||
| __device__ T operator()(size_type i) const | ||
| { | ||
| return has_nulls && super_t::col.is_null_nocheck(i) ? init : super_t::value(i); | ||
| } | ||
| }; | ||
|
|
||
| template <aggregation::Kind K> | ||
| struct reduce_functor { | ||
| template <typename T> | ||
|
|
@@ -61,51 +174,71 @@ struct reduce_functor { | |
| rmm::cuda_stream_view stream, | ||
| rmm::mr::device_memory_resource* mr) | ||
| { | ||
| using DeviceType = device_storage_type_t<T>; | ||
| using OpType = cudf::detail::corresponding_operator_t<K>; | ||
| using ResultType = cudf::detail::target_type_t<T, K>; | ||
| using ResultDType = device_storage_type_t<ResultType>; | ||
|
|
||
| auto result_type = is_fixed_point<ResultType>() | ||
| ? data_type{type_to_id<ResultType>(), values.type().scale()} | ||
| : data_type{type_to_id<ResultType>()}; | ||
|
|
||
| std::unique_ptr<column> result = | ||
| make_fixed_width_column(result_type, num_groups, mask_state::UNALLOCATED, stream, mr); | ||
|
|
||
| if (values.is_empty()) { return result; } | ||
|
|
||
| auto resultview = mutable_column_device_view::create(result->mutable_view(), stream); | ||
| auto valuesview = column_device_view::create(values, stream); | ||
| if constexpr (K == aggregation::ARGMAX || K == aggregation::ARGMIN) { | ||
| constexpr auto SENTINEL = | ||
| (K == aggregation::ARGMAX ? cudf::detail::ARGMAX_SENTINEL : cudf::detail::ARGMIN_SENTINEL); | ||
| auto idx_begin = | ||
| cudf::detail::make_counting_transform_iterator(0, null_as_sentinel{*valuesview, SENTINEL}); | ||
| // dictionary keys are sorted, so dictionary32 index comparison is enough. | ||
| auto column_begin = valuesview->begin<DeviceType>(); | ||
| auto begin = thrust::make_zip_iterator(thrust::make_tuple(column_begin, idx_begin)); | ||
| auto result_begin = thrust::make_transform_output_iterator(resultview->begin<ResultDType>(), | ||
| get_tuple_second_element{}); | ||
| using OpType = | ||
| std::conditional_t<(K == aggregation::ARGMAX), ArgMax<DeviceType>, ArgMin<DeviceType>>; | ||
| thrust::reduce_by_key(rmm::exec_policy(stream), | ||
| group_labels.data(), | ||
| group_labels.data() + group_labels.size(), | ||
| begin, | ||
| thrust::make_discard_iterator(), | ||
| result_begin, | ||
| thrust::equal_to<size_type>{}, | ||
| OpType{}); | ||
| } else { | ||
| auto init = OpType::template identity<DeviceType>(); | ||
| auto begin = cudf::detail::make_counting_transform_iterator( | ||
| 0, null_replaced_value_accessor{*valuesview, init, values.has_nulls()}); | ||
| thrust::reduce_by_key(rmm::exec_policy(stream), | ||
| group_labels.data(), | ||
| group_labels.data() + group_labels.size(), | ||
| begin, | ||
| thrust::make_discard_iterator(), | ||
| resultview->begin<ResultDType>(), | ||
| thrust::equal_to<size_type>{}, | ||
| OpType{}); | ||
| } | ||
|
|
||
| if (values.has_nulls()) { | ||
| rmm::device_uvector<bool> validity(num_groups, stream); | ||
| thrust::reduce_by_key(rmm::exec_policy(stream), | ||
| group_labels.data(), | ||
| group_labels.data() + group_labels.size(), | ||
| cudf::detail::make_validity_iterator(*valuesview), | ||
| thrust::make_discard_iterator(), | ||
| validity.begin(), | ||
| thrust::equal_to<size_type>{}, | ||
| thrust::logical_or<bool>{}); | ||
| auto [null_mask, null_count] = cudf::detail::valid_if( | ||
| validity.begin(), validity.end(), thrust::identity<bool>{}, stream, mr); | ||
| result->set_null_mask(std::move(null_mask)); | ||
| result->set_null_count(null_count); | ||
| } | ||
| return result; | ||
| } | ||
|
|
||
|
|
||
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Can't
null_replacement_iteratorbe used instead?cudf/cpp/include/cudf/detail/iterator.cuh
Line 162 in dfe0a03
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Can't.
null_replacement_iteratorreturns values of the column. Here, indices are needed.