Skip to content

Commit

Permalink
Add cudnnGetConvolutionDescriptor function
Browse files Browse the repository at this point in the history
There is already `cudnnGetTensorDescriptor` and
`cudnnGetFilterDescriptor`, so now we have everything to cache algorithm performances.
  • Loading branch information
RomeoV committed Jun 12, 2023
1 parent 2dd4f28 commit 2fd8986
Showing 1 changed file with 22 additions and 0 deletions.
22 changes: 22 additions & 0 deletions lib/cudnn/src/convolution.jl
Original file line number Diff line number Diff line change
Expand Up @@ -120,6 +120,28 @@ function cudnnConvolutionForwardAD(w, x, bias, z; y, activation, convDesc, wDesc
return y
end

function cudnnGetConvolutionDescriptor(d::cudnnConvolutionDescriptor)
# we don't know the dimension of the convolution, so we start by
# allocating the maximum size it can be.
nbDimsRequested = CUDNN_DIM_MAX - 2
# later, here we get the actual dimensionality of the convolution
arrlen = Ref{Cint}(nbDimsRequested)
padding = Array{Cint}(undef, nbDimsRequested)
stride = Array{Cint}(undef, nbDimsRequested)
dilation = Array{Cint}(undef, nbDimsRequested)
mode = Ref{cuDNN.cudnnConvolutionMode_t}(CUDNN_CONVOLUTION)
dataType = Ref{cuDNN.cudnnDataType_t}(cuDNN.CUDNN_DATA_FLOAT)

cudnnGetConvolutionNdDescriptor(d, nbDimsRequested, arrlen, padding, stride, dilation,
mode, dataType)
T = juliaDataType(dataType[])
SZ = arrlen[]
P = (padding[1:SZ]..., )
S = (stride[1:SZ]..., )
D = (dilation[1:SZ]..., )
return T, mode[], SZ, P, S, D
end

# Helper for cudnnConvolutionDescriptor
function cudnnSetConvolutionDescriptor(
ptr::cudnnConvolutionDescriptor_t,
Expand Down

0 comments on commit 2fd8986

Please sign in to comment.