[SYCL] Implement sycl_khr_max_work_group_queries

sycl/include/sycl/khr/max_work_group_queries.hpp

@@ -0,0 +1,52 @@
+//==-- max_work_group_queries.hpp - oneapi max work groups info traits
+//------------==//
+//
+// 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/detail/info_desc_traits.hpp>
+#include <sycl/id.hpp>
+#include <unified-runtime/ur_api.h>
+
+#include <cstddef>
+
+namespace sycl {
+inline namespace _V1 {
+namespace khr::info::device {
+
+template <int Dimensions = 3> struct max_work_group_range;
+
+// max_work_group_range<1> and <2> are RT-only; only <3> dispatches via UR.
+template <>
+struct max_work_group_range<1>
+    : sycl::detail::rt_traits_base<sycl::detail::info_class::device> {
+  using return_type = sycl::range<1>;
+};
+
+template <>
+struct max_work_group_range<2>
+    : sycl::detail::rt_traits_base<sycl::detail::info_class::device> {
+  using return_type = sycl::range<2>;
+};
+
+template <>
+struct max_work_group_range<3>
+    : sycl::detail::ur_traits_base<sycl::detail::info_class::device,
+                                   UR_DEVICE_INFO_MAX_WORK_GROUPS_3D> {
+  using return_type = sycl::range<3>;
+};
+
+struct max_work_group_range_size
+    : sycl::detail::ur_traits_base<sycl::detail::info_class::device,
+                                   UR_DEVICE_INFO_MAX_WORK_GROUPS> {
+  using return_type = size_t;
+};
+
+} // namespace khr::info::device
+} // namespace _V1
+} // namespace sycl

sycl/include/sycl/sycl.hpp

@@ -165,5 +165,6 @@ can be disabled by setting SYCL_DISABLE_FSYCL_SYCLHPP_WARNING macro.")
 #include <sycl/khr/dynamic_addrspace_cast.hpp>
 #include <sycl/khr/free_function_commands.hpp>
 #include <sycl/khr/group_interface.hpp>
+#include <sycl/khr/max_work_group_queries.hpp>
 #include <sycl/khr/static_addrspace_cast.hpp>
 #include <sycl/khr/work_item_queries.hpp>

sycl/source/detail/device_impl.hpp

@@ -26,6 +26,7 @@
 #include <sycl/ext/oneapi/matrix/query-types.hpp>
 #include <sycl/info/device.hpp>
 #include <sycl/kernel_bundle.hpp>
+#include <sycl/khr/max_work_group_queries.hpp>
 #include <sycl/platform.hpp>
 
 #include <memory>
@@ -1219,6 +1220,29 @@ class device_impl {
                         "ext_intel_max_lanes_per_hw_thread aspect");
       return get_info_impl<UR_DEVICE_INFO_MAX_LANES_PER_HW_THREAD>();
     }
+
+    // khr device traits (defined under sycl/ext/oneapi/...).
+
+    CASE(khr::info::device::max_work_group_range_size) {
+      return get_info_impl<UR_DEVICE_INFO_MAX_WORK_GROUPS>();
+    }
+    CASE(khr::info::device::max_work_group_range<3>) {
+      size_t result[3] = {};
+      getAdapter().call<UrApiKind::urDeviceGetInfo>(
+          getHandleRef(), UR_DEVICE_INFO_MAX_WORK_GROUPS_3D, sizeof(result),
+          &result, nullptr);
+      return range<3>(result[2], result[1], result[0]);
+    }
+    CASE(khr::info::device::max_work_group_range<2>) {
+      range<3> max_3d = get_info<khr::info::device::max_work_group_range<3>,
+                                 DependentFalse>();
+      return range<2>{max_3d[1], max_3d[2]};
+    }
+    CASE(khr::info::device::max_work_group_range<1>) {
+      range<3> max_3d = get_info<khr::info::device::max_work_group_range<3>,
+                                 DependentFalse>();
+      return range<1>{max_3d[2]};
+    }
     else {
       constexpr auto Desc = UrInfoCode<Param>::value;
       return static_cast<typename Param::return_type>(get_info_impl<Desc>());

sycl/source/device.cpp

@@ -394,6 +394,15 @@ __SYCL_ONEAPI_DEVICE_INST(ext::oneapi::experimental::info::device,
 __SYCL_ONEAPI_DEVICE_INST(ext::oneapi::info::device, num_compute_units, size_t)
 #undef __SYCL_ONEAPI_DEVICE_INST
 
+#define __SYCL_KHR_DEVICE_INST(NS, NAME, RETURN_T)                             \
+  template __SYCL_EXPORT detail::ABINeutralT_t<RETURN_T>                       \
+  device::get_info_impl<NS::NAME>() const;
+__SYCL_KHR_DEVICE_INST(khr::info::device, max_work_group_range_size, size_t)
+__SYCL_KHR_DEVICE_INST(khr::info::device, max_work_group_range<1>, range<1>)
+__SYCL_KHR_DEVICE_INST(khr::info::device, max_work_group_range<2>, range<2>)
+__SYCL_KHR_DEVICE_INST(khr::info::device, max_work_group_range<3>, range<3>)
+#undef __SYCL_KHR_DEVICE_INST
+
 #define __SYCL_ONEAPI_PROGRESS_INST(NAME, SCOPE)                               \
   template __SYCL_EXPORT detail::ABINeutralT_t<                                \
       std::vector<ext::oneapi::experimental::forward_progress_guarantee>>      \

sycl/source/feature_test.hpp.in

@@ -107,6 +107,7 @@ inline namespace _V1 {
 #define SYCL_EXT_ONEAPI_CURRENT_DEVICE 1
 #define SYCL_KHR_QUEUE_EMPTY_QUERY 1
 #define SYCL_KHR_QUEUE_FLUSH 1
+#define SYCL_KHR_MAX_WORK_GROUP_QUERIES 1
 #define SYCL_EXT_ONEAPI_MEMORY_EXPORT 1
 #define SYCL_EXT_ONEAPI_CLOCK 1
 #define SYCL_EXT_ONEAPI_DEVICE_IS_INTEGRATED_GPU 1

sycl/test-e2e/Basic/max_work_group_queries.cpp

@@ -0,0 +1,50 @@
+// RUN: %{build} -o %t.out
+// REQUIRES: cuda || hip || level_zero
+// RUN: %{run} %t.out
+
+#include <sycl/detail/core.hpp>
+#include <sycl/khr/max_work_group_queries.hpp>
+
+#include <cassert>
+#include <iostream>
+
+using namespace sycl;
+
+int main() {
+  queue q;
+  device dev = q.get_device();
+
+#if !defined(SYCL_KHR_MAX_WORK_GROUP_QUERIES)
+#error SYCL_KHR_MAX_WORK_GROUP_QUERIES is not defined!
+#endif
+
+  sycl::id<1> groupD =
+      dev.get_info<sycl::khr::info::device::max_work_group_range<1>>();
+  std::cout << "Max work group size in 1D \n";
+  std::cout << "Dimension 1:" << groupD[0] << std::endl;
+
+  sycl::id<2> group2D =
+      dev.get_info<sycl::khr::info::device::max_work_group_range<2>>();
+  std::cout << "Max work group size in 2D \n";
+  std::cout << "Dimension 1:" << group2D[0] << "\n"
+            << "Dimension 2:" << group2D[1] << std::endl;
+
+  sycl::id<3> group3D =
+      dev.get_info<sycl::khr::info::device::max_work_group_range<3>>();
+  std::cout << "Max work group size in 3D \n";
+  std::cout << "Dimension 1:" << group3D[0] << "\n"
+            << "Dimension 2:" << group3D[1] << "\n"
+            << "Dimension 3:" << group3D[2] << std::endl;
+
+  size_t group_max =
+      dev.get_info<sycl::khr::info::device::max_work_group_range_size>();
+  std::cout << "Max global work group size:" << group_max << "\n";
+
+  assert((group3D[0] <= group_max && group3D[1] <= group_max &&
+          group3D[2] <= group_max) &&
+         "Max work-group size of each dimension must be smaller than "
+         "global work-group size");
+
+  std::cout << "Passed!" << std::endl;
+  return 0;
+}

sycl/test/basic_tests/Inputs/khr_all.hpp

@@ -21,6 +21,7 @@
 #include <sycl/khr/dynamic_addrspace_cast.hpp>
 #include <sycl/khr/free_function_commands.hpp>
 #include <sycl/khr/group_interface.hpp>
+#include <sycl/khr/max_work_group_queries.hpp>
 #include <sycl/khr/split_headers/accessor.hpp>
 #include <sycl/khr/split_headers/atomic.hpp>
 #include <sycl/khr/split_headers/backend.hpp>

unified-runtime/source/adapters/opencl/device.cpp

@@ -254,23 +254,37 @@ ur_result_t urDeviceGetInfo(ur_device_handle_t hDevice,
     return ReturnValue(URValue.data(), URValue.size());
   }
   case UR_DEVICE_INFO_MAX_WORK_GROUPS_3D: {
-    /* Returns the maximum sizes of a work group for each dimension one could
-     * use to submit a kernel. There is no such query defined in OpenCL. So
-     * we'll return the maximum value. */
-    static constexpr uint32_t MaxWorkItemDimensions = 3u;
-    static constexpr size_t Max = (std::numeric_limits<size_t>::max)();
+    constexpr size_t ReturnBufferSize = 3;
+    cl_uint MaxWorkItemDimensions = 0;
 
-    struct {
-      size_t sizes[MaxWorkItemDimensions];
-    } ReturnSizes;
+    CL_RETURN_ON_FAILURE(clGetDeviceInfo(
+        Device->CLDevice, CL_DEVICE_MAX_WORK_ITEM_DIMENSIONS,
+        sizeof(MaxWorkItemDimensions), &MaxWorkItemDimensions, nullptr));
+    assert(MaxWorkItemDimensions >= 3);
+
+    size_t ReturnBuffer[ReturnBufferSize];
+    size_t *ClCallBuffer = new size_t[MaxWorkItemDimensions];
+
+    oclv::OpenCLVersion DevVer;
+    UR_RETURN_ON_FAILURE(Device->getDeviceVersion(DevVer));
 
-    ReturnSizes.sizes[0] = Max;
-    ReturnSizes.sizes[1] = Max;
-    ReturnSizes.sizes[2] = Max;
-    return ReturnValue(ReturnSizes);
+    CL_RETURN_ON_FAILURE(
+        clGetDeviceInfo(Device->CLDevice, CL_DEVICE_MAX_WORK_ITEM_SIZES,
+                        sizeof(ClCallBuffer), &ClCallBuffer, nullptr));
+
+    ReturnBuffer[0] = ClCallBuffer[0];
+    ReturnBuffer[1] = ClCallBuffer[1];
+    ReturnBuffer[2] = ClCallBuffer[2];
+    delete[] ClCallBuffer;
+    return ReturnValue(ReturnBuffer);
   }
   case UR_DEVICE_INFO_MAX_WORK_GROUPS: {
-    return ReturnValue(std::numeric_limits<size_t>::max());
+    size_t Max = 0;
+    CL_RETURN_ON_FAILURE(clGetDeviceInfo(Device->CLDevice,
+                                         CL_DEVICE_MAX_WORK_GROUP_SIZE,
+                                         sizeof(Max), &Max, nullptr));
+    assert(Max >= 0);
+    return ReturnValue(Max);
   }
   case UR_DEVICE_INFO_MAX_COMPUTE_QUEUE_INDICES: {
     return ReturnValue(static_cast<uint32_t>(1u));