Skip to content

Add env RLE #7908

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
gonidelis merged 4 commits into from
Mar 10, 2026
Merged

Add env RLE #7908

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
623aec3
Add env RLE
gonidelis Mar 5, 2026
40a08c5
Add missing test, fix licenses
gonidelis Mar 9, 2026
1de3972
Remove any references to determinism
gonidelis Mar 10, 2026
c30d037
Fix licenses?
gonidelis Mar 10, 2026
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
228 changes: 228 additions & 0 deletions cub/cub/device/device_run_length_encode.cuh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -21,13 +21,18 @@
#endif // no system header

#include <cub/detail/choose_offset.cuh>
#include <cub/detail/env_dispatch.cuh>
#include <cub/device/dispatch/dispatch_reduce_by_key.cuh>
#include <cub/device/dispatch/dispatch_rle.cuh>
#include <cub/device/dispatch/dispatch_streaming_reduce_by_key.cuh>
#include <cub/device/dispatch/tuning/tuning_rle_encode.cuh>

#include <thrust/iterator/constant_iterator.h>

#include <cuda/std/__execution/env.h>
#include <cuda/std/__type_traits/enable_if.h>
#include <cuda/std/__type_traits/is_same.h>

CUB_NAMESPACE_BEGIN

//! @rst
Expand Down Expand Up @@ -203,6 +208,121 @@ struct DeviceRunLengthEncode
policy_selector_t{});
}

//! @rst
//! Computes a run-length encoding of the sequence ``d_in``.
//!
//! .. versionadded:: 3.4.0
//! First appears in CUDA Toolkit 13.4.
//!
//! This is an environment-based API that allows customization of:
//!
//! - Stream: Query via ``cuda::get_stream``
//! - Memory resource: Query via ``cuda::mr::get_memory_resource``
//!
//! - For the *i*\ :sup:`th` run encountered, the first key of the run and
//! its length are written to ``d_unique_out[i]`` and ``d_counts_out[i]``, respectively.
//! - The total number of runs encountered is written to ``d_num_runs_out``.
//! - The ``==`` equality operator is used to determine whether values are equivalent
//! - In-place operations are not supported. There must be no overlap between any of the provided ranges:
//!
//! - ``[d_unique_out, d_unique_out + *d_num_runs_out)``
//! - ``[d_counts_out, d_counts_out + *d_num_runs_out)``
//! - ``[d_num_runs_out, d_num_runs_out + 1)``
//! - ``[d_in, d_in + num_items)``
//!
//! Snippet
//! +++++++++++++++++++++++++++++++++++++++++++++
//!
//! The code snippet below illustrates the run-length encoding of a sequence of ``int`` values
//! using environment-based API:
//!
//! .. literalinclude:: ../../../cub/test/catch2_test_device_run_length_encode_env_api.cu
//! :language: c++
//! :dedent:
//! :start-after: example-begin encode-env
//! :end-before: example-end encode-env
//!
//! @endrst
//!
//! @tparam InputIteratorT
//! **[inferred]** Random-access input iterator type for reading input items @iterator
//!
//! @tparam UniqueOutputIteratorT
//! **[inferred]** Random-access output iterator type for writing unique output items @iterator
//!
//! @tparam LengthsOutputIteratorT
//! **[inferred]** Random-access output iterator type for writing output counts @iterator
//!
//! @tparam NumRunsOutputIteratorT
//! **[inferred]** Output iterator type for recording the number of runs encountered @iterator
//!
//! @tparam NumItemsT
//! **[inferred]** Type of num_items
//!
//! @tparam EnvT
//! **[inferred]** Environment type (e.g., `cuda::std::execution::env<...>`)
//!
//! @param[in] d_in
//! Pointer to the input sequence of keys
//!
//! @param[out] d_unique_out
//! Pointer to the output sequence of unique keys (one key per run)
//!
//! @param[out] d_counts_out
//! Pointer to the output sequence of run-lengths (one count per run)
//!
//! @param[out] d_num_runs_out
//! Pointer to total number of runs
//!
//! @param[in] num_items
//! Total number of input items (i.e., the length of `d_in`)
//!
//! @param[in] env
//! **[optional]** Execution environment. Default is ``cuda::std::execution::env{}``.
//! @endrst
template <typename InputIteratorT,
typename UniqueOutputIteratorT,
typename LengthsOutputIteratorT,
typename NumRunsOutputIteratorT,
typename NumItemsT,
typename EnvT = ::cuda::std::execution::env<>,
::cuda::std::enable_if_t<!::cuda::std::is_same_v<InputIteratorT, void*>, int> = 0>
[[nodiscard]] CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE static cudaError_t Encode(
InputIteratorT d_in,
UniqueOutputIteratorT d_unique_out,
LengthsOutputIteratorT d_counts_out,
NumRunsOutputIteratorT d_num_runs_out,
NumItemsT num_items,
EnvT env = {})
{
_CCCL_NVTX_RANGE_SCOPE("cub::DeviceRunLengthEncode::Encode");

using equality_op = ::cuda::std::equal_to<>;
using reduction_op = ::cuda::std::plus<>;
Copy link
Copy Markdown
Contributor

@miscco miscco Mar 10, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Ditto: Should this rather be

Suggested change
using reduction_op = ::cuda::std::plus<>;
using reduction_op = ::cuda::std::plus<length_t>;

Otherwise this will always promote offset_t to a larger integer

Copy link
Copy Markdown
Contributor

@bernhardmgruber bernhardmgruber Mar 10, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The non-env overload where I guess this implementation comes from also uses plus<>. I agree that that's maybe not what we want, since plus<length_t> does not promote and influence the accumulator type. But this should be addressed in a separate PR.

using offset_t = detail::choose_signed_offset_t<NumItemsT>;
using length_t = cub::detail::non_void_value_t<LengthsOutputIteratorT, offset_t>;
using lengths_input_iterator_t = THRUST_NS_QUALIFIER::constant_iterator<length_t, offset_t>;
using accum_t = ::cuda::std::__accumulator_t<reduction_op, length_t, length_t>;
using key_t = cub::detail::non_void_value_t<UniqueOutputIteratorT, cub::detail::it_value_t<InputIteratorT>>;
using policy_selector_t = detail::rle::encode::policy_selector_from_types<accum_t, key_t>;

return detail::dispatch_with_env(env, [&]([[maybe_unused]] auto tuning, void* storage, size_t& bytes, auto stream) {
return detail::reduce_by_key::dispatch_streaming_reduce_by_key<accum_t>(
storage,
bytes,
d_in,
d_unique_out,
lengths_input_iterator_t(length_t{1}),
d_counts_out,
d_num_runs_out,
equality_op{},
reduction_op{},
static_cast<offset_t>(num_items),
stream,
policy_selector_t{});
});
}

//! @rst
//! Enumerates the starting offsets and lengths of all non-trivial runs
//! (of ``length > 1``) of same-valued keys in the sequence ``d_in``.
Expand Down Expand Up @@ -332,6 +452,114 @@ struct DeviceRunLengthEncode
static_cast<global_offset_t>(num_items),
stream);
}

//! @rst
//! Enumerates the starting offsets and lengths of all non-trivial runs
//! (of ``length > 1``) of same-valued keys in the sequence ``d_in``.
//!
//! .. versionadded:: 3.4.0
//! First appears in CUDA Toolkit 13.4.
//!
//! This is an environment-based API that allows customization of:
//!
//! - Stream: Query via ``cuda::get_stream``
//! - Memory resource: Query via ``cuda::mr::get_memory_resource``
//!
//! - For the *i*\ :sup:`th` non-trivial run, the run's starting offset and
//! its length are written to ``d_offsets_out[i]`` and ``d_lengths_out[i]``, respectively.
//! - The total number of runs encountered is written to ``d_num_runs_out``.
//! - The ``==`` equality operator is used to determine whether values are equivalent
//! - In-place operations are not supported. There must be no overlap between any of the provided ranges:
//!
//! - ``[d_offsets_out, d_offsets_out + *d_num_runs_out)``
//! - ``[d_lengths_out, d_lengths_out + *d_num_runs_out)``
//! - ``[d_num_runs_out, d_num_runs_out + 1)``
//! - ``[d_in, d_in + num_items)``
//!
//! Snippet
//! +++++++++++++++++++++++++++++++++++++++++++++
//!
//! The code snippet below illustrates the identification of non-trivial runs
//! within a sequence of ``int`` values using environment-based API:
//!
//! .. literalinclude:: ../../../cub/test/catch2_test_device_run_length_encode_env_api.cu
//! :language: c++
//! :dedent:
//! :start-after: example-begin non-trivial-runs-env
//! :end-before: example-end non-trivial-runs-env
//!
//! @endrst
//!
//! @tparam InputIteratorT
//! **[inferred]** Random-access input iterator type for reading input items @iterator
//!
//! @tparam OffsetsOutputIteratorT
//! **[inferred]** Random-access output iterator type for writing run-offset values @iterator
//!
//! @tparam LengthsOutputIteratorT
//! **[inferred]** Random-access output iterator type for writing run-length values @iterator
//!
//! @tparam NumRunsOutputIteratorT
//! **[inferred]** Output iterator type for recording the number of runs encountered @iterator
//!
//! @tparam NumItemsT
//! **[inferred]** Type of num_items
//!
//! @tparam EnvT
//! **[inferred]** Environment type (e.g., `cuda::std::execution::env<...>`)
//!
//! @param[in] d_in
//! Pointer to input sequence of data items
//!
//! @param[out] d_offsets_out
//! Pointer to output sequence of run-offsets
//! (one offset per non-trivial run)
//!
//! @param[out] d_lengths_out
//! Pointer to output sequence of run-lengths (one count per non-trivial run)
//!
//! @param[out] d_num_runs_out
//! Pointer to total number of runs (i.e., length of `d_offsets_out`)
//!
//! @param[in] num_items
//! Total number of input items (i.e., the length of `d_in`)
//!
//! @param[in] env
//! **[optional]** Execution environment. Default is ``cuda::std::execution::env{}``.
//! @endrst
template <typename InputIteratorT,
typename OffsetsOutputIteratorT,
typename LengthsOutputIteratorT,
typename NumRunsOutputIteratorT,
typename NumItemsT,
typename EnvT = ::cuda::std::execution::env<>,
::cuda::std::enable_if_t<!::cuda::std::is_same_v<InputIteratorT, void*>, int> = 0>
[[nodiscard]] CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE static cudaError_t NonTrivialRuns(
InputIteratorT d_in,
OffsetsOutputIteratorT d_offsets_out,
LengthsOutputIteratorT d_lengths_out,
NumRunsOutputIteratorT d_num_runs_out,
NumItemsT num_items,
EnvT env = {})
{
_CCCL_NVTX_RANGE_SCOPE("cub::DeviceRunLengthEncode::NonTrivialRuns");

using global_offset_t = detail::choose_signed_offset_t<NumItemsT>;
using equality_op = ::cuda::std::equal_to<>;
Copy link
Copy Markdown
Contributor

@miscco miscco Mar 10, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I believe this is fine, because it always just returns a bool and does not promote integers


return detail::dispatch_with_env(env, [&]([[maybe_unused]] auto tuning, void* storage, size_t& bytes, auto stream) {
return detail::rle::dispatch(
storage,
bytes,
d_in,
d_offsets_out,
d_lengths_out,
d_num_runs_out,
equality_op{},
static_cast<global_offset_t>(num_items),
stream);
});
}
};

CUB_NAMESPACE_END
1 change: 0 additions & 1 deletion cub/cub/device/device_select.cuh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -23,7 +23,6 @@
#include <cub/device/dispatch/dispatch_select_if.cuh>
#include <cub/device/dispatch/dispatch_unique_by_key.cuh>

#include <cuda/__execution/determinism.h>
Copy link
Copy Markdown
Member Author

@gonidelis gonidelis Mar 10, 2026 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

slipped in while on a workaround. let it pass it's a byfix

#include <cuda/__execution/require.h>
#include <cuda/std/__execution/env.h>
#include <cuda/std/__functional/operations.h>
Expand Down
Loading
Loading