| 62 | typename conditional<is_same<T, double>::value, double, float>::type; |
| 63 | |
| 64 | pair<cudnnConvolutionFwdAlgo_t, size_t> getForwardAlgorithm( |
| 65 | cudnnHandle_t cudnn, cudnnTensorDescriptor_t input_descriptor, |
| 66 | cudnnFilterDescriptor_t filter_descriptor, |
| 67 | cudnnConvolutionDescriptor_t convolution_descriptor, |
| 68 | cudnnTensorDescriptor_t output_descriptor) { |
| 69 | cudnnConvolutionFwdAlgo_t convolution_algorithm; |
| 70 | size_t workspace_bytes = 0; |
| 71 | |
| 72 | auto version = getCudnnPlugin().getVersion(); |
| 73 | if (version.major() >= 8) { |
| 74 | int maxAlgoCount = 0; |
| 75 | CUDNN_CHECK(cuda::cudnnGetConvolutionForwardAlgorithmMaxCount( |
| 76 | cudnn, &maxAlgoCount)); |
| 77 | |
| 78 | vector<cudnnConvolutionFwdAlgoPerf_t> perfResults(maxAlgoCount); |
| 79 | int returnAlgoCount = 0; |
| 80 | CUDNN_CHECK(cuda::cudnnFindConvolutionForwardAlgorithm( |
| 81 | cudnn, input_descriptor, filter_descriptor, convolution_descriptor, |
| 82 | output_descriptor, maxAlgoCount, &returnAlgoCount, |
| 83 | perfResults.data())); |
| 84 | |
| 85 | for (int i = 0; i < returnAlgoCount; ++i) { |
| 86 | if (perfResults[i].status == CUDNN_STATUS_SUCCESS) { |
| 87 | convolution_algorithm = perfResults[i].algo; |
| 88 | workspace_bytes = perfResults[i].memory; |
| 89 | break; |
| 90 | } |
| 91 | } |
| 92 | } else { |
| 93 | const int memory_limit = |
| 94 | 0; // TODO: set to remaining space in memory manager? |
| 95 | CUDNN_CHECK(cuda::cudnnGetConvolutionForwardAlgorithm( |
| 96 | cudnn, input_descriptor, filter_descriptor, convolution_descriptor, |
| 97 | output_descriptor, CUDNN_CONVOLUTION_FWD_PREFER_FASTEST, |
| 98 | memory_limit, &convolution_algorithm)); |
| 99 | CUDNN_CHECK(cuda::cudnnGetConvolutionForwardWorkspaceSize( |
| 100 | cudnn, input_descriptor, filter_descriptor, convolution_descriptor, |
| 101 | output_descriptor, convolution_algorithm, &workspace_bytes)); |
| 102 | } |
| 103 | |
| 104 | return {convolution_algorithm, workspace_bytes}; |
| 105 | } |
| 106 | |
| 107 | template<typename T> |
| 108 | Array<T> convolve2_cudnn(const Array<T> &signal, const Array<T> &filter, |
no test coverage detected