Adding Resize CUDA kernel (#869)

* Adding resize cuda kernel

* Fix a bug.

* Fix CI errors

* Fix CI issues.

* update onnx

* fix a linux CI issue.

* fix cuda kernels
This commit is contained in:
Du Li 2019-04-29 13:35:38 -07:00 committed by GitHub
parent ded7eeb033
commit 1658efd953
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
9 changed files with 308 additions and 16 deletions

View file

@ -11,7 +11,6 @@ template <typename T>
class Resize : public Upsample<T> {
public:
Resize(const OpKernelInfo& info) : Upsample<T>(info) {
UpsampleBase::is_resize = true;
}
Status Compute(OpKernelContext* context) const override {

View file

@ -512,6 +512,11 @@ class ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kO
class ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 7, 9, MLFloat16, Upsample);
class ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 7, 9, int32_t, Upsample);
class ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 7, 9, uint8_t, Upsample);
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 10, float, Resize);
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 10, double, Resize);
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 10, MLFloat16, Resize);
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 10, int32_t, Resize);
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 10, uint8_t, Resize);
static void RegisterCudaKernels(KernelRegistry& kernel_registry) {
static const BuildKernelCreateInfoFn function_table[] = {
@ -787,6 +792,11 @@ static void RegisterCudaKernels(KernelRegistry& kernel_registry) {
BuildKernelCreateInfo<ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 7, 9, MLFloat16, Upsample)>,
BuildKernelCreateInfo<ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 7, 9, int32_t, Upsample)>,
BuildKernelCreateInfo<ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 7, 9, uint8_t, Upsample)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 10, float, Resize)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 10, double, Resize)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 10, MLFloat16, Resize)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 10, int32_t, Resize)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 10, uint8_t, Resize)>,
};
for (auto& function_table_entry : function_table) {

View file

@ -0,0 +1,27 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
#include "resize.h"
namespace onnxruntime {
namespace cuda {
#define REGISTER_KERNEL_TYPED(T) \
ONNX_OPERATOR_TYPED_KERNEL_EX( \
Resize, \
kOnnxDomain, \
10, \
T, \
kCudaExecutionProvider, \
KernelDefBuilder() \
.InputMemoryType<OrtMemTypeCPUInput>(1) \
.TypeConstraint("T", DataTypeImpl::GetTensorType<T>()), \
Resize<T>);
REGISTER_KERNEL_TYPED(float)
REGISTER_KERNEL_TYPED(double)
REGISTER_KERNEL_TYPED(MLFloat16)
REGISTER_KERNEL_TYPED(int32_t)
REGISTER_KERNEL_TYPED(uint8_t)
} // namespace cuda
} // namespace onnxruntime

View file

@ -0,0 +1,23 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
#pragma once
#include "core/providers/cuda/tensor/upsample.h"
namespace onnxruntime {
namespace cuda {
template <typename T>
class Resize : public Upsample<T> {
public:
Resize(OpKernelInfo info) : Upsample<T>(info) {
}
Status ComputeInternal(OpKernelContext* context) const override {
return Upsample<T>::ComputeInternal(context);
}
};
} // namespace cuda
} // namespace onnxruntime

View file

@ -0,0 +1,133 @@
#include "core/providers/cuda/cu_inc/common.cuh"
#include "core/providers/cuda/tensor/resize_impl.h"
namespace onnxruntime {
namespace cuda {
template <typename T>
__global__ void _ResizeNearestKernel(const size_t rank,
const int64_t* input_pitches,
const fast_divmod* output_div_pitches,
const float* scales,
const T* input_data,
T* output_data,
const size_t N) {
CALCULATE_ELEMENTWISE_INDEX_OR_EXIT(id, N);
CUDA_LONG input_index = 0;
CUDA_LONG output_index = id;
int div, mod;
for (int dim = 0; dim < rank; ++dim) {
output_div_pitches[dim].divmod(output_index, div, mod);
output_index = mod;
if (scales[dim] <= 1) { //downsample
div = std::ceil(div / scales[dim]);
} else { //upsample
div = div / scales[dim];
}
input_index += input_pitches[dim] * div;
}
output_data[id] = input_data[input_index];
}
template <typename T>
__global__ void _ResizeBilinearKernel(const int64_t input_dim2,
const int64_t* input_pitches,
const fast_divmod* output_div_pitches,
const float* scales,
const T* input_data,
T* output_data,
const size_t N) {
CALCULATE_ELEMENTWISE_INDEX_OR_EXIT(id, N);
CUDA_LONG input_index = 0;
// For bilinear mode, scales[0]=scales[1]=1
int mod;
int index_of_dim0, index_of_dim1, index_of_dim2, index_of_dim3;
output_div_pitches[0].divmod(id, index_of_dim0, mod);
output_div_pitches[1].divmod(mod, index_of_dim1, mod);
output_div_pitches[2].divmod(mod, index_of_dim2, mod);
index_of_dim3 = mod;
int index_of_input_dim2, index_of_input_dim3;
float x_offset_0, y_offset_0, x_offset_1, y_offset_1;
index_of_input_dim2 = static_cast<int64_t>(index_of_dim2 / scales[2]);
index_of_input_dim3 = static_cast<int64_t>(index_of_dim3 / scales[3]);
input_index = index_of_dim0 * input_pitches[0] +
index_of_dim1 * input_pitches[1] +
index_of_input_dim2 * input_pitches[2] +
index_of_input_dim3;
T x00 = input_data[input_index];
T x10, x01, x11;
bool end_of_dim2 = false, end_of_dim3 = false;
if (index_of_input_dim2 == (input_dim2 - 1)) {
// It's the end in dimension 2
x01 = x00;
end_of_dim2 = true;
} else {
x01 = input_data[input_index + input_pitches[2]];
}
if (index_of_input_dim3 == (input_pitches[2] - 1)) {
// It's the end in dimension 3
x10 = x00;
x11 = x01;
end_of_dim3 = true;
} else {
x10 = input_data[input_index + 1];
x11 = end_of_dim2 ? x10 : input_data[input_index + input_pitches[2] + 1];
}
y_offset_0 = end_of_dim2 ? 0.5f : index_of_dim2 / scales[2] - index_of_input_dim2;
y_offset_1 = 1.0f - y_offset_0;
x_offset_0 = end_of_dim3 ? 0.5f : index_of_dim3 / scales[3] - index_of_input_dim3;
x_offset_1 = 1.0f - x_offset_0;
output_data[id] =
x00 * static_cast<T>(y_offset_1 * x_offset_1) +
x01 * static_cast<T>(y_offset_0 * x_offset_1) +
x10 * static_cast<T>(y_offset_1 * x_offset_0) +
x11 * static_cast<T>(y_offset_0 * x_offset_0);
}
template <typename T>
void ResizeImpl(const onnxruntime::UpsampleMode upsample_mode,
const size_t rank,
const int64_t input_dim2,
const int64_t* input_pitches,
const fast_divmod* output_div_pitches,
const float* scales_vals,
const T* input_data,
T* output_data,
const size_t N) {
int blocksPerGrid = (int)(ceil(static_cast<float>(N) / GridDim::maxThreadsPerBlock));
if (onnxruntime::UpsampleMode::NN == upsample_mode) {
_ResizeNearestKernel<T><<<blocksPerGrid, GridDim::maxThreadsPerBlock, 0>>>(
rank, input_pitches, output_div_pitches, scales_vals,
input_data, output_data, N);
} else if (onnxruntime::UpsampleMode::LINEAR == upsample_mode) {
_ResizeBilinearKernel<T><<<blocksPerGrid, GridDim::maxThreadsPerBlock, 0>>>(
input_dim2, input_pitches, output_div_pitches, scales_vals,
input_data, output_data, N);
}
}
#define SPECIALIZED_IMPL(T) \
template void ResizeImpl<T>(const onnxruntime::UpsampleMode upsample_mode, \
const size_t rank, \
const int64_t input_dim2, \
const int64_t* input_pitches, \
const fast_divmod* output_div_pitches, \
const float* scales_vals, \
const T* input_data, \
T* output_data, \
const size_t N);
SPECIALIZED_IMPL(float)
SPECIALIZED_IMPL(double)
SPECIALIZED_IMPL(half)
SPECIALIZED_IMPL(int32_t)
SPECIALIZED_IMPL(uint8_t)
} // namespace cuda
} // namespace onnxruntime

View file

@ -0,0 +1,25 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
#pragma once
#include <stdint.h>
#include "core/providers/cuda/shared_inc/cuda_utils.h"
#include "core/common/common.h"
#include "core/providers/cpu/tensor/resize.h"
namespace onnxruntime {
namespace cuda {
template <typename T>
void ResizeImpl(const onnxruntime::UpsampleMode upsample_mode,
const size_t rank,
const int64_t input_dim2,
const int64_t* input_pitches,
const fast_divmod* output_div_pitches,
const float* scales_vals,
const T* input_data,
T* output_data,
const size_t N);
} // namespace cuda
} // namespace onnxruntime

View file

@ -3,6 +3,7 @@
#include "upsample.h"
#include "upsample_impl.h"
#include "core/providers/cuda/tensor/resize_impl.h"
#include "core/providers/cpu/tensor/utils.h"
using namespace onnxruntime::common;
@ -32,6 +33,7 @@ REGISTER_KERNEL_TYPED(uint8_t)
template <typename T>
Status Upsample<T>::BaseCompute(OpKernelContext* context, const std::vector<float>& scales) const {
const Tensor* X = context->Input<Tensor>(0);
ORT_ENFORCE(nullptr != X);
const std::vector<int64_t>& X_dims = X->Shape().GetDims();
auto rank = X_dims.size();
@ -58,34 +60,55 @@ Status Upsample<T>::BaseCompute(OpKernelContext* context, const std::vector<floa
CudaAsyncBuffer<fast_divmod> output_div_pitches(this, device_id, rank);
gsl::span<fast_divmod> div_strides_span = output_div_pitches.CpuSpan();
CudaAsyncBuffer<fast_divmod> scales_div(this, device_id, rank);
gsl::span<fast_divmod> scales_div_span = scales_div.CpuSpan();
for (int i = 0; i < rank; ++i) {
input_stride_span[i] = input_pitches[i];
div_strides_span[i] = fast_divmod(gsl::narrow_cast<int>(output_pitches[i]));
scales_div_span[i] = fast_divmod(gsl::narrow_cast<int>(ceil(scales[i])));
}
input_strides.CopyToGpu();
output_div_pitches.CopyToGpu();
scales_div.CopyToGpu();
size_t output_count = Y->Shape().Size();
if (UpsampleMode::LINEAR == mode_) {
if (rank != 4)
return Status(ONNXRUNTIME, FAIL, "Upsample: linear mode upsample only support 4-D tensor with NCHW layout");
if (is_resize) {
return Status(ONNXRUNTIME, FAIL, "Resize: linear mode only supports 4-D tensor with NCHW layout");
} else {
return Status(ONNXRUNTIME, FAIL, "Upsample: linear mode only supports 4-D tensor with NCHW layout");
}
}
UpampleImpl(mode_,
rank,
(UpsampleMode::LINEAR == mode_) ? X_dims[2] : 0,
input_strides.GpuPtr(),
output_div_pitches.GpuPtr(),
scales_div.GpuPtr(),
reinterpret_cast<const CudaT*>(X->template Data<T>()),
reinterpret_cast<CudaT*>(Y->template MutableData<T>()),
output_count);
if (is_resize) {
CudaAsyncBuffer<float> scales_vals(this, device_id, scales);
scales_vals.CopyToGpu();
ResizeImpl(mode_,
rank,
(UpsampleMode::LINEAR == mode_) ? X_dims[2] : 0,
input_strides.GpuPtr(),
output_div_pitches.GpuPtr(),
scales_vals.GpuPtr(),
reinterpret_cast<const CudaT*>(X->template Data<T>()),
reinterpret_cast<CudaT*>(Y->template MutableData<T>()),
output_count);
} else {
CudaAsyncBuffer<fast_divmod> scales_div(this, device_id, rank);
gsl::span<fast_divmod> scales_div_span = scales_div.CpuSpan();
for (int i = 0; i < rank; ++i) {
scales_div_span[i] = fast_divmod(gsl::narrow_cast<int>(ceil(scales[i])));
}
scales_div.CopyToGpu();
UpampleImpl(mode_,
rank,
(UpsampleMode::LINEAR == mode_) ? X_dims[2] : 0,
input_strides.GpuPtr(),
output_div_pitches.GpuPtr(),
scales_div.GpuPtr(),
reinterpret_cast<const CudaT*>(X->template Data<T>()),
reinterpret_cast<CudaT*>(Y->template MutableData<T>()),
output_count);
}
return Status::OK();
}

View file

@ -21,5 +21,17 @@ void UpampleImpl(const onnxruntime::UpsampleMode upsample_mode,
T* output_data,
const size_t N);
template <typename T>
void ResizeImpl(
const onnxruntime::UpsampleMode upsample_mode,
int64_t batch_size,
int64_t num_channels,
int64_t input_height,
int64_t input_width,
float height_scale,
float width_scale,
const T* Xdata,
T* Ydata,
const size_t N);
} // namespace cuda
} // namespace onnxruntime

View file

@ -7,6 +7,46 @@
namespace onnxruntime {
namespace test {
TEST(ResizeOpTest, ResizeOpLineartDownSampleTest) {
OpTester test("Resize", 10);
std::vector<float> scales{1.0f, 1.0f, 0.6f, 0.6f};
test.AddAttribute("mode", "linear");
const int64_t N = 1, C = 1, H = 2, W = 4;
std::vector<float> X = {
1.0f, 2.0f, 3.0f, 4.0f,
5.0f, 6.0f, 7.0f, 8.0f};
test.AddInput<float>("X", {N, C, H, W}, X);
test.AddInput<float>("scales", {4}, scales);
std::vector<float> Y = {1.0f, 2.66666651f};
test.AddOutput<float>("Y", {N, C, (int64_t)(H * scales[2]), (int64_t)(W * scales[3])}, Y);
test.Run();
}
TEST(ResizeOpTest, ResizeOpUpsampleNearestTest) {
OpTester test("Resize", 10);
std::vector<float> scales{1.0f, 1.0f, 2.0f, 3.0f};
test.AddAttribute("mode", "nearest");
const int64_t N = 1, C = 1, H = 2, W = 2;
std::vector<float> X = {1.0f, 2.0f, 3.0f, 4.0f};
test.AddInput<float>("X", {N, C, H, W}, X);
test.AddInput<float>("scales", {4}, scales);
std::vector<float> Y = {1.0f, 1.0f, 1.0f, 2.0f, 2.0f, 2.0f,
1.0f, 1.0f, 1.0f, 2.0f, 2.0f, 2.0f,
3.0f, 3.0f, 3.0f, 4.0f, 4.0f, 4.0f,
3.0f, 3.0f, 3.0f, 4.0f, 4.0f, 4.0f};
test.AddOutput<float>("Y", {N, C, (int64_t)(H * scales[2]), (int64_t)(W * scales[3])}, Y);
test.Run();
}
TEST(ResizeOpTest, ResizeOpNearestTest) {
OpTester test("Resize", 10);