Merge pull request #2614 from kurapov-peter/spills

Add UR_KERNEL_INFO_SPILL_MEM_SIZE kernel info prop
oneapi-src · Feb 4, 2025 · 08d36b7 · 08d36b7
2 parents 4c504db + e6b61c6
commit 08d36b7
Show file tree

Hide file tree

Showing 14 changed files with 77 additions and 4 deletions.
diff --git a/include/ur_api.h b/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

diff --git a/include/ur_print.hpp b/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;

diff --git a/scripts/core/kernel.yml b/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"

diff --git a/source/adapters/cuda/kernel.cpp b/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;
   }

diff --git a/source/adapters/hip/kernel.cpp b/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;
   }

diff --git a/source/adapters/level_zero/kernel.cpp b/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:

diff --git a/source/adapters/level_zero/v2/kernel.cpp b/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: {

diff --git a/source/adapters/native_cpu/kernel.cpp b/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;
   }

diff --git a/source/adapters/opencl/kernel.cpp b/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,

diff --git a/source/loader/layers/validation/ur_valddi.cpp b/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)

diff --git a/source/loader/ur_libapi.cpp b/source/loader/ur_libapi.cpp
@@ -3798,7 +3798,7 @@ ur_result_t UR_APICALL urKernelSetArgLocal(
 ///     - ::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

diff --git a/source/ur_api.cpp b/source/ur_api.cpp
@@ -3330,7 +3330,7 @@ ur_result_t UR_APICALL urKernelSetArgLocal(
 ///     - ::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

diff --git a/test/conformance/kernel/urKernelGetInfo.cpp b/test/conformance/kernel/urKernelGetInfo.cpp
@@ -136,6 +136,22 @@ TEST_P(urKernelGetInfoTest, SuccessNumRegs) {
                                  property_value.data(), nullptr));
 }
 
+TEST_P(urKernelGetInfoTest, SuccessSpillMemSize) {
+  UUR_KNOWN_FAILURE_ON(uur::HIP{}, uur::OpenCL{});
+
+  ur_kernel_info_t property_name = UR_KERNEL_INFO_SPILL_MEM_SIZE;
+  size_t property_size = 0;
+
+  ASSERT_SUCCESS_OR_OPTIONAL_QUERY(
+      urKernelGetInfo(kernel, property_name, 0, nullptr, &property_size),
+      property_name);
+  ASSERT_EQ(property_size, sizeof(uint32_t));
+
+  std::vector<uint32_t> property_value(property_size / sizeof(uint32_t));
+  ASSERT_SUCCESS(urKernelGetInfo(kernel, property_name, property_size,
+                                 property_value.data(), nullptr));
+}
+
 TEST_P(urKernelGetInfoTest, InvalidNullHandleKernel) {
   size_t kernel_name_length = 0;
   ASSERT_EQ_RESULT(UR_RESULT_ERROR_INVALID_NULL_HANDLE,

diff --git a/test/conformance/testing/include/uur/optional_queries.h b/test/conformance/testing/include/uur/optional_queries.h
@@ -86,6 +86,7 @@ template <> inline bool isQueryOptional(ur_program_info_t query) {
 
 constexpr std::array optional_ur_kernel_info_t = {
     UR_KERNEL_INFO_NUM_REGS,
+    UR_KERNEL_INFO_SPILL_MEM_SIZE,
 };
 
 template <> inline bool isQueryOptional(ur_kernel_info_t query) {