saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-30 20:18:08 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Misc fixes to Conv and ConvTranspose CUDA kernels (#4281)

Browse source
This commit is contained in:
Hariharan Seshadri 2020-06-29 16:07:42 -07:00 committed by GitHub
parent 35a048ef9b
commit 465140b384
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
2 changed files with 60 additions and 45 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

17
onnxruntime/core/providers/cuda/nn/conv.cc
View file

@ -95,10 +95,6 @@ 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>());
// special case when there is a dim value of 0 in the shape.
if (Y->Shape().Size() == 0)
return Status::OK();
std::vector<int64_t> x_dims_cudnn = x_dims;
std::vector<int64_t> y_dims_cudnn = y_dims;
if (rank < 2) {
@ -112,12 +108,21 @@ Status Conv<T>::ComputeInternal(OpKernelContext* context) const {
strides.push_back(1);
dilations.push_back(1);
}
ORT_RETURN_IF_ERROR(s_.x_tensor.Set(x_dims_cudnn, CudnnTensor::GetDataType<CudaT>()));
ORT_RETURN_IF_ERROR(s_.y_tensor.Set(y_dims_cudnn, CudnnTensor::GetDataType<CudaT>()));
if (w_dims_changed)
ORT_RETURN_IF_ERROR(s_.filter_desc.Set(w_dims, CudnnTensor::GetDataType<CudaT>()));
// Special case when there is a dim value of 0 in the shape.
// Return only after we have cached the following for subsequent runs :
// 1) `w_dims` in the `filter_desc`
// 2) `y_dims` in s_.y_dims
if (Y->Shape().Size() == 0) {
return Status::OK();
}
ORT_RETURN_IF_ERROR(s_.x_tensor.Set(x_dims_cudnn, CudnnTensor::GetDataType<CudaT>()));
ORT_RETURN_IF_ERROR(s_.y_tensor.Set(y_dims_cudnn, CudnnTensor::GetDataType<CudaT>()));
cudnnConvolutionMode_t mode = CUDNN_CROSS_CORRELATION;
ORT_RETURN_IF_ERROR(s_.conv_desc.Set(kernel_shape.size(), pads, strides, dilations,
mode, CudnnTensor::GetDataType<CudaT>()));

88
onnxruntime/core/providers/cuda/nn/conv_transpose.cc
View file

@ -67,7 +67,7 @@ Status ConvTranspose<T>::DoConvTranspose(OpKernelContext* context, bool dynamic_
{
std::lock_guard<OrtMutex> lock(s_.mutex);
// TODO: add a global cache if need to handle cases for multiple frames running simultaneuously with different batch_size
// TODO: add a global cache if need to handle cases for multiple frames running simultaneously with different batch_size
bool input_dims_changed = (s_.last_x_dims != x_dims);
bool w_dims_changed = (s_.last_w_dims != w_dims);
if (input_dims_changed || w_dims_changed) {
@ -82,11 +82,6 @@ Status ConvTranspose<T>::DoConvTranspose(OpKernelContext* context, bool dynamic_
ConvTransposeAttributes::Prepare p;
ORT_RETURN_IF_ERROR(conv_transpose_attrs_.PrepareForCompute(context, has_bias, p, dynamic_padding));
// Bail out early if one of the dimensions is zero.
if (p.Y->Shape().Size() == 0) {
return Status::OK();
}
auto y_dims = p.Y->Shape().GetDims();
if (x_dimensions == 3) {
y_dims.insert(y_dims.begin() + 2, 1);
@ -98,12 +93,20 @@ Status ConvTranspose<T>::DoConvTranspose(OpKernelContext* context, bool dynamic_
}
s_.y_dims = y_dims;
ORT_RETURN_IF_ERROR(s_.x_tensor.Set(x_dims, CudnnTensor::GetDataType<CudaT>()));
ORT_RETURN_IF_ERROR(s_.y_tensor.Set(y_dims, CudnnTensor::GetDataType<CudaT>()));
if (w_dims_changed)
ORT_RETURN_IF_ERROR(s_.filter_desc.Set(w_dims, CudnnTensor::GetDataType<CudaT>()));
// Special case when there is a dim value of 0 in the shape.
// Return only after we have cached the following for subsequent runs :
// 1) `w_dims` in the `filter_desc`
// 2) `y_dims` in s_.y_dims
if (p.Y->Shape().Size() == 0) {
return Status::OK();
}
ORT_RETURN_IF_ERROR(s_.x_tensor.Set(x_dims, CudnnTensor::GetDataType<CudaT>()));
ORT_RETURN_IF_ERROR(s_.y_tensor.Set(y_dims, CudnnTensor::GetDataType<CudaT>()));
cudnnConvolutionMode_t mode = CUDNN_CROSS_CORRELATION;
ORT_RETURN_IF_ERROR(s_.conv_desc.Set(p.kernel_shape.size(), p.pads, p.strides,
p.dilations, mode, CudnnTensor::GetDataType<CudaT>()));
@ -155,42 +158,49 @@ Status ConvTranspose<T>::DoConvTranspose(OpKernelContext* context, bool dynamic_
s_.algo = perf.algo;
s_.workspace_bytes = perf.memory;
}
}
if (!y_data) {
auto y_dims = s_.y_dims;
if (x_dimensions == 3) {
y_dims.erase(y_dims.begin() + 2);
// The following block will be executed in case there has been no change in the shapes of the
// input and the filter compared to the previous run
if (!y_data) {
auto y_dims = s_.y_dims;
if (x_dimensions == 3) {
y_dims.erase(y_dims.begin() + 2);
}
Tensor* Y = context->Output(0, TensorShape(y_dims));
y_data = reinterpret_cast<CudaT*>(Y->template MutableData<T>());
// Bail out early if one of the output dimensions is zero.
if (Y->Shape().Size() == 0) {
return Status::OK();
}
}
Tensor* Y = context->Output(0, TensorShape(y_dims));
y_data = reinterpret_cast<CudaT*>(Y->template MutableData<T>());
}
const auto alpha = Consts<CudaT>::One;
const auto beta = Consts<CudaT>::Zero;
const auto alpha = Consts<CudaT>::One;
const auto beta = Consts<CudaT>::Zero;
IAllocatorUniquePtr<void> workspace = GetScratchBuffer<void>(s_.workspace_bytes);
IAllocatorUniquePtr<void> workspace = GetScratchBuffer<void>(s_.workspace_bytes);
CUDNN_RETURN_IF_ERROR(
cudnnConvolutionBackwardData(
CudnnHandle(),
&alpha,
s_.filter_desc,
w_data,
s_.x_tensor,
x_data,
s_.conv_desc,
s_.algo,
workspace.get(),
s_.workspace_bytes,
&beta,
s_.y_tensor,
y_data));
CUDNN_RETURN_IF_ERROR(
cudnnConvolutionBackwardData(
CudnnHandle(),
&alpha,
s_.filter_desc,
w_data,
s_.x_tensor,
x_data,
s_.conv_desc,
s_.algo,
workspace.get(),
s_.workspace_bytes,
&beta,
s_.y_tensor,
y_data));
if (has_bias) {
const Tensor* B = dynamic_padding ? context->Input<Tensor>(3) : context->Input<Tensor>(2);
auto b_data = reinterpret_cast<const CudaT*>(B->template Data<T>());
CUDNN_RETURN_IF_ERROR(cudnnAddTensor(CudnnHandle(), &alpha, s_.b_tensor, b_data, &alpha, s_.y_tensor, y_data));
if (has_bias) {
const Tensor* B = dynamic_padding ? context->Input<Tensor>(3) : context->Input<Tensor>(2);
auto b_data = reinterpret_cast<const CudaT*>(B->template Data<T>());
CUDNN_RETURN_IF_ERROR(cudnnAddTensor(CudnnHandle(), &alpha, s_.b_tensor, b_data, &alpha, s_.y_tensor, y_data));
}
}
return Status::OK();