Merge pull request #2525 from dyniols/num_compute_units

Addition of num_compute_units query for device info

Author: kbenzie

include/ur_api.h

@@ -2198,9 +2198,11 @@ typedef enum ur_device_info_t {
   /// to the `USMPool` entry points and usage of the `pool` parameter of the
   /// USM alloc entry points.
   UR_DEVICE_INFO_USM_POOL_SUPPORT = 119,
+  /// [uint32_t] the number of compute units for specific backend.
+  UR_DEVICE_INFO_NUM_COMPUTE_UNITS = 120,
   /// [::ur_bool_t] support the ::urProgramSetSpecializationConstants entry
   /// point
-  UR_DEVICE_INFO_PROGRAM_SET_SPECIALIZATION_CONSTANTS = 120,
+  UR_DEVICE_INFO_PROGRAM_SET_SPECIALIZATION_CONSTANTS = 121,
   /// [::ur_bool_t] Returns true if the device supports the use of
   /// command-buffers.
   UR_DEVICE_INFO_COMMAND_BUFFER_SUPPORT_EXP = 0x1000,

include/ur_print.hpp

@@ -2917,6 +2917,9 @@ inline std::ostream &operator<<(std::ostream &os, enum ur_device_info_t value) {
   case UR_DEVICE_INFO_USM_POOL_SUPPORT:
     os << "UR_DEVICE_INFO_USM_POOL_SUPPORT";
     break;
+  case UR_DEVICE_INFO_NUM_COMPUTE_UNITS:
+    os << "UR_DEVICE_INFO_NUM_COMPUTE_UNITS";
+    break;
   case UR_DEVICE_INFO_PROGRAM_SET_SPECIALIZATION_CONSTANTS:
     os << "UR_DEVICE_INFO_PROGRAM_SET_SPECIALIZATION_CONSTANTS";
     break;
@@ -4551,6 +4554,19 @@ inline ur_result_t printTagged(std::ostream &os, const void *ptr,
 
     os << ")";
   } break;
+  case UR_DEVICE_INFO_NUM_COMPUTE_UNITS: {
+    const uint32_t *tptr = (const uint32_t *)ptr;
+    if (sizeof(uint32_t) > size) {
+      os << "invalid size (is: " << size << ", expected: >=" << sizeof(uint32_t)
+         << ")";
+      return UR_RESULT_ERROR_INVALID_SIZE;
+    }
+    os << (const void *)(tptr) << " (";
+
+    os << *tptr;
+
+    os << ")";
+  } break;
   case UR_DEVICE_INFO_PROGRAM_SET_SPECIALIZATION_CONSTANTS: {
     const ur_bool_t *tptr = (const ur_bool_t *)ptr;
     if (sizeof(ur_bool_t) > size) {

scripts/core/device.yml

@@ -443,6 +443,8 @@ etors:
       desc: "[$x_bool_t] return true if the device supports the `EnqueueDeviceGlobalVariableWrite` and `EnqueueDeviceGlobalVariableRead` entry points."
     - name: USM_POOL_SUPPORT
       desc: "[$x_bool_t] return true if the device supports USM pooling. Pertains to the `USMPool` entry points and usage of the `pool` parameter of the USM alloc entry points."
+    - name: NUM_COMPUTE_UNITS
+      desc: "[uint32_t] the number of compute units for specific backend."
     - name: PROGRAM_SET_SPECIALIZATION_CONSTANTS
       desc: "[$x_bool_t] support the $xProgramSetSpecializationConstants entry point"
 --- #--------------------------------------------------------------------------

source/adapters/cuda/device.cpp

@@ -56,6 +56,7 @@ UR_APIEXPORT ur_result_t UR_APICALL urDeviceGetInfo(ur_device_handle_t hDevice,
   case UR_DEVICE_INFO_VENDOR_ID: {
     return ReturnValue(4318u);
   }
+  case UR_DEVICE_INFO_NUM_COMPUTE_UNITS:
   case UR_DEVICE_INFO_MAX_COMPUTE_UNITS: {
     return ReturnValue(hDevice->getNumComputeUnits());
   }

source/adapters/hip/device.cpp

@@ -57,6 +57,7 @@ UR_APIEXPORT ur_result_t UR_APICALL urDeviceGetInfo(ur_device_handle_t hDevice,
 #endif
     return ReturnValue(VendorId);
   }
+  case UR_DEVICE_INFO_NUM_COMPUTE_UNITS:
   case UR_DEVICE_INFO_MAX_COMPUTE_UNITS: {
     int ComputeUnits = 0;
     UR_CHECK_ERROR(hipDeviceGetAttribute(

source/adapters/level_zero/device.cpp

@@ -324,6 +324,12 @@ ur_result_t urDeviceGetInfo(
 
     return ReturnValue(uint32_t{MaxComputeUnits});
   }
+  case UR_DEVICE_INFO_NUM_COMPUTE_UNITS: {
+    uint32_t NumComputeUnits =
+        Device->ZeDeviceProperties->numSubslicesPerSlice *
+        Device->ZeDeviceProperties->numSlices;
+    return ReturnValue(uint32_t{NumComputeUnits});
+  }
   case UR_DEVICE_INFO_MAX_WORK_ITEM_DIMENSIONS:
     // Level Zero spec defines only three dimensions
     return ReturnValue(uint32_t{3});

source/adapters/native_cpu/device.cpp

@@ -164,6 +164,7 @@ UR_APIEXPORT ur_result_t UR_APICALL urDeviceGetInfo(ur_device_handle_t hDevice,
     return ReturnValue(bool{true});
   case UR_DEVICE_INFO_LINKER_AVAILABLE:
     return ReturnValue(bool{true});
+  case UR_DEVICE_INFO_NUM_COMPUTE_UNITS:
   case UR_DEVICE_INFO_MAX_COMPUTE_UNITS:
     return ReturnValue(static_cast<uint32_t>(hDevice->tp.num_threads()));
   case UR_DEVICE_INFO_PARTITION_MAX_SUB_DEVICES:

source/adapters/opencl/device.cpp

@@ -828,6 +828,38 @@ UR_APIEXPORT ur_result_t UR_APICALL urDeviceGetInfo(ur_device_handle_t hDevice,
   case UR_DEVICE_INFO_VIRTUAL_MEMORY_SUPPORT: {
     return ReturnValue(false);
   }
+  case UR_DEVICE_INFO_NUM_COMPUTE_UNITS: {
+
+    bool ExtensionSupported = false;
+    UR_RETURN_ON_FAILURE(cl_adapter::checkDeviceExtensions(
+        cl_adapter::cast<cl_device_id>(hDevice),
+        {"cl_intel_device_attribute_query"}, ExtensionSupported));
+
+    cl_device_type CLType;
+    CL_RETURN_ON_FAILURE(
+        clGetDeviceInfo(cl_adapter::cast<cl_device_id>(hDevice), CL_DEVICE_TYPE,
+                        sizeof(cl_device_type), &CLType, nullptr));
+
+    cl_uint NumComputeUnits;
+    if (ExtensionSupported && (CLType & CL_DEVICE_TYPE_GPU)) {
+      cl_uint SliceCount = 0;
+      cl_uint SubSlicePerSliceCount = 0;
+      CL_RETURN_ON_FAILURE(clGetDeviceInfo(
+          cl_adapter::cast<cl_device_id>(hDevice), CL_DEVICE_NUM_SLICES_INTEL,
+          sizeof(cl_uint), &SliceCount, nullptr));
+      CL_RETURN_ON_FAILURE(
+          clGetDeviceInfo(cl_adapter::cast<cl_device_id>(hDevice),
+                          CL_DEVICE_NUM_SUB_SLICES_PER_SLICE_INTEL,
+                          sizeof(cl_uint), &SubSlicePerSliceCount, nullptr));
+      NumComputeUnits = SliceCount * SubSlicePerSliceCount;
+    } else {
+      CL_RETURN_ON_FAILURE(clGetDeviceInfo(
+          cl_adapter::cast<cl_device_id>(hDevice), CL_DEVICE_MAX_COMPUTE_UNITS,
+          sizeof(cl_uint), &NumComputeUnits, nullptr));
+    }
+
+    return ReturnValue(static_cast<uint32_t>(NumComputeUnits));
+  }
   case UR_DEVICE_INFO_TIMESTAMP_RECORDING_SUPPORT_EXP: {
     return ReturnValue(false);
   }

test/conformance/device/urDeviceGetInfo.cpp

@@ -128,7 +128,8 @@ static std::unordered_map<ur_device_info_t, size_t> device_info_size_map = {
      sizeof(ur_memory_scope_capability_flags_t)},
     {UR_DEVICE_INFO_ESIMD_SUPPORT, sizeof(ur_bool_t)},
     {UR_DEVICE_INFO_IP_VERSION, sizeof(uint32_t)},
-    {UR_DEVICE_INFO_VIRTUAL_MEMORY_SUPPORT, sizeof(ur_bool_t)}};
+    {UR_DEVICE_INFO_VIRTUAL_MEMORY_SUPPORT, sizeof(ur_bool_t)},
+    {UR_DEVICE_INFO_NUM_COMPUTE_UNITS, sizeof(uint32_t)}};
 
 using urDeviceGetInfoTest = uur::urDeviceTestWithParam<ur_device_info_t>;
 
@@ -253,6 +254,7 @@ UUR_DEVICE_TEST_SUITE_P(
         UR_DEVICE_INFO_PREFERRED_VECTOR_WIDTH_DOUBLE,          //
         UR_DEVICE_INFO_PREFERRED_VECTOR_WIDTH_HALF,            //
         UR_DEVICE_INFO_PREFERRED_VECTOR_WIDTH_INT,             //
+        UR_DEVICE_INFO_NUM_COMPUTE_UNITS,                      //
         UR_DEVICE_INFO_PROGRAM_SET_SPECIALIZATION_CONSTANTS    //
         ),
     uur::deviceTestWithParamPrinter<ur_device_info_t>);

tools/urinfo/urinfo.hpp

@@ -327,6 +327,8 @@ inline void printDeviceInfos(ur_device_handle_t hDevice,
   std::cout << prefix;
   printDeviceInfo<ur_bool_t>(hDevice, UR_DEVICE_INFO_USM_POOL_SUPPORT);
   std::cout << prefix;
+  printDeviceInfo<uint32_t>(hDevice, UR_DEVICE_INFO_NUM_COMPUTE_UNITS);
+  std::cout << prefix;
   printDeviceInfo<ur_bool_t>(
       hDevice, UR_DEVICE_INFO_PROGRAM_SET_SPECIALIZATION_CONSTANTS);
   std::cout << prefix;