Skip to content

Improve base step of scan algorithm #1942

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
oleksandr-pavlyk merged 6 commits into from
Dec 27, 2024
Merged

Improve base step of scan algorithm #1942

Changes from 1 commit
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
18c5ed7
Add implementation of base scan using striped load/store pattern
oleksandr-pavlyk Dec 18, 2024
ffd2609
Rename variable in inclusive_scan_base_step_blocked, add const qualifier
oleksandr-pavlyk Dec 19, 2024
101514f
Implement custom scan group function for generic binary operator
oleksandr-pavlyk Dec 19, 2024
ce4cf5d
Fixed hang on CUDA GPUs in custom scanning code
oleksandr-pavlyk Dec 20, 2024
0d59c8f
For HIP devices use smaller wg-size parameter
oleksandr-pavlyk Dec 22, 2024
93eba50
Changed type of wg_size to uint32_t everywhere
oleksandr-pavlyk Dec 26, 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
Next Next commit
Add implementation of base scan using striped load/store pattern
  • Loading branch information
@oleksandr-pavlyk
oleksandr-pavlyk committed Dec 18, 2024
commit 18c5ed773152306914bbcce9d4d4892f6032b3e2
292 changes: 273 additions & 19 deletions dpctl/tensor/libtensor/include/kernels/accumulators.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -145,7 +145,7 @@ template <typename T> class stack_strided_t

namespace su_ns = dpctl::tensor::sycl_utils;

using nwiT = std::uint16_t;
using nwiT = std::uint32_t;

template <typename inputT,
typename outputT,
Expand All @@ -156,7 +156,18 @@ template <typename inputT,
typename TransformerT,
typename ScanOpT,
bool include_initial>
class inclusive_scan_iter_local_scan_krn;
class inclusive_scan_iter_local_scan_blocked_krn;

template <typename inputT,
typename outputT,
nwiT n_wi,
typename IterIndexerT,
typename InpIndexerT,
typename OutIndexerT,
typename TransformerT,
typename ScanOpT,
bool include_initial>
class inclusive_scan_iter_local_scan_striped_krn;

template <typename inputT,
typename outputT,
Expand All @@ -177,22 +188,22 @@ template <typename inputT,
typename ScanOpT,
bool include_initial = false>
sycl::event
inclusive_scan_base_step(sycl::queue &exec_q,
const std::size_t wg_size,
const std::size_t iter_nelems,
const std::size_t acc_nelems,
const inputT *input,
outputT *output,
const std::size_t s0,
const std::size_t s1,
const IterIndexerT &iter_indexer,
const InpIndexerT &inp_indexer,
const OutIndexerT &out_indexer,
TransformerT transformer,
const ScanOpT &scan_op,
outputT identity,
std::size_t &acc_groups,
const std::vector<sycl::event> &depends = {})
inclusive_scan_base_step_blocked(sycl::queue &exec_q,
const std::uint32_t wg_size,
const std::size_t iter_nelems,
const std::size_t acc_nelems,
const inputT *input,
outputT *output,
const std::size_t s0,
const std::size_t s1,
const IterIndexerT &iter_indexer,
const InpIndexerT &inp_indexer,
const OutIndexerT &out_indexer,
TransformerT transformer,
const ScanOpT &scan_op,
outputT identity,
std::size_t &acc_groups,
const std::vector<sycl::event> &depends = {})
{
acc_groups = ceiling_quotient<std::size_t>(acc_nelems, n_wi * wg_size);

Expand All @@ -208,7 +219,7 @@ inclusive_scan_base_step(sycl::queue &exec_q,

slmT slm_iscan_tmp(lws, cgh);

using KernelName = inclusive_scan_iter_local_scan_krn<
using KernelName = inclusive_scan_iter_local_scan_blocked_krn<
inputT, outputT, n_wi, IterIndexerT, InpIndexerT, OutIndexerT,
TransformerT, ScanOpT, include_initial>;

Expand All @@ -218,6 +229,7 @@ inclusive_scan_base_step(sycl::queue &exec_q,
const std::size_t gid = it.get_global_id(0);
const std::size_t lid = it.get_local_id(0);

const std::uint32_t wg_size = it.get_local_range(0);
const std::size_t reduce_chunks = acc_groups * wg_size;
const std::size_t iter_gid = gid / reduce_chunks;
const std::size_t chunk_gid = gid - (iter_gid * reduce_chunks);
Expand Down Expand Up @@ -296,6 +308,248 @@ inclusive_scan_base_step(sycl::queue &exec_q,
return inc_scan_phase1_ev;
}

template <typename inputT,
typename outputT,
nwiT n_wi,
typename IterIndexerT,
typename InpIndexerT,
typename OutIndexerT,
typename TransformerT,
typename ScanOpT,
bool include_initial = false>
sycl::event
inclusive_scan_base_step_striped(sycl::queue &exec_q,
const std::uint32_t wg_size,
const std::size_t iter_nelems,
const std::size_t acc_nelems,
const inputT *input,
outputT *output,
const std::size_t s0,
const std::size_t s1,
const IterIndexerT &iter_indexer,
const InpIndexerT &inp_indexer,
const OutIndexerT &out_indexer,
TransformerT transformer,
const ScanOpT &scan_op,
outputT identity,
std::size_t &acc_groups,
const std::vector<sycl::event> &depends = {})
{
const std::uint32_t reduce_nelems_per_wg = n_wi * wg_size;
acc_groups =
ceiling_quotient<std::size_t>(acc_nelems, reduce_nelems_per_wg);

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

using slmT = sycl::local_accessor<outputT, 1>;

const auto &gRange = sycl::range<1>{iter_nelems * acc_groups * wg_size};
const auto &lRange = sycl::range<1>{wg_size};

const auto &ndRange = sycl::nd_range<1>{gRange, lRange};

slmT slm_iscan_tmp(reduce_nelems_per_wg, cgh);

using KernelName = inclusive_scan_iter_local_scan_striped_krn<
inputT, outputT, n_wi, IterIndexerT, InpIndexerT, OutIndexerT,
TransformerT, ScanOpT, include_initial>;

cgh.parallel_for<KernelName>(ndRange, [=, slm_iscan_tmp =
std::move(slm_iscan_tmp)](
sycl::nd_item<1> it) {
const std::uint32_t lid = it.get_local_linear_id();
const std::uint32_t wg_size = it.get_local_range(0);

const auto &sg = it.get_sub_group();
const std::uint32_t sgSize = sg.get_max_local_range()[0];
const std::size_t sgroup_id = sg.get_group_id()[0];
const std::uint32_t lane_id = sg.get_local_id()[0];

const std::size_t flat_group_id = it.get_group(0);
const std::size_t iter_gid = flat_group_id / acc_groups;
const std::size_t acc_group_id =
flat_group_id - (iter_gid * acc_groups);

const auto &iter_offsets = iter_indexer(iter_gid);
const auto &inp_iter_offset = iter_offsets.get_first_offset();
const auto &out_iter_offset = iter_offsets.get_second_offset();

std::array<outputT, n_wi> local_iscan{};

const std::size_t inp_id0 = acc_group_id * n_wi * wg_size +
sgroup_id * n_wi * sgSize + lane_id;

#pragma unroll
for (nwiT m_wi = 0; m_wi < n_wi; ++m_wi) {
const std::size_t inp_id = inp_id0 + m_wi * sgSize;
if constexpr (!include_initial) {
local_iscan[m_wi] =
(inp_id < acc_nelems)
? transformer(input[inp_iter_offset +
inp_indexer(s0 + s1 * inp_id)])
: identity;
}
else {
// shift input to the left by a single element relative to
// output
local_iscan[m_wi] =
(inp_id < acc_nelems && inp_id > 0)
? transformer(
input[inp_iter_offset +
inp_indexer((s0 + s1 * inp_id) - 1)])
: identity;
}
}

// change layout from striped to blocked
{
{
const std::uint32_t local_offset0 = lid * n_wi;
#pragma unroll
for (std::uint32_t i = 0; i < n_wi; ++i) {
slm_iscan_tmp[local_offset0 + i] = local_iscan[i];
}

it.barrier(sycl::access::fence_space::local_space);
}

{
const std::uint32_t block_offset =
sgroup_id * sgSize * n_wi;
const std::uint32_t disp0 = lane_id * n_wi;
#pragma unroll
for (nwiT i = 0; i < n_wi; ++i) {
const std::uint32_t disp = disp0 + i;

// disp == lane_id1 + i1 * sgSize;
const std::uint32_t i1 = disp / sgSize;
const std::uint32_t lane_id1 = disp - i1 * sgSize;

const std::uint32_t disp_exchanged =
(lane_id1 * n_wi + i1);

local_iscan[i] =
slm_iscan_tmp[block_offset + disp_exchanged];
}

it.barrier(sycl::access::fence_space::local_space);
}
}

#pragma unroll
for (nwiT m_wi = 1; m_wi < n_wi; ++m_wi) {
local_iscan[m_wi] =
scan_op(local_iscan[m_wi], local_iscan[m_wi - 1]);
}
// local_iscan is now result of
// inclusive scan of locally stored inputs

outputT wg_iscan_val;
if constexpr (can_use_inclusive_scan_over_group<ScanOpT,
outputT>::value)
{
wg_iscan_val = sycl::inclusive_scan_over_group(
it.get_group(), local_iscan.back(), scan_op, identity);
}
else {
wg_iscan_val = su_ns::custom_inclusive_scan_over_group(
it.get_group(), slm_iscan_tmp, local_iscan.back(), scan_op);
// ensure all finished reading from SLM, to avoid race condition
// with subsequent writes into SLM
it.barrier(sycl::access::fence_space::local_space);
}

slm_iscan_tmp[(lid + 1) % wg_size] = wg_iscan_val;
it.barrier(sycl::access::fence_space::local_space);
const outputT modifier = (lid == 0) ? identity : slm_iscan_tmp[lid];

#pragma unroll
for (nwiT m_wi = 0; m_wi < n_wi; ++m_wi) {
local_iscan[m_wi] = scan_op(local_iscan[m_wi], modifier);
}

it.barrier(sycl::access::fence_space::local_space);

// convert back to blocked layout
{
{
const std::uint32_t local_offset0 = lid * n_wi;
#pragma unroll
for (nwiT m_wi = 0; m_wi < n_wi; ++m_wi) {
slm_iscan_tmp[local_offset0 + m_wi] = local_iscan[m_wi];
}

it.barrier(sycl::access::fence_space::local_space);
}
}

{
const std::uint32_t block_offset =
sgroup_id * sgSize * n_wi + lane_id;
#pragma unroll
for (nwiT m_wi = 0; m_wi < n_wi; ++m_wi) {
const std::uint32_t m_wi_scaled = m_wi * sgSize;
const std::size_t out_id = inp_id0 + m_wi_scaled;
if (out_id < acc_nelems) {
output[out_iter_offset + out_indexer(out_id)] =
slm_iscan_tmp[block_offset + m_wi_scaled];
}
}
}
});
});

return inc_scan_phase1_ev;
}

template <typename inputT,
typename outputT,
nwiT n_wi,
typename IterIndexerT,
typename InpIndexerT,
typename OutIndexerT,
typename TransformerT,
typename ScanOpT,
bool include_initial = false>
sycl::event
inclusive_scan_base_step(sycl::queue &exec_q,
const std::uint32_t wg_size,
Comment thread
oleksandr-pavlyk marked this conversation as resolved.
const std::size_t iter_nelems,
const std::size_t acc_nelems,
const inputT *input,
outputT *output,
const std::size_t s0,
const std::size_t s1,
const IterIndexerT &iter_indexer,
const InpIndexerT &inp_indexer,
const OutIndexerT &out_indexer,
TransformerT transformer,
const ScanOpT &scan_op,
outputT identity,
std::size_t &acc_groups,
const std::vector<sycl::event> &depends = {})
{
// For small stride use striped load/store.
// Threshold value chosen experimentally.
if (s1 <= 16) {
return inclusive_scan_base_step_striped<
inputT, outputT, n_wi, IterIndexerT, InpIndexerT, OutIndexerT,
TransformerT, ScanOpT, include_initial>(
exec_q, wg_size, iter_nelems, acc_nelems, input, output, s0, s1,
iter_indexer, inp_indexer, out_indexer, transformer, scan_op,
identity, acc_groups, depends);
}
else {
return inclusive_scan_base_step_blocked<
inputT, outputT, n_wi, IterIndexerT, InpIndexerT, OutIndexerT,
TransformerT, ScanOpT, include_initial>(
exec_q, wg_size, iter_nelems, acc_nelems, input, output, s0, s1,
iter_indexer, inp_indexer, out_indexer, transformer, scan_op,
identity, acc_groups, depends);
}
}

template <typename inputT,
typename outputT,
nwiT n_wi,
Expand Down