temp fix opencl runtime for docker

apache · Mar 30, 2018 · 06480f5 · 06480f5
1 parent 972b587
commit 06480f5
Showing 1 changed file with 5 additions and 0 deletions.
diff --git a/src/runtime/opencl/opencl_device_api.cc b/src/runtime/opencl/opencl_device_api.cc
@@ -145,6 +145,11 @@ std::string GetDeviceInfo(
 }
 
 std::vector<cl_platform_id> GetPlatformIDs() {
+  // Trigger CUDA initialziation when it exists
+  // on some NV platform opencl depends on CUDA to be loaded first
+  if (const PackedFunc* query = runtime::Registry::Get("_GetDeviceAttr")) {
+    (*query)(static_cast<int>(kDLGPU), 0, static_cast<int>(kExist));
+  }
   cl_uint ret_size;
   cl_int code = clGetPlatformIDs(0, nullptr, &ret_size);
   std::vector<cl_platform_id> ret;