conditional choice of fast modes

JuliaGPU · Jul 27, 2020 · a73a069 · a73a069
1 parent 009143c
commit a73a069
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 3 deletions.
diff --git a/lib/cudnn/conv.jl b/lib/cudnn/conv.jl
@@ -38,8 +38,9 @@ function ConvDesc(T, N, padding, stride, dilation, mode)
     else
         cudnnSetConvolutionNdDescriptor(cd[],N,cdsize(padding,N),cdsize(stride,N),cdsize(dilation,N),mode)
     end
-    if version() >= v"7"
-        cudnnSetConvolutionMathType(cd[], cudnnMathType_t(1))
+    # enable tensor math mode if our device supports it, and fast math is enabled
+    if Base.JLOptions().fast_math == 1 && capability(device()) >= v"7.0" && version() >= v"9"
+        cudnnSetConvolutionMathType(cd[], CUDNN_TENSOR_OP_MATH)
     end
     this = ConvDesc(cd[])
     finalizer(unsafe_free!, this)

diff --git a/lib/cudnn/nnlib.jl b/lib/cudnn/nnlib.jl
@@ -20,7 +20,13 @@ function softmax(xs::CuVecOrMat{T}; dims=1) where T<:CUDNNFloat
 end
 
 function softmax!(out::CuVecOrMat{T}, xs::CuVecOrMat{T}; dims=1) where T<:CUDNNFloat
-  cudnnSoftmaxForward(reshape4D(xs), reshape4D(out), algorithm=CUDNN_SOFTMAX_FAST, mode=cudnnSoftmaxMode_t(dims-1))
+  # use fast over accurate algorithm if fast math is enabled
+  if Base.JLOptions().fast_math == 1
+    algorithm = CUDNN_SOFTMAX_FAST
+  else
+    algorithm = CUDNN_SOFTMAX_ACCURATE
+  end
+  cudnnSoftmaxForward(reshape4D(xs), reshape4D(out), algorithm=algorithm, mode=cudnnSoftmaxMode_t(dims-1))
   return out
 end