Limit time held by CUDNN locks.

They inhibit finalizers, resulting in OOMs.

Author: maleadt

lib/cudnn/convolution.jl

@@ -172,52 +172,70 @@ end
 const cudnnConvolutionFwdAlgoPerfCache = Dict{Tuple,cudnnConvolutionFwdAlgoPerf_t}()
 const cudnnConvolutionFwdAlgoPerfCacheLock = ReentrantLock()
 function cudnnConvolutionFwdAlgoPerf(xDesc, x, wDesc, w, convDesc, yDesc, y, biasDesc, activation)
-    lock(cudnnConvolutionFwdAlgoPerfCacheLock) do
-        get!(cudnnConvolutionFwdAlgoPerfCache, (xDesc, wDesc, convDesc, biasDesc, activation)) do
-            requestedAlgoCount = Int(CUDNN_CONVOLUTION_FWD_ALGO_COUNT)
-            returnedAlgoCount = Cint[0]
-            perfResults = Array{cudnnConvolutionFwdAlgoPerf_t}(undef,requestedAlgoCount)
-            workspaceSize() = cudnnFindConvolutionAlgorithmWorkspaceSize(x)
-            with_workspace(workspaceSize) do workspace
-                cudnnFindConvolutionForwardAlgorithmEx(handle(),xDesc,x,wDesc,w,convDesc,yDesc,y,requestedAlgoCount,returnedAlgoCount,perfResults,workspace,sizeof(workspace))
-            end
-            cudnnConvolutionAlgoPerfChoose(perfResults, returnedAlgoCount[1])
+    key = (xDesc, wDesc, convDesc, biasDesc, activation)
+    val = lock(cudnnConvolutionFwdAlgoPerfCacheLock) do
+         get(cudnnConvolutionFwdAlgoPerfCache, key, nothing)
+    end
+    if val === nothing
+        requestedAlgoCount = Int(CUDNN_CONVOLUTION_FWD_ALGO_COUNT)
+        returnedAlgoCount = Cint[0]
+        perfResults = Array{cudnnConvolutionFwdAlgoPerf_t}(undef,requestedAlgoCount)
+        workspaceSize() = cudnnFindConvolutionAlgorithmWorkspaceSize(x)
+        with_workspace(workspaceSize) do workspace
+            cudnnFindConvolutionForwardAlgorithmEx(handle(),xDesc,x,wDesc,w,convDesc,yDesc,y,requestedAlgoCount,returnedAlgoCount,perfResults,workspace,sizeof(workspace))
+        end
+        val = cudnnConvolutionAlgoPerfChoose(perfResults, returnedAlgoCount[1])
+        lock(cudnnConvolutionFwdAlgoPerfCacheLock) do
+            cudnnConvolutionFwdAlgoPerfCache[key] = val
         end
     end
+    return val
 end
 
 const cudnnConvolutionBwdDataAlgoPerfCache = Dict{Tuple,cudnnConvolutionBwdDataAlgoPerf_t}()
 const cudnnConvolutionBwdDataAlgoPerfCacheLock = ReentrantLock()
 function cudnnConvolutionBwdDataAlgoPerf(wDesc, w, dyDesc, dy, convDesc, dxDesc, dx)
-    lock(cudnnConvolutionBwdDataAlgoPerfCacheLock) do
-        get!(cudnnConvolutionBwdDataAlgoPerfCache, (wDesc, dyDesc, convDesc)) do
-            requestedAlgoCount = Int(CUDNN_CONVOLUTION_BWD_DATA_ALGO_COUNT)
-            returnedAlgoCount = Cint[0]
-            perfResults = Array{cudnnConvolutionBwdDataAlgoPerf_t}(undef,requestedAlgoCount)
-            workspaceSize() = cudnnFindConvolutionAlgorithmWorkspaceSize(dx)
-            with_workspace(workspaceSize) do workspace
-                cudnnFindConvolutionBackwardDataAlgorithmEx(handle(),wDesc,w,dyDesc,dy,convDesc,dxDesc,dx,requestedAlgoCount,returnedAlgoCount,perfResults,workspace,sizeof(workspace))
-            end
-            cudnnConvolutionAlgoPerfChoose(perfResults, returnedAlgoCount[1])
+    key = (wDesc, dyDesc, convDesc)
+    val = lock(cudnnConvolutionBwdDataAlgoPerfCacheLock) do
+        get(cudnnConvolutionBwdDataAlgoPerfCache, key, nothing)
+    end
+    if val === nothing
+        requestedAlgoCount = Int(CUDNN_CONVOLUTION_BWD_DATA_ALGO_COUNT)
+        returnedAlgoCount = Cint[0]
+        perfResults = Array{cudnnConvolutionBwdDataAlgoPerf_t}(undef,requestedAlgoCount)
+        workspaceSize() = cudnnFindConvolutionAlgorithmWorkspaceSize(dx)
+        with_workspace(workspaceSize) do workspace
+            cudnnFindConvolutionBackwardDataAlgorithmEx(handle(),wDesc,w,dyDesc,dy,convDesc,dxDesc,dx,requestedAlgoCount,returnedAlgoCount,perfResults,workspace,sizeof(workspace))
+        end
+        val = cudnnConvolutionAlgoPerfChoose(perfResults, returnedAlgoCount[1])
+        lock(cudnnConvolutionBwdDataAlgoPerfCacheLock) do
+            cudnnConvolutionBwdDataAlgoPerfCache[key] = val
         end
     end
+    val
 end
 
 const cudnnConvolutionBwdFilterAlgoPerfCache = Dict{Tuple,cudnnConvolutionBwdFilterAlgoPerf_t}()
 const cudnnConvolutionBwdFilterAlgoPerfCacheLock = ReentrantLock()
 function cudnnConvolutionBwdFilterAlgoPerf(xDesc, x, dyDesc, dy, convDesc, dwDesc, dw)
-    lock(cudnnConvolutionBwdFilterAlgoPerfCacheLock) do
-        get!(cudnnConvolutionBwdFilterAlgoPerfCache, (xDesc, dyDesc, convDesc)) do
-            requestedAlgoCount = Int(CUDNN_CONVOLUTION_BWD_FILTER_ALGO_COUNT)
-            returnedAlgoCount = Cint[0]
-            perfResults = Array{cudnnConvolutionBwdFilterAlgoPerf_t}(undef,requestedAlgoCount)
-            workspaceSize() = cudnnFindConvolutionAlgorithmWorkspaceSize(x)
-            with_workspace(workspaceSize) do workspace
-                cudnnFindConvolutionBackwardFilterAlgorithmEx(handle(),xDesc,x,dyDesc,dy,convDesc,dwDesc,dw,requestedAlgoCount,returnedAlgoCount,perfResults,workspace,sizeof(workspace))
-            end
-            cudnnConvolutionAlgoPerfChoose(perfResults, returnedAlgoCount[1])
+    key = (xDesc, dyDesc, convDesc)
+    val = lock(cudnnConvolutionBwdFilterAlgoPerfCacheLock) do
+        get(cudnnConvolutionBwdFilterAlgoPerfCache, (xDesc, dyDesc, convDesc), nothing)
+    end
+    if val === nothing
+        requestedAlgoCount = Int(CUDNN_CONVOLUTION_BWD_FILTER_ALGO_COUNT)
+        returnedAlgoCount = Cint[0]
+        perfResults = Array{cudnnConvolutionBwdFilterAlgoPerf_t}(undef,requestedAlgoCount)
+        workspaceSize() = cudnnFindConvolutionAlgorithmWorkspaceSize(x)
+        with_workspace(workspaceSize) do workspace
+            cudnnFindConvolutionBackwardFilterAlgorithmEx(handle(),xDesc,x,dyDesc,dy,convDesc,dwDesc,dw,requestedAlgoCount,returnedAlgoCount,perfResults,workspace,sizeof(workspace))
+        end
+        val = cudnnConvolutionAlgoPerfChoose(perfResults, returnedAlgoCount[1])
+        lock(cudnnConvolutionBwdFilterAlgoPerfCacheLock) do
+            cudnnConvolutionBwdFilterAlgoPerfCache[key] = val
         end
     end
+    val
 end
 
 

lib/cudnn/descriptors.jl

@@ -30,23 +30,27 @@ macro cudnnDescriptor(x, set = Symbol("cudnnSet$(x)Descriptor"))
         const $cache = Dict{Tuple,$sname}()                 # Dict is 3x faster than IdDict!
         const $cache_lock = ReentrantLock()
         function $sname(args...)
-            lock($cache_lock) do
-                get!($cache, args) do
-                    ptr = $tname[C_NULL]
-                    $create(ptr)
-                    $set(ptr[1], args...)
-                    d = $sname(ptr[1])
-                    finalizer(x->$destroy(x.ptr), d)
-                    return d
+            d = lock($cache_lock) do
+                get($cache, args, nothing)
+            end
+            if d === nothing
+                ptr = $tname[C_NULL]
+                $create(ptr)
+                $set(ptr[1], args...)
+                d = $sname(ptr[1])
+                finalizer(x->$destroy(x.ptr), d)
+                lock($cache_lock) do
+                    $cache[args] = d
                 end
             end
+            return d
         end
     end |> esc
 end
 
 
 """
-    cudnnActivationDescriptor(mode::cudnnActivationMode_t, 
+    cudnnActivationDescriptor(mode::cudnnActivationMode_t,
                               reluNanOpt::cudnnNanPropagation_t,
                               coef::Cfloat)
 """
@@ -116,8 +120,8 @@ cudnnConvolutionDescriptor(pad::Vector{Cint},
 
 """
     cudnnLRNDescriptor(lrnN::Cuint,
-                       lrnAlpha::Cdouble, 
-                       lrnBeta::Cdouble, 
+                       lrnAlpha::Cdouble,
+                       lrnBeta::Cdouble,
                        lrnK::Cdouble)
 """
 @cudnnDescriptor(LRN)