| 2649 | } |
| 2650 | |
| 2651 | port::Status MIOpenSupport::DoPrepareForConvolution( |
| 2652 | dnn::ConvolutionKind kind, dnn::DataType element_type, Stream* stream, |
| 2653 | const dnn::BatchDescriptor& input_descriptor, DeviceMemoryBase input_data, |
| 2654 | const dnn::FilterDescriptor& filter_descriptor, |
| 2655 | DeviceMemoryBase filter_data, const dnn::BatchDescriptor& output_descriptor, |
| 2656 | DeviceMemoryBase output_data, |
| 2657 | const dnn::ConvolutionDescriptor& convolution_descriptor, |
| 2658 | const dnn::AlgorithmConfig& algorithm_config, |
| 2659 | ScratchAllocator* scratch_allocator, dnn::AlgorithmDesc* algorithm_desc, |
| 2660 | DeviceMemory<uint8>* scratch_memory) { |
| 2661 | ScopedTensorDescriptor input_nd{ |
| 2662 | input_descriptor, |
| 2663 | ToMIOpenDataType(element_type, input_descriptor.layout())}; |
| 2664 | ScopedFilterDescriptor filter{ |
| 2665 | filter_descriptor, input_descriptor, |
| 2666 | ToMIOpenDataType(element_type, filter_descriptor.layout())}; |
| 2667 | ScopedTensorDescriptor output_nd{ |
| 2668 | output_descriptor, |
| 2669 | ToMIOpenDataType(element_type, output_descriptor.layout())}; |
| 2670 | ScopedConvolutionDescriptor conv{ |
| 2671 | convolution_descriptor, |
| 2672 | ToMIOpenDataType(GetConvAccumulatorType(element_type))}; |
| 2673 | |
| 2674 | auto miopen = miopen_->GetHandle(parent_, stream); |
| 2675 | |
| 2676 | absl::optional<dnn::AlgorithmDesc> algo_desc = algorithm_config.algorithm(); |
| 2677 | size_t scratch_memory_size; |
| 2678 | |
| 2679 | if (!algo_desc.has_value()) { |
| 2680 | // With the default algorithm, use MIOpen's heuristics. |
| 2681 | assert(scratch_allocator); |
| 2682 | |
| 2683 | DeviceMemory<uint8> scratch_memory_temp; |
| 2684 | MIOpenAllocatorContext mac(scratch_allocator, stream); |
| 2685 | wrap::miopenSetAllocator(miopen.handle(), MIOpenAllocatorCallback, |
| 2686 | MIOpenDeallocatorCallback, &mac); |
| 2687 | size_t size_in_bytes; |
| 2688 | miopenStatus_t status = miopenStatusSuccess; |
| 2689 | |
| 2690 | switch (kind) { |
| 2691 | case dnn::ConvolutionKind::FORWARD: { |
| 2692 | status = wrap::miopenConvolutionForwardGetWorkSpaceSize( |
| 2693 | miopen.handle(), /*filterDesc=*/filter.handle(), |
| 2694 | /*srcDesc=*/input_nd.handle(), /*convDesc=*/conv.handle(), |
| 2695 | /*destDesc=*/output_nd.handle(), /*sizeInBytes=*/&size_in_bytes); |
| 2696 | break; |
| 2697 | } |
| 2698 | case dnn::ConvolutionKind::BACKWARD_DATA: { |
| 2699 | status = wrap::miopenConvolutionBackwardDataGetWorkSpaceSize( |
| 2700 | miopen.handle(), /*diffDesc=*/output_nd.handle(), |
| 2701 | /*filterDesc=*/filter.handle(), /*convDesc=*/conv.handle(), |
| 2702 | /*gradDesc=*/input_nd.handle(), /*sizeInBytes=*/&size_in_bytes); |
| 2703 | break; |
| 2704 | } |
| 2705 | case dnn::ConvolutionKind::BACKWARD_FILTER: { |
| 2706 | status = wrap::miopenConvolutionBackwardWeightsGetWorkSpaceSize( |
| 2707 | miopen.handle(), /*diffDesc=*/output_nd.handle(), |
| 2708 | /*srcDesc=*/input_nd.handle(), /*convDesc=*/conv.handle(), |
nothing calls this directly
no test coverage detected