Skip to content

[SYCL] Add prototype of sycl_khr_group_interface #15617

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.

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

Draft
wants to merge 4 commits into
base: sycl
Choose a base branch
from
Draft
Show file tree
Hide file tree
Changes from all commits
Commits
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
229 changes: 229 additions & 0 deletions sycl/include/sycl/khr/group_interface.hpp
Original file line number Diff line number Diff line change
@@ -0,0 +1,229 @@
//==----- group_interface.hpp --- sycl_khr_group_interface extension -------==//
//
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
//
//===----------------------------------------------------------------------===//
#pragma once
#ifdef __DPCPP_ENABLE_UNFINISHED_KHR_EXTENSIONS

#include <sycl/ext/oneapi/free_function_queries.hpp>
#include <sycl/id.hpp>
#include <sycl/range.hpp>

#if __cplusplus >= 202302L && defined(__has_include)
#if __has_include(<mdspan>)
#include <mdspan>
#endif
#endif

namespace sycl {
inline namespace _V1 {

namespace khr {

// Forward declarations for traits.
template <int Dimensions> class work_group;
class sub_group;
template <typename ParentGroup> class work_item;

} // namespace khr

namespace detail {
#if defined(__cpp_lib_mdspan)
template <typename IndexType, int Dimensions> struct single_extents;

template <typename IndexType> single_extents<1> {
using type = std::extents<IndexType, 1>;
}

template <typename IndexType> single_extents<2> {
using type = std::extents<IndexType, 1, 1>;
}

template <typename IndexType> single_extents<3> {
using type = std::extents<IndexType, 1, 1, 1>;
}
#endif

template <typename T> struct is_khr_group : public std::false_type {};

template <int Dimensions>
struct is_khr_group<khr::work_group<Dimensions>> : public std::true_type {};

template <> struct is_khr_group<khr::sub_group> : public std::true_type {};

} // namespace detail

namespace khr {

// Forward declaration for friend function.
template <typename ParentGroup>
std::enable_if_t<detail::is_khr_group<ParentGroup>::value,
work_item<ParentGroup>>
get_item(ParentGroup g);

template <int Dimensions = 1> class work_group {
public:
using id_type = id<Dimensions>;
using linear_id_type = size_t;
using range_type = range<Dimensions>;
#if defined(__cpp_lib_mdspan)
using extents_type = std::dextents<size_t, Dimensions>;
#endif
using size_type = size_t;
static constexpr int dimensions = Dimensions;
static constexpr memory_scope fence_scope = memory_scope::work_group;

work_group(group<Dimensions> g) noexcept {}

operator group<Dimensions>() const noexcept { return legacy(); }

id_type id() const noexcept { return legacy().get_group_id(); }

linear_id_type linear_id() const noexcept {
return legacy().get_group_linear_id();
}

range_type range() const noexcept { return legacy().get_group_range(); }

#if defined(__cpp_lib_mdspan)
constexpr extents_type extents() const noexcept {
auto LocalRange = legacy().get_local_range();
if constexpr (dimensions == 1) {
return extents_type(LocalRange[0]);
} else if constexpr (dimensions == 2) {
return extents_type(LocalRange[0], LocalRange[1]);
} else if constexpr (dimensions == 3) {
return extents_type(LocalRange[0], LocalRange[1], LocalRange[2]);
}
}

constexpr index_type extent(rank_type r) const noexcept {
return extents().extent(r);
}
#endif

constexpr size_type size() const noexcept {
return legacy().get_local_range().size();
}

private:
group<Dimensions> legacy() const noexcept {
return ext::oneapi::this_work_item::get_work_group<Dimensions>();
}
};

class sub_group {
public:
using id_type = id<1>;
using linear_id_type = uint32_t;
using range_type = range<1>;
#if defined(__cpp_lib_mdspan)
using extents_type = std::dextents<uint32_t, 1>;
#endif
using size_type = uint32_t;
static constexpr int dimensions = 1;
static constexpr memory_scope fence_scope = memory_scope::sub_group;

sub_group(sycl::sub_group g) noexcept {}

operator sycl::sub_group() const noexcept { return legacy(); }

id_type id() const noexcept { return legacy().get_group_id(); }

linear_id_type linear_id() const noexcept {
return legacy().get_group_linear_id();
}

range_type range() const noexcept { return legacy().get_group_range(); }

#if defined(__cpp_lib_mdspan)
constexpr extents_type extents() const noexcept {
return extents_type(legacy().get_local_range()[0]);
}

constexpr index_type extent(rank_type r) const noexcept {
return extents().extent(r);
}
#endif

constexpr size_type size() const noexcept {
return legacy().get_local_range()[0];
}

constexpr size_type max_size() const noexcept {
return legacy().get_max_local_range()[0];
}

private:
sycl::sub_group legacy() const noexcept {
return ext::oneapi::this_work_item::get_sub_group();
}
};

template <typename ParentGroup> class work_item {
public:
using id_type = typename ParentGroup::id_type;
using linear_id_type = typename ParentGroup::linear_id_type;
using range_type = typename ParentGroup::range_type;
#if defined(__cpp_lib_mdspan)
using extents_type =
detail::single_extents<typename ParentGroup::extents_type::index_type,
ParentGroup::dimensions>;
#endif
using size_type = typename ParentGroup::size_type;
static constexpr int dimensions = ParentGroup::dimensions;
static constexpr memory_scope fence_scope = memory_scope::work_item;

id_type id() const noexcept { return legacy().get_local_id(); }

linear_id_type linear_id() const noexcept {
return legacy().get_local_linear_id();
}

range_type range() const noexcept { return legacy().get_local_range(); }

#if defined(__cpp_lib_mdspan)
constexpr extents_type extents() const noexcept { return extents_type(); }

constexpr index_type extent(rank_type r) const noexcept {
return extents().extent(r);
}
#endif

constexpr size_type size() const noexcept { return 1; }

private:
auto legacy() const noexcept {
if constexpr (std::is_same_v<ParentGroup, sub_group>) {
return ext::oneapi::this_work_item::get_sub_group();
} else {
return ext::oneapi::this_work_item::get_work_group<
ParentGroup::dimensions>();
}
}

protected:
work_item() {}

friend work_item<ParentGroup> get_item<ParentGroup>(ParentGroup);
};

template <typename ParentGroup>
std::enable_if_t<detail::is_khr_group<ParentGroup>::value,
work_item<ParentGroup>>
get_item(ParentGroup g) {
return work_item<ParentGroup>{};
}

template <typename Group> bool leader_of(Group g) {
return get_item(g).linear_id() == 0;
}

} // namespace khr
} // namespace _V1
} // namespace sycl

#endif // __DPCPP_ENABLE_UNFINISHED_KHR_EXTENSIONS
40 changes: 40 additions & 0 deletions sycl/test-e2e/GroupInterface/leader_of.cpp
Original file line number Diff line number Diff line change
@@ -0,0 +1,40 @@
// RUN: %{build} -o %t.out
// RUN: %{run} %t.out

#include <cassert>
#include <iostream>
#include <sycl/detail/core.hpp>
#include <sycl/group_algorithm.hpp>
#define __DPCPP_ENABLE_UNFINISHED_KHR_EXTENSIONS
#include <sycl/khr/group_interface.hpp>

using namespace sycl;

void test(queue q) {
int out = 0;
size_t G = 4;

range<2> R(G, G);
{
buffer<int> out_buf(&out, 1);

q.submit([&](handler &cgh) {
auto out = out_buf.template get_access<access::mode::read_write>(cgh);
cgh.parallel_for(nd_range<2>(R, R), [=](nd_item<2> it) {
khr::work_group<2> g = it.get_group();
if (khr::leader_of(g)) {
out[0] += 1;
}
});
});
}
assert(out == 1);
}

int main() {
queue q;
test(q);

std::cout << "Test passed." << std::endl;
return 0;
}
111 changes: 111 additions & 0 deletions sycl/test-e2e/GroupInterface/sub_group.cpp
Original file line number Diff line number Diff line change
@@ -0,0 +1,111 @@
// REQUIRES: cpu

// RUN: %{build} %cxx_std_optionc++23 -o %t.out
// RUN: %{run} %t.out

#include <sycl/detail/core.hpp>
#define __DPCPP_ENABLE_UNFINISHED_KHR_EXTENSIONS
#include <sycl/khr/group_interface.hpp>

#include <sycl/builtins.hpp>

#include <type_traits>

using namespace sycl;

static_assert(std::is_same_v<khr::sub_group::id_type, id<1>>);
static_assert(std::is_same_v<khr::sub_group::linear_id_type, uint32_t>);
static_assert(std::is_same_v<khr::sub_group::range_type, range<1>>);
#if defined(__cpp_lib_mdspan)
static_assert(
std::is_same_v<khr::sub_group::extents_type, std::dextents<uint32_t, 1>>);
#endif
static_assert(std::is_same_v<khr::sub_group::size_type, uint32_t>);
static_assert(khr::sub_group::dimensions == 1);
static_assert(khr::sub_group::fence_scope == memory_scope::sub_group);

int main() {
queue q(cpu_selector_v);

const int sz = 16;
q.submit([&](handler &h) {
h.parallel_for(nd_range<1>{sz, sz}, [=](nd_item<1> item) {
sub_group g = item.get_sub_group();

khr::sub_group sg = g;
assert(sg.id() == g.get_group_id());
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'm not very confident in using device asserts for our tests... Does anybody know for sure it's stable enough?

assert(sg.linear_id() == g.get_group_linear_id());
assert(sg.range() == g.get_group_range());
#if defined(__cpp_lib_mdspan)
assert(sg.extents().rank() == 1);
assert(sg.extent(0) == g.get_local_range()[0]);
#endif
assert(sg.size() == g.get_local_linear_range());
assert(sg.max_size() == g.get_max_local_range()[0]);

khr::work_item wi = get_item(sg);
assert(wi.id() == g.get_local_id());
assert(wi.linear_id() == g.get_local_linear_id());
assert(wi.range() == g.get_local_range());
#if defined(__cpp_lib_mdspan)
assert(wi.extents().rank() == 1);
assert(wi.extent(0) == 1);
#endif
assert(wi.size() == 1);
});
});
q.submit([&](handler &h) {
h.parallel_for(nd_range<2>{range<2>{sz, sz}, range<2>{sz, sz}},
[=](nd_item<2> item) {
sub_group g = item.get_sub_group();

khr::sub_group sg = g;
assert(sg.id() == g.get_group_id());
assert(sg.linear_id() == g.get_group_linear_id());
assert(sg.range() == g.get_group_range());
#if defined(__cpp_lib_mdspan)
assert(sg.extents().rank() == 1);
assert(sg.extent(0) == g.get_local_range()[0]);
#endif
assert(sg.size() == g.get_local_linear_range());
assert(sg.max_size() == g.get_max_local_range()[0]);

khr::work_item wi = get_item(sg);
assert(wi.id() == g.get_local_id());
assert(wi.linear_id() == g.get_local_linear_id());
assert(wi.range() == g.get_local_range());
#if defined(__cpp_lib_mdspan)
assert(wi.extents().rank() == 1);
assert(wi.extent(0) == 1);
#endif
assert(wi.size() == 1);
});
});
q.submit([&](handler &h) {
h.parallel_for(nd_range<3>{range<3>{sz, sz, sz}, range<3>{sz, sz, sz}},
[=](nd_item<3> item) {
sub_group g = item.get_sub_group();

khr::sub_group sg = g;
assert(sg.id() == g.get_group_id());
assert(sg.linear_id() == g.get_group_linear_id());
assert(sg.range() == g.get_group_range());
#if defined(__cpp_lib_mdspan)
assert(sg.extents().rank() == 1);
assert(sg.extent(0) == g.get_local_range()[0]);
#endif
assert(sg.size() == g.get_local_linear_range());
assert(sg.max_size() == g.get_max_local_range()[0]);

khr::work_item wi = get_item(sg);
assert(wi.id() == g.get_local_id());
assert(wi.linear_id() == g.get_local_linear_id());
assert(wi.range() == g.get_local_range());
#if defined(__cpp_lib_mdspan)
assert(wi.extents().rank() == 1);
assert(wi.extent(0) == 1);
#endif
});
});
q.wait();
}
Loading
Loading