[SYCL] Support sycl_ext_oneapi_platform_device_index

sycl/include/sycl/device.hpp

@@ -366,6 +366,15 @@ class __SYCL_STANDALONE_DEBUG __SYCL_EXPORT device {
   /// \return the default context
   context ext_oneapi_get_default_context();
 
+  /// If this device is a root device as defined by the core SYCL specification,
+  /// returns the index that it has in the std::vector that is returned when
+  /// calling platform::get_devices() on the platform that contains this device,
+  /// otherwise throws an exception.
+  ///
+  /// \return the index that it has in the std::vector that is returned when
+  /// calling platform::get_devices() on the platform that contains this device.
+  size_t ext_oneapi_index_within_platform() const;
+
   // Definitions are in `<sycl/ext/oneapi/weak_object.hpp>` to avoid circular
   // dependencies:
   inline bool ext_oneapi_owner_before(const device &Other) const noexcept;

sycl/include/sycl/platform.hpp

@@ -204,6 +204,13 @@ class __SYCL_EXPORT platform : public detail::OwnerLessBase<platform> {
   /// \return the default context
   context khr_get_default_context() const;
 
+  /// If the index is within range of the std::vector that is returned when
+  /// calling platform::get_devices(), returns a copy of the device object which
+  /// has that index.
+  ///
+  /// \return a copy of the device object which has that index.
+  device ext_oneapi_device_at_index(size_t index) const;
+
 private:
   ur_native_handle_t getNative() const;
 

sycl/source/detail/device_impl.cpp

@@ -23,14 +23,14 @@ namespace detail {
 /// Constructs a SYCL device instance using the provided
 /// UR device instance.
 device_impl::device_impl(ur_device_handle_t Device, platform_impl &Platform,
-                         device_impl::private_tag)
+                         device_impl::private_tag, size_t idx)
     : MDevice(Device), MPlatform(Platform),
       // No need to set MRootDevice when MAlwaysRootDevice is true
       MRootDevice(Platform.MAlwaysRootDevice
                       ? nullptr
                       : get_info_impl<UR_DEVICE_INFO_PARENT_DEVICE>()),
       // TODO catch an exception and put it to list of asynchronous exceptions:
-      MCache{*this} {
+      MCache{*this}, MIndexWithinPlatform(idx) {
   // Interoperability Constructor already calls DeviceRetain in
   // urDeviceCreateWithNativeHandle.
   getAdapter().call<UrApiKind::urDeviceRetain>(MDevice);

sycl/source/detail/device_impl.hpp

@@ -424,10 +424,15 @@ class device_impl {
   // Must be called through `platform_impl::getOrMakeDeviceImpl` only.
   // `private_tag` ensures that is true.
   explicit device_impl(ur_device_handle_t Device, platform_impl &Platform,
-                       private_tag);
+                       private_tag, size_t idx);
 
   ~device_impl();
 
+  /// Get the index of the device within the device vector of its platform.
+  ///
+  /// \return the index if the device
+  size_t getIndexWithinPlatform() const { return MIndexWithinPlatform; }
+
   /// Get instance of OpenCL device
   ///
   /// \return a valid cl_device_id instance in accordance with the
@@ -2292,6 +2297,8 @@ class device_impl {
           aspect::ext_oneapi_is_composite, aspect::ext_oneapi_is_component>>
       MCache;
 
+  const size_t MIndexWithinPlatform = 0;
+
 }; // class device_impl
 
 using devices_iterator = variadic_iterator<

sycl/source/detail/platform_impl.cpp

@@ -281,7 +281,7 @@ device_impl &platform_impl::getOrMakeDeviceImpl(ur_device_handle_t UrDevice) {
 
   // Otherwise make the impl
   MDevices.emplace_back(std::make_unique<device_impl>(
-      UrDevice, *this, device_impl::private_tag{}));
+      UrDevice, *this, device_impl::private_tag{}, MDevices.size()));
 
   return *MDevices.back();
 }

sycl/source/detail/platform_impl.hpp

@@ -160,8 +160,6 @@ class platform_impl : public std::enable_shared_from_this<platform_impl> {
   ///
   /// \param UrDevice is the UrDevice whose impl is requested
   ///
-  /// \param PlatormImpl is the Platform for that Device
-  ///
   /// \return a device_impl* corresponding to the device
   device_impl &getOrMakeDeviceImpl(ur_device_handle_t UrDevice);
 

sycl/source/device.cpp

@@ -351,5 +351,12 @@ void device::ext_oneapi_wait() {
 
 void device::ext_oneapi_throw_asynchronous() { impl->throwAsynchronous(); }
 
+size_t device::ext_oneapi_index_within_platform() const {
+  if (!impl->isRootDevice())
+    throw sycl::exception(sycl::make_error_code(errc::invalid),
+                          "this device is not a root device");
+  return impl->getIndexWithinPlatform();
+}
+
 } // namespace _V1
 } // namespace sycl

sycl/source/feature_test.hpp.in

@@ -111,6 +111,7 @@ inline namespace _V1 {
 #define SYCL_EXT_ONEAPI_CLOCK 1
 #define SYCL_EXT_ONEAPI_DEVICE_IS_INTEGRATED_GPU 1
 #define SYCL_EXT_ONEAPI_DEVICE_DEFAULT_CONTEXT 1
+#define SYCL_EXT_ONEAPI_PLATFORM_DEVICE_INDEX 1
 // In progress yet
 #define SYCL_EXT_ONEAPI_ATOMIC16 0
 #define SYCL_KHR_DEFAULT_CONTEXT 1

sycl/source/platform.cpp

@@ -130,6 +130,15 @@ std::vector<device> platform::ext_oneapi_get_composite_devices() const {
   return Result;
 }
 
+device platform::ext_oneapi_device_at_index(size_t index) const {
+  auto devices = get_devices();
+  if (index < devices.size())
+    return devices[index];
+  else
+    throw sycl::exception(sycl::make_error_code(errc::invalid),
+                          "index is out of range");
+}
+
 namespace detail {
 
 void enable_ext_oneapi_default_context(bool Val) {

sycl/test-e2e/Basic/sycl_ext_oneapi_platform_device_index.cpp

@@ -0,0 +1,122 @@
+// RUN: %{build} %level_zero_options %opencl_lib -o %t.out
+// RUN: %{run} %t.out
+
+#include "../helpers.hpp"
+
+#include <level_zero/ze_api.h>
+#include <sycl/detail/core.hpp>
+#include <sycl/ext/oneapi/backend/level_zero.hpp>
+#include <sycl/platform.hpp>
+
+#include <algorithm>
+#include <cassert>
+#include <cstdlib>
+#include <iostream>
+#include <iterator>
+#include <vector>
+
+int main() {
+  sycl::device orig_dev;
+  auto plt = orig_dev.get_platform();
+  auto devices = plt.get_devices();
+  auto it = std::find(devices.begin(), devices.end(), orig_dev);
+  auto orig_dev_index_within_plt = std::distance(devices.begin(), it);
+
+  // ext_oneapi_index_within_platform
+  size_t ext_oneapi_index_within_platform =
+      orig_dev.ext_oneapi_index_within_platform();
+
+  // sycl_ext_oneapi_platform_device_index guarantees:
+  // The device index returned from device::ext_oneapi_index_within_platform is
+  // compatible with the index of the underlying backend device when the
+  // ONEAPI_DEVICE_SELECTOR environment variable is not set.
+  //
+  // When the platform’s backend is backend::ext_oneapi_level_zero, the index
+  // returned from device::ext_oneapi_index_within_platform matches the index of
+  // the device’s underlying ze_device_handle_t within the list of handles
+  // returned from zeDeviceGet.
+
+  // When the platform’s backend is backend::opencl, the index returned from
+  // device::ext_oneapi_index_within_platform matches the index of the device’s
+  // underlying cl_device_id within the list of IDs returned from
+  // clGetDeviceIDs.
+
+  // Check if the index matches the index of the device’s underlying handle
+  // within the list of handles returned from zeDeviceGet/clGetDeviceIDs.
+  std::string selector = env::getVal("ONEAPI_DEVICE_SELECTOR");
+  if (!selector.empty()) {
+    assert(orig_dev_index_within_plt == ext_oneapi_index_within_platform &&
+           "The index returned from device::ext_oneapi_index_within_platform "
+           "doesn't match the index that the device has in the std::vector "
+           "that is returned when calling platform::get_devices() on the "
+           "platform that contains this device");
+  } else {
+    if (orig_dev.get_backend() == sycl::backend::ext_oneapi_level_zero) {
+      auto l0_plt = sycl::get_native<sycl::backend::ext_oneapi_level_zero>(plt);
+      auto l0_dev =
+          sycl::get_native<sycl::backend::ext_oneapi_level_zero>(orig_dev);
+
+      uint32_t num_devices = 0;
+      zeDeviceGet(l0_plt, &num_devices, nullptr);
+
+      std::vector<ze_device_handle_t> l0_devices(num_devices);
+      zeDeviceGet(l0_plt, &num_devices, l0_devices.data());
+
+      auto it = std::find(l0_devices.begin(), l0_devices.end(), l0_dev);
+      assert(ext_oneapi_index_within_platform ==
+                 (std::distance(l0_devices.begin(), it)) &&
+             "The index returned from device::ext_oneapi_index_within_platform "
+             "doesn't match the index of the device’s underlying cl_device_id "
+             "within the list of IDs returned from clGetDeviceIDs");
+
+    } else if (orig_dev.get_backend() == sycl::backend::opencl) {
+      auto cl_plt = sycl::get_native<sycl::backend::opencl>(plt);
+      auto cl_dev = sycl::get_native<sycl::backend::opencl>(orig_dev);
+
+      cl_uint num_devices = 0;
+      clGetDeviceIDs(cl_plt, CL_DEVICE_TYPE_ALL, 0, nullptr, &num_devices);
+
+      std::vector<cl_device_id> cl_devices(num_devices);
+      clGetDeviceIDs(cl_plt, CL_DEVICE_TYPE_ALL, num_devices, cl_devices.data(),
+                     nullptr);
+
+      auto it = std::find(cl_devices.begin(), cl_devices.end(), cl_dev);
+      assert(orig_dev_index_within_plt ==
+                 (std::distance(cl_devices.begin(), it)) &&
+             "The index returned from device::ext_oneapi_index_within_platform "
+             "doesn't match the index of the device’s underlying cl_device_id "
+             "within the list of IDs returned from clGetDeviceIDs");
+    }
+  }
+  // Test non-root device exception (if partition is supported)
+  auto partition_properties =
+      orig_dev.get_info<sycl::info::device::partition_properties>();
+  if (std::find(partition_properties.begin(), partition_properties.end(),
+                sycl::info::partition_property::partition_equally) !=
+      partition_properties.end()) {
+    std::vector<sycl::device> sub_devices = orig_dev.create_sub_devices<
+        sycl::info::partition_property::partition_equally>(2);
+    try {
+      sub_devices[0].ext_oneapi_index_within_platform();
+      assert(false && "Missing an exception");
+    } catch (sycl::exception &e) {
+      std::cout << e.what() << std::endl;
+    }
+  }
+
+  // ext_oneapi_device_at_index
+  auto ext_oneapi_device_at_index = plt.ext_oneapi_device_at_index(
+      static_cast<size_t>(orig_dev_index_within_plt));
+  assert(orig_dev == ext_oneapi_device_at_index &&
+         "A copy of the device object which has that index doesn't match the "
+         "original device");
+  // Test out-of-range exception
+  try {
+    plt.ext_oneapi_device_at_index(devices.size());
+    assert(false && "Missing an exception");
+  } catch (sycl::exception &e) {
+    std::cout << e.what() << std::endl;
+  }
+
+  return 0;
+}

sycl/test/abi/sycl_symbols_linux.dump

@@ -3907,6 +3907,7 @@ _ZNK4sycl3_V16device18create_sub_devicesILNS0_4info18partition_propertyE4231EEES
 _ZNK4sycl3_V16device18create_sub_devicesILNS0_4info18partition_propertyE4232EEESt6vectorIS1_SaIS1_EENS3_25partition_affinity_domainE
 _ZNK4sycl3_V16device18create_sub_devicesILNS0_4info18partition_propertyE4233EEESt6vectorIS1_SaIS1_EEv
 _ZNK4sycl3_V16device26ext_oneapi_cl_profile_implEv
+_ZNK4sycl3_V16device32ext_oneapi_index_within_platformEv
 _ZNK4sycl3_V16device32ext_oneapi_supports_cl_c_versionERKNS0_3ext6oneapi12experimental10cl_versionE
 _ZNK4sycl3_V16device32ext_oneapi_supports_cl_extensionENS0_6detail11string_viewEPNS0_3ext6oneapi12experimental10cl_versionE
 _ZNK4sycl3_V16device3getEv
@@ -3995,6 +3996,7 @@ _ZNK4sycl3_V18platform13get_info_implINS0_4info8platform7profileEEENS0_6detail11
 _ZNK4sycl3_V18platform13get_info_implINS0_4info8platform7versionEEENS0_6detail11ABINeutralTINS6_21is_platform_info_descIT_E11return_typeEE4typeEv
 _ZNK4sycl3_V18platform13has_extensionENS0_6detail11string_viewE
 _ZNK4sycl3_V18platform23khr_get_default_contextEv
+_ZNK4sycl3_V18platform26ext_oneapi_device_at_indexEm
 _ZNK4sycl3_V18platform30ext_oneapi_get_default_contextEv
 _ZNK4sycl3_V18platform32ext_oneapi_get_composite_devicesEv
 _ZNK4sycl3_V18platform3getEv

sycl/test/abi/sycl_symbols_windows.dump

@@ -3832,6 +3832,7 @@
 ?ext_oneapi_copy@queue@_V1@sycl@@QEAA?AVevent@23@Uimage_mem_handle@experimental@oneapi@ext@23@V?$range@$02@23@AEBUimage_descriptor@67823@U567823@121AEBUcode_location@detail@23@@Z
 ?ext_oneapi_copy@queue@_V1@sycl@@QEAA?AVevent@23@Uimage_mem_handle@experimental@oneapi@ext@23@V?$range@$02@23@AEBUimage_descriptor@67823@U567823@121AEBV?$vector@Vevent@_V1@sycl@@V?$allocator@Vevent@_V1@sycl@@@std@@@std@@AEBUcode_location@detail@23@@Z
 ?ext_oneapi_copy@queue@_V1@sycl@@QEAA?AVevent@23@Uimage_mem_handle@experimental@oneapi@ext@23@V?$range@$02@23@AEBUimage_descriptor@67823@U567823@121V423@AEBUcode_location@detail@23@@Z
+?ext_oneapi_device_at_index@platform@_V1@sycl@@QEBA?AVdevice@23@_K@Z
 ?ext_oneapi_disable_peer_access@device@_V1@sycl@@QEAAXAEBV123@@Z
 ?ext_oneapi_empty@queue@_V1@sycl@@QEBA_NXZ
 ?ext_oneapi_enable_peer_access@device@_V1@sycl@@QEAAXAEBV123@@Z
@@ -3862,6 +3863,7 @@
 ?ext_oneapi_has_device_global@kernel_bundle_plain@detail@_V1@sycl@@QEAA_NAEBV?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@@Z
 ?ext_oneapi_has_kernel@kernel_bundle_plain@detail@_V1@sycl@@AEAA_NVstring_view@234@@Z
 ?ext_oneapi_has_kernel@kernel_bundle_plain@detail@_V1@sycl@@QEAA_NAEBV?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@@Z
+?ext_oneapi_index_within_platform@device@_V1@sycl@@QEBA_KXZ
 ?ext_oneapi_memcpy2d_impl@handler@_V1@sycl@@AEAAXPEAX_KPEBX111@Z
 ?ext_oneapi_memset2d_impl@handler@_V1@sycl@@AEAAXPEAX_KH11@Z
 ?ext_oneapi_owner_before@?$OwnerLessBase@Vcontext@_V1@sycl@@@detail@_V1@sycl@@QEBA_NAEBV?$weak_object_base@Vcontext@_V1@sycl@@@2oneapi@ext@34@@Z

sycl/unittests/Extensions/CMakeLists.txt

@@ -35,3 +35,4 @@ add_subdirectory(NumComputeUnits)
 add_subdirectory(FreeFunctionCommands)
 add_subdirectory(KernelQueries)
 add_subdirectory(InterProcessCommunication)
+add_subdirectory(DeviceIndex)

sycl/unittests/Extensions/DeviceIndex/CMakeLists.txt

@@ -0,0 +1,3 @@
+add_sycl_unittest(DeviceIndexExtensionTests OBJECT
+  DeviceIndex.cpp
+)

sycl/unittests/Extensions/DeviceIndex/DeviceIndex.cpp

@@ -0,0 +1,46 @@
+//==- DeviceIndex.cpp -- sycl_ext_oneapi_platform_device_index 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>
+
+namespace {
+const auto DEVICE1 = reinterpret_cast<ur_device_handle_t>(1u);
+const auto DEVICE2 = reinterpret_cast<ur_device_handle_t>(2u);
+const auto DEVICE3 = reinterpret_cast<ur_device_handle_t>(3u);
+
+ur_result_t redefine_urDeviceGet(void *pParams) {
+  auto params = *static_cast<ur_device_get_params_t *>(pParams);
+  if (*params.ppNumDevices)
+    **params.ppNumDevices = 3;
+  if (*params.pphDevices && *params.pNumEntries > 0) {
+    (*params.pphDevices)[0] = DEVICE1;
+    (*params.pphDevices)[1] = DEVICE2;
+    (*params.pphDevices)[2] = DEVICE3;
+  }
+  return UR_RESULT_SUCCESS;
+}
+
+} // namespace
+
+TEST(sycl_ext_oneapi_platform_device_index, CheckDeviceIndexes) {
+  sycl::unittest::UrMock<> Mock;
+  mock::getCallbacks().set_replace_callback("urDeviceGet",
+                                            &redefine_urDeviceGet);
+
+  sycl::platform plt = sycl::platform();
+  auto devs = plt.get_devices();
+
+  ASSERT_EQ(devs.size(), 3ull);
+
+  for (size_t i = 0; i < devs.size(); i++)
+    ASSERT_EQ(devs[i].ext_oneapi_index_within_platform(), i);
+}