| 106 | |
| 107 | template<typename T> |
| 108 | Array<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 | |
| 165 | template<typename T> |
nothing calls this directly
no test coverage detected