From 2fd8986d59a7732e08ecd2e09031de4c25e9a817 Mon Sep 17 00:00:00 2001 From: Romeo Valentin Date: Mon, 12 Jun 2023 02:49:01 -0700 Subject: [PATCH] Add `cudnnGetConvolutionDescriptor` function There is already `cudnnGetTensorDescriptor` and `cudnnGetFilterDescriptor`, so now we have everything to cache algorithm performances. --- lib/cudnn/src/convolution.jl | 22 ++++++++++++++++++++++ 1 file changed, 22 insertions(+) diff --git a/lib/cudnn/src/convolution.jl b/lib/cudnn/src/convolution.jl index 724dac233c..ec581f6511 100644 --- a/lib/cudnn/src/convolution.jl +++ b/lib/cudnn/src/convolution.jl @@ -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,