Merge pull request #404 from JuliaGPU/tb/flux

WIP: CUDNN improvements

Author: MikeInnes

.gitlab-ci.yml

@@ -47,7 +47,7 @@ flux:
                 Pkg.develop(PackageSpec(path=pwd()));
                 Pkg.build();'
     - julia -e 'using Pkg;
-                Pkg.add("Flux");
+                Pkg.add(PackageSpec(name="Flux", rev="tb/cuarrays_dnn"));
                 Pkg.test("Flux");'
   allow_failure: true
 

src/CuArrays.jl

@@ -64,6 +64,7 @@ include("broadcast.jl")
 include("matmul.jl")
 include("mapreduce.jl")
 include("accumulate.jl")
+include("linalg.jl")
 
 include("gpuarray_interface.jl")
 

src/array.jl

@@ -239,45 +239,6 @@ function Base.fill!(A::CuArray{T}, x) where T <: MemsetCompatTypes
 end
 
 
-## generic linear algebra routines
-
-function LinearAlgebra.tril!(A::CuMatrix{T}, d::Integer = 0) where T
-  function kernel!(_A, _d)
-    li = (blockIdx().x - 1) * blockDim().x + threadIdx().x
-    m, n = size(_A)
-    if 0 < li <= m*n
-      i, j = Tuple(CartesianIndices(_A)[li])
-      if i < j - _d
-        _A[i, j] = 0
-      end
-    end
-    return nothing
-  end
-
-  blk, thr = cudims(A)
-  @cuda blocks=blk threads=thr kernel!(A, d)
-  return A
-end
-
-function LinearAlgebra.triu!(A::CuMatrix{T}, d::Integer = 0) where T
-  function kernel!(_A, _d)
-    li = (blockIdx().x - 1) * blockDim().x + threadIdx().x
-    m, n = size(_A)
-    if 0 < li <= m*n
-      i, j = Tuple(CartesianIndices(_A)[li])
-      if j < i + _d
-        _A[i, j] = 0
-      end
-    end
-    return nothing
-  end
-
-  blk, thr = cudims(A)
-  @cuda blocks=blk threads=thr kernel!(A, d)
-  return A
-end
-
-
 ## reversing
 
 # the kernel works by treating the array as 1d. after reversing by dimension x an element at

src/blas/CUBLAS.jl

@@ -2,7 +2,7 @@ module CUBLAS
 
 using CUDAapi
 
-import CUDAdrv: CUDAdrv, CuContext, CuStream_t, CuPtr, PtrOrCuPtr, CU_NULL, devices
+using CUDAdrv: CUDAdrv, CuContext, CuStream_t, CuPtr, PtrOrCuPtr, CU_NULL, devices
 
 import CUDAnative
 

src/deprecated.jl

@@ -1,6 +1,6 @@
 # Deprecated functionality
 
-import Base: @deprecate_binding
+using Base: @deprecate_binding
 
 @deprecate_binding BLAS CUBLAS
 @deprecate_binding FFT CUFFT

src/dnn/CUDNN.jl

@@ -1,21 +1,22 @@
 module CUDNN
 
-import CUDAapi
+using CUDAapi
+using CUDAapi: libraryPropertyType
 
-import CUDAdrv: CUDAdrv, CuContext, CuPtr, CU_NULL
+using CUDAdrv
+using CUDAdrv: CuContext, CuPtr, PtrOrCuPtr, CU_NULL, CuStream_t
 
 import CUDAnative
 
+using CEnum
+
 using ..CuArrays
-using ..CuArrays: @libcudnn, active_context, unsafe_free!
-using ..CuArrays: CuVecOrMat, CuVector
+using ..CuArrays: @libcudnn, active_context, CuVecOrMat, CuVector
+import ..CuArrays.unsafe_free!
 
-using NNlib
-import NNlib: conv!, ∇conv_filter!, ∇conv_data!, stride, dilation, flipkernel,
-  maxpool!, meanpool!, ∇maxpool!, ∇meanpool!, spatial_dims, padding, kernel_size,
-  softmax, softmax!, ∇softmax!, logsoftmax, logsoftmax!, ∇logsoftmax
+import NNlib
 
-include("libcudnn_types.jl")
+include("libcudnn_common.jl")
 include("error.jl")
 
 const _handles = Dict{CuContext,cudnnHandle_t}()
@@ -35,13 +36,23 @@ function handle()
     return _handle[]
 end
 
+include("base.jl")
 include("libcudnn.jl")
+
 include("helpers.jl")
+include("tensor.jl")
+include("conv.jl")
+include("pooling.jl")
+include("activation.jl")
+include("filter.jl")
+include("softmax.jl")
+include("batchnorm.jl")
+include("dropout.jl")
+include("rnn.jl")
+
+# interfaces with other software
 include("nnlib.jl")
-include("compat.jl")
 
-version() = VersionNumber(cudnnGetProperty(CUDAapi.MAJOR_VERSION),
-                          cudnnGetProperty(CUDAapi.MINOR_VERSION),
-                          cudnnGetProperty(CUDAapi.PATCH_LEVEL))
+include("compat.jl")
 
 end

src/dnn/activation.jl

@@ -0,0 +1,36 @@
+# descriptor
+
+mutable struct ActivationDesc
+    ptr::cudnnActivationDescriptor_t
+end
+
+unsafe_free!(ad::ActivationDesc)=cudnnDestroyActivationDescriptor(ad.ptr)
+
+Base.unsafe_convert(::Type{cudnnActivationDescriptor_t}, ad::ActivationDesc)=ad.ptr
+
+function ActivationDesc(mode, coeff, reluNanOpt=CUDNN_NOT_PROPAGATE_NAN)
+    ad = Ref{cudnnActivationDescriptor_t}()
+    cudnnCreateActivationDescriptor(ad)
+    cudnnSetActivationDescriptor(ad[],mode,reluNanOpt,coeff)
+    this = ActivationDesc(ad[])
+    finalizer(unsafe_free!, this)
+    return this
+end
+
+
+# wrappers
+
+function cudnnActivationForward(y::CuArray{T,N}, x::CuArray{T,N}; mode=CUDNN_ACTIVATION_RELU, #CUDNN_ACTIVATION_IDENTITY will not work
+                                coeff=0.0, reluNanOpt=CUDNN_NOT_PROPAGATE_NAN, alpha=1, beta=0) where {T,N}
+    ad = ActivationDesc(mode, T(coeff), reluNanOpt)
+    cudnnActivationForward(handle(), ad, Ref(T(alpha)), TensorDesc(x), x, Ref(T(beta)), TensorDesc(y), y)
+    return y
+end
+
+function cudnnActivationBackward(dx::CuArray{T,N}, x::CuArray{T,N}, y::CuArray{T,N}, dy::CuArray{T,N};
+                                 mode=CUDNN_ACTIVATION_RELU, #CUDNN_ACTIVATION_IDENTITY will not work
+                                 coeff=0.0, reluNanOpt=CUDNN_NOT_PROPAGATE_NAN, alpha=1, beta=0) where {T,N}
+    ad = ActivationDesc(mode, T(coeff), reluNanOpt)
+    cudnnActivationBackward(handle(), ad, Ref(T(alpha)), TensorDesc(y), y, TensorDesc(dy), dy, TensorDesc(x), x, Ref(T(beta)), TensorDesc(dx), dx)
+    return dx
+end

src/dnn/base.jl

@@ -0,0 +1,15 @@
+function cudnnCreate()
+    handle = Ref{cudnnHandle_t}()
+    cudnnCreate(handle)
+    return handle[]
+end
+
+function cudnnGetProperty(property::CUDAapi.libraryPropertyType)
+  value_ref = Ref{Cint}()
+  cudnnGetProperty(property, value_ref)
+  value_ref[]
+end
+
+version() = VersionNumber(cudnnGetProperty(CUDAapi.MAJOR_VERSION),
+                          cudnnGetProperty(CUDAapi.MINOR_VERSION),
+                          cudnnGetProperty(CUDAapi.PATCH_LEVEL))

src/dnn/batchnorm.jl

@@ -0,0 +1,117 @@
+mutable struct BNCache
+  mean
+  ivar
+end
+
+BNCache() = BNCache(nothing, nothing)
+
+@inline _wsize(y) = (map(_ -> 1, size(y)[1:end-2])..., size(y)[end-1], 1)
+
+@inline _reddims(y) = (collect(1:ndims(y)-2)..., ndims(y))
+
+# NOTE: CuDNN supports only 4D and 5D Tensors for BatchNorm Operations
+# so reshape a 2D Tensor into 4D
+batchnorm(g::CuArray{T}, b::CuArray{T}, x::CuArray{T, 2},
+          running_mean::CuArray{T}, running_var::CuArray{T}, momentum;
+          cache = nothing, alpha = T(1), beta = T(0),
+          eps = T(1e-5), training = true) where T<:Union{Float32, Float64} =
+  dropdims(batchnorm(g, b, reshape(x, 1, 1, size(x, 1), size(x, 2)), running_mean, running_var, momentum,
+            cache = cache, alpha = alpha, beta = beta, eps = eps, training = training), dims = (1, 2))
+
+function batchnorm(g::CuArray{T}, b::CuArray{T}, x::Union{CuArray{T, 4},CuArray{T,5}},
+                   running_mean::CuArray{T}, running_var::CuArray{T}, momentum;
+                   cache = nothing, alpha = T(1), beta = T(0),
+                   eps = T(1e-5), training = true) where T<:Union{Float32, Float64}
+  y = similar(x)
+  cudnnBNForward!(y, g, b, x, running_mean, running_var, momentum, cache = cache,
+      alpha = alpha, beta = beta, eps = eps, training = training)
+  y
+end
+
+function cudnnBNForward!(y::CuArray{T}, g::CuArray{T}, b::CuArray{T}, x::CuArray{T},
+                        running_mean::CuArray{T}, running_var::CuArray{T},
+                        momentum; cache = nothing,
+                        alpha = T(1), beta = T(0),
+                        eps = T(1e-5), training = true) where T<:Union{Float32, Float64}
+  dims = _wsize(x)
+  if eps < CUDNN_BN_MIN_EPSILON
+    # warn("eps ",eps," is too small for CuDNN so eps has been assigned the value ", CUDNN_BN_MIN_EPSILON)
+    eps = CUDNN_BN_MIN_EPSILON
+  end
+  xd = TensorDesc(x)
+  yd = TensorDesc(y)
+  gd = TensorDesc(T, dims)
+
+  if training
+
+    if cache !== nothing
+      mean = zeros(CuArray{T}, dims...)
+      ivar = ones(CuArray{T}, dims...)
+    else
+      mean = CU_NULL
+      ivar = CU_NULL
+    end
+
+    cudnnBatchNormalizationForwardTraining(handle(), CUDNN_BATCHNORM_SPATIAL, Ref(T(alpha)), Ref(T(beta)), xd, x, yd, y, gd, g, b, momentum, running_mean, running_var, eps, mean, ivar)
+
+    if cache !== nothing
+      cache.mean = mean
+      cache.ivar = ivar
+    end
+  else
+    cudnnBatchNormalizationForwardInference(handle(), CUDNN_BATCHNORM_SPATIAL, Ref(T(alpha)), Ref(T(beta)), xd, x, yd, y, gd, g, b, running_mean, running_var, eps)
+  end
+end
+
+function ∇batchnorm(g::CuArray{T}, b::CuArray{T}, x::CuArray{T, 2}, dy::CuArray{T, 2},
+           running_mean::CuArray{T}, running_var::CuArray{T}, momentum;
+           cache = nothing, eps = T(1e-5), alpha = T(1),
+           beta = T(0), training = true) where T<:Union{Float32, Float64}
+  dg, db, dx = ∇batchnorm(g, b, reshape(x, 1, 1, size(x, 1), size(x, 2)), reshape(dy, 1, 1, size(dy, 1),
+                          size(dy, 2)), running_mean, running_var, momentum, cache = cache, eps = eps,
+                          alpha = alpha, beta = beta, training = training)
+  (dg, db, dropdims(dx, dims = (1, 2)))
+end
+
+function ∇batchnorm(g::CuArray{T}, b::CuArray{T}, x::CuArray{T}, dy::CuArray{T},
+                    running_mean::CuArray{T}, running_var::CuArray{T}, momentum;
+                    cache = nothing, eps = T(1e-5), alpha = T(1),
+                    beta = T(0), training = true) where T<:Union{Float32, Float64}
+  dg = similar(g)
+  db = similar(b)
+  dx = similar(x)
+  cudnnBNBackward!(dg, g, db, dx, x, dy, running_mean, running_var, T(momentum),
+    training = training, cache = cache, eps = eps, alpha = alpha, beta = beta)
+  (dg, db, dx)
+end
+
+function cudnnBNBackward!(dg::CuArray{T}, g::CuArray{T}, db::CuArray{T},
+                          dx::CuArray{T}, x::CuArray{T}, dy::CuArray{T},
+                          running_mean::CuArray{T}, running_var::CuArray{T},
+                          momentum; cache = nothing, eps = T(1e-5),
+                          alpha = T(1), beta = T(0),
+                          dalpha = T(1), dbeta = T(0), training = true) where T<:Union{Float32, Float64}
+  if training
+    xd = TensorDesc(x)
+    dyd = TensorDesc(dy)
+    dxd = TensorDesc(dx)
+    gd = TensorDesc(T, _wsize(x))
+    if cache !== nothing
+      mean, ivar = cache.mean, cache.ivar
+      info("mean and ivar are fetched from the cache")
+    else
+      mean, ivar = CU_NULL, CU_NULL
+    end
+
+    if eps < CUDNN_BN_MIN_EPSILON
+      eps = CUDNN_BN_MIN_EPSILON
+    end
+
+    cudnnBatchNormalizationBackward(handle(), CUDNN_BATCHNORM_SPATIAL, Ref(T(alpha)), Ref(T(beta)), Ref(T(dalpha)), Ref(T(dbeta)), xd, x, dyd, dy, dxd, dx, gd, g, dg, db, eps, mean, ivar)
+  else
+    ivar = 1 ./ sqrt.(reshape(running_var, _wsize(x)) .+ eps)
+    dx .= dy .* reshape(g, _wsize(x)) .* ivar
+    dg .= squeeze(sum(dy .* (x .- reshape(running_mean, _wsize(x))) .* ivar, _reddims(dy)), dims = (1,2,4))
+    db .= squeeze(sum(dy, _reddims(dy)), dims = (1,2,4))
+  end
+end