Add OpenMP segmented prefix sum #1837
Conversation
ginkgo-bot
added
reg:build
| { | ||
| for (int i = 0; i < nthreads - 1; i++) { | ||
| // the next block carries over the previous partial sum | ||
| // if it starts and ends with the same key as the next one |
There was a problem hiding this comment.
| // if it starts and ends with the same key as the next one | |
| // if it starts and ends with the same key as the previous one |
| if (num_entries < 2) { | ||
| if (num_entries == 0) { | ||
| return; | ||
| } else { | ||
| *it = init; | ||
| return; | ||
| } | ||
| } |
There was a problem hiding this comment.
| if (num_entries < 2) { | |
| if (num_entries == 0) { | |
| return; | |
| } else { | |
| *it = init; | |
| return; | |
| } | |
| } | |
| if (num_entries == 0) { | |
| return; | |
| } else if (num_entries == 1) { | |
| *it = init; | |
| return; | |
| } | |
| void segmented_prefix_sum( | ||
| std::shared_ptr<const OmpExecutor> exec, KeyIterator key, Iterator it, | ||
| const size_type num_entries, | ||
| typename std::iterator_traits<KeyIterator>::value_type key_init = {}, |
There was a problem hiding this comment.
does exclusive_scan_by_key have key_init?
There was a problem hiding this comment.
Is it for the situation when KeyIterator is float or double?
I think we only need to consider it is integer (or bool/char) type because float and double as the key having the same value sounds unusual to me
| const size_type num_entries, | ||
| typename std::iterator_traits<KeyIterator>::value_type key_init = {}, | ||
| typename std::iterator_traits<Iterator>::value_type init = {}, | ||
| ScanOp op = {}) |
There was a problem hiding this comment.
As it call prefix_sum, we do not need the ScanOp.
| std::min(num_entries, (thread_id + 1) * def_num_witems); | ||
|
|
||
| auto partial_sum = init; | ||
| auto cur_key = startidx < num_entries ? key[startidx] : key_init; |
There was a problem hiding this comment.
I think key_init can be the last key.
It may introduce one or two operation in proc_sums[i + 1] = op(proc_sums[i], proc_sums[i + 1]);, but I think it should be okay.
| * Similar to prefix_sum, only reduces within runs of the same key value (each | ||
| * key run must only occur once, otherwise the scan operation is not necessarily | ||
| * associaive). It also doesn't ignore the last value! | ||
| * Similar to thrust::exclusive_scan_by_key |
There was a problem hiding this comment.
I do not get it. if it does not ignore the last value, should it be the inclusive one?
| gko::kernels::omp::components::segmented_prefix_sum( | ||
| this->exec, keys.cbegin(), input.begin(), keys.size()); |
There was a problem hiding this comment.
| gko::kernels::omp::components::segmented_prefix_sum( | |
| this->exec, keys.cbegin(), input.begin(), keys.size()); | |
| gko::kernels::omp::components::segmented_prefix_sum( | |
| this->exec, keys.begin(), input.begin(), keys.size()); |
| TYPED_TEST_SUITE(PrefixSum, gko::test::IndexTypes, TypenameNameGenerator); | ||
|
|
||
|
|
||
| TYPED_TEST(PrefixSum, SegmentedPrefixSumWorks) |
There was a problem hiding this comment.
another test for checking the init is not 0
| template <typename KeyIterator, typename Iterator, | ||
| typename ScanOp = | ||
| std::plus<typename std::iterator_traits<Iterator>::value_type>> | ||
| void segmented_prefix_sum( |
Again part of #1758, this provides the equivalent of
thrust::exclusive_scan_by_key