From ce7ff27bac622aedb284ff3b57832055d715ae08 Mon Sep 17 00:00:00 2001 From: Sherlock Date: Mon, 19 Apr 2021 23:37:05 -0700 Subject: [PATCH] Fix perf issue in Conv CUDA kernel (#7348) * Fix perf issue in Conv CUDA kernel * Read avaiable memory from device * assuming 10% fragmentation Co-authored-by: Sherlock Huang --- onnxruntime/core/providers/cuda/nn/conv.cc | 104 +++++++++++++----- onnxruntime/core/providers/cuda/nn/conv.h | 5 + .../test/gradient/gradient_ops_test.cc | 4 +- 3 files changed, 83 insertions(+), 30 deletions(-) diff --git a/onnxruntime/core/providers/cuda/nn/conv.cc b/onnxruntime/core/providers/cuda/nn/conv.cc index e17ffe676e..4333d7fffc 100644 --- a/onnxruntime/core/providers/cuda/nn/conv.cc +++ b/onnxruntime/core/providers/cuda/nn/conv.cc @@ -34,6 +34,55 @@ REGISTER_KERNEL_TYPED(float) REGISTER_KERNEL_TYPED(double) REGISTER_KERNEL_TYPED(MLFloat16) +template +const cudnnConvolutionFwdAlgo_t Conv::kAllAlgos[] = { + CUDNN_CONVOLUTION_FWD_ALGO_GEMM, + CUDNN_CONVOLUTION_FWD_ALGO_FFT, + CUDNN_CONVOLUTION_FWD_ALGO_FFT_TILING, + CUDNN_CONVOLUTION_FWD_ALGO_IMPLICIT_GEMM, + CUDNN_CONVOLUTION_FWD_ALGO_IMPLICIT_PRECOMP_GEMM, + CUDNN_CONVOLUTION_FWD_ALGO_DIRECT, + CUDNN_CONVOLUTION_FWD_ALGO_WINOGRAD, + CUDNN_CONVOLUTION_FWD_ALGO_WINOGRAD_NONFUSED, +}; + +cudnnStatus_t getWorkspaceSize(const CudnnConvState& s, + cudnnConvolutionFwdAlgo_t algo, size_t* sz) { + return cudnnGetConvolutionForwardWorkspaceSize( + s.handle, + s.x_tensor, + s.w_desc, + s.conv_desc, + s.y_tensor, + algo, + sz); +} + +template +size_t getMaxWorkspaceSize(const CudnnConvState& s, + const algo_t* algo, int n_algo) { + size_t max_ws_size = 0; + + // TODO: get maximum available size from memory areana + + size_t free, total; + CUDA_CALL_THROW(cudaMemGetInfo(&free, &total)); + // Assuming 10% of fragmentation + free = static_cast(static_cast(free) * 0.9); + + std::cout << "free: " << free << " total: " << total << std::endl; + + for (int i = 0; i < n_algo; i++) { + cudnnStatus_t err; + size_t sz; + err = getWorkspaceSize(s, algo[i], &sz); + if (CUDNN_STATUS_SUCCESS != err || sz == 0 || sz < max_ws_size || sz > free) + continue; + max_ws_size = sz; + } + return max_ws_size; +} + Status SliceOutUnwantedOutputSection(cudaStream_t stream, const void* input_data, const std::vector& input_dims, @@ -159,14 +208,16 @@ Status Conv::UpdateState(OpKernelContext* context, bool bias_expected) const std::vector y_dims_cudnn = !post_slicing_required ? y_dims : y_dims_with_adjusted_pads; if (rank < 2) { // cudnn only takes 4D or 5D input, so pad dimensions if needed - x_dims_cudnn.push_back(1); - y_dims_cudnn.push_back(1); - w_dims.push_back(1); - pads.insert(pads.begin() + rank, 0); - pads.insert(pads.end(), 0); - kernel_shape.push_back(1); - strides.push_back(1); - dilations.push_back(1); + // If input shape is [N, C, D], we pad the shape to [N, C, 1, D], as it results in + // more efficient algorithm than padding to [N, C, D, 1] + x_dims_cudnn.insert(x_dims_cudnn.begin() + 2, 1); + y_dims_cudnn.insert(y_dims_cudnn.begin() + 2, 1); + w_dims.insert(w_dims.begin() + 2, 1); + pads.insert(pads.begin(), 0); + pads.insert(pads.begin() + 2, 0); + kernel_shape.insert(kernel_shape.begin(), 1); + strides.insert(strides.begin(), 1); + dilations.insert(dilations.begin(), 1); } if (w_dims_changed) { @@ -200,8 +251,6 @@ Status Conv::UpdateState(OpKernelContext* context, bool bias_expected) const } if (!s_.cached_benchmark_results.contains(x_dims_cudnn)) { - IAllocatorUniquePtr algo_search_workspace = GetScratchBuffer(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)); @@ -212,9 +261,13 @@ Status Conv::UpdateState(OpKernelContext* context, bool bias_expected) const int cudnn_conv_algo = cuda_ep->GetCudnnConvAlgo(); ORT_ENFORCE(cudnn_conv_algo > -1 && cudnn_conv_algo < 3, "cudnn_conv_algo should be 0, 1 or 2, but got ", cudnn_conv_algo); switch (cudnn_conv_algo) { - case 0: + case 0: { + static constexpr int num_algos = CUDNN_CONVOLUTION_FWD_ALGO_COUNT; + size_t max_ws_size = getMaxWorkspaceSize(s_, kAllAlgos, num_algos); + IAllocatorUniquePtr algo_search_workspace = GetScratchBuffer(max_ws_size); + CUDNN_RETURN_IF_ERROR(cudnnFindConvolutionForwardAlgorithmEx( - CudnnHandle(), + s_.handle, s_.x_tensor, s_.x_data, s_.w_desc, @@ -222,35 +275,28 @@ Status Conv::UpdateState(OpKernelContext* context, bool bias_expected) const s_.conv_desc, s_.y_tensor, s_.y_data, - 1, - &algo_count, + 1, // requestedAlgoCount + &algo_count, // returnedAlgoCount &perf, algo_search_workspace.get(), - AlgoSearchWorkspaceSize)); + max_ws_size)); break; - + } case 1: CUDNN_RETURN_IF_ERROR(cudnnGetConvolutionForwardAlgorithm_v7( - CudnnHandle(), + s_.handle, s_.x_tensor, s_.w_desc, s_.conv_desc, s_.y_tensor, - 1, - &algo_count, + 1, // requestedAlgoCount + &algo_count, // returnedAlgoCount &perf)); break; default: perf.algo = kDefaultConvAlgo; - CUDNN_RETURN_IF_ERROR(cudnnGetConvolutionForwardWorkspaceSize( - CudnnHandle(), - s_.x_tensor, - s_.w_desc, - s_.conv_desc, - s_.y_tensor, - perf.algo, - &perf.memory)); + CUDNN_RETURN_IF_ERROR(getWorkspaceSize(s_, perf.algo, &perf.memory)); if (std::is_same::value) { perf.mathType = CUDNN_TENSOR_OP_MATH; } else { @@ -289,7 +335,7 @@ Status Conv::ComputeInternal(OpKernelContext* context) const { const auto alpha = Consts::One; const auto beta = Consts::Zero; IAllocatorUniquePtr workspace = GetWorkSpace(); - CUDNN_RETURN_IF_ERROR(cudnnConvolutionForward(CudnnHandle(), + CUDNN_RETURN_IF_ERROR(cudnnConvolutionForward(s_.handle, &alpha, s_.x_tensor, s_.x_data, @@ -303,7 +349,7 @@ Status Conv::ComputeInternal(OpKernelContext* context) const { s_.y_tensor, s_.y_data)); if (nullptr != s_.b_data) { - CUDNN_RETURN_IF_ERROR(cudnnAddTensor(CudnnHandle(), &alpha, s_.b_tensor, s_.b_data, + CUDNN_RETURN_IF_ERROR(cudnnAddTensor(s_.handle, &alpha, s_.b_tensor, s_.b_data, &alpha, s_.y_tensor, s_.y_data)); } // To deal with asymmetric padding, we may have over-padded on one or both sides of the spatial dimensions diff --git a/onnxruntime/core/providers/cuda/nn/conv.h b/onnxruntime/core/providers/cuda/nn/conv.h index 8f27265e37..8e54c1fde2 100644 --- a/onnxruntime/core/providers/cuda/nn/conv.h +++ b/onnxruntime/core/providers/cuda/nn/conv.h @@ -112,6 +112,8 @@ constexpr size_t MAX_CACHED_ALGO_PERF_RESULTS = 10000; template struct CudnnConvState { + cudnnHandle_t handle; + // if x/w dims changed, update algo and cudnnTensors std::vector last_x_dims; std::vector last_w_dims; @@ -174,6 +176,8 @@ class Conv : public CudaKernel { Conv(const OpKernelInfo& info) : CudaKernel(info), conv_attrs_(info) { auto pads_size = conv_attrs_.pads.size(); ORT_ENFORCE(pads_size % 2 == 0); + + s_.handle = CudnnHandle(); } Status ComputeInternal(OpKernelContext* context) const override; @@ -187,6 +191,7 @@ class Conv : public CudaKernel { ConvAttributes conv_attrs_; mutable CudnnConvState s_; constexpr static auto kDefaultConvAlgo = CUDNN_CONVOLUTION_FWD_ALGO_IMPLICIT_PRECOMP_GEMM; + static const cudnnConvolutionFwdAlgo_t kAllAlgos[]; }; Status SliceOutUnwantedOutputSection(cudaStream_t stream, diff --git a/orttraining/orttraining/test/gradient/gradient_ops_test.cc b/orttraining/orttraining/test/gradient/gradient_ops_test.cc index 114ce159f5..1ea29d8110 100755 --- a/orttraining/orttraining/test/gradient/gradient_ops_test.cc +++ b/orttraining/orttraining/test/gradient/gradient_ops_test.cc @@ -811,7 +811,9 @@ void ConvGradientCheckerTest(std::vector>* e float max_error; GradientChecker gradient_checker; OpDef op_def{"Conv"}; - float error_tolerance = 1e-1f; + + // TODO: revisit the tol when ConvGrad impl is completed + float error_tolerance = 2e-1f; // 1D convolution {