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 <bahuang@OrtTrainingDev3.af05slrtruoetgaxwwjv5nsq5e.px.internal.cloudapp.net>
This commit is contained in:
Sherlock 2021-04-19 23:37:05 -07:00 committed by GitHub
parent ac346a1b90
commit ce7ff27bac
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
3 changed files with 83 additions and 30 deletions

View file

@ -34,6 +34,55 @@ REGISTER_KERNEL_TYPED(float)
REGISTER_KERNEL_TYPED(double)
REGISTER_KERNEL_TYPED(MLFloat16)
template <typename T>
const cudnnConvolutionFwdAlgo_t Conv<T>::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<cudnnConvolutionFwdAlgoPerf_t>& 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 <typename algo_t>
size_t getMaxWorkspaceSize(const CudnnConvState<cudnnConvolutionFwdAlgoPerf_t>& 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<size_t>(static_cast<double>(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<int64_t>& input_dims,
@ -159,14 +208,16 @@ Status Conv<T>::UpdateState(OpKernelContext* context, bool bias_expected) const
std::vector<int64_t> 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<T>::UpdateState(OpKernelContext* context, bool bias_expected) const
}
if (!s_.cached_benchmark_results.contains(x_dims_cudnn)) {
IAllocatorUniquePtr<void> algo_search_workspace = GetScratchBuffer<void>(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));
@ -212,9 +261,13 @@ Status Conv<T>::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<void> algo_search_workspace = GetScratchBuffer<void>(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<T>::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<T, MLFloat16>::value) {
perf.mathType = CUDNN_TENSOR_OP_MATH;
} else {
@ -289,7 +335,7 @@ Status Conv<T>::ComputeInternal(OpKernelContext* context) const {
const auto alpha = Consts<CudaT>::One;
const auto beta = Consts<CudaT>::Zero;
IAllocatorUniquePtr<void> 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<T>::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

View file

@ -112,6 +112,8 @@ constexpr size_t MAX_CACHED_ALGO_PERF_RESULTS = 10000;
template <typename AlgoPerfType>
struct CudnnConvState {
cudnnHandle_t handle;
// if x/w dims changed, update algo and cudnnTensors
std::vector<int64_t> last_x_dims;
std::vector<int64_t> 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<cudnnConvolutionFwdAlgoPerf_t> s_;
constexpr static auto kDefaultConvAlgo = CUDNN_CONVOLUTION_FWD_ALGO_IMPLICIT_PRECOMP_GEMM;
static const cudnnConvolutionFwdAlgo_t kAllAlgos[];
};
Status SliceOutUnwantedOutputSection(cudaStream_t stream,

View file

@ -811,7 +811,9 @@ void ConvGradientCheckerTest(std::vector<std::unique_ptr<IExecutionProvider>>* e
float max_error;
GradientChecker<float, float, float> 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
{