mirror of
https://github.com/saymrwulf/onnxruntime.git
synced 2026-07-23 19:32:23 +00:00
Upsample opset 9 cuda implementation (#330)
This commit is contained in:
parent
0efc48a11a
commit
779123cf55
5 changed files with 57 additions and 40 deletions
|
|
@ -17,19 +17,32 @@ enum UpsampleMode {
|
|||
|
||||
class UpsampleBase {
|
||||
protected:
|
||||
UpsampleBase(OpKernelInfo info) {
|
||||
UpsampleBase(OpKernelInfo info): scales_cached_(false) {
|
||||
std::string mode;
|
||||
ORT_ENFORCE(info.GetAttr<std::string>("mode", &mode).IsOK());
|
||||
mode_ = StringToUpsampleMode(mode);
|
||||
|
||||
if (info.GetInputCount() == 1) {
|
||||
auto input_count = info.GetInputCount();
|
||||
if (input_count == 1) {
|
||||
ORT_ENFORCE(info.GetAttrs<float>("scales", scales_).IsOK());
|
||||
ScalesValidation(scales_, mode_);
|
||||
}
|
||||
|
||||
// opset 9
|
||||
if (input_count > 1) {
|
||||
const Tensor* scale;
|
||||
bool get_scale = info.TryGetConstantInput(1, &scale);
|
||||
|
||||
if (get_scale) {
|
||||
ParseScalesData(scale, scales_);
|
||||
scales_cached_ = true;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
UpsampleMode mode_;
|
||||
std::vector<float> scales_;
|
||||
bool scales_cached_;
|
||||
|
||||
UpsampleMode StringToUpsampleMode(const std::string& mode) {
|
||||
if (strcmp(mode.c_str(), UpsampleModeNN) == 0) {
|
||||
|
|
@ -53,28 +66,7 @@ class UpsampleBase {
|
|||
"Upsample: linear mode upsample only support bilinear, the first 2 scales should be 1.");
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
template <typename T>
|
||||
class Upsample : public UpsampleBase, public OpKernel {
|
||||
public:
|
||||
Upsample(OpKernelInfo info) : UpsampleBase(info), OpKernel(info), scales_cached_(false) {
|
||||
if (info.GetInputCount() > 1) {
|
||||
const Tensor* scale;
|
||||
bool get_scale = info.TryGetConstantInput(1, &scale);
|
||||
|
||||
if (get_scale) {
|
||||
ParseScalesData(scale, scales_);
|
||||
scales_cached_ = true;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Status Compute(OpKernelContext* context) const override;
|
||||
|
||||
Status BaseCompute(OpKernelContext* context, const std::vector<float>& scales) const;
|
||||
|
||||
private:
|
||||
void ParseScalesData(const Tensor* scale, std::vector<float>& scales) const {
|
||||
const float* scale_data = scale->template Data<float>();
|
||||
int64_t scales_size = scale->Shape().Size();
|
||||
|
|
@ -82,9 +74,17 @@ private:
|
|||
memcpy(scales.data(), scale_data, scales_size * sizeof(float));
|
||||
ScalesValidation(scales, mode_);
|
||||
}
|
||||
};
|
||||
|
||||
private:
|
||||
bool scales_cached_;
|
||||
template <typename T>
|
||||
class Upsample : public UpsampleBase, public OpKernel {
|
||||
public:
|
||||
Upsample(OpKernelInfo info) : UpsampleBase(info), OpKernel(info) {
|
||||
}
|
||||
|
||||
Status Compute(OpKernelContext* context) const override;
|
||||
|
||||
Status BaseCompute(OpKernelContext* context, const std::vector<float>& scales) const;
|
||||
};
|
||||
|
||||
} // namespace onnxruntime
|
||||
|
|
|
|||
|
|
@ -498,10 +498,10 @@ class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain,
|
|||
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 1, int32_t, DynamicSlice);
|
||||
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 1, int64_t, DynamicSlice);
|
||||
class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 9, Compress);
|
||||
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 7, float, Upsample);
|
||||
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 7, double, Upsample);
|
||||
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 7, MLFloat16, Upsample);
|
||||
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 7, int32_t, Upsample);
|
||||
class ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 7, 9, float, Upsample);
|
||||
class ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 7, 9, double, Upsample);
|
||||
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);
|
||||
|
||||
static void RegisterCudaKernels(std::function<void(KernelCreateInfo&&)> fn) {
|
||||
fn(BuildKernel<ONNX_OPERATOR_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 1, MemcpyFromHost)>());
|
||||
|
|
@ -761,10 +761,10 @@ static void RegisterCudaKernels(std::function<void(KernelCreateInfo&&)> fn) {
|
|||
fn(BuildKernel<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 1, int32_t, DynamicSlice)>());
|
||||
fn(BuildKernel<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 1, int64_t, DynamicSlice)>());
|
||||
fn(BuildKernel<ONNX_OPERATOR_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 9, Compress)>());
|
||||
fn(BuildKernel<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 7, float, Upsample)>());
|
||||
fn(BuildKernel<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 7, double, Upsample)>());
|
||||
fn(BuildKernel<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 7, MLFloat16, Upsample)>());
|
||||
fn(BuildKernel<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 7, int32_t, Upsample)>());
|
||||
fn(BuildKernel<ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 7, 9, float, Upsample)>());
|
||||
fn(BuildKernel<ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 7, 9, double, Upsample)>());
|
||||
fn(BuildKernel<ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 7, 9, MLFloat16, Upsample)>());
|
||||
fn(BuildKernel<ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 7, 9, int32_t, Upsample)>());
|
||||
}
|
||||
|
||||
} // namespace cuda
|
||||
|
|
|
|||
|
|
@ -11,13 +11,15 @@ namespace onnxruntime {
|
|||
namespace cuda {
|
||||
|
||||
#define REGISTER_KERNEL_TYPED(T) \
|
||||
ONNX_OPERATOR_TYPED_KERNEL_EX( \
|
||||
ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_EX( \
|
||||
Upsample, \
|
||||
kOnnxDomain, \
|
||||
7, \
|
||||
9, \
|
||||
T, \
|
||||
kCudaExecutionProvider, \
|
||||
KernelDefBuilder() \
|
||||
.InputMemoryType<OrtMemTypeCPUInput>(1) \
|
||||
.TypeConstraint("T", DataTypeImpl::GetTensorType<T>()), \
|
||||
Upsample<T>);
|
||||
|
||||
|
|
@ -27,7 +29,7 @@ REGISTER_KERNEL_TYPED(MLFloat16)
|
|||
REGISTER_KERNEL_TYPED(int32_t)
|
||||
|
||||
template <typename T>
|
||||
Status Upsample<T>::ComputeInternal(OpKernelContext* context) const {
|
||||
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();
|
||||
|
|
@ -35,12 +37,12 @@ Status Upsample<T>::ComputeInternal(OpKernelContext* context) const {
|
|||
if (rank == 0)
|
||||
return Status(ONNXRUNTIME, INVALID_ARGUMENT, "Upsample: input tensor cannot be scalar.");
|
||||
|
||||
if (rank != scales_.size())
|
||||
if (rank != scales.size())
|
||||
return Status(ONNXRUNTIME, INVALID_ARGUMENT, "Upsample: input tensor's dimension does not match the scales.");
|
||||
|
||||
std::vector<int64_t> Y_dims;
|
||||
for (std::size_t i = 0; i < rank; i++) {
|
||||
Y_dims.push_back(static_cast<int64_t>(scales_[i] * X_dims[i]));
|
||||
Y_dims.push_back(static_cast<int64_t>(scales[i] * X_dims[i]));
|
||||
}
|
||||
Tensor* Y = context->Output(0, Y_dims);
|
||||
typedef typename ToCudaType<T>::MappedType CudaT;
|
||||
|
|
@ -55,14 +57,13 @@ Status Upsample<T>::ComputeInternal(OpKernelContext* context) const {
|
|||
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])));
|
||||
scales_div_span[i] = fast_divmod(gsl::narrow_cast<int>(ceil(scales[i])));
|
||||
}
|
||||
input_strides.CopyToGpu();
|
||||
output_div_pitches.CopyToGpu();
|
||||
|
|
@ -88,5 +89,21 @@ Status Upsample<T>::ComputeInternal(OpKernelContext* context) const {
|
|||
return Status::OK();
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
Status Upsample<T>::ComputeInternal(OpKernelContext* context) const {
|
||||
// Opset 7
|
||||
if (OpKernel::Node().InputDefs().size() == 1 || scales_cached_) {
|
||||
return BaseCompute(context, scales_);
|
||||
}
|
||||
|
||||
// Opset 9
|
||||
const Tensor* scales = context->Input<Tensor>(1);
|
||||
ORT_ENFORCE(scales != nullptr);
|
||||
int64_t scales_size = scales->Shape().Size();
|
||||
std::vector<float> scales_arrary(scales_size);
|
||||
ParseScalesData(scales, scales_arrary);
|
||||
return BaseCompute(context, scales_arrary);
|
||||
}
|
||||
|
||||
} // namespace cuda
|
||||
} // namespace onnxruntime
|
||||
|
|
|
|||
|
|
@ -17,6 +17,7 @@ class Upsample : public UpsampleBase, public CudaKernel {
|
|||
}
|
||||
|
||||
Status ComputeInternal(OpKernelContext* context) const override;
|
||||
Status BaseCompute(OpKernelContext* context, const std::vector<float>& scales) const;
|
||||
};
|
||||
|
||||
} // namespace cuda
|
||||
|
|
|
|||
|
|
@ -304,7 +304,6 @@ int real_main(int argc, char* argv[]) {
|
|||
{"operator_rnn_single_layer", "disable reason"},
|
||||
{"prelu_broadcast", "disable reason"},
|
||||
{"prelu_example", "disable reason"},
|
||||
{"upsample_nearest", "opset 9 not supported yet"},
|
||||
{"sinh_example", "opset 9 not supported yet"},
|
||||
{"cosh_example", "opset 9 not supported yet"},
|
||||
{"asinh_example", "opset 9 not supported yet"},
|
||||
|
|
|
|||
Loading…
Reference in a new issue