Implement SYCL_KHR_WORK_ITEM_QUERIES

.github/workflows/extension_ci.yml

@@ -36,6 +36,7 @@ jobs:
         -S ${{ github.workspace }}
         -DCMAKE_INSTALL_PREFIX=${{ steps.strings.outputs.install-dir }}
         -DSIMSYCL_ENABLE_SYCL_KHR_QUEUE_FLUSH=OFF
+        -DSIMSYCL_ENABLE_SYCL_KHR_WORK_ITEM_QUERIES=OFF
 
     - name: Build SimSYCL (no extensions)
       run: >
@@ -71,6 +72,7 @@ jobs:
         -S ${{ github.workspace }}
         -DCMAKE_INSTALL_PREFIX=${{ steps.strings.outputs.install-dir }}
         -DSIMSYCL_ENABLE_SYCL_KHR_QUEUE_FLUSH=ON
+        -DSIMSYCL_ENABLE_SYCL_KHR_WORK_ITEM_QUERIES=ON
 
     - name: Build SimSYCL (with extensions)
       run: >

CMakeLists.txt

@@ -69,6 +69,7 @@ set(SIMSYCL_CHECK_MODE "ABORT" CACHE STRING "Runtime assertion handling NONE|LOG
 
 # Extension options
 option(SIMSYCL_ENABLE_SYCL_KHR_QUEUE_FLUSH "Enable the SYCL_KHR_QUEUE_FLUSH extension" ON)
+option(SIMSYCL_ENABLE_SYCL_KHR_WORK_ITEM_QUERIES "Enable the SYCL_KHR_WORK_ITEM_QUERIES extension" ON)
 
 set(CONFIG_PATH "${CMAKE_CURRENT_BINARY_DIR}/include/simsycl/config.hh")
 configure_file(

cmake/simsycl-config.cmake.in

@@ -28,5 +28,6 @@ set(SIMSYCL_CHECK_MODE "@SIMSYCL_CHECK_MODE@")
 set(SIMSYCL_ENABLE_ASAN "@SIMSYCL_ENABLE_ASAN@")
 
 set(SIMSYCL_ENABLE_SYCL_KHR_QUEUE_FLUSH "@SIMSYCL_ENABLE_SYCL_KHR_QUEUE_FLUSH@")
+set(SIMSYCL_ENABLE_SYCL_KHR_WORK_ITEM_QUERIES "@SIMSYCL_ENABLE_SYCL_KHR_WORK_ITEM_QUERIES@")
 
 include("${CMAKE_CURRENT_LIST_DIR}/AddToTarget.cmake")

include/simsycl/config.hh.in

@@ -9,6 +9,7 @@
 #cmakedefine01 SIMSYCL_FEATURE_HALF_TYPE
 
 #cmakedefine01 SIMSYCL_ENABLE_SYCL_KHR_QUEUE_FLUSH
+#cmakedefine01 SIMSYCL_ENABLE_SYCL_KHR_WORK_ITEM_QUERIES
 
 #ifndef SIMSYCL_CHECK_MODE
 #define SIMSYCL_CHECK_MODE SIMSYCL_CHECK_@SIMSYCL_CHECK_MODE@

include/simsycl/sycl.hh

@@ -44,4 +44,6 @@
 #include "sycl/type_traits.hh"
 #include "sycl/usm.hh"
 #include "sycl/vec.hh"
+
+#include "sycl/khr/sub_group_queries.hh"
 // IWYU pragma: end_keep

include/simsycl/sycl/khr/sub_group_queries.hh

@@ -0,0 +1,44 @@
+#include <simsycl/sycl/group.hh>
+#include <simsycl/sycl/nd_item.hh>
+#include <simsycl/sycl/sub_group.hh>
+
+namespace simsycl::sycl::khr {
+
+#if SIMSYCL_ENABLE_SYCL_KHR_WORK_ITEM_QUERIES
+
+namespace detail {
+template<int Dimensions>
+std::optional<simsycl::sycl::nd_item<Dimensions>> g_khr_wi_query_this_nd_item;
+
+template<int Dimensions>
+std::optional<simsycl::sycl::group<Dimensions>> g_khr_wi_query_this_group;
+
+inline std::optional<simsycl::sycl::sub_group> g_khr_wi_query_this_sub_group;
+} // namespace detail
+
+template<int Dimensions>
+simsycl::sycl::nd_item<Dimensions> this_nd_item() {
+    SIMSYCL_CHECK_MSG(!!detail::g_khr_wi_query_this_nd_item<Dimensions>,
+        "Work item query state 'this_nd_item' is not available.\n"
+        "Make sure that the query originated from a kernel launched with a sycl::nd_range argument");
+    return detail::g_khr_wi_query_this_nd_item<Dimensions>.value();
+}
+
+template<int Dimensions>
+simsycl::sycl::group<Dimensions> this_group() {
+    SIMSYCL_CHECK_MSG(!!detail::g_khr_wi_query_this_group<Dimensions>,
+        "Work item query state 'this_group' is not available.\n"
+        "Make sure that the query originated from a kernel launched with a sycl::nd_range argument");
+    return detail::g_khr_wi_query_this_group<Dimensions>.value();
+}
+
+inline simsycl::sycl::sub_group this_sub_group() {
+    SIMSYCL_CHECK_MSG(!!detail::g_khr_wi_query_this_sub_group,
+        "Work item query state 'this_sub_group' is not available.\n"
+        "Make sure that the query originated from a kernel launched with a sycl::nd_range argument");
+    return detail::g_khr_wi_query_this_sub_group.value();
+}
+
+#endif // SIMSYCL_ENABLE_SYCL_KHR_WORK_ITEM_QUERIES
+
+} // namespace simsycl::sycl::khr

src/simsycl/schedule.cc

@@ -1,3 +1,7 @@
+#include "simsycl/sycl/group.hh"
+#include "simsycl/sycl/khr/sub_group_queries.hh"
+#include "simsycl/sycl/nd_item.hh"
+
 #include <simsycl/detail/utils.hh>
 #include <simsycl/schedule.hh>
 #include <simsycl/sycl/device.hh>
@@ -181,6 +185,23 @@ void cooperative_for_nd_range(const sycl::device &device, const sycl::nd_range<D
     std::vector<detail::concurrent_sub_group> concurrent_sub_groups(num_concurrent_sub_groups);
     std::vector<detail::concurrent_nd_item> num_concurrent_nd_items(num_concurrent_items);
 
+#if SIMSYCL_ENABLE_SYCL_KHR_WORK_ITEM_QUERIES
+    std::vector<const sycl::nd_item<Dimensions> *> concurrent_khr_wi_query_nd_item_ptrs(num_concurrent_items, nullptr);
+
+    auto update_global_khr_wi_query_data = [&](int cc_g_idx = -1) {
+        if(cc_g_idx != -1 && concurrent_khr_wi_query_nd_item_ptrs[cc_g_idx] != nullptr) {
+            const auto nd_item = *concurrent_khr_wi_query_nd_item_ptrs[cc_g_idx];
+            sycl::khr::detail::g_khr_wi_query_this_nd_item<Dimensions> = nd_item;
+            sycl::khr::detail::g_khr_wi_query_this_group<Dimensions> = nd_item.get_group();
+            sycl::khr::detail::g_khr_wi_query_this_sub_group = nd_item.get_sub_group();
+        } else {
+            sycl::khr::detail::g_khr_wi_query_this_nd_item<Dimensions> = std::nullopt;
+            sycl::khr::detail::g_khr_wi_query_this_group<Dimensions> = std::nullopt;
+            sycl::khr::detail::g_khr_wi_query_this_sub_group = std::nullopt;
+        }
+    };
+#endif // SIMSYCL_ENABLE_SYCL_KHR_WORK_ITEM_QUERIES
+
     for(auto &cgroup : concurrent_groups) {
         cgroup.local_memory_allocations.resize(local_memory.size());
         for(size_t i = 0; i < local_memory.size(); ++i) {
@@ -220,8 +241,13 @@ void cooperative_for_nd_range(const sycl::device &device, const sycl::nd_range<D
                 group_linear_range, sub_group_linear_id_in_group, sub_group_linear_range_in_group,
                 sub_group_max_local_linear_range, sub_group_max_local_range, thread_id_in_sub_group,
                 sub_group_id_in_group, sub_group_range_in_group, &concurrent_nd_item, &concurrent_group,
-                &concurrent_sub_group, &kernel, &concurrent_items_exited, &caught_exceptions,
-                &range](boost::context::continuation &&scheduler) //
+                &concurrent_sub_group, &kernel, &concurrent_items_exited, &caught_exceptions, &range
+#if SIMSYCL_ENABLE_SYCL_KHR_WORK_ITEM_QUERIES
+                ,
+                concurrent_global_idx, &concurrent_khr_wi_query_nd_item_ptrs,
+                &update_global_khr_wi_query_data
+#endif
+        ](boost::context::continuation &&scheduler) //
             {
                 // yield immediately to allow the scheduling loop to set up local memory pointers
                 enter_kernel_fiber(std::move(scheduler));
@@ -245,7 +271,8 @@ void cooperative_for_nd_range(const sycl::device &device, const sycl::nd_range<D
 
                     SIMSYCL_START_IGNORING_DEPRECATIONS;
                     const auto group_id = linear_index_to_id(group_range, group_linear_id);
-                    const auto global_id = range.get_offset() + (group_id * sycl::id<Dimensions>(local_range)) + local_id;
+                    const auto global_id
+                        = range.get_offset() + (group_id * sycl::id<Dimensions>(local_range)) + local_id;
 
                     // if sub-group range is not divisible by local range, the last sub-group will be smaller
                     const auto sub_group_local_linear_range = std::min(sub_group_max_local_linear_range,
@@ -265,6 +292,12 @@ void cooperative_for_nd_range(const sycl::device &device, const sycl::nd_range<D
                     const auto nd_item
                         = detail::make_nd_item(global_item, local_item, group, sub_group, &concurrent_nd_item);
 
+#if SIMSYCL_ENABLE_SYCL_KHR_WORK_ITEM_QUERIES
+                    concurrent_khr_wi_query_nd_item_ptrs[concurrent_global_idx] = &nd_item;
+                    // adjust the globals now that the data is available, before starting the kernel
+                    update_global_khr_wi_query_data(concurrent_global_idx);
+#endif // SIMSYCL_ENABLE_SYCL_KHR_WORK_ITEM_QUERIES
+
                     try {
                         kernel(nd_item);
                         // Add an implicit "exit" operations to groups and sub-groups to catch potential divergence on
@@ -311,11 +344,21 @@ void cooperative_for_nd_range(const sycl::device &device, const sycl::nd_range<D
                 *local_memory[i].ptr = concurrent_groups[concurrent_group_idx].local_memory_allocations[i].get();
             }
 
+#if SIMSYCL_ENABLE_SYCL_KHR_WORK_ITEM_QUERIES
+            // adjust globals before switching fibers
+            update_global_khr_wi_query_data(concurrent_global_idx);
+#endif // SIMSYCL_ENABLE_SYCL_KHR_WORK_ITEM_QUERIES
+
             fibers[concurrent_global_idx] = fibers[concurrent_global_idx].resume();
         }
         schedule_state = schedule.update(schedule_state, order);
     }
 
+#if SIMSYCL_ENABLE_SYCL_KHR_WORK_ITEM_QUERIES
+    // reset globals
+    update_global_khr_wi_query_data();
+#endif // SIMSYCL_ENABLE_SYCL_KHR_WORK_ITEM_QUERIES
+
     // rethrow any encountered exceptions
     for(auto &exception : caught_exceptions) { std::rethrow_exception(exception); }
 }

test/CMakeLists.txt

@@ -22,6 +22,7 @@ add_executable(tests
   simulation_tests.cc
   alloc_tests.cc
   vec_tests.cc
+  extensions/work_item_queries_test.cc
 )
 
 add_sycl_to_target(TARGET tests SIMSYCL_ALL_WARNINGS)

test/extensions/extensions_test.cc

@@ -11,4 +11,13 @@ int main() {
 
     // SIMSYCL_ENABLE_SYCL_KHR_QUEUE_FLUSH
     queue.khr_flush();
+
+    // SIMSYCL_ENABLE_SYCL_KHR_WORK_ITEM_QUERIES
+    queue.submit([&](sycl::handler &cgh) {
+        cgh.parallel_for(sycl::nd_range<1>(1024, 64), [=](sycl::nd_item<1>) {
+            [[maybe_unused]] const auto item = sycl::khr::this_nd_item<1>();
+            [[maybe_unused]] const auto group = sycl::khr::this_group<1>();
+            [[maybe_unused]] const auto sub_group = sycl::khr::this_sub_group();
+        });
+    });
 }

test/extensions/work_item_queries_test.cc

@@ -0,0 +1,51 @@
+#include <simsycl/sycl.hh>
+
+#include <catch2/catch_template_test_macros.hpp>
+#include <catch2/catch_test_macros.hpp>
+#include <catch2/generators/catch_generators.hpp>
+
+
+using namespace simsycl;
+
+TEMPLATE_TEST_CASE_SIG(
+    "work item queries are correct if supported", "[khr][work_item_queries]", ((int Dims), Dims), 1, 2, 3) {
+
+    #if SIMSYCL_ENABLE_SYCL_KHR_WORK_ITEM_QUERIES
+
+    sycl::range<Dims> global_range;
+    sycl::range<Dims> local_range;
+    for(int d = 0; d < Dims; ++d) {
+        const int s = d+1;
+        global_range[d] = s * (2 + s);
+        local_range[d] = 2 + s;
+    }
+
+    std::vector<bool> visited(global_range.size(), false);
+    sycl::queue()
+        .submit([&](sycl::handler &cgh) {
+            cgh.parallel_for(sycl::nd_range(global_range, local_range), [=, &visited](sycl::nd_item<Dims> it) {
+                const auto global_linear_id = it.get_global_linear_id();
+                CHECK(global_linear_id < global_range.size());
+                CHECK(!visited[global_linear_id]);
+                visited[global_linear_id] = true;
+
+                CHECK(sycl::khr::this_nd_item<Dims>() == it);
+                CHECK(sycl::khr::this_group<Dims>() == it.get_group());
+                CHECK(sycl::khr::this_sub_group() == it.get_sub_group());
+
+                group_barrier(it.get_group());
+
+                // check again after scheduling through group_barrier
+                CHECK(sycl::khr::this_nd_item<Dims>() == it);
+                CHECK(sycl::khr::this_group<Dims>() == it.get_group());
+                CHECK(sycl::khr::this_sub_group() == it.get_sub_group());
+            });
+        })
+        .wait();
+
+    for(size_t i = 0; i < global_range.size(); ++i) { CAPTURE(i); CHECK(visited[i]); }
+
+    #else // SIMSYCL_ENABLE_SYCL_KHR_WORK_ITEM_QUERIES
+    SKIP("SYCL_KHR_WORK_ITEM_QUERIES not enabled");
+    #endif // SIMSYCL_ENABLE_SYCL_KHR_WORK_ITEM_QUERIES
+}