Skip to content

Contribution to fix gh 1887 #1891

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 2 commits into from
Nov 13, 2024
Merged

Contribution to fix gh 1887 #1891

Changes from 1 commit
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
04fd35c
Tweaks to copy_as_contiguous CContig functor
oleksandr-pavlyk Nov 13, 2024
18f7ea0
Added static assert to verify that n_vecs * vec_sz fits in uint8_t
oleksandr-pavlyk Nov 13, 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
Tweaks to copy_as_contiguous CContig functor 
1. Save common subexpressions to variables
2. Sub-group size type changed to uint16 (from uint32)
3. sg.get_local_range() replaced with sg.get_max_local_range()

   This is safe to do since work-group size is chosen to be
   a multiple of sub-group size for all possile choices
   of sub-group size (1, 8, 16, 32, 64)

4. Simplified computation of base value in generic branch for
   complex types, or when sg_load is disabled, to avoid a
   division (and left a comment)
  • Loading branch information
@oleksandr-pavlyk
oleksandr-pavlyk committed Nov 13, 2024
commit 04fd35ccf66bbcd3f89dad094be38a7bb55bbe7f
70 changes: 41 additions & 29 deletions dpctl/tensor/libtensor/include/kernels/copy_as_contiguous.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -44,8 +44,8 @@ namespace copy_as_contig

template <typename T,
typename IndexerT,
int vec_sz = 4,
int n_vecs = 2,
std::uint32_t vec_sz = 4u,
std::uint32_t n_vecs = 2u,
bool enable_sg_loadstore = true>
class CopyAsCContigFunctor
{
Expand All @@ -66,53 +66,63 @@ class CopyAsCContigFunctor

void operator()(sycl::nd_item<1> ndit) const
{
static_assert(vec_sz > 0);
static_assert(n_vecs > 0);
static_assert(vec_sz * n_vecs < (std::uint32_t(1) << 8));

constexpr std::uint8_t elems_per_wi =
static_cast<std::uint8_t>(vec_sz * n_vecs);

using dpctl::tensor::type_utils::is_complex;
if constexpr (!enable_sg_loadstore || is_complex<T>::value) {
const std::uint32_t sgSize =
const std::uint16_t sgSize =
ndit.get_sub_group().get_local_range()[0];
const std::size_t gid = ndit.get_global_linear_id();

const std::size_t base =
(gid / sgSize) * sgSize * n_vecs * vec_sz + (gid % sgSize);
for (size_t offset = base;
offset < std::min(nelems, base + sgSize * (n_vecs * vec_sz));
offset += sgSize)
{
// base = (gid / sgSize) * sgSize * elems_per_wi + (gid % sgSize)
// gid % sgSize == gid - (gid / sgSize) * sgSize
const std::size_t elems_per_sg = sgSize * (elems_per_wi - 1);
const std::size_t base = (gid / sgSize) * elems_per_sg + gid;
const std::size_t offset_max =
std::min(nelems, base + sgSize * elems_per_wi);

for (size_t offset = base; offset < offset_max; offset += sgSize) {
auto src_offset = src_indexer(offset);
dst_p[offset] = src_p[src_offset];
}
}
else {
auto sg = ndit.get_sub_group();
const std::uint32_t sgSize = sg.get_local_range()[0];
const size_t base = n_vecs * vec_sz *
(ndit.get_group(0) * ndit.get_local_range(0) +
sg.get_group_id()[0] * sgSize);
const std::uint16_t sgSize = sg.get_max_local_range()[0];
const size_t base =
elems_per_wi * (ndit.get_group(0) * ndit.get_local_range(0) +
sg.get_group_id()[0] * sgSize);

if (base + n_vecs * vec_sz * sgSize < nelems) {
if (base + elems_per_wi * sgSize < nelems) {
sycl::vec<T, vec_sz> dst_vec;

#pragma unroll
for (std::uint32_t it = 0; it < n_vecs * vec_sz; it += vec_sz) {
for (std::uint8_t it = 0; it < elems_per_wi; it += vec_sz) {
const size_t block_start_id = base + it * sgSize;
auto dst_multi_ptr = sycl::address_space_cast<
sycl::access::address_space::global_space,
sycl::access::decorated::yes>(
&dst_p[base + it * sgSize]);
sycl::access::decorated::yes>(&dst_p[block_start_id]);

const size_t elem_id0 = block_start_id + sg.get_local_id();
#pragma unroll
for (std::uint32_t k = 0; k < vec_sz; k++) {
ssize_t src_offset = src_indexer(
base + (it + k) * sgSize + sg.get_local_id());
for (std::uint8_t k = 0; k < vec_sz; k++) {
const size_t elem_id = elem_id0 + k * sgSize;
const ssize_t src_offset = src_indexer(elem_id);
dst_vec[k] = src_p[src_offset];
}
sg.store<vec_sz>(dst_multi_ptr, dst_vec);
}
}
else {
for (size_t k = base + sg.get_local_id()[0]; k < nelems;
k += sgSize)
{
ssize_t src_offset = src_indexer(k);
const size_t lane_id = sg.get_local_id()[0];
const size_t k0 = base + lane_id;
for (size_t k = k0; k < nelems; k += sgSize) {
const ssize_t src_offset = src_indexer(k);
dst_p[k] = src_p[src_offset];
}
}
Expand Down Expand Up @@ -150,9 +160,11 @@ sycl::event submit_c_contiguous_copy(sycl::queue &exec_q,
const std::size_t lws =
((preferred_lws + max_sg_size - 1) / max_sg_size) * max_sg_size;

constexpr std::uint32_t nelems_per_wi = n_vecs * vec_sz;
size_t n_groups =
(nelems + nelems_per_wi * lws - 1) / (nelems_per_wi * lws);
constexpr std::uint8_t nelems_per_wi = n_vecs * vec_sz;
Comment thread
oleksandr-pavlyk marked this conversation as resolved.

const size_t nelems_per_group = nelems_per_wi * lws;
const size_t n_groups =
(nelems + nelems_per_group - 1) / (nelems_per_group);

sycl::event copy_ev = exec_q.submit([&](sycl::handler &cgh) {
cgh.depends_on(depends);
Expand All @@ -171,8 +183,8 @@ sycl::event submit_c_contiguous_copy(sycl::queue &exec_q,

template <typename T,
typename IndexT,
int vec_sz,
int n_vecs,
std::uint32_t vec_sz,
std::uint32_t n_vecs,
bool enable_sgload>
class as_contig_krn;

Expand Down