Clear OpenCL function pointer cache when context is released.

source/adapters/opencl/common.hpp

@@ -319,45 +319,33 @@ cl_int(CL_API_CALL *)(cl_command_buffer_khr command_buffer,
 template <typename T> struct FuncPtrCache {
   std::map<cl_context, T> Map;
   std::mutex Mutex;
+
+  void clear(cl_context context) {
+    std::lock_guard<std::mutex> CacheLock{Mutex};
+    Map.erase(context);
+  }
 };
 
-// FIXME: There's currently no mechanism for cleaning up this cache, meaning
-// that it is invalidated whenever a context is destroyed. This could lead to
-// reusing an invalid function pointer if another context happens to have the
-// same native handle.
 struct ExtFuncPtrCacheT {
-  FuncPtrCache<clHostMemAllocINTEL_fn> clHostMemAllocINTELCache;
-  FuncPtrCache<clDeviceMemAllocINTEL_fn> clDeviceMemAllocINTELCache;
-  FuncPtrCache<clSharedMemAllocINTEL_fn> clSharedMemAllocINTELCache;
-  FuncPtrCache<clGetDeviceFunctionPointer_fn> clGetDeviceFunctionPointerCache;
-  FuncPtrCache<clGetDeviceGlobalVariablePointer_fn>
-      clGetDeviceGlobalVariablePointerCache;
-  FuncPtrCache<clCreateBufferWithPropertiesINTEL_fn>
-      clCreateBufferWithPropertiesINTELCache;
-  FuncPtrCache<clMemBlockingFreeINTEL_fn> clMemBlockingFreeINTELCache;
-  FuncPtrCache<clSetKernelArgMemPointerINTEL_fn>
-      clSetKernelArgMemPointerINTELCache;
-  FuncPtrCache<clEnqueueMemFillINTEL_fn> clEnqueueMemFillINTELCache;
-  FuncPtrCache<clEnqueueMemcpyINTEL_fn> clEnqueueMemcpyINTELCache;
-  FuncPtrCache<clGetMemAllocInfoINTEL_fn> clGetMemAllocInfoINTELCache;
-  FuncPtrCache<clEnqueueWriteGlobalVariable_fn>
-      clEnqueueWriteGlobalVariableCache;
-  FuncPtrCache<clEnqueueReadGlobalVariable_fn> clEnqueueReadGlobalVariableCache;
-  FuncPtrCache<clEnqueueReadHostPipeINTEL_fn> clEnqueueReadHostPipeINTELCache;
-  FuncPtrCache<clEnqueueWriteHostPipeINTEL_fn> clEnqueueWriteHostPipeINTELCache;
-  FuncPtrCache<clSetProgramSpecializationConstant_fn>
-      clSetProgramSpecializationConstantCache;
-  FuncPtrCache<clCreateCommandBufferKHR_fn> clCreateCommandBufferKHRCache;
-  FuncPtrCache<clRetainCommandBufferKHR_fn> clRetainCommandBufferKHRCache;
-  FuncPtrCache<clReleaseCommandBufferKHR_fn> clReleaseCommandBufferKHRCache;
-  FuncPtrCache<clFinalizeCommandBufferKHR_fn> clFinalizeCommandBufferKHRCache;
-  FuncPtrCache<clCommandNDRangeKernelKHR_fn> clCommandNDRangeKernelKHRCache;
-  FuncPtrCache<clCommandCopyBufferKHR_fn> clCommandCopyBufferKHRCache;
-  FuncPtrCache<clCommandCopyBufferRectKHR_fn> clCommandCopyBufferRectKHRCache;
-  FuncPtrCache<clCommandFillBufferKHR_fn> clCommandFillBufferKHRCache;
-  FuncPtrCache<clEnqueueCommandBufferKHR_fn> clEnqueueCommandBufferKHRCache;
-  FuncPtrCache<clGetCommandBufferInfoKHR_fn> clGetCommandBufferInfoKHRCache;
-  FuncPtrCache<clUpdateMutableCommandsKHR_fn> clUpdateMutableCommandsKHRCache;
+#define CL_EXTENSION_FUNC(func) FuncPtrCache<func##_fn> func##Cache;
+
+#include "extension_functions.def"
+
+#undef CL_EXTENSION_FUNC
+
+  // If a context stored in the current caching mechanism is destroyed by the
+  // CL driver all of its function pointers are invalidated. This can lead to a
+  // pathological case where a subsequently created context gets returned with
+  // a coincidentally identical handle to the destroyed one and ends up being
+  // used to retrieve bad function pointers. To avoid this we clear the cache
+  // when contexts are released.
+  void clearCache(cl_context context) {
+#define CL_EXTENSION_FUNC(func) func##Cache.clear(context);
+
+#include "extension_functions.def"
+
+#undef CL_EXTENSION_FUNC
+  }
 };
 // A raw pointer is used here since the lifetime of this map has to be tied to
 // piTeardown to avoid issues with static destruction order (a user application

source/adapters/opencl/context.cpp

@@ -113,9 +113,30 @@ urContextGetInfo(ur_context_handle_t hContext, ur_context_info_t propName,
 
 UR_APIEXPORT ur_result_t UR_APICALL
 urContextRelease(ur_context_handle_t hContext) {
+  // If we're reasonably sure this context is about to be detroyed we should
+  // clear the ext function pointer cache. This isn't foolproof sadly but it
+  // should drastically reduce the chances of the pathological case described
+  // in the comments in common.hpp.
+  static std::mutex contextReleaseMutex;
+  auto clContext = cl_adapter::cast<cl_context>(hContext);
 
-  cl_int Ret = clReleaseContext(cl_adapter::cast<cl_context>(hContext));
-  return mapCLErrorToUR(Ret);
+  {
+    std::lock_guard<std::mutex> lock(contextReleaseMutex);
+    size_t refCount = 0;
+    CL_RETURN_ON_FAILURE(clGetContextInfo(clContext, CL_CONTEXT_REFERENCE_COUNT,
+                                          sizeof(size_t), &refCount, nullptr));
+
+    // ExtFuncPtrCache is destroyed in an atexit() callback, so it doesn't
+    // necessarily outlive the adapter (or all the contexts).
+    if (refCount == 1 && cl_ext::ExtFuncPtrCache) {
+      cl_ext::ExtFuncPtrCache->clearCache(clContext);
+    }
+  }
+
+  CL_RETURN_ON_FAILURE(
+      clReleaseContext(cl_adapter::cast<cl_context>(hContext)));
+
+  return UR_RESULT_SUCCESS;
 }
 
 UR_APIEXPORT ur_result_t UR_APICALL

source/adapters/opencl/extension_functions.def

@@ -0,0 +1,27 @@
+CL_EXTENSION_FUNC(clHostMemAllocINTEL)
+CL_EXTENSION_FUNC(clDeviceMemAllocINTEL)
+CL_EXTENSION_FUNC(clSharedMemAllocINTEL)
+CL_EXTENSION_FUNC(clGetDeviceFunctionPointer)
+CL_EXTENSION_FUNC(clGetDeviceGlobalVariablePointer)
+CL_EXTENSION_FUNC(clCreateBufferWithPropertiesINTEL)
+CL_EXTENSION_FUNC(clMemBlockingFreeINTEL)
+CL_EXTENSION_FUNC(clSetKernelArgMemPointerINTEL)
+CL_EXTENSION_FUNC(clEnqueueMemFillINTEL)
+CL_EXTENSION_FUNC(clEnqueueMemcpyINTEL)
+CL_EXTENSION_FUNC(clGetMemAllocInfoINTEL)
+CL_EXTENSION_FUNC(clEnqueueWriteGlobalVariable)
+CL_EXTENSION_FUNC(clEnqueueReadGlobalVariable)
+CL_EXTENSION_FUNC(clEnqueueReadHostPipeINTEL)
+CL_EXTENSION_FUNC(clEnqueueWriteHostPipeINTEL)
+CL_EXTENSION_FUNC(clSetProgramSpecializationConstant)
+CL_EXTENSION_FUNC(clCreateCommandBufferKHR)
+CL_EXTENSION_FUNC(clRetainCommandBufferKHR)
+CL_EXTENSION_FUNC(clReleaseCommandBufferKHR)
+CL_EXTENSION_FUNC(clFinalizeCommandBufferKHR)
+CL_EXTENSION_FUNC(clCommandNDRangeKernelKHR)
+CL_EXTENSION_FUNC(clCommandCopyBufferKHR)
+CL_EXTENSION_FUNC(clCommandCopyBufferRectKHR)
+CL_EXTENSION_FUNC(clCommandFillBufferKHR)
+CL_EXTENSION_FUNC(clEnqueueCommandBufferKHR)
+CL_EXTENSION_FUNC(clGetCommandBufferInfoKHR)
+CL_EXTENSION_FUNC(clUpdateMutableCommandsKHR)