Skip to content

Avoid copies in sorting #1967

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
merged 3 commits into from
Jan 15, 2025
Merged

Avoid copies in sorting #1967

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
082fca1
Use forwarding reference in leaf_sort_impl to avoid copies
oleksandr-pavlyk Jan 14, 2025
9a7868d
Do not invoke map_back kernel if iter_nelems is 1
oleksandr-pavlyk Jan 14, 2025
4847101
Update dpctl/tensor/libtensor/include/kernels/sorting/merge_sort.hpp
oleksandr-pavlyk Jan 15, 2025
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
62 changes: 23 additions & 39 deletions dpctl/tensor/libtensor/include/kernels/sorting/merge_sort.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -212,10 +212,10 @@ void merge_impl(const std::size_t offset,
}

template <typename Iter, typename Compare>
void insertion_sort_impl(Iter first,
const std::size_t begin,
const std::size_t end,
Compare comp)
void insertion_sort_impl(Iter &&first,
std::size_t begin,
std::size_t end,
Compare &&comp)
{
for (std::size_t i = begin + 1; i < end; ++i) {
const auto val_i = first[i];
Expand All @@ -231,31 +231,14 @@ void insertion_sort_impl(Iter first,
}

template <typename Iter, typename Compare>
void bubble_sort_impl(Iter first,
const std::size_t begin,
const std::size_t end,
Compare comp)
void leaf_sort_impl(Iter &&first,
std::size_t begin,
std::size_t end,
Compare &&comp)
{
if (begin < end) {
for (std::size_t i = begin; i < end; ++i) {
// Handle intermediate items
for (std::size_t idx = i + 1; idx < end; ++idx) {
if (comp(first[idx], first[i])) {
std::swap(first[i], first[idx]);
}
}
}
}
}

template <typename Iter, typename Compare>
void leaf_sort_impl(Iter first,
const std::size_t begin,
const std::size_t end,
Compare comp)
{
return insertion_sort_impl<Iter, Compare>(
std::move(first), std::move(begin), std::move(end), std::move(comp));
return insertion_sort_impl<Iter, Compare>(std::forward<Iter>(first),
std::move(begin), std::move(end),
std::forward<Compare>(comp));
}

template <typename Iter> struct GetValueType
Expand Down Expand Up @@ -356,7 +339,7 @@ sort_base_step_contig_impl(sycl::queue &q,
using KernelName = sort_base_step_contig_krn<inpT, outT, Comp>;

const std::size_t n_segments =
quotient_ceil<std::size_t>(sort_nelems, conseq_nelems_sorted);
quotient_ceil(sort_nelems, conseq_nelems_sorted);

sycl::event base_sort = q.submit([&](sycl::handler &cgh) {
cgh.depends_on(depends);
Expand All @@ -375,8 +358,7 @@ sort_base_step_contig_impl(sycl::queue &q,
iter_offset + segment_id * conseq_nelems_sorted;
const std::size_t end_id =
iter_offset +
std::min<std::size_t>((segment_id + 1) * conseq_nelems_sorted,
sort_nelems);
std::min((segment_id + 1) * conseq_nelems_sorted, sort_nelems);
for (std::size_t i = beg_id; i < end_id; ++i) {
output_acc[i] = input_acc[i];
}
Expand Down Expand Up @@ -444,8 +426,7 @@ sort_over_work_group_contig_impl(sycl::queue &q,
// This assumption permits doing away with using a loop
assert(nelems_wg_sorts % lws == 0);

const std::size_t n_segments =
quotient_ceil<std::size_t>(sort_nelems, nelems_wg_sorts);
const std::size_t n_segments = quotient_ceil(sort_nelems, nelems_wg_sorts);

sycl::event base_sort_ev = q.submit([&](sycl::handler &cgh) {
cgh.depends_on(depends);
Expand All @@ -471,8 +452,8 @@ sort_over_work_group_contig_impl(sycl::queue &q,
const std::size_t lid = it.get_local_linear_id();

const std::size_t segment_start_idx = segment_id * nelems_wg_sorts;
const std::size_t segment_end_idx = std::min<std::size_t>(
segment_start_idx + nelems_wg_sorts, sort_nelems);
const std::size_t segment_end_idx =
std::min(segment_start_idx + nelems_wg_sorts, sort_nelems);
const std::size_t wg_chunk_size =
segment_end_idx - segment_start_idx;

Expand All @@ -487,8 +468,7 @@ sort_over_work_group_contig_impl(sycl::queue &q,
}
sycl::group_barrier(it.get_group());

const std::size_t chunk =
quotient_ceil<std::size_t>(nelems_wg_sorts, lws);
const std::size_t chunk = quotient_ceil(nelems_wg_sorts, lws);

const std::size_t chunk_start_idx = lid * chunk;
const std::size_t chunk_end_idx =
Expand Down Expand Up @@ -620,8 +600,7 @@ merge_sorted_block_contig_impl(sycl::queue &q,
used_depends = true;
}

const std::size_t n_chunks =
quotient_ceil<std::size_t>(sort_nelems, chunk_size);
const std::size_t n_chunks = quotient_ceil(sort_nelems, chunk_size);

if (needs_copy) {
sycl::accessor temp_acc{temp_buf, cgh, sycl::write_only,
Expand Down Expand Up @@ -835,6 +814,11 @@ sycl::event stable_argsort_axis1_contig_impl(
exec_q, iter_nelems, sort_nelems, res_tp, index_comp, sorted_block_size,
{base_sort_ev});

// no need to map back if iter_nelems == 1
if (iter_nelems == 1u) {
return merges_ev;
}

using MapBackKernelName = index_map_to_rows_krn<argTy, IndexTy>;
using dpctl::tensor::kernels::sort_utils_detail::map_back_impl;

Expand Down
12 changes: 9 additions & 3 deletions dpctl/tensor/libtensor/include/kernels/sorting/radix_sort.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1867,11 +1867,17 @@ radix_argsort_axis1_contig_impl(sycl::queue &exec_q,
using MapBackKernelName = radix_argsort_index_write_out_krn<argTy, IndexTy>;
using dpctl::tensor::kernels::sort_utils_detail::map_back_impl;

sycl::event map_back_ev = map_back_impl<MapBackKernelName, IndexTy>(
exec_q, total_nelems, res_tp, res_tp, sort_nelems, {radix_sort_ev});
sycl::event dep = radix_sort_ev;

// no need to perform map_back ( id % sort_nelems)
// if total_nelems == sort_nelems
if (iter_nelems > 1u) {
dep = map_back_impl<MapBackKernelName, IndexTy>(
exec_q, total_nelems, res_tp, res_tp, sort_nelems, {dep});
}

sycl::event cleanup_ev = dpctl::tensor::alloc_utils::async_smart_free(
exec_q, {map_back_ev}, workspace_owner);
exec_q, {dep}, workspace_owner);

return cleanup_ev;
}
Expand Down
Loading