Merge pull request oneapi-src#2614 from kurapov-peter/spills

Add UR_KERNEL_INFO_SPILL_MEM_SIZE kernel info prop

Author: kbenzie

include/ur_api.h

@@ -5993,6 +5993,13 @@ typedef enum ur_kernel_info_t {
   /// [uint32_t][optional-query] Return the number of registers used by the
   /// compiled kernel.
   UR_KERNEL_INFO_NUM_REGS = 6,
+  /// [uint32_t[]][optional-query] Return the spill memory size allocated by
+  /// the compiler.
+  /// The returned values correspond to the associated devices.
+  /// The order of the devices is guaranteed (i.e., the same as queried by
+  /// `urDeviceGet`) by the UR within a single application even if the runtime
+  /// is reinitialized.
+  UR_KERNEL_INFO_SPILL_MEM_SIZE = 7,
   /// @cond
   UR_KERNEL_INFO_FORCE_UINT32 = 0x7fffffff
   /// @endcond
@@ -6093,7 +6100,7 @@ typedef enum ur_kernel_exec_info_t {
 ///     - ::UR_RESULT_ERROR_INVALID_NULL_HANDLE
 ///         + `NULL == hKernel`
 ///     - ::UR_RESULT_ERROR_INVALID_ENUMERATION
-///         + `::UR_KERNEL_INFO_NUM_REGS < propName`
+///         + `::UR_KERNEL_INFO_SPILL_MEM_SIZE < propName`
 ///     - ::UR_RESULT_ERROR_UNSUPPORTED_ENUMERATION
 ///         + If `propName` is not supported by the adapter.
 ///     - ::UR_RESULT_ERROR_INVALID_SIZE

include/ur_print.hpp

@@ -8781,6 +8781,9 @@ inline std::ostream &operator<<(std::ostream &os, enum ur_kernel_info_t value) {
   case UR_KERNEL_INFO_NUM_REGS:
     os << "UR_KERNEL_INFO_NUM_REGS";
     break;
+  case UR_KERNEL_INFO_SPILL_MEM_SIZE:
+    os << "UR_KERNEL_INFO_SPILL_MEM_SIZE";
+    break;
   default:
     os << "unknown enumerator";
     break;
@@ -8873,6 +8876,20 @@ inline ur_result_t printTagged(std::ostream &os, const void *ptr,
 
     os << ")";
   } break;
+  case UR_KERNEL_INFO_SPILL_MEM_SIZE: {
+
+    const uint32_t *tptr = (const uint32_t *)ptr;
+    os << "{";
+    size_t nelems = size / sizeof(uint32_t);
+    for (size_t i = 0; i < nelems; ++i) {
+      if (i != 0) {
+        os << ", ";
+      }
+
+      os << tptr[i];
+    }
+    os << "}";
+  } break;
   default:
     os << "unknown enumerator";
     return UR_RESULT_ERROR_INVALID_ENUMERATION;

scripts/core/kernel.yml

@@ -125,6 +125,12 @@ etors:
       desc: "[char[]] Return null-terminated kernel attributes string."
     - name: NUM_REGS
       desc: "[uint32_t][optional-query] Return the number of registers used by the compiled kernel."
+    - name: SPILL_MEM_SIZE
+      desc: | 
+            [uint32_t[]][optional-query] Return the spill memory size allocated by the compiler.
+            The returned values correspond to the associated devices.
+            The order of the devices is guaranteed (i.e., the same as queried by `urDeviceGet`)
+            by the UR within a single application even if the runtime is reinitialized.
 --- #--------------------------------------------------------------------------
 type: enum
 desc: "Get Kernel Work Group information"

source/adapters/cuda/kernel.cpp

@@ -299,6 +299,8 @@ UR_APIEXPORT ur_result_t UR_APICALL urKernelGetInfo(ur_kernel_handle_t hKernel,
                                       hKernel->get()));
     return ReturnValue(static_cast<uint32_t>(NumRegs));
   }
+  case UR_KERNEL_INFO_SPILL_MEM_SIZE:
+    return UR_RESULT_ERROR_UNSUPPORTED_ENUMERATION;
   default:
     break;
   }

source/adapters/hip/kernel.cpp

@@ -234,6 +234,8 @@ UR_APIEXPORT ur_result_t UR_APICALL urKernelGetInfo(ur_kernel_handle_t hKernel,
                                        hKernel->get()));
     return ReturnValue(static_cast<uint32_t>(NumRegs));
   }
+  case UR_KERNEL_INFO_SPILL_MEM_SIZE:
+    return UR_RESULT_ERROR_UNSUPPORTED_ENUMERATION;
   default:
     break;
   }

source/adapters/level_zero/kernel.cpp

@@ -751,6 +751,12 @@ ur_result_t urKernelGetInfo(
   case UR_KERNEL_INFO_NUM_REGS:
   case UR_KERNEL_INFO_NUM_ARGS:
     return ReturnValue(uint32_t{Kernel->ZeKernelProperties->numKernelArgs});
+  case UR_KERNEL_INFO_SPILL_MEM_SIZE: {
+    std::vector<uint32_t> spills = {
+        uint32_t{Kernel->ZeKernelProperties->spillMemSize}};
+    return ReturnValue(static_cast<const uint32_t *>(spills.data()),
+                       spills.size());
+  }
   case UR_KERNEL_INFO_REFERENCE_COUNT:
     return ReturnValue(uint32_t{Kernel->RefCount.load()});
   case UR_KERNEL_INFO_ATTRIBUTES:

source/adapters/level_zero/v2/kernel.cpp

@@ -625,6 +625,16 @@ ur_result_t urKernelGetInfo(ur_kernel_handle_t hKernel,
   case UR_KERNEL_INFO_NUM_REGS:
   case UR_KERNEL_INFO_NUM_ARGS:
     return ReturnValue(uint32_t{hKernel->getCommonProperties().numKernelArgs});
+  case UR_KERNEL_INFO_SPILL_MEM_SIZE: {
+    std::shared_lock<ur_shared_mutex> Guard(hKernel->getProgramHandle()->Mutex);
+    auto devices = hKernel->getProgramHandle()->AssociatedDevices;
+    std::vector<uint32_t> spills(devices.size());
+    for (size_t i = 0; i < spills.size(); ++i) {
+      spills[i] = uint32_t{hKernel->getProperties(devices[i]).spillMemSize};
+    }
+    return ReturnValue(static_cast<const uint32_t *>(spills.data()),
+                       spills.size());
+  }
   case UR_KERNEL_INFO_REFERENCE_COUNT:
     return ReturnValue(uint32_t{hKernel->RefCount.load()});
   case UR_KERNEL_INFO_ATTRIBUTES: {

source/adapters/native_cpu/kernel.cpp

@@ -105,6 +105,8 @@ UR_APIEXPORT ur_result_t UR_APICALL urKernelGetInfo(ur_kernel_handle_t hKernel,
     return ReturnValue(uint32_t{hKernel->getReferenceCount()});
   case UR_KERNEL_INFO_ATTRIBUTES:
     return ReturnValue("");
+  case UR_KERNEL_INFO_SPILL_MEM_SIZE:
+    return UR_RESULT_ERROR_UNSUPPORTED_ENUMERATION;
   default:
     return UR_RESULT_ERROR_INVALID_VALUE;
   }

source/adapters/opencl/kernel.cpp

@@ -64,6 +64,7 @@ static cl_int mapURKernelInfoToCL(ur_kernel_info_t URPropName) {
     return CL_KERNEL_ATTRIBUTES;
   // NUM_REGS doesn't have a CL equivalent
   case UR_KERNEL_INFO_NUM_REGS:
+  case UR_KERNEL_INFO_SPILL_MEM_SIZE:
   default:
     return -1;
   }
@@ -78,6 +79,9 @@ UR_APIEXPORT ur_result_t UR_APICALL urKernelGetInfo(ur_kernel_handle_t hKernel,
   if (propName == UR_KERNEL_INFO_NUM_REGS) {
     return UR_RESULT_ERROR_UNSUPPORTED_ENUMERATION;
   }
+  if (propName == UR_KERNEL_INFO_SPILL_MEM_SIZE) {
+    return UR_RESULT_ERROR_UNSUPPORTED_ENUMERATION;
+  }
   size_t CheckPropSize = 0;
   cl_int ClResult = clGetKernelInfo(cl_adapter::cast<cl_kernel>(hKernel),
                                     mapURKernelInfoToCL(propName), propSize,

source/loader/layers/validation/ur_valddi.cpp

@@ -3315,7 +3315,7 @@ __urdlllocal ur_result_t UR_APICALL urKernelGetInfo(
     if (pPropValue == NULL && pPropSizeRet == NULL)
       return UR_RESULT_ERROR_INVALID_NULL_POINTER;
 
-    if (UR_KERNEL_INFO_NUM_REGS < propName)
+    if (UR_KERNEL_INFO_SPILL_MEM_SIZE < propName)
       return UR_RESULT_ERROR_INVALID_ENUMERATION;
 
     if (propSize == 0 && pPropValue != NULL)