MCPcopy Create free account
hub / github.com/arrayfire/arrayfire / convolve2_cudnn

Function convolve2_cudnn

src/backend/cuda/convolveNN.cpp:108–163  ·  view source on GitHub ↗

Source from the content-addressed store, hash-verified

106
107template<typename T>
108Array<T> convolve2_cudnn(const Array<T> &signal, const Array<T> &filter,
109 const dim4 &stride, const dim4 &padding,
110 const dim4 &dilation) {
111 cudnnHandle_t cudnn = nnHandle();
112
113 cudnnDataType_t cudnn_dtype = getCudnnDataType<T>();
114 auto input_descriptor = toCudnn<cudnnTensorDescriptor_t>(signal);
115 auto filter_descriptor = toCudnn<cudnnFilterDescriptor_t>(filter);
116
117 // create convolution descriptor
118 auto convolution_descriptor = make_handle<cudnnConvolutionDescriptor_t>();
119
120 CUDNN_CHECK(cuda::cudnnSetConvolution2dDescriptor(
121 convolution_descriptor, padding[1], padding[0], stride[1], stride[0],
122 dilation[1], dilation[0], CUDNN_CONVOLUTION, cudnn_dtype));
123
124 // get output dimensions
125 const int tensorDims = 4;
126 int convolved_output_dim[tensorDims];
127 CUDNN_CHECK(cuda::cudnnGetConvolutionNdForwardOutputDim(
128 convolution_descriptor, input_descriptor, filter_descriptor, tensorDims,
129 convolved_output_dim));
130
131 // create output descriptor
132 const int n_out = convolved_output_dim[0];
133 const int c_out = convolved_output_dim[1];
134 const int h_out = convolved_output_dim[2];
135 const int w_out = convolved_output_dim[3];
136
137 // prepare output array and scratch space
138 dim4 odims(w_out, h_out, c_out, n_out);
139 Array<T> out = createEmptyArray<T>(odims);
140
141 auto output_descriptor = toCudnn<cudnnTensorDescriptor_t>(out);
142
143 // get convolution algorithm
144 cudnnConvolutionFwdAlgo_t convolution_algorithm;
145 size_t workspace_bytes = 0;
146
147 tie(convolution_algorithm, workspace_bytes) =
148 getForwardAlgorithm(cudnn, input_descriptor, filter_descriptor,
149 convolution_descriptor, output_descriptor);
150
151 auto workspace_buffer = memAlloc<char>(workspace_bytes);
152
153 // perform convolution
154 auto alpha = scalar<scale_type<T>>(1.0);
155 auto beta = scalar<scale_type<T>>(0.0);
156 CUDNN_CHECK(cuda::cudnnConvolutionForward(
157 cudnn, &alpha, input_descriptor, signal.device(), filter_descriptor,
158 filter.device(), convolution_descriptor, convolution_algorithm,
159 (void *)workspace_buffer.get(), workspace_bytes, &beta,
160 output_descriptor, out.device()));
161
162 return out;
163}
164
165template<typename T>

Callers

nothing calls this directly

Calls 7

nnHandleFunction · 0.85
getForwardAlgorithmFunction · 0.85
cudnnConvolutionForwardFunction · 0.85
deviceMethod · 0.45
getMethod · 0.45

Tested by

no test coverage detected