[SYCL] Implement sycl_ext_oneapi_current_device (#16970)

Spec:
[link](https://github.com/intel/llvm/blob/75a5e284a9b228f056ccfb9defb8cfb18d8f515a/sycl/doc/extensions/proposed/sycl_ext_oneapi_current_device.asciidoc)

Author: dm-vodopyanov

sycl/doc/extensions/proposed/sycl_ext_oneapi_current_device.asciidoc renamed to sycl/doc/extensions/experimental/sycl_ext_oneapi_current_device.asciidoc

@@ -43,11 +43,12 @@ SYCL specification refer to that revision.
 
 == Status
 
-This is a proposed extension specification, intended to gather community
-feedback.  Interfaces defined in this specification may not be implemented yet
-or may be in a preliminary state.  The specification itself may also change in
-incompatible ways before it is finalized.  *Shipping software products should
-not rely on APIs defined in this specification.*
+This is an experimental extension specification, intended to provide early
+access to features and gather community feedback.  Interfaces defined in this
+specification are implemented in {dpcpp}, but they are not finalized and may
+change incompatibly in future versions of {dpcpp} without prior notice.
+*Shipping software products should not rely on APIs defined in this
+specification.*
 
 == Overview
 

sycl/include/sycl/ext/oneapi/experimental/current_device.hpp

@@ -0,0 +1,44 @@
+//==--------- current_device.hpp - sycl_ext_oneapi_current_device ----------==//
+//
+// 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
+
+#include <sycl/device.hpp>
+
+namespace sycl {
+inline namespace _V1 {
+namespace ext::oneapi::experimental::this_thread {
+
+namespace detail {
+inline sycl::device &get_current_device_ref() {
+  static thread_local sycl::device current_device{sycl::default_selector_v};
+  return current_device;
+}
+} // namespace detail
+
+/// @return The current default device for the calling host thread. If
+/// `set_current_device()` has not been called by this thread, returns the
+/// device selected by the default device selector.
+///
+/// @pre The function is called from a host thread, executing outside of a host
+/// task or an asynchronous error handler.
+inline sycl::device get_current_device() {
+  return detail::get_current_device_ref();
+}
+
+/// @brief Sets the current default device to `dev` for the calling host thread.
+///
+/// @pre The function is called from a host thread, executing outside of a host
+/// task or an asynchronous error handler.
+inline void set_current_device(sycl::device dev) {
+  detail::get_current_device_ref() = dev;
+}
+
+} // namespace ext::oneapi::experimental::this_thread
+} // namespace _V1
+} // namespace sycl

sycl/include/sycl/sycl.hpp

@@ -93,6 +93,7 @@
 #include <sycl/ext/oneapi/experimental/cluster_group_prop.hpp>
 #include <sycl/ext/oneapi/experimental/composite_device.hpp>
 #include <sycl/ext/oneapi/experimental/cuda/barrier.hpp>
+#include <sycl/ext/oneapi/experimental/current_device.hpp>
 #include <sycl/ext/oneapi/experimental/enqueue_functions.hpp>
 #include <sycl/ext/oneapi/experimental/fixed_size_group.hpp>
 #include <sycl/ext/oneapi/experimental/forward_progress.hpp>

sycl/source/feature_test.hpp.in

@@ -114,6 +114,7 @@ inline namespace _V1 {
 #define SYCL_EXT_ONEAPI_WORK_GROUP_STATIC 1 
 #define SYCL_EXT_ONEAPI_NUM_COMPUTE_UNITS 1
 #define SYCL_EXT_ONEAPI_DEVICE_IMAGE_BACKEND_CONTENT 1
+#define SYCL_EXT_ONEAPI_CURRENT_DEVICE 1
 #define SYCL_KHR_FREE_FUNCTION_COMMANDS 1
 // In progress yet
 #define SYCL_EXT_ONEAPI_ATOMIC16 0

sycl/test-e2e/CurrentDevice/current_device.cpp

@@ -0,0 +1,73 @@
+// REQUIRES: any-device-is-cpu, any-device-is-gpu
+
+// RUN: %{build} -o %t.out %threads_lib
+// RUN: %{run-unfiltered-devices} %t.out
+
+// E2E tests for sycl_ext_oneapi_current_device
+
+#include <sycl/ext/oneapi/experimental/current_device.hpp>
+
+#include <thread>
+
+void check_get_eq(sycl::device dev) {
+  auto device =
+      sycl::ext::oneapi::experimental::this_thread::get_current_device();
+  if (device != dev)
+    assert(false && "check_get_eq failed.");
+}
+
+void check_get_ne(sycl::device dev) {
+  auto device =
+      sycl::ext::oneapi::experimental::this_thread::get_current_device();
+  if (device == dev)
+    assert(false && "check_get_ne failed.");
+}
+
+void check_set_get_eq(sycl::device dev) {
+  sycl::ext::oneapi::experimental::this_thread::set_current_device(dev);
+  auto device =
+      sycl::ext::oneapi::experimental::this_thread::get_current_device();
+  if (device != dev)
+    assert(false && "check_set_get_eq failed.");
+}
+
+int main() {
+  // Test 1
+  std::thread t1(check_get_eq, sycl::device{sycl::default_selector_v});
+  std::thread t2(check_get_eq, sycl::device{sycl::default_selector_v});
+
+  t1.join();
+  t2.join();
+
+  // Test 2
+  // As GPU device is required, it is always has higher score than CPU device,
+  // so test must not fail.
+  t1 = std::thread(check_get_ne, sycl::device{sycl::cpu_selector_v});
+  t2 = std::thread(check_get_ne, sycl::device{sycl::cpu_selector_v});
+
+  t1.join();
+  t2.join();
+
+  // Test 3
+  t1 = std::thread(check_set_get_eq, sycl::device{sycl::cpu_selector_v});
+  t2 = std::thread(check_set_get_eq, sycl::device{sycl::gpu_selector_v});
+
+  t1.join();
+  t2.join();
+
+  // Test 4
+  auto device_1 =
+      sycl::ext::oneapi::experimental::this_thread::get_current_device();
+  if (device_1 != sycl::device{sycl::default_selector_v})
+    assert(false && "get_current_device check failed.");
+
+  // Test 5
+  sycl::ext::oneapi::experimental::this_thread::set_current_device(
+      sycl::device{sycl::cpu_selector_v});
+  auto device_2 =
+      sycl::ext::oneapi::experimental::this_thread::get_current_device();
+  if (device_2 != sycl::device{sycl::cpu_selector_v})
+    assert(false && "set/get_current_device check failed.");
+
+  return 0;
+}

sycl/unittests/Extensions/CMakeLists.txt

@@ -1,6 +1,7 @@
 set(CMAKE_CXX_EXTENSIONS OFF)
 
 add_sycl_unittest(ExtensionsTests OBJECT
+  CurrentDevice.cpp
   DefaultContext.cpp
   FPGADeviceSelectors.cpp
   DeviceArchitecture.cpp

sycl/unittests/Extensions/CurrentDevice.cpp

@@ -0,0 +1,90 @@
+//==---- CurrentDevice.cpp -- sycl_ext_oneapi_current_device unit tests ----==//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+#include <sycl/sycl.hpp>
+
+#include <helpers/UrMock.hpp>
+
+#include <gtest/gtest.h>
+#include <thread>
+
+namespace {
+const auto DEVICE_CPU = reinterpret_cast<ur_device_handle_t>(1u);
+const auto DEVICE_GPU = reinterpret_cast<ur_device_handle_t>(2u);
+
+ur_result_t redefine_urDeviceGet(void *pParams) {
+  auto params = *static_cast<ur_device_get_params_t *>(pParams);
+  if (*params.ppNumDevices)
+    **params.ppNumDevices = 2;
+  if (*params.pphDevices && *params.pNumEntries > 0) {
+    (*params.pphDevices)[0] = DEVICE_CPU;
+    (*params.pphDevices)[1] = DEVICE_GPU;
+  }
+  return UR_RESULT_SUCCESS;
+}
+
+ur_result_t after_urDeviceGetInfo(void *pParams) {
+  auto params = *static_cast<ur_device_get_info_params_t *>(pParams);
+  switch (*params.ppropName) {
+  case UR_DEVICE_INFO_TYPE: {
+    ur_device_type_t UrDeviceType = UR_DEVICE_TYPE_DEFAULT;
+    if (*params.phDevice == DEVICE_CPU) {
+      UrDeviceType = UR_DEVICE_TYPE_CPU;
+    } else if (*params.phDevice == DEVICE_GPU) {
+      UrDeviceType = UR_DEVICE_TYPE_GPU;
+    }
+    if (*params.ppPropValue)
+      *static_cast<ur_device_type_t *>(*params.ppPropValue) = UrDeviceType;
+    if (*params.ppPropSizeRet)
+      **params.ppPropSizeRet = sizeof(UrDeviceType);
+    return UR_RESULT_SUCCESS;
+  }
+  default:
+    return UR_RESULT_SUCCESS;
+  }
+}
+
+void callable_set_get_eq(sycl::device dev) {
+  sycl::ext::oneapi::experimental::this_thread::set_current_device(dev);
+  ASSERT_NO_FATAL_FAILURE(
+      sycl::ext::oneapi::experimental::this_thread::get_current_device() = dev);
+}
+} // namespace
+
+TEST(CurrentDeviceTest, CheckGetCurrentDeviceReturnDefaultDeviceInHostThread) {
+  sycl::unittest::UrMock<> Mock;
+  mock::getCallbacks().set_replace_callback("urDeviceGet",
+                                            &redefine_urDeviceGet);
+  ASSERT_EQ(sycl::ext::oneapi::experimental::this_thread::get_current_device(),
+            sycl::device{sycl::default_selector_v});
+}
+
+TEST(CurrentDeviceTest,
+     CheckGetCurrentDeviceReturnDefaultSelectorByDefaultInTwoThreads) {
+  sycl::unittest::UrMock<> Mock;
+  mock::getCallbacks().set_replace_callback("urDeviceGet",
+                                            &redefine_urDeviceGet);
+  mock::getCallbacks().set_after_callback("urDeviceGetInfo",
+                                          &after_urDeviceGetInfo);
+
+  sycl::platform Plt = sycl::platform();
+
+  ASSERT_EQ(Plt.get_devices().size(), 2ull);
+
+  sycl::device cpu_device = Plt.get_devices()[0];
+  sycl::device gpu_device = Plt.get_devices()[1];
+
+  ASSERT_TRUE(cpu_device.is_cpu());
+  ASSERT_TRUE(gpu_device.is_gpu());
+
+  std::thread t1(callable_set_get_eq, cpu_device);
+  std::thread t2(callable_set_get_eq, gpu_device);
+
+  t1.join();
+  t2.join();
+}