| 3712 | } // namespace |
| 3713 | |
| 3714 | port::Status CudnnSupport::DoPrepareForConvolution( |
| 3715 | dnn::ConvolutionKind kind, dnn::DataType element_type, Stream* stream, |
| 3716 | const dnn::BatchDescriptor& input_descriptor, DeviceMemoryBase input_data, |
| 3717 | const dnn::FilterDescriptor& filter_descriptor, |
| 3718 | DeviceMemoryBase filter_data, const dnn::BatchDescriptor& output_descriptor, |
| 3719 | DeviceMemoryBase output_data, |
| 3720 | const dnn::ConvolutionDescriptor& convolution_descriptor, |
| 3721 | const dnn::AlgorithmConfig& algorithm_config, |
| 3722 | ScratchAllocator* scratch_allocator, dnn::AlgorithmDesc* algorithm_desc, |
| 3723 | DeviceMemory<uint8>* scratch_memory) { |
| 3724 | CudnnTensorDescriptor input_nd( |
| 3725 | input_descriptor, |
| 3726 | ToCudnnDataType(element_type, input_descriptor.layout())); |
| 3727 | CudnnFilterDescriptor filter_nd( |
| 3728 | filter_descriptor, |
| 3729 | ToCudnnDataType(element_type, filter_descriptor.layout())); |
| 3730 | CudnnTensorDescriptor output_nd( |
| 3731 | output_descriptor, |
| 3732 | ToCudnnDataType(element_type, output_descriptor.layout())); |
| 3733 | CudnnConvolutionDescriptor conv( |
| 3734 | convolution_descriptor, |
| 3735 | ToCudnnDataType(GetConvAccumulatorType(element_type))); |
| 3736 | |
| 3737 | auto cudnn = cudnn_->GetHandle(parent_, stream); |
| 3738 | |
| 3739 | switch (kind) { |
| 3740 | case dnn::ConvolutionKind::FORWARD: { |
| 3741 | SE_ASSIGN_OR_RETURN( |
| 3742 | *algorithm_desc, |
| 3743 | GetCudnnConvolutionForwardAlgorithm( |
| 3744 | stream, cudnn, algorithm_config, input_nd, filter_nd, conv, |
| 3745 | output_nd, scratch_allocator, scratch_memory)); |
| 3746 | break; |
| 3747 | } |
| 3748 | case dnn::ConvolutionKind::BACKWARD_DATA: { |
| 3749 | SE_ASSIGN_OR_RETURN( |
| 3750 | *algorithm_desc, |
| 3751 | GetCudnnConvolutionBackwardDataAlgorithm( |
| 3752 | stream, cudnn, algorithm_config, input_nd, filter_nd, conv, |
| 3753 | output_nd, scratch_allocator, scratch_memory)); |
| 3754 | break; |
| 3755 | } |
| 3756 | case dnn::ConvolutionKind::BACKWARD_FILTER: { |
| 3757 | SE_ASSIGN_OR_RETURN( |
| 3758 | *algorithm_desc, |
| 3759 | GetCudnnConvolutionBackwardFilterAlgorithm( |
| 3760 | stream, cudnn, algorithm_config, input_nd, filter_nd, conv, |
| 3761 | output_nd, scratch_allocator, scratch_memory)); |
| 3762 | break; |
| 3763 | } |
| 3764 | default: |
| 3765 | return port::InternalError( |
| 3766 | absl::StrCat("Unexpected convolution kind ", static_cast<int>(kind))); |
| 3767 | } |
| 3768 | |
| 3769 | return port::Status::OK(); |
| 3770 | } |
| 3771 |
nothing calls this directly
no test coverage detected