[SYCL] Implement sycl_khr_work_item_queries extension (#18519)

Implements the extension defined in
KhronosGroup/SYCL-Docs#682.

---------

Signed-off-by: Michael Aziz <[email protected]>

Author: 0x12CC

sycl/include/sycl/khr/work_item_queries.hpp

@@ -0,0 +1,34 @@
+//===-- work_item_queries.hpp --- KHR work item queries 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>
+
+namespace sycl {
+inline namespace _V1 {
+namespace khr {
+
+template <int Dimensions> nd_item<Dimensions> this_nd_item() {
+  return ext::oneapi::experimental::this_nd_item<Dimensions>();
+}
+
+template <int Dimensions> group<Dimensions> this_group() {
+  return ext::oneapi::this_work_item::get_work_group<Dimensions>();
+}
+
+inline sub_group this_sub_group() {
+  return ext::oneapi::this_work_item::get_sub_group();
+}
+
+} // namespace khr
+} // namespace _V1
+} // namespace sycl
+
+#endif

sycl/include/sycl/sycl.hpp

@@ -124,3 +124,4 @@
 #include <sycl/ext/oneapi/weak_object.hpp>
 #include <sycl/khr/free_function_commands.hpp>
 #include <sycl/khr/group_interface.hpp>
+#include <sycl/khr/work_item_queries.hpp>

sycl/source/feature_test.hpp.in

@@ -123,6 +123,12 @@ inline namespace _V1 {
 #define SYCL_KHR_DEFAULT_CONTEXT 1
 #define SYCL_EXT_INTEL_EVENT_MODE 1
 
+// Unfinished KHR extensions. These extensions are only available if the
+// __DPCPP_ENABLE_UNFINISHED_KHR_EXTENSIONS macro is defined.
+#ifdef __DPCPP_ENABLE_UNFINISHED_KHR_EXTENSIONS
+#define SYCL_KHR_WORK_ITEM_QUERIES 1
+#endif
+
 #ifndef __has_include
 #define __has_include(x) 0
 #endif

sycl/test-e2e/Basic/work_item_queries/work_item_queries.cpp

@@ -0,0 +1,120 @@
+// RUN: %{build} -o %t.out
+// RUN: %{run} %t.out
+
+//===- work_item_queries.cpp - KHR work item queries test -----------------===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+#define __DPCPP_ENABLE_UNFINISHED_KHR_EXTENSIONS
+
+#include <cassert>
+#include <iostream>
+#include <sycl/detail/core.hpp>
+#include <sycl/khr/work_item_queries.hpp>
+
+template <size_t... Dims> static int check_this_nd_item_api() {
+  // Define the kernel ranges.
+  constexpr int Dimensions = sizeof...(Dims);
+  const sycl::range<Dimensions> local_range{Dims...};
+  const sycl::range<Dimensions> global_range = local_range;
+  const sycl::nd_range<Dimensions> nd_range{global_range, local_range};
+  // Launch an ND-range kernel.
+  sycl::queue q;
+  sycl::buffer<bool, Dimensions> results{global_range};
+  q.submit([&](sycl::handler &cgh) {
+    sycl::accessor acc{results, cgh, sycl::write_only};
+    cgh.parallel_for(nd_range, [=](sycl::nd_item<Dimensions> it) {
+      // Compare it to this_nd_item<Dimensions>().
+      acc[it.get_global_id()] = (it == sycl::khr::this_nd_item<Dimensions>());
+    });
+  });
+  // Check the test results.
+  sycl::host_accessor acc{results};
+  for (const auto &result : acc) {
+    if (!result) {
+      std::cerr << "check_this_nd_item_api failed for dimensionality "
+                << Dimensions << ".\n";
+      return 1;
+    }
+  }
+  return 0;
+}
+
+template <size_t... Dims> static int check_this_group_api() {
+  // Define the kernel ranges.
+  constexpr int Dimensions = sizeof...(Dims);
+  const sycl::range<Dimensions> local_range{Dims...};
+  const sycl::range<Dimensions> global_range = local_range;
+  const sycl::nd_range<Dimensions> nd_range{global_range, local_range};
+  // Launch an ND-range kernel.
+  sycl::queue q;
+  sycl::buffer<bool, Dimensions> results{global_range};
+  q.submit([&](sycl::handler &cgh) {
+    sycl::accessor acc{results, cgh, sycl::write_only};
+    cgh.parallel_for(nd_range, [=](sycl::nd_item<Dimensions> it) {
+      // Compare it.get_group() to this_group<Dimensions>().
+      acc[it.get_global_id()] =
+          (it.get_group() == sycl::khr::this_group<Dimensions>());
+    });
+  });
+  // Check the test results.
+  sycl::host_accessor acc{results};
+  for (const auto &result : acc) {
+    if (!result) {
+      std::cerr << "check_this_group_api failed for dimensionality "
+                << Dimensions << ".\n";
+      return 1;
+    }
+  }
+  return 0;
+}
+
+template <size_t... Dims> static int check_this_sub_group_api() {
+  // Define the kernel ranges.
+  constexpr int Dimensions = sizeof...(Dims);
+  const sycl::range<Dimensions> local_range{Dims...};
+  const sycl::range<Dimensions> global_range = local_range;
+  const sycl::nd_range<Dimensions> nd_range{global_range, local_range};
+  // Launch an ND-range kernel.
+  sycl::queue q;
+  sycl::buffer<bool, Dimensions> results{global_range};
+  q.submit([&](sycl::handler &cgh) {
+    sycl::accessor acc{results, cgh, sycl::write_only};
+    cgh.parallel_for(nd_range, [=](sycl::nd_item<Dimensions> it) {
+      // Compare it.get_sub_group() to this_sub_group().
+      acc[it.get_global_id()] =
+          (it.get_sub_group() == sycl::khr::this_sub_group());
+    });
+  });
+  // Check the test results.
+  sycl::host_accessor acc{results};
+  for (const auto &result : acc) {
+    if (!result) {
+      std::cerr << "check_this_sub_group_api failed for dimensionality "
+                << Dimensions << ".\n";
+      return 1;
+    }
+  }
+  return 0;
+}
+
+int main() {
+  int failed = 0;
+  // nd_item
+  failed += check_this_nd_item_api<2>();
+  failed += check_this_nd_item_api<2, 3>();
+  failed += check_this_nd_item_api<2, 3, 4>();
+  // group
+  failed += check_this_group_api<2>();
+  failed += check_this_group_api<2, 3>();
+  failed += check_this_group_api<2, 3, 4>();
+  // sub_group
+  failed += check_this_sub_group_api<2>();
+  failed += check_this_sub_group_api<2, 3>();
+  failed += check_this_sub_group_api<2, 3, 4>();
+  return failed;
+}