Cache CUDNN convolution benchmark results in cuda::Conv kernels (#712)

* Cache CUDNN convolution benchmark results in cuda::Conv kernels

Previously, the best convolution algorithm was determined by running
cudnnFindConvolutionForwardAlgorithmEx and cudnnFindConvolutionBackwardDataAlgorithmEx
on every shape change.

This is very detrimental for variable input shapes, such as variable batch
sizes.

This change adds a map to cache previously determined benchmark results.

The caching results in significant speedups for variable input shapes.

* Use LRU to limit cached benchmark results

* Only cache benchmark results for a fixed weight shape

In case the weight shape changes, all cached results are discarded.

* Use padded shape as key for cached benchmarks

* Add constant for max number of cached benchmark results

* Use unordered_map to store cached benchmark results

* Only store the parameters that are actuallt needed

onnxruntime/core/providers/cuda/nn/conv.cc

@@ -51,8 +51,10 @@ Status Conv<T>::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<T>::ComputeInternal(OpKernelContext* context) const {
       ORT_RETURN_IF_ERROR(s_.conv_desc.Set(kernel_shape.size(), pads, strides, dilations, mode, CudnnTensor::GetDataType<CudaT>()));
       CUDNN_RETURN_IF_ERROR(cudnnSetConvolutionGroupCount(s_.conv_desc, gsl::narrow_cast<int>(group_)));
 
-      IAllocatorUniquePtr<void> algo_search_workspace = GetScratchBuffer<void>(AlgoSearchWorkspaceSize);
-
       if (has_bias) {
         const Tensor* B = context->Input<Tensor>(2);
         const auto& b_shape = B->Shape();
@@ -121,26 +121,33 @@ Status Conv<T>::ComputeInternal(OpKernelContext* context) const {
       Tensor* Y = context->Output(0, TensorShape(s_.y_dims));
       y_data = reinterpret_cast<CudaT*>(Y->template MutableData<T>());
 
-      // set math type to tensor core before algorithm search
-      if (std::is_same<T, MLFloat16>::value)
-        CUDNN_RETURN_IF_ERROR(cudnnSetConvolutionMathType(s_.conv_desc, CUDNN_TENSOR_OP_MATH));
+      if (!s_.cached_benchmark_results.contains(x_dims_cudnn)) {
+        IAllocatorUniquePtr<void> algo_search_workspace = GetScratchBuffer<void>(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<T, MLFloat16>::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;

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 <list>
 
 namespace onnxruntime {
 namespace cuda {
@@ -30,8 +31,85 @@ class CudnnConvolutionDescriptor final {
   cudnnConvolutionDescriptor_t desc_;
 };
 
+template <typename T>
+struct vector_hash {
+  std::size_t operator()(const std::vector<T>& values) const {
+    std::size_t seed = values.size();
+    for (auto& val : values)
+      seed ^= std::hash<T>()(val) + 0x9e3779b9 + (seed << 6) + (seed >> 2);
+    return seed;
+  }
+};
+
+template <typename Key, typename T,
+          typename Hash = std::hash<Key>,
+          typename KeyEqual = std::equal_to<Key>,
+          typename Allocator = std::allocator<std::pair<const Key, T>>,
+          typename ListAllocator = std::allocator<Key>>
+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<Key, ListAllocator>;
+  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<Key, value_type, Hash, KeyEqual, Allocator> items_;
+  list_type lru_list_;
+};
+
 // cached cudnn descriptors
-template <typename AlgoType>
+constexpr size_t MAX_CACHED_ALGO_PERF_RESULTS = 10000;
+
+template <typename AlgoPerfType>
 struct CudnnConvState {
   // if x/w dims changed, update algo and cudnnTensors
   std::vector<int64_t> last_x_dims;
@@ -40,13 +118,21 @@ struct CudnnConvState {
   // these would be recomputed if x/w dims change
   std::vector<int64_t> 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<std::vector<int64_t>, PerfResultParams, vector_hash<int64_t>> 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<cudnnConvolutionFwdAlgo_t> s_;
+  mutable CudnnConvState<cudnnConvolutionFwdAlgoPerf_t> s_;
 };
 
 }  // namespace cuda

onnxruntime/core/providers/cuda/nn/conv_transpose.cc

@@ -48,8 +48,10 @@ Status ConvTranspose<T>::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<T>::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<CudaT>()));
       CUDNN_RETURN_IF_ERROR(cudnnSetConvolutionGroupCount(s_.conv_desc, gsl::narrow_cast<int>(group_)));
 
-      IAllocatorUniquePtr<void> algo_search_workspace = GetScratchBuffer<void>(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<T>::ComputeInternal(OpKernelContext* context) const {
 
       y_data = reinterpret_cast<CudaT*>(p.Y->template MutableData<T>());
 
-      // set math type to tensor core before algorithm search
-      if (std::is_same<T, MLFloat16>::value)
-        CUDNN_RETURN_IF_ERROR(cudnnSetConvolutionMathType(s_.conv_desc, CUDNN_TENSOR_OP_MATH));
+      if (!s_.cached_benchmark_results.contains(x_dims)) {
+        IAllocatorUniquePtr<void> algo_search_workspace = GetScratchBuffer<void>(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<T, MLFloat16>::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;

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<cudnnConvolutionBwdDataAlgo_t> s_;
+  mutable CudnnConvState<cudnnConvolutionBwdDataAlgoPerf_t> s_;
 };
 
 }  // namespace cuda