mirror of
https://github.com/saymrwulf/onnxruntime.git
synced 2026-07-29 20:14:01 +00:00
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:
parent
ded7eeb033
commit
1658efd953
9 changed files with 308 additions and 16 deletions
|
|
@ -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 {
|
||||
|
|
|
|||
|
|
@ -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) {
|
||||
|
|
|
|||
27
onnxruntime/core/providers/cuda/tensor/resize.cc
Normal file
27
onnxruntime/core/providers/cuda/tensor/resize.cc
Normal 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
|
||||
23
onnxruntime/core/providers/cuda/tensor/resize.h
Normal file
23
onnxruntime/core/providers/cuda/tensor/resize.h
Normal 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
|
||||
133
onnxruntime/core/providers/cuda/tensor/resize_impl.cu
Normal file
133
onnxruntime/core/providers/cuda/tensor/resize_impl.cu
Normal 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
|
||||
25
onnxruntime/core/providers/cuda/tensor/resize_impl.h
Normal file
25
onnxruntime/core/providers/cuda/tensor/resize_impl.h
Normal 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
|
||||
|
|
@ -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();
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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);
|
||||
|
|
|
|||
Loading…
Reference in a new issue