Skip to content
This repository has been archived by the owner on Mar 12, 2021. It is now read-only.

Commit

Permalink
Merge #315
Browse files Browse the repository at this point in the history 
315: Dup of nnlib_overhaul [DO NOT MERGE] r=maleadt a=kshyatt

opened a new PR to test the rebased on top of master version without corrupting Elliot's PR.

Co-authored-by: Elliot Saba <[email protected]>
Co-authored-by: Katharine Hyatt <[email protected]>
Co-authored-by: Tim Besard <[email protected]>
  • Loading branch information
@bors @staticfloat
@kshyatt @maleadt
4 people committed Apr 5, 2019
2 parents 86e1ea6 + 1181599 commit 291cc82
Show file tree
Hide file tree
Showing 7 changed files with 175 additions and 119 deletions.
2 changes: 1 addition & 1 deletion Project.toml
View file
Original file line number Diff line number Diff line change
Expand Up @@ -31,6 +31,6 @@ julia = "1.0"
CUDAnative = "2.0"
CUDAdrv = "2.0"
CUDAapi = "0.5.3, 0.6"
NNlib = "0.5"
NNlib = "0.5, 0.6"
GPUArrays = "0.5, 0.6"
Adapt = "0.4"
8 changes: 7 additions & 1 deletion src/dnn/CUDNN.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,12 @@ import CUDAdrv: CUDAdrv, CuContext, CuPtr, CU_NULL

using ..CuArrays
using ..CuArrays: libcudnn, active_context, configured, unsafe_free!

using ..CuArrays: CuVecOrMat, CuVector
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
using CUDAnative
include("libcudnn_types.jl")
include("error.jl")

Expand All @@ -30,6 +35,7 @@ end
include("libcudnn.jl")
include("helpers.jl")
include("nnlib.jl")
include("compat.jl")

version() = VersionNumber(cudnnGetProperty(CUDAapi.MAJOR_VERSION),
cudnnGetProperty(CUDAapi.MINOR_VERSION),
Expand Down
21 changes: 21 additions & 0 deletions src/dnn/compat.jl
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,21 @@
# Compatibility shims until users upgrade to new NNlib format
function conv!(y::CuArray{T}, x::CuArray{T}, w::CuArray{T}; pad=0, stride=1, flipkernel=0, dilation=1, kwargs...) where {T<:CUDNNFloat}
cdims = DenseConvDims(x, w; padding=pad, stride=stride, flipkernel=flipkernel, dilation=dilation)
return conv!(y, x, w, cdims; kwargs...)
end

function ∇conv_filter!(dw::CuArray{T}, dy::CuArray{T}, x::CuArray{T}; pad=0, stride=1, flipkernel=0, dilation=1, kwargs...) where {T<:CUDNNFloat}
cdims = DenseConvDims(x, dw; padding=pad, stride=stride, flipkernel=flipkernel, dilation=dilation)
# NOTE!!! This compat shim re-arranges the argument order!
return ∇conv_filter!(dw, x, dy, cdims; kwargs...)
end

function maxpool!(y::CuArray{T}, x::CuArray{T}, k; pad=map(_->0,k), stride=k) where {T<:CUDNNFloat}
pdims = PoolDims(x, k; padding=pad, stride=stride)
return maxpool!(y, x, pdims)
end

function meanpool!(y::CuArray{T}, x::CuArray{T}, k; pad=map(_->0,k), stride=k) where {T<:CUDNNFloat}
pdims = PoolDims(x, k; padding=pad, stride=stride)
return meanpool!(y, x, pdims)
end
32 changes: 28 additions & 4 deletions src/dnn/helpers.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -83,13 +83,19 @@ free(cd::ConvDesc) = cudnnDestroyConvolutionDescriptor(cd.ptr)
Base.unsafe_convert(::Type{cudnnConvolutionDescriptor_t}, cd::ConvDesc)=cd.ptr

function cdsize(w, nd)
isa(w, Integer) ? Cint[fill(w,nd)...] :
length(w)!=nd ? error("Dimension mismatch") :
Cint[reverse(w)...]
isa(w, Integer) && return Cint[fill(w,nd)...]
length(w) == nd && return Cint[reverse(w)...]
length(w) == 2*nd && return Cint[reverse(w[nd+1:end])...]
throw(DimensionMismatch())
end

pdsize(w, nd)=Cint[reverse(psize(w,nd))...]
psize(w, nd)=(isa(w,Integer) ? fill(w,nd) : length(w) != nd ? error("Dimension mismatch") : w)
function psize(w, nd)
isa(w, Integer) && return Cint[fill(w,nd)...]
length(w) == nd && return w
length(w) == 2*nd && return w[1:nd]
throw(DimensionMismatch())
end

function ConvDesc(T, N, padding, stride, dilation, mode)
cd = Ref{cudnnConvolutionDescriptor_t}()
Expand All @@ -102,6 +108,15 @@ function ConvDesc(T, N, padding, stride, dilation, mode)
return this
end

function ConvDesc(T, cdims::DenseConvDims)
pd = NNlib.padding(cdims)
if !all(pd[1:2:end] .== pd[2:2:end])
@warn("CuDNN does not support asymmetric padding; defaulting to symmetric choice")
end
return ConvDesc(T, NNlib.spatial_dims(cdims), pd[1:2:end], NNlib.stride(cdims),
NNlib.dilation(cdims), NNlib.flipkernel(cdims))
end

mutable struct PoolDesc; ptr; end
free(pd::PoolDesc)=cudnnDestroyPoolingDescriptor(pd.ptr)
Base.unsafe_convert(::Type{cudnnPoolingDescriptor_t}, pd::PoolDesc)=pd.ptr
Expand All @@ -115,6 +130,15 @@ function PoolDesc(nd, window, padding, stride, mode, maxpoolingNanOpt=CUDNN_NOT_
return this
end

function PoolDesc(pdims::PoolDims, mode, maxpoolingNanOpt=CUDNN_NOT_PROPAGATE_NAN)
pd = NNlib.padding(pdims)
if !all(pd[1:2:end] .== pd[2:2:end])
@warn("CuDNN does not support asymmetric padding; defaulting to symmetric choice")
end
return PoolDesc(NNlib.spatial_dims(pdims), NNlib.kernel_size(pdims), pd[1:2:end],
NNlib.stride(pdims), mode, maxpoolingNanOpt)
end

mutable struct ActivationDesc; ptr; end
free(ad::ActivationDesc)=cudnnDestroyActivationDescriptor(ad.ptr)
Base.unsafe_convert(::Type{cudnnActivationDescriptor_t}, ad::ActivationDesc)=ad.ptr
Expand Down
88 changes: 45 additions & 43 deletions src/dnn/libcudnn.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -262,13 +262,13 @@ function cudnnConvolutionForward(alpha, xDesc, x, wDesc, w, convDesc, algo, work
workspace_size, beta, yDesc, y)
end

function cudnnConvolutionForward(y::CuArray{T,N}, x::CuArray{T,N}, w::CuArray{T,N};
algo=0, workspace=CU_NULL, workspace_size=0,
alpha=1, beta=0, padding=0, stride=1, dilation=1, mode=0) where {T,N}
cd = ConvDesc(T, N-2, padding, stride, dilation, mode)
function cudnnConvolutionForward(y::CuArray{T,N}, x::CuArray{T,N}, w::CuArray{T,N},
cdims::DenseConvDims; algo=0, workspace=CU_NULL,
workspace_size=0, alpha=1, beta=0) where {T,N}
cudnnConvolutionForward(
Ref(T(alpha)),TensorDesc(x),x,FilterDesc(w),w,cd,algo,workspace,
workspace_size,Ref(T(beta)),TensorDesc(y),y)
Ref(T(alpha)), TensorDesc(x), x, FilterDesc(w), w, ConvDesc(T,cdims),
algo, workspace, workspace_size, Ref(T(beta)), TensorDesc(y), y
)
return y
end

Expand All @@ -281,12 +281,13 @@ function cudnnGetConvolutionForwardWorkspaceSize(xDesc, wDesc, convDesc, yDesc,
handle(), xDesc, wDesc, convDesc, yDesc, algo, workspace_size)
end

function cudnnGetConvolutionForwardWorkspaceSize(y::CuArray{T,N}, x::CuArray{T,N}, w::CuArray{T,N};
algo=0, padding=0, stride=1,
dilation=1, mode=0) where {T,N}
cd = ConvDesc(T, N-2, padding, stride, dilation, mode)
function cudnnGetConvolutionForwardWorkspaceSize(y::CuArray{T,N}, x::CuArray{T,N}, w::CuArray{T,N},
cdims::DenseConvDims; algo=0) where {T,N}
workspace_size = Ref{Cint}()
cudnnGetConvolutionForwardWorkspaceSize(TensorDesc(x), FilterDesc(w), cd, TensorDesc(y), algo, workspace_size)
cudnnGetConvolutionForwardWorkspaceSize(
TensorDesc(x), FilterDesc(w), ConvDesc(T, cdims),
TensorDesc(y), algo, workspace_size
)
return Int(workspace_size[])
end

Expand All @@ -301,13 +302,13 @@ function cudnnConvolutionBackwardData(alpha, wDesc, w, dyDesc, dy, convDesc, alg
workspace_size, beta, dxDesc, dx)
end

function cudnnConvolutionBackwardData(dx::CuArray{T,N}, w::CuArray{T,N}, dy::CuArray{T,N};
algo=0, workspace=CU_NULL, workspace_size=0,
alpha=1, beta=0, padding=0, stride=1, dilation=1, mode=0) where {T,N}
cd = ConvDesc(T, N-2, padding, stride, dilation, mode)
function cudnnConvolutionBackwardData(dx::CuArray{T,N}, w::CuArray{T,N}, dy::CuArray{T,N},
cdims::DenseConvDims; algo=0, workspace=CU_NULL,
workspace_size=0, alpha=1, beta=0) where {T,N}
cudnnConvolutionBackwardData(
Ref(T(alpha)),FilterDesc(w),w,TensorDesc(dy),dy,cd,algo,workspace,
workspace_size,Ref(T(beta)),TensorDesc(dx),dx)
Ref(T(alpha)), FilterDesc(w), w, TensorDesc(dy), dy, ConvDesc(T, cdims),
algo, workspace, workspace_size, Ref(T(beta)), TensorDesc(dx), dx
)
return dx
end

Expand All @@ -320,12 +321,13 @@ function cudnnGetConvolutionBackwardDataWorkspaceSize(wDesc, dyDesc, convDesc, d
handle(), wDesc, dyDesc, convDesc, dxDesc, algo, workspace_size)
end

function cudnnGetConvolutionBackwardDataWorkspaceSize(dx::CuArray{T,N}, w::CuArray{T,N}, dy::CuArray{T,N};
algo=0, padding=0, stride=1,
dilation=1, mode=0) where {T,N}
cd = ConvDesc(T, N-2, padding, stride, dilation, mode)
function cudnnGetConvolutionBackwardDataWorkspaceSize(dx::CuArray{T,N}, w::CuArray{T,N}, dy::CuArray{T,N},
cdims::DenseConvDims; algo=0) where {T,N}
workspace_size = Ref{Cint}()
cudnnGetConvolutionBackwardDataWorkspaceSize(FilterDesc(w), TensorDesc(dy), cd, TensorDesc(dx), algo, workspace_size)
cudnnGetConvolutionBackwardDataWorkspaceSize(
FilterDesc(w), TensorDesc(dy), ConvDesc(T, cdims),
TensorDesc(dx), algo, workspace_size
)
return Int(workspace_size[])
end

Expand All @@ -340,13 +342,13 @@ function cudnnConvolutionBackwardFilter(alpha, xDesc, x, dyDesc, dy, convDesc, a
workspace_size, beta, dwDesc, dw)
end

function cudnnConvolutionBackwardFilter(dw::CuArray{T,N}, x::CuArray{T,N}, dy::CuArray{T,N};
algo=0, workspace=CU_NULL, workspace_size=0,
alpha=1, beta=0, padding=0, stride=1, dilation=1, mode=0) where {T,N}
cd = ConvDesc(T, N-2, padding, stride, dilation, mode)
function cudnnConvolutionBackwardFilter(dw::CuArray{T,N}, x::CuArray{T,N}, dy::CuArray{T,N},
cdims::DenseConvDims; algo=0, workspace=CU_NULL,
workspace_size=0, alpha=1, beta=0) where {T,N}
cudnnConvolutionBackwardFilter(
Ref(T(alpha)),TensorDesc(x),x,TensorDesc(dy),dy,cd,algo,workspace,
workspace_size,Ref(T(beta)),FilterDesc(dw),dw)
Ref(T(alpha)), TensorDesc(x), x, TensorDesc(dy), dy, ConvDesc(T, cdims),
algo, workspace, workspace_size, Ref(T(beta)), FilterDesc(dw), dw
)
return dw
end

Expand All @@ -359,12 +361,13 @@ function cudnnGetConvolutionBackwardFilterWorkspaceSize(xDesc, dyDesc, convDesc,
handle(), xDesc, dyDesc, convDesc, dwDesc, algo, workspace_size)
end

function cudnnGetConvolutionBackwardFilterWorkspaceSize(dw::CuArray{T,N}, x::CuArray{T,N}, dy::CuArray{T,N};
algo=0, padding=0, stride=1,
dilation=1, mode=0) where {T,N}
cd = ConvDesc(T, N-2, padding, stride, dilation, mode)
function cudnnGetConvolutionBackwardFilterWorkspaceSize(dw::CuArray{T,N}, x::CuArray{T,N}, dy::CuArray{T,N},
cdims::DenseConvDims; algo=0) where {T,N}
workspace_size = Ref{Cint}()
cudnnGetConvolutionBackwardFilterWorkspaceSize(TensorDesc(x), TensorDesc(dy), cd, FilterDesc(dw), algo, workspace_size)
cudnnGetConvolutionBackwardFilterWorkspaceSize(
TensorDesc(x), TensorDesc(dy), ConvDesc(T, cdims),
FilterDesc(dw), algo, workspace_size
)
return Int(workspace_size[])
end

Expand Down Expand Up @@ -398,22 +401,21 @@ function cudnnPoolingBackward(poolingDesc,alpha,yDesc,y,dyDesc,dy,xDesc,x,beta,d
handle(),poolingDesc,alpha,yDesc,y,dyDesc,dy,xDesc,x,beta,dxDesc,dx)
end

function cudnnPoolingForward(y::CuArray{T,N}, x::CuArray{T,N}; alpha=1,
window=2, padding=0, stride=window, mode=0) where {T,N}
function cudnnPoolingForward(y::CuArray{T,N}, x::CuArray{T,N}, pdims::PoolDims;
alpha=1, mode=0) where {T,N}
beta = 0
pd = PoolDesc(N-2, window, padding, stride, mode)
cudnnPoolingForward(pd, Ref(T(alpha)), TensorDesc(x), x, Ref(T(beta)), TensorDesc(y), y)
cudnnPoolingForward(PoolDesc(pdims, mode), Ref(T(alpha)), TensorDesc(x), x, Ref(T(beta)), TensorDesc(y), y)
return y
end

function cudnnPoolingBackward(dx::CuArray{T,N}, dy::CuArray{T,N}, x::CuArray{T,N}, y::CuArray{T,N};
alpha=1,
window=2, padding=0, stride=window, mode=0) where {T,N}
function cudnnPoolingBackward(dx::CuArray{T,N}, dy::CuArray{T,N}, x::CuArray{T,N}, y::CuArray{T,N},
pdims::PoolDims; alpha=1, mode=0) where {T,N}
if alpha!=1 && mode==0; error("Gradient of pool(alpha!=1,mode=0) broken in CUDNN"); end
beta = 0
pd = PoolDesc(N-2, window, padding, stride, mode)
cudnnPoolingBackward(pd, Ref(T(alpha)), TensorDesc(y), y,
TensorDesc(dy), dy, TensorDesc(x), x, Ref(T(beta)), TensorDesc(dx), dx)
cudnnPoolingBackward(
PoolDesc(pdims, mode), Ref(T(alpha)), TensorDesc(y), y,
TensorDesc(dy), dy, TensorDesc(x), x, Ref(T(beta)), TensorDesc(dx), dx
)
return dx
end

Expand Down
64 changes: 28 additions & 36 deletions src/dnn/nnlib.jl
View file
Original file line number Diff line number Diff line change
@@ -1,4 +1,3 @@
using NNlib
import NNlib: conv!, ∇conv_filter!, ∇conv_data!,
maxpool!, meanpool!, ∇maxpool!, ∇meanpool!,
softmax, softmax!, ∇softmax!, logsoftmax, logsoftmax!, ∇logsoftmax
Expand Down Expand Up @@ -40,68 +39,61 @@ end

# Convolution

function conv!(y::CuArray{T}, x::CuArray{T}, w::CuArray{T};
pad=0, stride=1, flipkernel=0, alpha=1, dilation=1, algo=0) where T<:CUDNNFloat
function conv!(y::CuArray{T}, x::CuArray{T}, w::CuArray{T}, cdims::DenseConvDims;
alpha=1, algo=0) where T<:CUDNNFloat
if version() < v"6"
all(x -> x == 1, dilation) || error("Only dilation = 1 is supported in cuDNN version < 6")
all(x -> x == 1, dilation(cdims)) || error("Only dilation = 1 is supported in cuDNN version < 6")
end

workspace_size =
cudnnGetConvolutionForwardWorkspaceSize(y, x, w, padding=pad, stride=stride, dilation=dilation,
algo=algo, mode=flipkernel)
workspace_size = cudnnGetConvolutionForwardWorkspaceSize(y, x, w, cdims, algo=algo)

CuVector{UInt8}(undef, workspace_size) do workspace
cudnnConvolutionForward(y, x, w, padding=pad, stride=stride, dilation=dilation, mode=flipkernel,
alpha=alpha, algo=algo, workspace=workspace, workspace_size=workspace_size)
cudnnConvolutionForward(y, x, w, cdims, alpha=alpha, algo=algo,
workspace=workspace, workspace_size=workspace_size)
end
end

function ∇conv_filter!(dw::CuArray{T}, dy::CuArray{T}, x::CuArray{T};
pad=0, stride=1, flipkernel=0, alpha=1, dilation=1, algo=0) where T<:CUDNNFloat
function ∇conv_filter!(dw::CuArray{T}, x::CuArray{T}, dy::CuArray{T},
cdims::DenseConvDims; alpha=1, algo=0) where T<:CUDNNFloat
if version() < v"6"
all(x -> x == 1, dilation) || error("Only dilation = 1 is supported in cuDNN version < 6")
all(x -> x == 1, dilation(cdims)) || error("Only dilation = 1 is supported in cuDNN version < 6")
end

workspace_size =
cudnnGetConvolutionBackwardFilterWorkspaceSize(dw, x, dy, padding=pad, stride=stride,
dilation=dilation, algo=algo, mode=flipkernel)
workspace_size = cudnnGetConvolutionBackwardFilterWorkspaceSize(dw, x, dy, cdims, algo=algo)

CuVector{UInt8}(undef, workspace_size) do workspace
cudnnConvolutionBackwardFilter(dw, x, dy, padding=pad, stride=stride, dilation=dilation,
mode=flipkernel, alpha=alpha, algo=algo, workspace=workspace,
workspace_size=workspace_size)
cudnnConvolutionBackwardFilter(dw, x, dy, cdims, alpha=alpha, algo=algo,
workspace=workspace, workspace_size=workspace_size)
end
end

function ∇conv_data!(dx::CuArray{T}, dy::CuArray{T}, w::CuArray{T};
pad=0, stride=1, flipkernel=0, alpha=1, dilation=1, algo=0) where T<:CUDNNFloat
function ∇conv_data!(dx::CuArray{T}, dy::CuArray{T}, w::CuArray{T},
cdims::DenseConvDims; alpha=1, algo=0) where T<:CUDNNFloat
if version() < v"6"
all(x -> x == 1, dilation) || error("Only dilation = 1 is supported in cuDNN version < 6")
all(x -> x == 1, dilation(cdims)) || error("Only dilation = 1 is supported in cuDNN version < 6")
end

workspace_size =
cudnnGetConvolutionBackwardDataWorkspaceSize(dx, w, dy, padding=pad, stride=stride,
dilation=dilation, algo=algo, mode=flipkernel)
cudnnGetConvolutionBackwardDataWorkspaceSize(dx, w, dy, cdims; algo=algo)
CuVector{UInt8}(undef, workspace_size) do workspace
cudnnConvolutionBackwardData(dx, w, dy, padding=pad, stride=stride, dilation=dilation,
mode=flipkernel, alpha=alpha, algo=algo, workspace=workspace,
workspace_size=workspace_size)
cudnnConvolutionBackwardData(dx, w, dy, cdims, alpha=alpha, algo=algo,
workspace=workspace, workspace_size=workspace_size)
end
end

∇conv_bias!(db::CuArray{T}, dy::CuArray{T}; alpha=1, beta=0) where T<:CUDNNFloat =
cudnnConvolutionBackwardBias(db, dy, alpha=alpha, beta=beta)

maxpool!(y::CuArray{T}, x::CuArray{T}, k; pad=map(_->0,k), stride=k) where T<:CUDNNFloat =
cudnnPoolingForward(y, x, window=k, padding=pad, stride=stride, mode=0)
maxpool!(y::CuArray{T}, x::CuArray{T}, pdims::PoolDims) where T<:CUDNNFloat =
cudnnPoolingForward(y, x, pdims; mode=0)

∇maxpool!(dx::CuArray{T}, dy::CuArray{T}, y::CuArray{T}, x::CuArray{T}, k;
pad=map(_->0,k), stride=k) where T<:CUDNNFloat =
cudnnPoolingBackward(dx, dy, x, y, window=k, padding=pad, stride=stride, mode=0)
∇maxpool!(dx::CuArray{T}, dy::CuArray{T}, y::CuArray{T}, x::CuArray{T},
pdims::PoolDims) where T<:CUDNNFloat =
cudnnPoolingBackward(dx, dy, x, y, pdims, mode=0)

meanpool!(y::CuArray{T}, x::CuArray{T}, k; pad=map(_->0,k), stride=k) where T<:CUDNNFloat =
cudnnPoolingForward(y, x, window=k, padding=pad, stride=stride, mode=1)
meanpool!(y::CuArray{T}, x::CuArray{T}, pdims::PoolDims) where T<:CUDNNFloat =
cudnnPoolingForward(y, x, pdims, mode=1)

∇meanpool!(dx::CuArray{T}, dy::CuArray{T}, y::CuArray{T}, x::CuArray{T}, k;
pad=map(_->0,k), stride=k) where T<:CUDNNFloat =
cudnnPoolingBackward(dx, dy, x, y, window=k, padding=pad, stride=stride, mode=1)
∇meanpool!(dx::CuArray{T}, dy::CuArray{T}, y::CuArray{T}, x::CuArray{T},
pdims::PoolDims) where T<:CUDNNFloat =
cudnnPoolingBackward(dx, dy, x, y, pdims, mode=1)
Loading

0 comments on commit 291cc82

Please sign in to comment.