From 95238df3119dc140bc90b6b0df26c1b5521c236f Mon Sep 17 00:00:00 2001
From: Peter Thoman <petert@dps.uibk.ac.at>
Date: Wed, 30 Apr 2025 13:44:04 +0200
Subject: [PATCH] SYCL_KHR_GROUP_INTERFACE prototype

---
 .../simsycl/extensions/khr_group_interface.hh | 140 +++++++++
 test/CMakeLists.txt                           |   1 +
 test/extensions/khr_group_interface_tests.cc  | 272 ++++++++++++++++++
 3 files changed, 413 insertions(+)
 create mode 100644 include/simsycl/extensions/khr_group_interface.hh
 create mode 100644 test/extensions/khr_group_interface_tests.cc

diff --git a/include/simsycl/extensions/khr_group_interface.hh b/include/simsycl/extensions/khr_group_interface.hh
new file mode 100644
index 0000000..5513bf9
--- /dev/null
+++ b/include/simsycl/extensions/khr_group_interface.hh
@@ -0,0 +1,140 @@
+#include "sycl/sycl.hpp" // IWYU pragma: keep
+
+#define SYCL_KHR_GROUP_INTERFACE 1
+
+namespace simsycl::sycl::khr {
+
+template<typename ParentGroup>
+class member_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;
+    // using extents_type = /* extents of all 1s with ParentGroup's index type */; // C++23
+    using size_type = typename ParentGroup::size_type;
+    static constexpr int dimensions = ParentGroup::dimensions;
+    static constexpr memory_scope fence_scope = memory_scope::work_item;
+
+    /* -- common by-value interface members -- */
+
+    id_type id() const noexcept { return m_parent_group.get_local_id(); }
+    linear_id_type linear_id() const noexcept { return m_parent_group.get_local_linear_id(); }
+
+    range_type range() const noexcept { return m_parent_group.get_local_range(); }
+
+    // constexpr extents_type extents() const noexcept;                                     // C++23
+    // constexpr extents_type::index_type extent(extents_type::rank_type r) const noexcept; // C++23
+
+    // static constexpr extents_type::rank_type rank() noexcept;         // C++23
+    // static constexpr extents_type::rank_type rank_dynamic() noexcept; // C++23
+    // static constexpr size_t static_extent(rank_type r) noexcept;      // C++23
+
+    constexpr size_type size() const noexcept { return 1; }
+
+  private:
+    ParentGroup m_parent_group;
+    member_item(ParentGroup g) noexcept : m_parent_group(g) {}
+
+    linear_id_type get_local_linear_id() const noexcept { return m_parent_group.get_local_linear_id(); }
+
+    template<typename Group>
+    friend member_item<Group> get_member_item(Group g) noexcept;
+    template<typename Group>
+    friend bool leader_of(Group g) noexcept;
+};
+
+template<int Dimensions = 1>
+class work_group {
+  public:
+    using id_type = sycl::id<Dimensions>;
+    using linear_id_type = size_t;
+    using range_type = sycl::range<Dimensions>;
+    // using extents_type = std::dextents<size_t, Dimensions>; // C++23
+    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 : m_group(g) {}
+
+    operator group<Dimensions>() const noexcept { return m_group; }
+
+    /* -- common by-value interface members -- */
+
+    id_type id() const noexcept { return m_group.get_group_id(); }
+    linear_id_type linear_id() const noexcept { return m_group.get_group_linear_id(); }
+
+    range_type range() const noexcept { return m_group.get_group_range(); }
+
+    // extents_type extents() const noexcept;                                     // C++23
+    // extents_type::index_type extent(extents_type::rank_type r) const noexcept; // C++23
+
+    // static constexpr extents_type::rank_type rank() noexcept;         // C++23
+    // static constexpr extents_type::rank_type rank_dynamic() noexcept; // C++23
+    // static constexpr size_t static_extent(rank_type r) noexcept;      // C++23
+
+    size_type size() const noexcept { return m_group.get_local_range().size(); }
+
+  private:
+    group<Dimensions> m_group;
+
+    id_type get_local_id() const noexcept { return m_group.get_local_id(); }
+    linear_id_type get_local_linear_id() const noexcept { return m_group.get_local_linear_id(); }
+    range_type get_local_range() const noexcept { return m_group.get_local_range(); }
+    friend class member_item<work_group>;
+    template<typename Group>
+    friend bool leader_of(Group g) noexcept;
+};
+
+class sub_group {
+  public:
+    using id_type = sycl::id<1>;
+    using linear_id_type = uint32_t;
+    using range_type = sycl::range<1>;
+    // using extents_type = std::dextents<uint32_t, 1>; // C++23
+    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 sg) noexcept : m_sub_group(sg) {}
+
+    operator sycl::sub_group() const noexcept { return m_sub_group; }
+
+    /* -- common by-value interface members -- */
+
+    id_type id() const noexcept { return m_sub_group.get_group_id(); }
+    linear_id_type linear_id() const noexcept { return m_sub_group.get_group_linear_id(); }
+
+    range_type range() const noexcept { return m_sub_group.get_group_range(); }
+
+    // extents_type extents() const noexcept;                                     // C++23
+    // extents_type::index_type extent(extents_type::rank_type r) const noexcept; // C++23
+
+    // static constexpr extents_type::rank_type rank() noexcept;         // C++23
+    // static constexpr extents_type::rank_type rank_dynamic() noexcept; // C++23
+    // static constexpr size_t static_extent(rank_type r) noexcept;      // C++23
+
+    size_type size() const noexcept { return m_sub_group.get_local_range().size(); }
+    size_type max_size() const noexcept { return m_sub_group.get_max_local_range().size(); }
+
+  private:
+    sycl::sub_group m_sub_group;
+
+    id_type get_local_id() const noexcept { return m_sub_group.get_local_id(); }
+    linear_id_type get_local_linear_id() const noexcept { return m_sub_group.get_local_linear_id(); }
+    range_type get_local_range() const noexcept { return m_sub_group.get_local_range(); }
+    friend class member_item<sub_group>;
+    template<typename Group>
+    friend bool leader_of(Group g) noexcept;
+};
+
+template<typename Group>
+member_item<Group> get_member_item(Group g) noexcept {
+    return member_item<Group>(g);
+}
+
+template<typename Group>
+bool leader_of(Group g) noexcept {
+    return g.get_local_linear_id() == 0;
+}
+
+} // namespace simsycl::sycl::khr
diff --git a/test/CMakeLists.txt b/test/CMakeLists.txt
index 3b95dd7..457b23c 100644
--- a/test/CMakeLists.txt
+++ b/test/CMakeLists.txt
@@ -22,6 +22,7 @@ add_executable(tests
   simulation_tests.cc
   alloc_tests.cc
   vec_tests.cc
+  extensions/khr_group_interface_tests.cc
 )
 
 add_sycl_to_target(TARGET tests SIMSYCL_ALL_WARNINGS)
diff --git a/test/extensions/khr_group_interface_tests.cc b/test/extensions/khr_group_interface_tests.cc
new file mode 100644
index 0000000..5d5fa2b
--- /dev/null
+++ b/test/extensions/khr_group_interface_tests.cc
@@ -0,0 +1,272 @@
+#include <simsycl/extensions/khr_group_interface.hh>
+#include <sycl/sycl.hpp>
+
+#include <catch2/catch_test_macros.hpp>
+
+namespace group_interface::tests {
+
+template<int Dimensions>
+static void test_work_group(sycl::nd_item<Dimensions> it) {
+    sycl::khr::work_group<Dimensions> work_group{it.get_group()};
+    sycl::group<Dimensions> group{it.get_group()};
+
+    // id
+    static_assert(std::is_same_v<decltype(work_group.id()), typename sycl::khr::work_group<Dimensions>::id_type>);
+    CHECK(group.get_group_id() == work_group.id());
+
+    // linear_id
+    static_assert(
+        std::is_same_v<decltype(work_group.linear_id()), typename sycl::khr::work_group<Dimensions>::linear_id_type>);
+    CHECK(group.get_group_linear_id() == work_group.linear_id());
+
+    // range
+    static_assert(std::is_same_v<decltype(work_group.range()), typename sycl::khr::work_group<Dimensions>::range_type>);
+    CHECK(group.get_group_range() == work_group.range());
+
+#if __cplusplus >= 202302L
+    // extents
+    static_assert(
+        std::is_same_v<decltype(work_group.extents()), typename sycl::khr::work_group<Dimensions>::extents_type>);
+    {
+        const sycl::range<Dimensions> localRange = group.get_local_range();
+        if constexpr(Dimensions == 1)
+            CHECK(work_group.extents() == std::dextents<std::size_t, Dimensions>(localRange[0]));
+        else if constexpr(Dimensions == 2)
+            CHECK(work_group.extents() == std::dextents<std::size_t, Dimensions>(localRange[0], localRange[1]));
+        else if constexpr(Dimensions == 3)
+            CHECK(work_group.extents()
+                == std::dextents<std::size_t, Dimensions>(localRange[0], localRange[1], localRange[2]));
+    }
+
+    // extent
+    static_assert(std::is_same_v<decltype(work_group.extent(0)),
+        typename sycl::khr::work_group<Dimensions>::extents_type::index_type>);
+    for(int i = 0; i < Dimensions; i++) CHECK(work_group.extent(i) == work_group.extents().extent(i));
+
+    // rank
+    static_assert(std::is_same_v<decltype(sycl::khr::work_group<Dimensions>::rank()),
+        typename sycl::khr::work_group<Dimensions>::extents_type::rank_type>);
+    static_assert(decltype(work_group)::rank() == decltype(work_group.extents())::rank());
+
+    // rank_dynamic
+    static_assert(std::is_same_v<decltype(sycl::khr::work_group<Dimensions>::rank_dynamic()),
+        typename sycl::khr::work_group<Dimensions>::extents_type::rank_type>);
+    static_assert(decltype(work_group)::rank_dynamic() == decltype(work_group.extents())::rank_dynamic());
+
+    // static_extent
+    static_assert(std::is_same_v<decltype(sycl::khr::work_group<Dimensions>::static_extent(0)), std::size_t>);
+    static_assert(decltype(work_group)::static_extent(0) == decltype(work_group.extents())::static_extent(0));
+    if constexpr(Dimensions >= 2)
+        static_assert(decltype(work_group)::static_extent(1) == decltype(work_group.extents())::static_extent(1));
+    if constexpr(Dimensions == 3)
+        static_assert(decltype(work_group)::static_extent(2) == decltype(work_group.extents())::static_extent(2));
+#endif
+
+    // size
+    static_assert(std::is_same_v<decltype(work_group.size()), typename sycl::khr::work_group<Dimensions>::size_type>);
+    CHECK(group.get_local_linear_range() == work_group.size());
+
+    // leader_of
+    static_assert(std::is_same_v<decltype(sycl::khr::leader_of(work_group)), bool>);
+    CHECK(group.leader() == sycl::khr::leader_of(work_group));
+}
+
+template<int Dimensions>
+static void test_sub_group(sycl::nd_item<Dimensions> it) {
+    sycl::khr::sub_group sub_group{it.get_sub_group()};
+    sycl::sub_group group{it.get_sub_group()};
+
+    // id
+    static_assert(std::is_same_v<decltype(sub_group.id()), typename sycl::khr::sub_group::id_type>);
+    CHECK(group.get_group_id() == sub_group.id());
+
+    // linear_id
+    static_assert(std::is_same_v<decltype(sub_group.linear_id()), typename sycl::khr::sub_group::linear_id_type>);
+    CHECK(group.get_group_linear_id() == sub_group.linear_id());
+
+    // range
+    static_assert(std::is_same_v<decltype(sub_group.range()), typename sycl::khr::sub_group::range_type>);
+    CHECK(group.get_group_range() == sub_group.range());
+
+#if __cplusplus >= 202302L
+    // extents
+    static_assert(std::is_same_v<decltype(sub_group.extents()), typename sycl::khr::sub_group::extents_type>);
+    CHECK(sub_group.extents() == std::dextents<std::uint32_t, 1>(group.get_local_linear_range()));
+
+    // extent
+    static_assert(
+        std::is_same_v<decltype(sub_group.extent(0)), typename sycl::khr::sub_group ::extents_type::index_type>);
+    CHECK(sub_group.extent(0) == sub_group.extents().extent(0));
+
+    // rank
+    static_assert(
+        std::is_same_v<decltype(sycl::khr::sub_group::rank()), typename sycl::khr::sub_group::extents_type::rank_type>);
+    static_assert(decltype(sub_group)::rank() == decltype(sub_group.extents())::rank());
+
+    // rank_dynamic
+    static_assert(std::is_same_v<decltype(sycl::khr::sub_group::rank_dynamic()),
+        typename sycl::khr::sub_group::extents_type::rank_type>);
+    static_assert(decltype(sub_group)::rank_dynamic() == decltype(sub_group.extents())::rank_dynamic());
+
+    // static_extent
+    static_assert(std::is_same_v<decltype(sycl::khr::sub_group::static_extent(0)), std::size_t>);
+    static_assert(decltype(sub_group)::static_extent(0) == decltype(sub_group.extents())::static_extent(0));
+#endif
+
+    // size
+    static_assert(std::is_same_v<decltype(sub_group.size()), typename sycl::khr::sub_group::size_type>);
+    CHECK(group.get_local_range()[0] == sub_group.size());
+
+    // max_size
+    static_assert(std::is_same_v<decltype(sub_group.max_size()), typename sycl::khr::sub_group::size_type>);
+    CHECK(group.get_max_local_range()[0] == sub_group.max_size());
+
+    // leader_of
+    static_assert(std::is_same_v<decltype(sycl::khr::leader_of(sub_group)), bool>);
+    CHECK(group.leader() == sycl::khr::leader_of(sub_group));
+}
+
+template<int Dimensions>
+static void test_work_item_group(sycl::nd_item<Dimensions> it) {
+    sycl::group<Dimensions> group{it.get_group()};
+    sycl::khr::work_group<Dimensions> work_group{group};
+    sycl::khr::member_item item{sycl::khr::get_member_item(work_group)};
+
+    // id
+    static_assert(std::is_same_v<decltype(item.id()),
+        typename sycl::khr::member_item<sycl::khr::work_group<Dimensions>>::id_type>);
+    CHECK(group.get_local_id() == item.id());
+
+    // linear_id
+    static_assert(std::is_same_v<decltype(item.linear_id()),
+        typename sycl::khr::member_item<sycl::khr::work_group<Dimensions>>::linear_id_type>);
+    CHECK(group.get_local_linear_id() == item.linear_id());
+
+    // range
+    static_assert(std::is_same_v<decltype(item.range()),
+        typename sycl::khr::member_item<sycl::khr::work_group<Dimensions>>::range_type>);
+    CHECK(group.get_local_range() == item.range());
+
+#if __cplusplus >= 202302L
+    // extents
+    static_assert(std::is_same_v<decltype(item.extents()),
+        typename sycl::khr::member_item<sycl::khr::work_group<Dimensions>>::extents_type>);
+    if constexpr(Dimensions == 1)
+        CHECK(item.extents() == std::extents<std::size_t, 1>());
+    else if constexpr(Dimensions == 2)
+        CHECK(item.extents() == std::extents<std::size_t, 1, 1>());
+    else if constexpr(Dimensions == 3)
+        CHECK(item.extents() == std::extents<std::size_t, 1, 1, 1>());
+
+    // extent
+    static_assert(std::is_same_v<decltype(item.extent(0)),
+        typename sycl::khr::member_item<sycl::khr::work_group<Dimensions>>::extents_type::index_type>);
+    for(int i = 0; i < Dimensions; i++) CHECK(item.extent(i) == item.extents().extent(i));
+
+    // rank
+    static_assert(std::is_same_v<decltype(sycl::khr::member_item<sycl::khr::work_group<Dimensions>>::rank()),
+        typename sycl::khr::member_item<sycl::khr::work_group<Dimensions>>::extents_type::rank_type>);
+    static_assert(decltype(item)::rank() == decltype(item.extents())::rank());
+
+    // rank_dynamic
+    static_assert(std::is_same_v<decltype(sycl::khr::member_item<sycl::khr::work_group<Dimensions>>::rank_dynamic()),
+        typename sycl::khr::member_item<sycl::khr::work_group<Dimensions>>::extents_type::rank_type>);
+    static_assert(decltype(item)::rank_dynamic() == decltype(item.extents())::rank_dynamic());
+
+    // static_extent
+    static_assert(std::is_same_v<decltype(sycl::khr::member_item<sycl::khr::work_group<Dimensions>>::static_extent(0)),
+        std::size_t>);
+    static_assert(decltype(item)::static_extent(0) == decltype(item.extents())::static_extent(0));
+    if constexpr(Dimensions >= 2)
+        static_assert(decltype(item)::static_extent(1) == decltype(item.extents())::static_extent(1));
+    if constexpr(Dimensions == 3)
+        static_assert(decltype(item)::static_extent(2) == decltype(item.extents())::static_extent(2));
+#endif
+
+    // size
+    static_assert(std::is_same_v<decltype(item.size()),
+        typename sycl::khr::member_item<sycl::khr::work_group<Dimensions>>::size_type>);
+    CHECK(1 == item.size());
+}
+
+template<int Dimensions>
+static void test_work_item_subgroup(sycl::nd_item<Dimensions> it) {
+    sycl::sub_group group{it.get_sub_group()};
+    sycl::khr::sub_group sub_group{group};
+    sycl::khr::member_item item{sycl::khr::get_member_item(sub_group)};
+
+    // id
+    static_assert(std::is_same_v<decltype(item.id()), typename sycl::khr::member_item<sycl::khr::sub_group>::id_type>);
+    CHECK(group.get_local_id() == item.id());
+
+    // linear_id
+    static_assert(std::is_same_v<decltype(item.linear_id()),
+        typename sycl::khr::member_item<sycl::khr::sub_group>::linear_id_type>);
+    CHECK(group.get_local_linear_id() == item.linear_id());
+
+    // range
+    static_assert(
+        std::is_same_v<decltype(item.range()), typename sycl::khr::member_item<sycl::khr::sub_group>::range_type>);
+    CHECK(group.get_local_range() == item.range());
+
+#if __cplusplus >= 202302L
+    // extents
+    static_assert(
+        std::is_same_v<decltype(item.extents()), typename sycl::khr::member_item<sycl::khr::sub_group>::extents_type>);
+    CHECK(item.extents() == std::extents<std::uint32_t, 1>());
+
+    // extent
+    static_assert(std::is_same_v<decltype(item.extent(0)),
+        typename sycl::khr::member_item<sycl::khr::sub_group>::extents_type::index_type>);
+    CHECK(item.extent(0) == item.extents().extent(0));
+
+    // rank
+    static_assert(std::is_same_v<decltype(sycl::khr::member_item<sycl::khr::sub_group>::rank()),
+        typename sycl::khr::member_item<sycl::khr::sub_group>::extents_type::rank_type>);
+    static_assert(decltype(item)::rank() == decltype(item.extents())::rank());
+
+    // rank_dynamic
+    static_assert(std::is_same_v<decltype(sycl::khr::member_item<sycl::khr::sub_group>::rank_dynamic()),
+        typename sycl::khr::member_item<sycl::khr::sub_group>::extents_type::rank_type>);
+    static_assert(decltype(item)::rank_dynamic() == decltype(item.extents())::rank_dynamic());
+
+    // static_extent
+    static_assert(std::is_same_v<decltype(sycl::khr::member_item<sycl::khr::work_group<Dimensions>>::static_extent(0)),
+        std::size_t>);
+    static_assert(decltype(item)::static_extent(0) == decltype(item.extents())::static_extent(0));
+#endif
+
+    // size
+    static_assert(
+        std::is_same_v<decltype(item.size()), typename sycl::khr::member_item<sycl::khr::sub_group>::size_type>);
+    CHECK(1 == item.size());
+}
+
+template<int Dimensions>
+static void test_group_interface(sycl::nd_range<Dimensions> nd_range) {
+    sycl::queue q{};
+    q.submit([&](sycl::handler &cgh) {
+        cgh.parallel_for(nd_range, [=](sycl::nd_item<Dimensions> it) {
+            test_work_group(it);
+            test_sub_group(it);
+            test_work_item_group(it);
+            test_work_item_subgroup(it);
+        });
+    });
+}
+
+TEST_CASE("the group interface extension defines the SYCL_KHR_GROUP_INTERFACE "
+          "macro",
+    "[khr_group_interface]") {
+#ifndef SYCL_KHR_GROUP_INTERFACE
+    static_assert(false, "SYCL_KHR_GROUP_INTERFACE is not defined");
+#endif
+}
+
+TEST_CASE("khr group interface extension work_group, sub_group, and work_item classes", "[khr_group_interface]") {
+    test_group_interface(sycl::nd_range<1>(10, 2));
+    test_group_interface(sycl::nd_range<2>({6, 6}, {2, 3}));
+    test_group_interface(sycl::nd_range<3>({2, 4, 6}, {1, 2, 3}));
+}
+
+} // namespace group_interface::tests