Skip to content

Do more internal refactoring in reduce-then-scan #2830

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

Open
akukanov wants to merge 8 commits into from
+332 −364
Open

Do more internal refactoring in reduce-then-scan #2830

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
8 commits
Select commit Hold shift + click to select a range
c20e9f9
move block carry routines out of scan submitter
akukanov Sep 16, 2026
20280d5
lower tag dispatching to __scan_through_elements_helper_impl
akukanov Sep 16, 2026
1dc6718
simplify scan-through-elements for the reduce stage
akukanov Sep 16, 2026
bd4d7da
integrate __scan_through_elements into the scan submitter
akukanov Sep 18, 2026
6f44a26
Move more code into __scan_through_elements
akukanov Sep 18, 2026
9605e58
flatten __scan_through_elements
akukanov Sep 18, 2026
c658233
improve file structure
akukanov Sep 11, 2026
bd47d73
simplify _require_access_args
akukanov Sep 11, 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
295 changes: 152 additions & 143 deletions include/oneapi/dpl/pstl/hetero/dpcpp/parallel_backend_sycl.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -100,7 +100,7 @@ __parallel_copy_impl(sycl::queue& __q, _Index __count, _Range1&& __rng1, _Range2
}

//------------------------------------------------------------------------
// parallel_transform_scan single group - async pattern
// parallel_transform_scan - async pattern
//------------------------------------------------------------------------

template <typename _ValueType, bool _Inclusive, typename _Group, typename _InPtr, typename _OutPtr,
Expand Down Expand Up @@ -216,109 +216,6 @@ struct __parallel_transform_scan_static_single_group_submitter<_Inclusive, _Elem
}
};

struct __parallel_copy_if_single_group_base
{
using _ValueType = std::uint16_t;

template <typename _Size>
static std::pair<std::make_unsigned_t<_Size>, std::make_unsigned_t<_Size>>
__local_memory_needed(_Size __n)
{
// Next power of 2 greater than or equal to __n
std::make_unsigned_t<_Size> __n_uniform =
oneapi::dpl::__internal::__dpl_bit_ceil(static_cast<std::make_unsigned_t<_Size>>(__n));
// The kernel needs memory for: N predicate evaluations, N output offsets, and the input stop position
return {__n_uniform * 2 + 1, __n_uniform};
}

template <typename _Size>
static bool
__enough_local_memory(sycl::queue __q, _Size __n)
{
// Pessimistically expect only half of local memory to account for possible memory use by the compiled code
std::size_t __available_size = __q.get_device().template get_info<sycl::info::device::local_mem_size>() / 2;
return __available_size >= __local_memory_needed(__n).first * sizeof(_ValueType);
}
};

template <typename _KernelName>
struct __parallel_copy_if_single_group_functor;

template <typename... _ScanKernelName>
struct __parallel_copy_if_single_group_functor<__internal::__optional_kernel_name<_ScanKernelName...>>
: __parallel_copy_if_single_group_base
{
template <typename _InRng, typename _OutRng, typename _Size, typename _IndexPred, typename _Assign>
std::array<_Size, 2>
operator()(sycl::queue& __q, _InRng&& __in_rng, _OutRng&& __out_rng, _Size __n, _Size __n_out, _IndexPred __pred,
_Assign __assign, std::size_t __max_wg_size)
{
assert(__max_wg_size <= std::numeric_limits<std::uint16_t>::max());
// This type is used as a workaround for when an internal tuple is assigned to std::tuple, such as
// with zip_iterator
using __tuple_type = typename oneapi::dpl::__internal::__get_tuple_type<
std::decay_t<decltype(__in_rng[0])>, std::decay_t<decltype(__out_rng[0])>>::__type;

__result_storage<_Size> __result{__q, 2};

__q.submit([&](sycl::handler& __hdl) {
oneapi::dpl::__ranges::__require_access(__hdl, __in_rng, __out_rng);

std::make_unsigned_t<_Size> __lsize, __n_uniform;
// Since __n_uniform is captured into a lambda, structured binding cannot be used here till C++20
std::tie(__lsize, __n_uniform) = __local_memory_needed(__n);
auto __lacc = __dpl_sycl::__local_accessor<_ValueType>(sycl::range<1>(__lsize), __hdl);
auto __res_acc = __get_accessor(sycl::write_only, __result, __hdl, __dpl_sycl::__no_init{});
const auto __wg_size = static_cast<std::uint16_t>(std::min<std::size_t>(__n_uniform, __max_wg_size));

__hdl.parallel_for<_ScanKernelName...>(sycl::nd_range<1>(__wg_size, __wg_size),
[=](sycl::nd_item<1> __self_item) {
sycl::group __group = __self_item.get_group();
// This kernel is only launched for sizes less than 2^16
const std::uint16_t __item_id = __self_item.get_local_linear_id();
_ValueType* __lacc_ptr = __dpl_sycl::__get_accessor_ptr(__lacc);
for (std::uint16_t __idx = __item_id; __idx < __n; __idx += __wg_size)
{
__lacc[__idx] = __pred(__in_rng, __idx);
}
if (__item_id == 0)
{
// Store the input size as the expected stop position
__lacc[2 * __n_uniform] = __n;
}

__scan_work_group<_ValueType, /* _Inclusive */ false>(
__group, __lacc_ptr, __lacc_ptr + __n, __lacc_ptr + __n_uniform, sycl::plus<_ValueType>{});

for (std::uint16_t __idx = __item_id; __idx < __n; __idx += __wg_size)
{
if (__lacc[__idx]) {
_ValueType __out_idx = __lacc[__idx + __n_uniform];
if (__out_idx < __n_out)
__assign(static_cast<__tuple_type>(__in_rng[__idx]), __out_rng[__out_idx]);
if (__out_idx == __n_out)
__lacc[2 * __n_uniform] = __idx; // the actual stop position in the input
}
}
sycl::group_barrier(__group);

if (__item_id == 0)
{
_Size* __res_ptr = __res_acc.__data();
_ValueType __stop_in = __lacc[2 * __n_uniform];
__res_ptr[1] = __stop_in;
// Add predicate of last element to account for the scan's exclusivity
__res_ptr[0] = (__stop_in == __n) ? __lacc[__n_uniform + __n - 1] + __lacc[__n - 1] : __n_out;
}
});
}).wait_and_throw();

std::array<_Size, 2> __ret;
__result.__copy_result(__ret.data(), __ret.size());
return __ret;
}
};

template <typename _CustomName, typename _InRng, typename _OutRng, typename _UnaryOperation, typename _InitType,
typename _BinaryOperation, typename _Inclusive>
sycl::event
Expand Down Expand Up @@ -448,6 +345,113 @@ __parallel_transform_scan(oneapi::dpl::__internal::__device_backend_tag, _Execut
return __future(std::move(__event), std::move(__holder).__extract());
}

//------------------------------------------------------------------------
// Filtering patterns: copy_if, unique_copy, etc.; also partition_copy
//------------------------------------------------------------------------

struct __parallel_copy_if_single_group_base
{
using _ValueType = std::uint16_t;

template <typename _Size>
static std::pair<std::make_unsigned_t<_Size>, std::make_unsigned_t<_Size>>
__local_memory_needed(_Size __n)
{
// Next power of 2 greater than or equal to __n
std::make_unsigned_t<_Size> __n_uniform =
oneapi::dpl::__internal::__dpl_bit_ceil(static_cast<std::make_unsigned_t<_Size>>(__n));
// The kernel needs memory for: N predicate evaluations, N output offsets, and the input stop position
return {__n_uniform * 2 + 1, __n_uniform};
}

template <typename _Size>
static bool
__enough_local_memory(sycl::queue __q, _Size __n)
{
// Pessimistically expect only half of local memory to account for possible memory use by the compiled code
std::size_t __available_size = __q.get_device().template get_info<sycl::info::device::local_mem_size>() / 2;
return __available_size >= __local_memory_needed(__n).first * sizeof(_ValueType);
}
};

template <typename _KernelName>
struct __parallel_copy_if_single_group_functor;

template <typename... _ScanKernelName>
struct __parallel_copy_if_single_group_functor<__internal::__optional_kernel_name<_ScanKernelName...>>
: __parallel_copy_if_single_group_base
{
template <typename _InRng, typename _OutRng, typename _Size, typename _IndexPred, typename _Assign>
std::array<_Size, 2>
operator()(sycl::queue& __q, _InRng&& __in_rng, _OutRng&& __out_rng, _Size __n, _Size __n_out, _IndexPred __pred,
_Assign __assign, std::size_t __max_wg_size)
{
assert(__max_wg_size <= std::numeric_limits<std::uint16_t>::max());
// This type is used as a workaround for when an internal tuple is assigned to std::tuple, such as
// with zip_iterator
using __tuple_type = typename oneapi::dpl::__internal::__get_tuple_type<
std::decay_t<decltype(__in_rng[0])>, std::decay_t<decltype(__out_rng[0])>>::__type;

__result_storage<_Size> __result{__q, 2};

__q.submit([&](sycl::handler& __hdl) {
oneapi::dpl::__ranges::__require_access(__hdl, __in_rng, __out_rng);

std::make_unsigned_t<_Size> __lsize, __n_uniform;
// Since __n_uniform is captured into a lambda, structured binding cannot be used here till C++20
std::tie(__lsize, __n_uniform) = __local_memory_needed(__n);
auto __lacc = __dpl_sycl::__local_accessor<_ValueType>(sycl::range<1>(__lsize), __hdl);
auto __res_acc = __get_accessor(sycl::write_only, __result, __hdl, __dpl_sycl::__no_init{});
const auto __wg_size = static_cast<std::uint16_t>(std::min<std::size_t>(__n_uniform, __max_wg_size));

__hdl.parallel_for<_ScanKernelName...>(sycl::nd_range<1>(__wg_size, __wg_size),
[=](sycl::nd_item<1> __self_item) {
sycl::group __group = __self_item.get_group();
// This kernel is only launched for sizes less than 2^16
const std::uint16_t __item_id = __self_item.get_local_linear_id();
_ValueType* __lacc_ptr = __dpl_sycl::__get_accessor_ptr(__lacc);
for (std::uint16_t __idx = __item_id; __idx < __n; __idx += __wg_size)
{
__lacc[__idx] = __pred(__in_rng, __idx);
}
if (__item_id == 0)
{
// Store the input size as the expected stop position
__lacc[2 * __n_uniform] = __n;
}

__scan_work_group<_ValueType, /* _Inclusive */ false>(
__group, __lacc_ptr, __lacc_ptr + __n, __lacc_ptr + __n_uniform, sycl::plus<_ValueType>{});

for (std::uint16_t __idx = __item_id; __idx < __n; __idx += __wg_size)
{
if (__lacc[__idx]) {
_ValueType __out_idx = __lacc[__idx + __n_uniform];
if (__out_idx < __n_out)
__assign(static_cast<__tuple_type>(__in_rng[__idx]), __out_rng[__out_idx]);
if (__out_idx == __n_out)
__lacc[2 * __n_uniform] = __idx; // the actual stop position in the input
}
}
sycl::group_barrier(__group);

if (__item_id == 0)
{
_Size* __res_ptr = __res_acc.__data();
_ValueType __stop_in = __lacc[2 * __n_uniform];
__res_ptr[1] = __stop_in;
// Add predicate of last element to account for the scan's exclusivity
__res_ptr[0] = (__stop_in == __n) ? __lacc[__n_uniform + __n - 1] + __lacc[__n - 1] : __n_out;
}
});
}).wait_and_throw();

std::array<_Size, 2> __ret;
__result.__copy_result(__ret.data(), __ret.size());
return __ret;
}
};

template <bool _Bounded, typename _CustomName, typename _InRng, typename _OutRng, typename _Size, typename _GenMask,
typename _WriteOp, typename _IsUniquePattern>
std::array<_Size, 2>
Expand Down Expand Up @@ -522,45 +526,6 @@ __parallel_unique_copy(oneapi::dpl::__internal::__device_backend_tag, _Execution
return __ret;
}

template <typename _CustomName, typename _Range1, typename _Range2, typename _Range3, typename _Range4,
typename _BinaryPredicate, typename _BinaryOperator>
auto /*__future<sycl::event, ...>*/
__parallel_reduce_by_segment_reduce_then_scan(sycl::queue& __q, _Range1&& __keys, _Range2&& __values,
_Range3&& __out_keys, _Range4&& __out_values,
_BinaryPredicate __binary_pred, _BinaryOperator __binary_op)
{
// Flags new segments and passes input value through a 2-tuple
using _GenReduceInput = __gen_red_by_seg_reduce_input<_BinaryPredicate>;
// Operation that computes output indices and output reduction values per segment
using _ReduceOp = __red_by_seg_op<_BinaryOperator>;
// Returns 4-component tuple which contains flags, keys, value, and a flag to write output
using _GenScanInput = __gen_red_by_seg_scan_input<_BinaryPredicate>;
// Returns the first component from scan input which is scanned over
using _ScanInputTransform = __get_zeroth_element;
// Writes current segment's output reduction and the next segment's output key
using _WriteOp = __write_red_by_seg<_BinaryPredicate>;
using _KeyType = oneapi::dpl::__internal::__value_t<_Range1>;
using _ValueType = oneapi::dpl::__internal::__value_t<_Range2>;
using _ResultType = oneapi::dpl::__internal::tuple<std::size_t, _ValueType>;

std::size_t __n = oneapi::dpl::__ranges::__size(__keys);
// __gen_red_by_seg_scan_input requires that __n > 1
assert(__n > 1);
__transform_scan_storage_holder_simple<_ResultType> __holder(__q);
// Each work-item iteration reads one key and one value from the zipped input. The comparison against the previous
// key is not counted separately, as that key is read by the adjacent index's iteration.
constexpr std::uint32_t __bytes_per_work_item_iter = sizeof(_KeyType) + sizeof(_ValueType);

sycl::event __event = __parallel_transform_reduce_then_scan<
/*_Bounded*/ false, __bytes_per_work_item_iter, _CustomName>(
__q, __n, oneapi::dpl::__ranges::make_zip_view(std::forward<_Range1>(__keys), std::forward<_Range2>(__values)),
oneapi::dpl::__ranges::make_zip_view(std::forward<_Range3>(__out_keys), std::forward<_Range4>(__out_values)),
_GenReduceInput{__binary_pred}, _ReduceOp{__binary_op}, _GenScanInput{__binary_pred, __n},
_ScanInputTransform{}, _WriteOp{__binary_pred, __n}, oneapi::dpl::unseq_backend::__no_init_value<_ResultType>{},
__holder, /*Inclusive*/ std::true_type{}, /*_IsUniquePattern=*/std::false_type{});
return __future(std::move(__event), std::move(__holder).__extract());
}

template <bool _Bounded, typename _ExecutionPolicy, typename _Range1, typename _Range2, typename _Range3,
typename _UnaryPredicate>
std::array<std::ptrdiff_t, 2>
Expand Down Expand Up @@ -655,6 +620,10 @@ __parallel_copy_if(oneapi::dpl::__internal::__device_backend_tag, _ExecutionPoli
return __ret;
}

//------------------------------------------------------------------------
// Set operations
//------------------------------------------------------------------------

// balanced path
template <bool _Bounded, typename _CustomName, typename _SetTag, typename _Range1, typename _Range2, typename _Range3,
typename _Compare, typename _Proj1, typename _Proj2>
Expand Down Expand Up @@ -1497,6 +1466,46 @@ __parallel_partial_sort(oneapi::dpl::__internal::__device_backend_tag, _Executio
// inability to create event dependency chains across separate parallel pattern calls. If we ever add support for
// cross parallel pattern dependencies, then we can implement this as an async pattern.
//------------------------------------------------------------------------

template <typename _CustomName, typename _Range1, typename _Range2, typename _Range3, typename _Range4,
typename _BinaryPredicate, typename _BinaryOperator>
auto /*__future<sycl::event, ...>*/
__parallel_reduce_by_segment_reduce_then_scan(sycl::queue& __q, _Range1&& __keys, _Range2&& __values,
_Range3&& __out_keys, _Range4&& __out_values,
_BinaryPredicate __binary_pred, _BinaryOperator __binary_op)
{
// Flags new segments and passes input value through a 2-tuple
using _GenReduceInput = __gen_red_by_seg_reduce_input<_BinaryPredicate>;
// Operation that computes output indices and output reduction values per segment
using _ReduceOp = __red_by_seg_op<_BinaryOperator>;
// Returns 4-component tuple which contains flags, keys, value, and a flag to write output
using _GenScanInput = __gen_red_by_seg_scan_input<_BinaryPredicate>;
// Returns the first component from scan input which is scanned over
using _ScanInputTransform = __get_zeroth_element;
// Writes current segment's output reduction and the next segment's output key
using _WriteOp = __write_red_by_seg<_BinaryPredicate>;
using _KeyType = oneapi::dpl::__internal::__value_t<_Range1>;
using _ValueType = oneapi::dpl::__internal::__value_t<_Range2>;
using _ResultType = oneapi::dpl::__internal::tuple<std::size_t, _ValueType>;

std::size_t __n = oneapi::dpl::__ranges::__size(__keys);
// __gen_red_by_seg_scan_input requires that __n > 1
assert(__n > 1);
__transform_scan_storage_holder_simple<_ResultType> __holder(__q);
// Each work-item iteration reads one key and one value from the zipped input. The comparison against the previous
// key is not counted separately, as that key is read by the adjacent index's iteration.
constexpr std::uint32_t __bytes_per_work_item_iter = sizeof(_KeyType) + sizeof(_ValueType);

sycl::event __event = __parallel_transform_reduce_then_scan<
/*_Bounded*/ false, __bytes_per_work_item_iter, _CustomName>(
__q, __n, oneapi::dpl::__ranges::make_zip_view(std::forward<_Range1>(__keys), std::forward<_Range2>(__values)),
oneapi::dpl::__ranges::make_zip_view(std::forward<_Range3>(__out_keys), std::forward<_Range4>(__out_values)),
_GenReduceInput{__binary_pred}, _ReduceOp{__binary_op}, _GenScanInput{__binary_pred, __n},
_ScanInputTransform{}, _WriteOp{__binary_pred, __n}, oneapi::dpl::unseq_backend::__no_init_value<_ResultType>{},
__holder, /*Inclusive*/ std::true_type{}, /*_IsUniquePattern=*/std::false_type{});
return __future(std::move(__event), std::move(__holder).__extract());
}

template <typename _Name>
struct __reduce1_wrapper;

Expand Down
Loading
Loading