Move max_global_work_groups query to UR (#21840)

This PR makes the following changes:
1. Moves implementation of `max_global_work_groups` to UR. This query
has been implemented in SYCL RT because there's no backend support for
`max_global_work_groups` query. However, it was recently decided that
OpenCL will add a corresponding query. See CMPLRLLVM-73572 for more
info.
2. Changes `max_global_work_groups` from `INT_MAX` to `SIZE_MAX` for all
backends. For CUDA, HIP, OFFLOAD, and L0 adapter, we calculate the value
of `max_global_work_groups` by taking minimum of `SIZE_MAX` and
multiplication of per-dimension max group size.
3. Changed `max_work_groups<3>` so that `max_global_work_groups` no
longer limits per-dimension max work group size.

---------

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
Co-authored-by: iclsrc <iclsrc@intel.com>

Author: 3 people

include/unified-runtime/ur_api.h

@@ -2380,24 +2380,26 @@ typedef enum ur_device_info_t {
   UR_DEVICE_INFO_PREFERRED_VECTOR_WIDTH_LONG_LONG = 131,
   /// [uint32_t] native vector width for long long
   UR_DEVICE_INFO_NATIVE_VECTOR_WIDTH_LONG_LONG = 132,
+  /// [size_t] return max total number of work groups
+  UR_DEVICE_INFO_MAX_WORK_GROUPS = 133,
   /// [uint32_t][optional-query] return Intel GPU number of
   /// stacks/chiplets/tiles
-  UR_DEVICE_INFO_XE_STACK_COUNT = 133,
+  UR_DEVICE_INFO_XE_STACK_COUNT = 134,
   /// [uint32_t][optional-query] return Intel GPU number of regions sharing
   /// local L2/L3 (XE_CU) per stack
-  UR_DEVICE_INFO_XE_REGIONS_PER_STACK = 134,
+  UR_DEVICE_INFO_XE_REGIONS_PER_STACK = 135,
   /// [uint32_t][optional-query] return Intel GPU number of clusters
   /// (slices) per region
-  UR_DEVICE_INFO_XE_CLUSTERS_PER_REGION = 135,
+  UR_DEVICE_INFO_XE_CLUSTERS_PER_REGION = 136,
   /// [uint32_t][optional-query] return Intel GPU number of XE cores per
   /// cluster
-  UR_DEVICE_INFO_XE_CORES_PER_CLUSTER = 136,
+  UR_DEVICE_INFO_XE_CORES_PER_CLUSTER = 137,
   /// [uint32_t][optional-query] return Intel GPU number of execution
   /// engines (EUs) per XE Core
-  UR_DEVICE_INFO_EUS_PER_XE_CORE = 137,
+  UR_DEVICE_INFO_EUS_PER_XE_CORE = 138,
   /// [uint32_t][optional-query] return Intel GPU maximal number of lanes
   /// (virtual SIMD size) per hardware thread
-  UR_DEVICE_INFO_MAX_LANES_PER_HW_THREAD = 138,
+  UR_DEVICE_INFO_MAX_LANES_PER_HW_THREAD = 139,
   /// [::ur_bool_t] Returns true if the device supports the use of
   /// command-buffers.
   UR_DEVICE_INFO_COMMAND_BUFFER_SUPPORT_EXP = 0x1000,

include/unified-runtime/ur_print.hpp

@@ -3180,6 +3180,9 @@ inline std::ostream &operator<<(std::ostream &os, enum ur_device_info_t value) {
   case UR_DEVICE_INFO_NATIVE_VECTOR_WIDTH_LONG_LONG:
     os << "UR_DEVICE_INFO_NATIVE_VECTOR_WIDTH_LONG_LONG";
     break;
+  case UR_DEVICE_INFO_MAX_WORK_GROUPS:
+    os << "UR_DEVICE_INFO_MAX_WORK_GROUPS";
+    break;
   case UR_DEVICE_INFO_XE_STACK_COUNT:
     os << "UR_DEVICE_INFO_XE_STACK_COUNT";
     break;
@@ -5038,6 +5041,19 @@ inline ur_result_t printTagged(std::ostream &os, const void *ptr,
 
     os << ")";
   } break;
+  case UR_DEVICE_INFO_MAX_WORK_GROUPS: {
+    const size_t *tptr = (const size_t *)ptr;
+    if (sizeof(size_t) > size) {
+      os << "invalid size (is: " << size << ", expected: >=" << sizeof(size_t)
+         << ")";
+      return UR_RESULT_ERROR_INVALID_SIZE;
+    }
+    os << (const void *)(tptr) << " (";
+
+    os << *tptr;
+
+    os << ")";
+  } break;
   case UR_DEVICE_INFO_XE_STACK_COUNT: {
     const uint32_t *tptr = (const uint32_t *)ptr;
     if (sizeof(uint32_t) > size) {

scripts/core/device.yml

@@ -473,6 +473,8 @@ etors:
       desc: "[uint32_t] preferred vector width for long long"
     - name: NATIVE_VECTOR_WIDTH_LONG_LONG
       desc: "[uint32_t] native vector width for long long"
+    - name: MAX_WORK_GROUPS
+      desc: "[size_t] return max total number of work groups"
     - name: XE_STACK_COUNT
       desc: "[uint32_t][optional-query] return Intel GPU number of stacks/chiplets/tiles"
     - name: XE_REGIONS_PER_STACK

source/adapters/cuda/device.cpp

@@ -110,6 +110,25 @@ UR_APIEXPORT ur_result_t UR_APICALL urDeviceGetInfo(ur_device_handle_t hDevice,
     return ReturnValue(ReturnSizes);
   }
 
+  case UR_DEVICE_INFO_MAX_WORK_GROUPS: {
+    int MaxX = 0, MaxY = 0, MaxZ = 0;
+    UR_CHECK_ERROR(cuDeviceGetAttribute(
+        &MaxX, CU_DEVICE_ATTRIBUTE_MAX_GRID_DIM_X, hDevice->get()));
+    assert(MaxX >= 0);
+
+    UR_CHECK_ERROR(cuDeviceGetAttribute(
+        &MaxY, CU_DEVICE_ATTRIBUTE_MAX_GRID_DIM_Y, hDevice->get()));
+    assert(MaxY >= 0);
+
+    UR_CHECK_ERROR(cuDeviceGetAttribute(
+        &MaxZ, CU_DEVICE_ATTRIBUTE_MAX_GRID_DIM_Z, hDevice->get()));
+    assert(MaxZ >= 0);
+
+    return ReturnValue(multiplyWithOverflowCheck(static_cast<size_t>(MaxX),
+                                                 static_cast<size_t>(MaxY),
+                                                 static_cast<size_t>(MaxZ)));
+  }
+
   case UR_DEVICE_INFO_MAX_WORK_GROUP_SIZE: {
     int MaxWorkGroupSize = 0;
     UR_CHECK_ERROR(cuDeviceGetAttribute(

source/adapters/hip/device.cpp

@@ -117,6 +117,25 @@ UR_APIEXPORT ur_result_t UR_APICALL urDeviceGetInfo(ur_device_handle_t hDevice,
     return ReturnValue(return_sizes);
   }
 
+  case UR_DEVICE_INFO_MAX_WORK_GROUPS: {
+    int MaxX = 0, MaxY = 0, MaxZ = 0;
+    UR_CHECK_ERROR(hipDeviceGetAttribute(&MaxX, hipDeviceAttributeMaxGridDimX,
+                                         hDevice->get()));
+    assert(MaxX >= 0);
+
+    UR_CHECK_ERROR(hipDeviceGetAttribute(&MaxY, hipDeviceAttributeMaxGridDimY,
+                                         hDevice->get()));
+    assert(MaxY >= 0);
+
+    UR_CHECK_ERROR(hipDeviceGetAttribute(&MaxZ, hipDeviceAttributeMaxGridDimZ,
+                                         hDevice->get()));
+    assert(MaxZ >= 0);
+
+    return ReturnValue(multiplyWithOverflowCheck(static_cast<size_t>(MaxX),
+                                                 static_cast<size_t>(MaxY),
+                                                 static_cast<size_t>(MaxZ)));
+  }
+
   case UR_DEVICE_INFO_MAX_WORK_GROUP_SIZE: {
     int MaxWorkGroupSize = 0;
     UR_CHECK_ERROR(hipDeviceGetAttribute(&MaxWorkGroupSize,

source/adapters/level_zero/device.cpp

@@ -471,6 +471,14 @@ ur_result_t urDeviceGetInfo(
                          Device->ZeDeviceComputeProperties->maxGroupCountZ}};
     return ReturnValue(MaxGroupCounts);
   }
+  case UR_DEVICE_INFO_MAX_WORK_GROUPS: {
+    // Multiply the max group counts in each dimension to get the total max
+    // number of work groups. Prevent overflow.
+    return ReturnValue(multiplyWithOverflowCheck(
+        Device->ZeDeviceComputeProperties->maxGroupCountX,
+        Device->ZeDeviceComputeProperties->maxGroupCountY,
+        Device->ZeDeviceComputeProperties->maxGroupCountZ));
+  }
   case UR_DEVICE_INFO_MAX_CLOCK_FREQUENCY:
     return ReturnValue(uint32_t{Device->ZeDeviceProperties->coreClockRate});
   case UR_DEVICE_INFO_ADDRESS_BITS: {

source/adapters/native_cpu/device.cpp

@@ -363,6 +363,7 @@ UR_APIEXPORT ur_result_t UR_APICALL urDeviceGetInfo(ur_device_handle_t hDevice,
   case UR_DEVICE_INFO_SUB_GROUP_INDEPENDENT_FORWARD_PROGRESS:
   case UR_DEVICE_INFO_IL_VERSION:
   case UR_DEVICE_INFO_MAX_WORK_GROUPS_3D:
+  case UR_DEVICE_INFO_MAX_WORK_GROUPS:
   case UR_DEVICE_INFO_MEMORY_CLOCK_RATE:
   case UR_DEVICE_INFO_MEMORY_BUS_WIDTH:
   case UR_DEVICE_INFO_GLOBAL_MEM_FREE:

source/adapters/offload/device.cpp

@@ -8,6 +8,7 @@
 //===----------------------------------------------------------------------===//
 
 #include <OffloadAPI.h>
+#include <limits>
 #include <unified-runtime/ur_api.h>
 #include <ur/ur.hpp>
 
@@ -211,6 +212,27 @@ UR_APIEXPORT ur_result_t UR_APICALL urDeviceGetInfo(ur_device_handle_t hDevice,
 
     return UR_RESULT_SUCCESS;
   }
+  case UR_DEVICE_INFO_MAX_WORK_GROUPS: {
+    // OL dimensions are uint32_t while UR is size_t, so they need to be mapped.
+    if (pPropSizeRet) {
+      *pPropSizeRet = sizeof(size_t);
+    }
+
+    if (pPropValue) {
+      ol_dimensions_t olVec;
+      OL_RETURN_ON_ERR(olGetDeviceInfo(
+          hDevice->OffloadDevice, OL_DEVICE_INFO_MAX_WORK_SIZE_PER_DIMENSION,
+          sizeof(olVec), &olVec));
+
+      // Multiply the max group counts in each dimension to get the total max
+      // number of work groups. Prevent overflow.
+      *reinterpret_cast<size_t *>(pPropValue) = multiplyWithOverflowCheck(
+          static_cast<size_t>(olVec.x), static_cast<size_t>(olVec.y),
+          static_cast<size_t>(olVec.z));
+    }
+
+    return UR_RESULT_SUCCESS;
+  }
 
   // Unimplemented features
   case UR_DEVICE_INFO_PROGRAM_SET_SPECIALIZATION_CONSTANTS:

source/adapters/opencl/device.cpp

@@ -269,6 +269,9 @@ UR_APIEXPORT ur_result_t UR_APICALL urDeviceGetInfo(ur_device_handle_t hDevice,
     ReturnSizes.sizes[2] = Max;
     return ReturnValue(ReturnSizes);
   }
+  case UR_DEVICE_INFO_MAX_WORK_GROUPS: {
+    return ReturnValue(std::numeric_limits<size_t>::max());
+  }
   case UR_DEVICE_INFO_MAX_COMPUTE_QUEUE_INDICES: {
     return ReturnValue(static_cast<uint32_t>(1u));
   }

source/common/ur_util.hpp

@@ -322,6 +322,37 @@ inline ur_result_t exceptionToResult(std::exception_ptr eptr) {
   }
 }
 
+// Multiply a, b and c, and check for overflow. If overflow occurs, return
+// MAX_SIZE_T.
+inline size_t multiplyWithOverflowCheck(size_t a, size_t b, size_t c) {
+
+  size_t Product = 0;
+  size_t MaxSizeTVal = std::numeric_limits<size_t>::max();
+
+  if (a == 0 || b == 0 || c == 0) {
+    return 0;
+  }
+
+#ifndef _MSC_VER
+  if (__builtin_mul_overflow(a, b, &Product) ||
+      __builtin_mul_overflow(Product, c, &Product)) {
+    return MaxSizeTVal; // Overflow occurred, return max possible value.
+  }
+#else
+  if (b > MaxSizeTVal / a) {
+    return MaxSizeTVal; // Overflow occurred, return max possible value.
+  }
+  Product = a * b;
+
+  if (c > MaxSizeTVal / Product) {
+    return MaxSizeTVal; // Overflow occurred, return max possible value.
+  }
+  Product *= c;
+#endif
+
+  return Product;
+}
+
 template <class> inline constexpr bool ur_always_false_t = false;
 
 namespace {