Skip to content
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

Add iota kernel #1946

Merged
merged 15 commits into from
Dec 27, 2024
Merged

Add iota kernel #1946

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
15 commits
Select commit Hold shift + click to select a range
be9f5ce
Add test file for top_k functionality
oleksandr-pavlyk Dec 14, 2024
ff81365
Adds test to triage AMD CPU top_k failure
ndgrigorian Dec 17, 2024
58da416
radix sort triage test now parametrizes for identical sizes to topk_1…
ndgrigorian Dec 18, 2024
fd4bab3
Add sort_utils.hpp with iota_impl
oleksandr-pavlyk Dec 20, 2024
d9cc919
Simplify write-out kernels in topk implementation (avoid recomputing …
oleksandr-pavlyk Dec 22, 2024
eb28d1a
Add explicit read into registers
oleksandr-pavlyk Dec 22, 2024
da3fbcc
Use unique_ptr as temporary owner of USM allocation
oleksandr-pavlyk Dec 20, 2024
79feecd
Factor out write-out kernel into separate detail function
oleksandr-pavlyk Dec 25, 2024
83b1771
Typo: smart_malloc_jost -> smart_malloc_host
oleksandr-pavlyk Dec 25, 2024
1c25d0b
Reimplemented map_back_impl to process few elements per work-item
oleksandr-pavlyk Dec 25, 2024
a74c788
Replace use of std::log2(size_t_value)
oleksandr-pavlyk Dec 25, 2024
9294095
Remove dead code disable via preprocessor conditional
oleksandr-pavlyk Dec 27, 2024
1bb83bc
Add information displayed on failure, renamed variables
oleksandr-pavlyk Dec 25, 2024
fde7408
Add comment to explain ceil_log2 algo
oleksandr-pavlyk Dec 27, 2024
4c8d7d7
Add static assers to async_smart_free
oleksandr-pavlyk Dec 27, 2024
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
35 changes: 9 additions & 26 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 @@ -33,6 +33,7 @@

#include "kernels/dpctl_tensor_types.hpp"
#include "kernels/sorting/search_sorted_detail.hpp"
#include "kernels/sorting/sort_utils.hpp"

namespace dpctl
{
Expand Down Expand Up @@ -811,20 +812,12 @@ sycl::event stable_argsort_axis1_contig_impl(

const size_t total_nelems = iter_nelems * sort_nelems;

sycl::event populate_indexed_data_ev =
exec_q.submit([&](sycl::handler &cgh) {
cgh.depends_on(depends);
using dpctl::tensor::kernels::sort_utils_detail::iota_impl;

const sycl::range<1> range{total_nelems};
using IotaKernelName = populate_index_data_krn<argTy, IndexTy, ValueComp>;

using KernelName =
populate_index_data_krn<argTy, IndexTy, ValueComp>;

cgh.parallel_for<KernelName>(range, [=](sycl::id<1> id) {
size_t i = id[0];
res_tp[i] = static_cast<IndexTy>(i);
});
});
sycl::event populate_indexed_data_ev = iota_impl<IotaKernelName, IndexTy>(
exec_q, res_tp, total_nelems, depends);

// Sort segments of the array
sycl::event base_sort_ev =
Expand All @@ -839,21 +832,11 @@ sycl::event stable_argsort_axis1_contig_impl(
exec_q, iter_nelems, sort_nelems, res_tp, index_comp, sorted_block_size,
{base_sort_ev});

sycl::event write_out_ev = exec_q.submit([&](sycl::handler &cgh) {
cgh.depends_on(merges_ev);

auto temp_acc =
merge_sort_detail::GetReadOnlyAccess<decltype(res_tp)>{}(res_tp,
cgh);

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

const sycl::range<1> range{total_nelems};

cgh.parallel_for<KernelName>(range, [=](sycl::id<1> id) {
res_tp[id] = (temp_acc[id] % sort_nelems);
});
});
sycl::event write_out_ev = map_back_impl<MapBackKernelName, IndexTy>(
exec_q, total_nelems, res_tp, res_tp, sort_nelems, {merges_ev});

return write_out_ev;
}
Expand Down
174 changes: 83 additions & 91 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 @@ -38,6 +38,7 @@
#include <sycl/sycl.hpp>

#include "kernels/dpctl_tensor_types.hpp"
#include "kernels/sorting/sort_utils.hpp"
#include "utils/sycl_alloc_utils.hpp"

namespace dpctl
Expand All @@ -62,6 +63,47 @@ class radix_sort_reorder_peer_kernel;
template <std::uint32_t, bool, typename... TrailingNames>
class radix_sort_reorder_kernel;

/*! @brief Computes smallest exponent such that `n <= (1 << exponent)` */
template <typename SizeT,
std::enable_if_t<std::is_unsigned_v<SizeT> &&
sizeof(SizeT) == sizeof(std::uint64_t),
int> = 0>
std::uint32_t ceil_log2(SizeT n)
{
if (n <= 1)
return std::uint32_t{1};

std::uint32_t exp{1};
--n;
// if n > 2^b, n = q * 2^b + r for q > 0 and 0 <= r < 2^b
// ceil_log2(q * 2^b + r) == ceil_log2(q * 2^b) == q + ceil_log2(n1)
if (n >= (SizeT{1} << 32)) {
n >>= 32;
exp += 32;
}
if (n >= (SizeT{1} << 16)) {
n >>= 16;
exp += 16;
}
if (n >= (SizeT{1} << 8)) {
n >>= 8;
exp += 8;
}
if (n >= (SizeT{1} << 4)) {
n >>= 4;
exp += 4;
}
if (n >= (SizeT{1} << 2)) {
n >>= 2;
exp += 2;
}
if (n >= (SizeT{1} << 1)) {
n >>= 1;
++exp;
}
return exp;
}

//----------------------------------------------------------
// bitwise order-preserving conversions to unsigned integers
//----------------------------------------------------------
Expand Down Expand Up @@ -1144,7 +1186,7 @@ struct subgroup_radix_sort
const std::size_t max_slm_size =
dev.template get_info<sycl::info::device::local_mem_size>() / 2;

const auto n_uniform = 1 << (std::uint32_t(std::log2(n - 1)) + 1);
const auto n_uniform = 1 << ceil_log2(n);
const auto req_slm_size_val = sizeof(T) * n_uniform;

return ((req_slm_size_val + req_slm_size_counters) <= max_slm_size)
Expand Down Expand Up @@ -1256,9 +1298,7 @@ struct subgroup_radix_sort
const uint16_t id = wi * block_size + i;
if (id < n)
values[i] = std::move(
this_input_arr[iter_val_offset +
static_cast<std::size_t>(
id)]);
this_input_arr[iter_val_offset + id]);
}

while (true) {
Expand All @@ -1272,8 +1312,7 @@ struct subgroup_radix_sort
// counting phase
auto pcounter =
get_accessor_pointer(counter_acc) +
static_cast<std::size_t>(wi) +
iter_counter_offset;
(wi + iter_counter_offset);

// initialize counters
#pragma unroll
Expand Down Expand Up @@ -1348,19 +1387,15 @@ struct subgroup_radix_sort

// scan contiguous numbers
uint16_t bin_sum[bin_count];
bin_sum[0] =
counter_acc[iter_counter_offset +
static_cast<std::size_t>(
wi * bin_count)];
const std::size_t counter_offset0 =
iter_counter_offset + wi * bin_count;
bin_sum[0] = counter_acc[counter_offset0];

#pragma unroll
for (uint16_t i = 1; i < bin_count; ++i)
bin_sum[i] =
bin_sum[i - 1] +
counter_acc
[iter_counter_offset +
static_cast<std::size_t>(
wi * bin_count + i)];
counter_acc[counter_offset0 + i];

sycl::group_barrier(ndit.get_group());

Expand All @@ -1374,10 +1409,7 @@ struct subgroup_radix_sort
// add to local sum, generate exclusive scan result
#pragma unroll
for (uint16_t i = 0; i < bin_count; ++i)
counter_acc[iter_counter_offset +
static_cast<std::size_t>(
wi * bin_count + i +
1)] =
counter_acc[counter_offset0 + i + 1] =
sum_scan + bin_sum[i];

if (wi == 0)
Expand Down Expand Up @@ -1407,10 +1439,8 @@ struct subgroup_radix_sort
if (r < n) {
// move the values to source range and
// destroy the values
this_output_arr
[iter_val_offset +
static_cast<std::size_t>(r)] =
std::move(values[i]);
this_output_arr[iter_val_offset + r] =
std::move(values[i]);
}
}

Expand All @@ -1422,8 +1452,7 @@ struct subgroup_radix_sort
for (uint16_t i = 0; i < block_size; ++i) {
const uint16_t r = indices[i];
if (r < n)
exchange_acc[iter_exchange_offset +
static_cast<std::size_t>(r)] =
exchange_acc[iter_exchange_offset + r] =
std::move(values[i]);
}

Expand All @@ -1435,8 +1464,7 @@ struct subgroup_radix_sort
if (id < n)
values[i] = std::move(
exchange_acc[iter_exchange_offset +
static_cast<std::size_t>(
id)]);
id]);
}

sycl::group_barrier(ndit.get_group());
Expand Down Expand Up @@ -1601,11 +1629,11 @@ sycl::event parallel_radix_sort_impl(sycl::queue &exec_q,
using CountT = std::uint32_t;

// memory for storing count and offset values
CountT *count_ptr =
sycl::malloc_device<CountT>(n_iters * n_counts, exec_q);
if (nullptr == count_ptr) {
throw std::runtime_error("Could not allocate USM-device memory");
}
auto count_owner =
dpctl::tensor::alloc_utils::smart_malloc_device<CountT>(
n_iters * n_counts, exec_q);

CountT *count_ptr = count_owner.get();

constexpr std::uint32_t zero_radix_iter{0};

Expand All @@ -1618,25 +1646,17 @@ sycl::event parallel_radix_sort_impl(sycl::queue &exec_q,
n_counts, count_ptr, proj_op,
is_ascending, depends);

sort_ev = exec_q.submit([=](sycl::handler &cgh) {
cgh.depends_on(sort_ev);
const sycl::context &ctx = exec_q.get_context();

using dpctl::tensor::alloc_utils::sycl_free_noexcept;
cgh.host_task(
[ctx, count_ptr]() { sycl_free_noexcept(count_ptr, ctx); });
});
sort_ev = dpctl::tensor::alloc_utils::async_smart_free(
exec_q, {sort_ev}, count_owner);

return sort_ev;
}

ValueT *tmp_arr =
sycl::malloc_device<ValueT>(n_iters * n_to_sort, exec_q);
if (nullptr == tmp_arr) {
using dpctl::tensor::alloc_utils::sycl_free_noexcept;
sycl_free_noexcept(count_ptr, exec_q);
throw std::runtime_error("Could not allocate USM-device memory");
}
auto tmp_arr_owner =
dpctl::tensor::alloc_utils::smart_malloc_device<ValueT>(
n_iters * n_to_sort, exec_q);

ValueT *tmp_arr = tmp_arr_owner.get();

// iterations per each bucket
assert("Number of iterations must be even" && radix_iters % 2 == 0);
Expand Down Expand Up @@ -1670,17 +1690,8 @@ sycl::event parallel_radix_sort_impl(sycl::queue &exec_q,
}
}

sort_ev = exec_q.submit([=](sycl::handler &cgh) {
cgh.depends_on(sort_ev);

const sycl::context &ctx = exec_q.get_context();

using dpctl::tensor::alloc_utils::sycl_free_noexcept;
cgh.host_task([ctx, count_ptr, tmp_arr]() {
sycl_free_noexcept(tmp_arr, ctx);
sycl_free_noexcept(count_ptr, ctx);
});
});
sort_ev = dpctl::tensor::alloc_utils::async_smart_free(
exec_q, {sort_ev}, tmp_arr_owner, count_owner);
}

return sort_ev;
Expand Down Expand Up @@ -1782,57 +1793,38 @@ radix_argsort_axis1_contig_impl(sycl::queue &exec_q,
reinterpret_cast<IndexTy *>(res_cp) + iter_res_offset + sort_res_offset;

const std::size_t total_nelems = iter_nelems * sort_nelems;
const std::size_t padded_total_nelems = ((total_nelems + 63) / 64) * 64;
IndexTy *workspace = sycl::malloc_device<IndexTy>(
padded_total_nelems + total_nelems, exec_q);
auto workspace_owner =
dpctl::tensor::alloc_utils::smart_malloc_device<IndexTy>(total_nelems,
exec_q);

if (nullptr == workspace) {
throw std::runtime_error("Could not allocate workspace on device");
}
// get raw USM pointer
IndexTy *workspace = workspace_owner.get();

using IdentityProjT = radix_sort_details::IdentityProj;
using IndexedProjT =
radix_sort_details::IndexedProj<IndexTy, argTy, IdentityProjT>;
const IndexedProjT proj_op{arg_tp, IdentityProjT{}};

sycl::event iota_ev = exec_q.submit([&](sycl::handler &cgh) {
cgh.depends_on(depends);
using IotaKernelName = radix_argsort_iota_krn<argTy, IndexTy>;

using KernelName = radix_argsort_iota_krn<argTy, IndexTy>;
using dpctl::tensor::kernels::sort_utils_detail::iota_impl;

cgh.parallel_for<KernelName>(
sycl::range<1>(total_nelems), [=](sycl::id<1> id) {
size_t i = id[0];
IndexTy sort_id = static_cast<IndexTy>(i);
workspace[i] = sort_id;
});
});
sycl::event iota_ev = iota_impl<IotaKernelName, IndexTy>(
exec_q, workspace, total_nelems, depends);

sycl::event radix_sort_ev =
radix_sort_details::parallel_radix_sort_impl<IndexTy, IndexedProjT>(
exec_q, iter_nelems, sort_nelems, workspace, res_tp, proj_op,
sort_ascending, {iota_ev});

sycl::event map_back_ev = exec_q.submit([&](sycl::handler &cgh) {
cgh.depends_on(radix_sort_ev);

using KernelName = radix_argsort_index_write_out_krn<argTy, IndexTy>;

cgh.parallel_for<KernelName>(
sycl::range<1>(total_nelems), [=](sycl::id<1> id) {
IndexTy linear_index = res_tp[id];
res_tp[id] = (linear_index % sort_nelems);
});
});

sycl::event cleanup_ev = exec_q.submit([&](sycl::handler &cgh) {
cgh.depends_on(map_back_ev);
using MapBackKernelName = radix_argsort_index_write_out_krn<argTy, IndexTy>;
using dpctl::tensor::kernels::sort_utils_detail::map_back_impl;

const sycl::context &ctx = exec_q.get_context();
sycl::event map_back_ev = map_back_impl<MapBackKernelName, IndexTy>(
exec_q, total_nelems, res_tp, res_tp, sort_nelems, {radix_sort_ev});

using dpctl::tensor::alloc_utils::sycl_free_noexcept;
cgh.host_task([ctx, workspace] { sycl_free_noexcept(workspace, ctx); });
});
sycl::event cleanup_ev = dpctl::tensor::alloc_utils::async_smart_free(
exec_q, {map_back_ev}, workspace_owner);

return cleanup_ev;
}
Expand Down
Loading
Loading