From d71376010b5a7d511550abcd6d953c11eeced8a7 Mon Sep 17 00:00:00 2001 From: Zhang Lei Date: Mon, 10 Dec 2018 10:12:53 -0800 Subject: [PATCH] Implement Dynamic Range Operator with shape inference Implement dynamic range operator in custom ops, with shape inference. --- onnxruntime/contrib_ops/contrib_ops.cc | 4 + onnxruntime/contrib_ops/cpu/range.cc | 88 ++++++++++ onnxruntime/contrib_ops/cpu/range.h | 21 +++ .../contrib_ops/cpu/range_schema_defs.cc | 160 ++++++++++++++++++ .../contrib_ops/cpu/range_schema_defs.h | 22 +++ onnxruntime/test/contrib_ops/range_test.cc | 94 ++++++++++ 6 files changed, 389 insertions(+) create mode 100644 onnxruntime/contrib_ops/cpu/range.cc create mode 100644 onnxruntime/contrib_ops/cpu/range.h create mode 100644 onnxruntime/contrib_ops/cpu/range_schema_defs.cc create mode 100644 onnxruntime/contrib_ops/cpu/range_schema_defs.h create mode 100644 onnxruntime/test/contrib_ops/range_test.cc diff --git a/onnxruntime/contrib_ops/contrib_ops.cc b/onnxruntime/contrib_ops/contrib_ops.cc index c2c8d3318d..43f0c3f8d7 100644 --- a/onnxruntime/contrib_ops/contrib_ops.cc +++ b/onnxruntime/contrib_ops/contrib_ops.cc @@ -8,6 +8,7 @@ #include "onnx/defs/schema.h" #include "./cpu/attnlstm/attn_lstm_schema_defs.h" +#include "./cpu/range_schema_defs.h" namespace onnxruntime { namespace contrib { @@ -75,6 +76,7 @@ Sample echo operator.)DOC"); .SetDoc(R"DOC(ExpandDims echo operator.)DOC"); ONNX_CONTRIB_OPERATOR_SCHEMA_ELSEWHERE(AttnLSTM, RegisterAttnLSTMContribOpSchema); + ONNX_CONTRIB_OPERATOR_SCHEMA_ELSEWHERE(Range, RegisterRangeOpSchema); ONNX_CONTRIB_OPERATOR_SCHEMA(IsNaN) .SetDomain(kMSDomain) @@ -527,6 +529,7 @@ class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, float, QuantizeLinear); class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, string, StringNormalizer); class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, float, NonMaxSuppression); +class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, Range); void RegisterContribKernels(std::function fn) { fn(BuildKernel()); @@ -542,6 +545,7 @@ void RegisterContribKernels(std::function fn) { fn(BuildKernel()); fn(BuildKernel()); fn(BuildKernel()); + fn(BuildKernel()); } } // namespace contrib } // namespace onnxruntime diff --git a/onnxruntime/contrib_ops/cpu/range.cc b/onnxruntime/contrib_ops/cpu/range.cc new file mode 100644 index 0000000000..ec0fc8478b --- /dev/null +++ b/onnxruntime/contrib_ops/cpu/range.cc @@ -0,0 +1,88 @@ +#include "range.h" +#include "onnx/defs/schema.h" + +#include + +namespace onnxruntime { +namespace contrib { + +template +static Status ComputeRange(OpKernelContext* ctx) { + auto& start_tensor = *ctx->Input(0); + auto& limit_tensor = *ctx->Input(1); + auto delta_tensor_ptr = ctx->Input(2); + + if (!start_tensor.Shape().IsScalar()) { + return ONNXRUNTIME_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "start in Range operator should be scalar like tensor, yet got shape:", + start_tensor.Shape()); + } + if (!limit_tensor.Shape().IsScalar()) { + return ONNXRUNTIME_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "limit in Range operator should be scalar like tensor, yet got shape:", + limit_tensor.Shape()); + } + if (delta_tensor_ptr != nullptr && !delta_tensor_ptr->Shape().IsScalar()) { + return ONNXRUNTIME_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "delta in Range operator should be scalar like tensor, yet got shape:", + delta_tensor_ptr->Shape()); + } + + T start = *start_tensor.template Data(); + T limit = *limit_tensor.template Data(); + T delta = (delta_tensor_ptr == nullptr) ? T{1} : *(delta_tensor_ptr->template Data()); + + if (delta == T{0}) { + return ONNXRUNTIME_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "delta in Range operator can not be zero!"); + } + int64_t n = static_cast(ceil((1.0 * (limit - start)) / delta)); + if (n <= 0) n = 1; + TensorShape shape = {n}; + T* y = ctx->Output(0, shape)->template MutableData(); + for (int64_t i = 0; i < n; ++i) { + *y++ = start; + start += delta; + } + + return Status::OK(); +} + +Status Range::Compute(OpKernelContext* ctx) const { + auto data_type = ctx->Input(0)->DataType(); + if (data_type == DataTypeImpl::GetType()) { + return ComputeRange(ctx); + } + else if (data_type == DataTypeImpl::GetType()) { + return ComputeRange(ctx); + } + else if (data_type == DataTypeImpl::GetType()) { + return ComputeRange(ctx); + } + else if (data_type == DataTypeImpl::GetType()) { + return ComputeRange(ctx); + } + else if (data_type == DataTypeImpl::GetType()) { + return ComputeRange(ctx); + } + return ONNXRUNTIME_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "Unsupportted tensor data type:", + data_type); +} + +/* Range operator */ +ONNX_OPERATOR_KERNEL_EX( + Range, //name + kMSDomain, + 1, + kCpuExecutionProvider, + KernelDefBuilder().TypeConstraint("T", { + DataTypeImpl::GetTensorType(), + DataTypeImpl::GetTensorType(), + DataTypeImpl::GetTensorType(), + DataTypeImpl::GetTensorType(), + DataTypeImpl::GetTensorType()}), + Range); + + +} // namespace contrib +} // namespace onnxruntime diff --git a/onnxruntime/contrib_ops/cpu/range.h b/onnxruntime/contrib_ops/cpu/range.h new file mode 100644 index 0000000000..5d7ef25238 --- /dev/null +++ b/onnxruntime/contrib_ops/cpu/range.h @@ -0,0 +1,21 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once + +#include "core/common/common.h" +#include "core/framework/op_kernel.h" +#include "core/framework/tensor.h" + +namespace onnxruntime { +namespace contrib { + +class Range : public OpKernel { + public: + explicit Range(const OpKernelInfo& info) : OpKernel(info) {} + + Status Compute(OpKernelContext* context) const override; +}; + +} // namespace contrib +} // namespace onnxruntime diff --git a/onnxruntime/contrib_ops/cpu/range_schema_defs.cc b/onnxruntime/contrib_ops/cpu/range_schema_defs.cc new file mode 100644 index 0000000000..cd9f96e99c --- /dev/null +++ b/onnxruntime/contrib_ops/cpu/range_schema_defs.cc @@ -0,0 +1,160 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "range_schema_defs.h" + +#include "core/graph/constants.h" +#include "core/graph/op.h" + +#include + +namespace onnxruntime { +namespace contrib { + +using ::ONNX_NAMESPACE::OPTIONAL; +using ::ONNX_NAMESPACE::OpSchema; +using ::ONNX_NAMESPACE::InferenceContext; +using ::ONNX_NAMESPACE::TensorShapeProto; +using ::ONNX_NAMESPACE::TensorProto; +using ::ONNX_NAMESPACE::TensorProto_DataType; + +// This Doc based on LSTM_ver7, and modification +static const char* Range_ver1_doc = R"DOC( +Creates a sequence of numbers that begins at `start` and extends by increments of `delta` +up to but not including `limit`. +)DOC"; + +template +static T get_data(const TensorProto*) { + fail_shape_inference("Unsupported non-raw-data data type!"); +} + +template <> +int32_t get_data(const TensorProto* shapeInitializer) { + if (shapeInitializer->int32_data_size() > 0) return shapeInitializer->int32_data(0); + fail_shape_inference("Can not get shape initializer data!"); +} + +template <> +int64_t get_data(const TensorProto* shapeInitializer) { + if (shapeInitializer->int64_data_size() > 0) return shapeInitializer->int64_data(0); + fail_shape_inference("Can not get shape initializer data!"); +} + +template <> +float get_data(const TensorProto* shapeInitializer) { + if (shapeInitializer->float_data_size() > 0) return shapeInitializer->float_data(0); + fail_shape_inference("Can not get shape initializer data!"); +} + +template <> +double get_data(const TensorProto* shapeInitializer) { + if (shapeInitializer->double_data_size() > 0) return shapeInitializer->double_data(0); + fail_shape_inference("Can not get shape initializer data!"); +} + +template +static T GetFirstElement(const TensorProto* shapeInitializer) { + if (shapeInitializer == nullptr) return T{1}; + + if (shapeInitializer->has_raw_data()) { + const std::string& bytes = shapeInitializer->raw_data(); + return *reinterpret_cast(bytes.c_str()); + } else { + return get_data(shapeInitializer); + } +} + +template +static int64_t CalcRangeDim(const TensorProto* startShapeInitializer, + const TensorProto* limitShapeInitializer, + const TensorProto* deltaShapeInitializer) { + T start = GetFirstElement(startShapeInitializer); + T limit = GetFirstElement(limitShapeInitializer); + T delta = GetFirstElement(deltaShapeInitializer); + if (delta == T{0}) { + fail_shape_inference("delta in Range operator can not be zero!"); + } + return static_cast(ceil((1.0 * (limit - start)) / delta)); +} + +static int64_t CalcResultDim(const TensorProto* startShapeInitializer, + const TensorProto* limitShapeInitializer, + const TensorProto* deltaShapeInitializer, + TensorProto_DataType dtype) { + int64_t dim = -1LL; + if (dtype == TensorProto::FLOAT) { + dim = CalcRangeDim(startShapeInitializer, limitShapeInitializer, deltaShapeInitializer); + } + else if (dtype == TensorProto::INT32) { + dim = CalcRangeDim(startShapeInitializer, limitShapeInitializer, deltaShapeInitializer); + } + else if (dtype == TensorProto::INT64) { + dim = CalcRangeDim(startShapeInitializer, limitShapeInitializer, deltaShapeInitializer); + } + else if (dtype == TensorProto::INT16) { + dim = CalcRangeDim(startShapeInitializer, limitShapeInitializer, deltaShapeInitializer); + } + else if (dtype == TensorProto::DOUBLE) { + dim = CalcRangeDim(startShapeInitializer, limitShapeInitializer, deltaShapeInitializer); + } + else { + fail_shape_inference("Unsupported type:", dtype); + } + return dim; +} + +OpSchema& RegisterRangeOpSchema(OpSchema&& op_schema){ + return op_schema + .SetDomain(kMSDomain) + .SinceVersion(1) + .TypeConstraint( + "T", + {"tensor(float)", "tensor(double)", "tensor(int16)", "tensor(int32)", "tensor(int64)" }, + "Constrain input and output types.") + .Input( + 0, + "start", + "Tensor(scalar, or dims=[1]). First entry in the range.", + "T") + .Input( + 1, + "limit", + "Tensor(scalar, or dims=[1]). Upper limit of sequence, exclusive.", + "T") + .Input( + 2, + "delta", + "Tensor(scalar, or dims=[1]). Number that increments start. Defaults to 1.", + "T", + OpSchema::Optional) + .Output( + 0, + "Y", + "1-D Tensor of the range.", + "T") + .SetDoc(Range_ver1_doc) + .TypeAndShapeInferenceFunction([](InferenceContext& ctx) { + propagateElemTypeFromInputToOutput(ctx, 0, 0); + + TensorShapeProto::Dimension dim; + bool enoughShapeInfo = ctx.getInputData(0) != nullptr + && ctx.getInputData(1) != nullptr + && (ctx.getNumInputs() == 2 || ctx.getInputData(2) != nullptr); + if (enoughShapeInfo) { + const TensorProto* startShapeInitializer = ctx.getInputData(0); + const TensorProto* limitShapeInitializer = ctx.getInputData(1); + const TensorProto* deltaShapeInitializer = (ctx.getNumInputs() > 2) ? ctx.getInputData(2) : nullptr; + const auto& startTensorType = ctx.getInputType(0)->tensor_type(); + TensorProto_DataType dtype = startTensorType.elem_type(); + + int64_t n = CalcResultDim(startShapeInitializer, limitShapeInitializer, deltaShapeInitializer, dtype); + dim.set_dim_value(n); + } // else unknown 1-D tensor + + updateOutputShape(ctx, 0, {dim}); + }); +} + +} // namespace contrib +} // namespace onnxruntime diff --git a/onnxruntime/contrib_ops/cpu/range_schema_defs.h b/onnxruntime/contrib_ops/cpu/range_schema_defs.h new file mode 100644 index 0000000000..3029882756 --- /dev/null +++ b/onnxruntime/contrib_ops/cpu/range_schema_defs.h @@ -0,0 +1,22 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once + +#ifdef __GNUC__ +#pragma GCC diagnostic push +#pragma GCC diagnostic ignored "-Wignored-qualifiers" +#pragma GCC diagnostic ignored "-Wunused-parameter" +#endif +#include "onnx/defs/schema.h" +#ifdef __GNUC__ +#pragma GCC diagnostic pop +#endif + +namespace onnxruntime { +namespace contrib { + +::ONNX_NAMESPACE::OpSchema& RegisterRangeOpSchema(::ONNX_NAMESPACE::OpSchema&& op_schema); + +} // namespace contrib +} // namespace onnxruntime diff --git a/onnxruntime/test/contrib_ops/range_test.cc b/onnxruntime/test/contrib_ops/range_test.cc new file mode 100644 index 0000000000..9c09faf83b --- /dev/null +++ b/onnxruntime/test/contrib_ops/range_test.cc @@ -0,0 +1,94 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "gtest/gtest.h" +#include "test/providers/provider_test_utils.h" + +namespace onnxruntime { +namespace test { + +TEST(RangeTest, PositiveInt32DeltaDefault) { + OpTester test("Range", 1, onnxruntime::kMSDomain); + std::vector start = {0}; + std::vector limit = {5}; + std::vector expected_output = {0, 1, 2, 3, 4}; + + test.AddInput("start", {1}, start); + test.AddInput("limit", {1}, limit); + test.AddOutput("Y", {5LL}, expected_output); + test.Run(); +} + +TEST(RangeTest, PositiveInt32Delta_0) { + OpTester test("Range", 1, onnxruntime::kMSDomain); + std::vector start = {0}; + std::vector limit = {10}; + std::vector delta = {2}; + std::vector expected_output = {0, 2, 4, 6, 8}; + + test.AddInput("start", {1}, start); + test.AddInput("limit", {1}, limit); + test.AddInput("delta", {1}, delta); + test.AddOutput("Y", {5LL}, expected_output); + test.Run(); +} + +TEST(RangeTest, PositiveInt32Delta_1) { + OpTester test("Range", 1, onnxruntime::kMSDomain); + std::vector start = {0}; + std::vector limit = {9}; + std::vector delta = {2}; + std::vector expected_output = {0, 2, 4, 6, 8}; + + test.AddInput("start", {1}, start); + test.AddInput("limit", {1}, limit); + test.AddInput("delta", {1}, delta); + test.AddOutput("Y", {5LL}, expected_output); + test.Run(); +} + +TEST(RangeTest, Int32ScalarNegativeDelta_0) { + OpTester test("Range", 1, onnxruntime::kMSDomain); + std::vector start = {2}; + std::vector limit = {-9}; + std::vector delta = {-2}; + std::vector expected_output = {2, 0, -2, -4, -6, -8}; + + test.AddInput("start", {}, start); + test.AddInput("limit", {}, limit); + test.AddInput("delta", {}, delta); + test.AddOutput("Y", {6LL}, expected_output); + test.Run(); +} + + +TEST(RangeTest, Int32ScalarNegativeDelta_1) { + OpTester test("Range", 1, onnxruntime::kMSDomain); + std::vector start = {2}; + std::vector limit = {-9}; + std::vector delta = {-2}; + std::vector expected_output = {2, 0, -2, -4, -6, -8}; + + test.AddInput("start", {}, start); + test.AddInput("limit", {}, limit); + test.AddInput("delta", {}, delta); + test.AddOutput("Y", {6LL}, expected_output); + test.Run(); +} + +TEST(RangeTest, ScalarFloatNegativeDelta) { + OpTester test("Range", 1, onnxruntime::kMSDomain); + std::vector start = {2.0f}; + std::vector limit = {-8.1f}; + std::vector delta = {-2.0f}; + std::vector expected_output = {2.0f, 0.0f, -2.0f, -4.0f, -6.0f, -8.0f}; + + test.AddInput("start", {}, start); + test.AddInput("limit", {}, limit); + test.AddInput("delta", {}, delta); + test.AddOutput("Y", {6LL}, expected_output); + test.Run(); +} + +} // namespace test +} // namespace onnxruntime