Save the function pointer for clSetKernelArgMemPointerINTEL per kernel

This eliminates the need to query the context for the kernel and
look up the function pointer in the extension function pointer cache
each time a pointer is set as a kernel argument.

Signed-off-by: Ben Ashbaugh <[email protected]>

Author: bashbaug

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

@@ -100,8 +100,13 @@ urKernelSetArgLocal(ur_kernel_handle_t hKernel, uint32_t argIndex,
   return UR_RESULT_SUCCESS;
 }
 
-static cl_int mapURKernelInfoToCL(ur_kernel_info_t URPropName) {
+// Querying the number of registers that a kernel uses is supported unofficially
+// on some devices.
+#ifndef CL_KERNEL_REGISTER_COUNT_INTEL
+#define CL_KERNEL_REGISTER_COUNT_INTEL 0x425B
+#endif
 
+static cl_int mapURKernelInfoToCL(ur_kernel_info_t URPropName) {
   switch (static_cast<uint32_t>(URPropName)) {
   case UR_KERNEL_INFO_FUNCTION_NAME:
     return CL_KERNEL_FUNCTION_NAME;
@@ -115,9 +120,10 @@ static cl_int mapURKernelInfoToCL(ur_kernel_info_t URPropName) {
     return CL_KERNEL_PROGRAM;
   case UR_KERNEL_INFO_ATTRIBUTES:
     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:
+    return CL_KERNEL_SPILL_MEM_SIZE_INTEL;
+  case UR_KERNEL_INFO_NUM_REGS:
+    return CL_KERNEL_REGISTER_COUNT_INTEL;
   default:
     return -1;
   }
@@ -132,10 +138,6 @@ UR_APIEXPORT ur_result_t UR_APICALL urKernelGetInfo(ur_kernel_handle_t hKernel,
   UrReturnHelper ReturnValue(propSize, pPropValue, pPropSizeRet);
 
   switch (propName) {
-  // OpenCL doesn't have a way to support this.
-  case UR_KERNEL_INFO_NUM_REGS: {
-    return UR_RESULT_ERROR_UNSUPPORTED_ENUMERATION;
-  }
   case UR_KERNEL_INFO_PROGRAM: {
     return ReturnValue(hKernel->Program);
   }
@@ -145,14 +147,14 @@ UR_APIEXPORT ur_result_t UR_APICALL urKernelGetInfo(ur_kernel_handle_t hKernel,
   case UR_KERNEL_INFO_REFERENCE_COUNT: {
     return ReturnValue(hKernel->getReferenceCount());
   }
-  case UR_KERNEL_INFO_SPILL_MEM_SIZE: {
-    return UR_RESULT_ERROR_UNSUPPORTED_ENUMERATION;
-  }
   default: {
     size_t CheckPropSize = 0;
     cl_int ClResult =
         clGetKernelInfo(hKernel->CLKernel, mapURKernelInfoToCL(propName),
                         propSize, pPropValue, &CheckPropSize);
+    if (ClResult == CL_INVALID_VALUE) {
+      return UR_RESULT_ERROR_UNSUPPORTED_ENUMERATION;
+    }
     if (pPropValue && CheckPropSize != propSize) {
       return UR_RESULT_ERROR_INVALID_SIZE;
     }
@@ -334,12 +336,18 @@ urKernelGetSubGroupInfo(ur_kernel_handle_t hKernel, ur_device_handle_t hDevice,
 }
 
 UR_APIEXPORT ur_result_t UR_APICALL urKernelRetain(ur_kernel_handle_t hKernel) {
+
+  UR_ASSERT(hKernel->getReferenceCount() > 0u, UR_RESULT_ERROR_INVALID_KERNEL);
+
   hKernel->incrementReferenceCount();
   return UR_RESULT_SUCCESS;
 }
 
 UR_APIEXPORT ur_result_t UR_APICALL
 urKernelRelease(ur_kernel_handle_t hKernel) {
+
+  UR_ASSERT(hKernel->getReferenceCount() != 0, UR_RESULT_ERROR_INVALID_KERNEL);
+
   if (hKernel->decrementReferenceCount() == 0) {
     delete hKernel;
   }
@@ -428,25 +436,16 @@ UR_APIEXPORT ur_result_t UR_APICALL urKernelSetArgPointer(
     ur_kernel_handle_t hKernel, uint32_t argIndex,
     const ur_kernel_arg_pointer_properties_t *, const void *pArgValue) {
 
-  cl_context CLContext;
-  CL_RETURN_ON_FAILURE(clGetKernelInfo(hKernel->CLKernel, CL_KERNEL_CONTEXT,
-                                       sizeof(cl_context), &CLContext,
-                                       nullptr));
-
-  clSetKernelArgMemPointerINTEL_fn FuncPtr = nullptr;
-  UR_RETURN_ON_FAILURE(
-      cl_ext::getExtFuncFromContext<clSetKernelArgMemPointerINTEL_fn>(
-          CLContext,
-          ur::cl::getAdapter()->fnCache.clSetKernelArgMemPointerINTELCache,
-          cl_ext::SetKernelArgMemPointerName, &FuncPtr));
-
-  if (FuncPtr) {
-    CL_RETURN_ON_FAILURE(
-        FuncPtr(hKernel->CLKernel, static_cast<cl_uint>(argIndex), pArgValue));
+  if (hKernel->clSetKernelArgMemPointerINTEL == nullptr) {
+    return UR_RESULT_ERROR_UNSUPPORTED_FEATURE;
   }
 
+  CL_RETURN_ON_FAILURE(hKernel->clSetKernelArgMemPointerINTEL(
+      hKernel->CLKernel, static_cast<cl_uint>(argIndex), pArgValue));
+
   return UR_RESULT_SUCCESS;
 }
+
 UR_APIEXPORT ur_result_t UR_APICALL urKernelGetNativeHandle(
     ur_kernel_handle_t hKernel, ur_native_handle_t *phNativeKernel) {
 

unified-runtime/source/adapters/opencl/kernel.hpp

@@ -9,6 +9,7 @@
 //===----------------------------------------------------------------------===//
 #pragma once
 
+#include "adapter.hpp"
 #include "common.hpp"
 #include "context.hpp"
 #include "program.hpp"
@@ -22,13 +23,19 @@ struct ur_kernel_handle_t_ {
   ur_context_handle_t Context;
   std::atomic<uint32_t> RefCount = 0;
   bool IsNativeHandleOwned = true;
+  clSetKernelArgMemPointerINTEL_fn clSetKernelArgMemPointerINTEL = nullptr;
 
   ur_kernel_handle_t_(native_type Kernel, ur_program_handle_t Program,
                       ur_context_handle_t Context)
       : CLKernel(Kernel), Program(Program), Context(Context) {
     RefCount = 1;
     urProgramRetain(Program);
     urContextRetain(Context);
+
+    cl_ext::getExtFuncFromContext<clSetKernelArgMemPointerINTEL_fn>(
+        Context->CLContext,
+        ur::cl::getAdapter()->fnCache.clSetKernelArgMemPointerINTELCache,
+        cl_ext::SetKernelArgMemPointerName, &clSetKernelArgMemPointerINTEL);
   }
 
   ~ur_kernel_handle_t_() {