diff --git a/onnxruntime/core/providers/cuda/nn/conv.cc b/onnxruntime/core/providers/cuda/nn/conv.cc index 41e068a6fc..be1c24ebed 100644 --- a/onnxruntime/core/providers/cuda/nn/conv.cc +++ b/onnxruntime/core/providers/cuda/nn/conv.cc @@ -51,8 +51,10 @@ Status Conv::ComputeInternal(OpKernelContext* context) const { if (input_dims_changed) s_.last_x_dims = x_dims; - if (w_dims_changed) + if (w_dims_changed) { s_.last_w_dims = w_dims; + s_.cached_benchmark_results.clear(); + } const int64_t N = X->Shape()[0]; const int64_t M = W->Shape()[0]; @@ -103,8 +105,6 @@ Status Conv::ComputeInternal(OpKernelContext* context) const { ORT_RETURN_IF_ERROR(s_.conv_desc.Set(kernel_shape.size(), pads, strides, dilations, mode, CudnnTensor::GetDataType())); CUDNN_RETURN_IF_ERROR(cudnnSetConvolutionGroupCount(s_.conv_desc, gsl::narrow_cast(group_))); - IAllocatorUniquePtr algo_search_workspace = GetScratchBuffer(AlgoSearchWorkspaceSize); - if (has_bias) { const Tensor* B = context->Input(2); const auto& b_shape = B->Shape(); @@ -121,26 +121,33 @@ Status Conv::ComputeInternal(OpKernelContext* context) const { Tensor* Y = context->Output(0, TensorShape(s_.y_dims)); y_data = reinterpret_cast(Y->template MutableData()); - // set math type to tensor core before algorithm search - if (std::is_same::value) - CUDNN_RETURN_IF_ERROR(cudnnSetConvolutionMathType(s_.conv_desc, CUDNN_TENSOR_OP_MATH)); + if (!s_.cached_benchmark_results.contains(x_dims_cudnn)) { + IAllocatorUniquePtr algo_search_workspace = GetScratchBuffer(AlgoSearchWorkspaceSize); - cudnnConvolutionFwdAlgoPerf_t perf; - int algo_count = 1; - CUDNN_RETURN_IF_ERROR(cudnnFindConvolutionForwardAlgorithmEx( - CudnnHandle(), - s_.x_tensor, - x_data, - s_.filter_desc, - w_data, - s_.conv_desc, - s_.y_tensor, - y_data, - 1, - &algo_count, - &perf, - algo_search_workspace.get(), - AlgoSearchWorkspaceSize)); + // set math type to tensor core before algorithm search + if (std::is_same::value) + CUDNN_RETURN_IF_ERROR(cudnnSetConvolutionMathType(s_.conv_desc, CUDNN_TENSOR_OP_MATH)); + + cudnnConvolutionFwdAlgoPerf_t perf; + int algo_count = 1; + CUDNN_RETURN_IF_ERROR(cudnnFindConvolutionForwardAlgorithmEx( + CudnnHandle(), + s_.x_tensor, + x_data, + s_.filter_desc, + w_data, + s_.conv_desc, + s_.y_tensor, + y_data, + 1, + &algo_count, + &perf, + algo_search_workspace.get(), + AlgoSearchWorkspaceSize)); + s_.cached_benchmark_results.insert(x_dims_cudnn, {perf.algo, perf.memory, perf.mathType}); + } + + const auto& perf = s_.cached_benchmark_results.at(x_dims_cudnn); CUDNN_RETURN_IF_ERROR(cudnnSetConvolutionMathType(s_.conv_desc, perf.mathType)); s_.algo = perf.algo; s_.workspace_bytes = perf.memory; diff --git a/onnxruntime/core/providers/cuda/nn/conv.h b/onnxruntime/core/providers/cuda/nn/conv.h index 9535477e43..0ae0ae5be7 100644 --- a/onnxruntime/core/providers/cuda/nn/conv.h +++ b/onnxruntime/core/providers/cuda/nn/conv.h @@ -8,6 +8,7 @@ #include "core/framework/op_kernel.h" #include "core/providers/cuda/cudnn_common.h" #include "core/providers/cpu/nn/conv_base.h" +#include namespace onnxruntime { namespace cuda { @@ -30,8 +31,85 @@ class CudnnConvolutionDescriptor final { cudnnConvolutionDescriptor_t desc_; }; +template +struct vector_hash { + std::size_t operator()(const std::vector& values) const { + std::size_t seed = values.size(); + for (auto& val : values) + seed ^= std::hash()(val) + 0x9e3779b9 + (seed << 6) + (seed >> 2); + return seed; + } +}; + +template , + typename KeyEqual = std::equal_to, + typename Allocator = std::allocator>, + typename ListAllocator = std::allocator> +class lru_unordered_map { + public: + lru_unordered_map(size_t max_size) : max_size_(max_size) {} + + void insert(const Key& key, const T& value) { + auto it = items_.find(key); + if (it != items_.end()) { + it->second.value = value; + move_to_front(it->second.lru_iterator); + return; + } + + while (size() + 1 > max_size_) { + items_.erase(lru_list_.back()); + lru_list_.pop_back(); + } + + lru_list_.emplace_front(key); + items_.emplace(key, value_type{value, lru_list_.begin()}); + } + + T& at(const Key& key) { + auto it = items_.find(key); + if (it == items_.end()) { + throw std::out_of_range("There is no such key in cache"); + } + move_to_front(it->second.lru_iterator); + return it->second.value; + } + + bool contains(const Key& key) const { + return items_.find(key) != items_.end(); + } + + size_t size() const { + return items_.size(); + } + + void clear() { + items_.clear(); + lru_list_.clear(); + } + +private: + using list_type = std::list; + using iterator_type = typename list_type::iterator; + struct value_type { + T value; + iterator_type lru_iterator; + }; + + void move_to_front(iterator_type it) { + lru_list_.splice(lru_list_.begin(), lru_list_, it); + } + + size_t max_size_; + std::unordered_map items_; + list_type lru_list_; +}; + // cached cudnn descriptors -template +constexpr size_t MAX_CACHED_ALGO_PERF_RESULTS = 10000; + +template struct CudnnConvState { // if x/w dims changed, update algo and cudnnTensors std::vector last_x_dims; @@ -40,13 +118,21 @@ struct CudnnConvState { // these would be recomputed if x/w dims change std::vector y_dims; size_t workspace_bytes; - AlgoType algo; + decltype(AlgoPerfType().algo) algo; CudnnTensor x_tensor; CudnnFilterDescriptor filter_desc; CudnnTensor b_tensor; CudnnTensor y_tensor; CudnnConvolutionDescriptor conv_desc; + struct PerfResultParams { + decltype(AlgoPerfType().algo) algo; + decltype(AlgoPerfType().memory) memory; + decltype(AlgoPerfType().mathType) mathType; + }; + + lru_unordered_map, PerfResultParams, vector_hash> cached_benchmark_results { MAX_CACHED_ALGO_PERF_RESULTS }; + // note that conv objects are shared between execution frames, and a lock is needed to avoid multi-thread racing OrtMutex mutex; }; @@ -70,7 +156,7 @@ class Conv : public CudaKernel, public ConvBase { Status ComputeInternal(OpKernelContext* context) const override; private: - mutable CudnnConvState s_; + mutable CudnnConvState s_; }; } // namespace cuda diff --git a/onnxruntime/core/providers/cuda/nn/conv_transpose.cc b/onnxruntime/core/providers/cuda/nn/conv_transpose.cc index 16406af5af..c15bd2978c 100644 --- a/onnxruntime/core/providers/cuda/nn/conv_transpose.cc +++ b/onnxruntime/core/providers/cuda/nn/conv_transpose.cc @@ -48,8 +48,10 @@ Status ConvTranspose::ComputeInternal(OpKernelContext* context) const { if (input_dims_changed) s_.last_x_dims = x_dims; - if (w_dims_changed) + if (w_dims_changed) { s_.last_w_dims = w_dims; + s_.cached_benchmark_results.clear(); + } Prepare p; ORT_RETURN_IF_ERROR(PrepareForCompute(context, has_bias, p)); @@ -67,8 +69,6 @@ Status ConvTranspose::ComputeInternal(OpKernelContext* context) const { ORT_RETURN_IF_ERROR(s_.conv_desc.Set(p.kernel_shape.size(), p.pads, p.strides, p.dilations, mode, CudnnTensor::GetDataType())); CUDNN_RETURN_IF_ERROR(cudnnSetConvolutionGroupCount(s_.conv_desc, gsl::narrow_cast(group_))); - IAllocatorUniquePtr algo_search_workspace = GetScratchBuffer(AlgoSearchWorkspaceSize); - if (has_bias) { const auto& b_shape = p.B->Shape(); ORT_RETURN_IF_NOT(b_shape.NumDimensions() == 1, "bias should be 1D"); @@ -83,26 +83,33 @@ Status ConvTranspose::ComputeInternal(OpKernelContext* context) const { y_data = reinterpret_cast(p.Y->template MutableData()); - // set math type to tensor core before algorithm search - if (std::is_same::value) - CUDNN_RETURN_IF_ERROR(cudnnSetConvolutionMathType(s_.conv_desc, CUDNN_TENSOR_OP_MATH)); + if (!s_.cached_benchmark_results.contains(x_dims)) { + IAllocatorUniquePtr algo_search_workspace = GetScratchBuffer(AlgoSearchWorkspaceSize); - cudnnConvolutionBwdDataAlgoPerf_t perf; - int algo_count = 1; - CUDNN_RETURN_IF_ERROR(cudnnFindConvolutionBackwardDataAlgorithmEx( - CudnnHandle(), - s_.filter_desc, - w_data, - s_.x_tensor, - x_data, - s_.conv_desc, - s_.y_tensor, - y_data, - 1, - &algo_count, - &perf, - algo_search_workspace.get(), - AlgoSearchWorkspaceSize)); + // set math type to tensor core before algorithm search + if (std::is_same::value) + CUDNN_RETURN_IF_ERROR(cudnnSetConvolutionMathType(s_.conv_desc, CUDNN_TENSOR_OP_MATH)); + + cudnnConvolutionBwdDataAlgoPerf_t perf; + int algo_count = 1; + CUDNN_RETURN_IF_ERROR(cudnnFindConvolutionBackwardDataAlgorithmEx( + CudnnHandle(), + s_.filter_desc, + w_data, + s_.x_tensor, + x_data, + s_.conv_desc, + s_.y_tensor, + y_data, + 1, + &algo_count, + &perf, + algo_search_workspace.get(), + AlgoSearchWorkspaceSize)); + s_.cached_benchmark_results.insert(x_dims, {perf.algo, perf.memory, perf.mathType}); + } + + const auto& perf = s_.cached_benchmark_results.at(x_dims); CUDNN_RETURN_IF_ERROR(cudnnSetConvolutionMathType(s_.conv_desc, perf.mathType)); s_.algo = perf.algo; s_.workspace_bytes = perf.memory; diff --git a/onnxruntime/core/providers/cuda/nn/conv_transpose.h b/onnxruntime/core/providers/cuda/nn/conv_transpose.h index 93846950b2..9b6abd46f9 100644 --- a/onnxruntime/core/providers/cuda/nn/conv_transpose.h +++ b/onnxruntime/core/providers/cuda/nn/conv_transpose.h @@ -17,7 +17,7 @@ class ConvTranspose : public CudaKernel, public ConvTransposeBase { Status ComputeInternal(OpKernelContext* context) const override; private: - mutable CudnnConvState s_; + mutable CudnnConvState s_; }; } // namespace cuda