Implement Dynamic Range Operator with shape inference

Implement dynamic range operator in custom ops, with shape inference.

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<void(KernelCreateInfo&&)> fn) {
   fn(BuildKernel<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, float, SampleOp)>());
@@ -542,6 +545,7 @@ void RegisterContribKernels(std::function<void(KernelCreateInfo&&)> fn) {
   fn(BuildKernel<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, float, QuantizeLinear)>());
   fn(BuildKernel<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, string, StringNormalizer)>());
   fn(BuildKernel<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, float, NonMaxSuppression)>());
+  fn(BuildKernel<ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, Range)>());
 }
 }  // namespace contrib
 }  // namespace onnxruntime

onnxruntime/contrib_ops/cpu/range.cc

@@ -0,0 +1,88 @@
+#include "range.h"
+#include "onnx/defs/schema.h"
+
+#include <cmath>
+
+namespace onnxruntime {
+namespace contrib {
+
+template <typename T>
+static Status ComputeRange(OpKernelContext* ctx) {
+  auto& start_tensor = *ctx->Input<Tensor>(0);
+  auto& limit_tensor = *ctx->Input<Tensor>(1);
+  auto  delta_tensor_ptr = ctx->Input<Tensor>(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>();
+  T limit = *limit_tensor.template Data<T>();
+  T delta = (delta_tensor_ptr == nullptr) ? T{1} : *(delta_tensor_ptr->template Data<T>());
+  
+  if (delta == T{0}) {
+    return ONNXRUNTIME_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,  "delta in Range operator can not be zero!");
+  }
+  int64_t n = static_cast<int64_t>(ceil((1.0 * (limit - start)) / delta));
+  if (n <= 0) n = 1;
+  TensorShape shape = {n};
+  T* y = ctx->Output(0, shape)->template MutableData<T>();
+  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<Tensor>(0)->DataType();
+  if (data_type == DataTypeImpl::GetType<int32_t>()) {
+      return ComputeRange<int32_t>(ctx);
+  }
+  else if (data_type == DataTypeImpl::GetType<int16_t>()) {
+      return ComputeRange<int16_t>(ctx);
+  }
+  else if (data_type == DataTypeImpl::GetType<int64_t>()) {
+      return ComputeRange<int64_t>(ctx);
+  }
+  else if (data_type == DataTypeImpl::GetType<float>()) {
+      return ComputeRange<float>(ctx);
+  }
+  else if (data_type == DataTypeImpl::GetType<double>()) {
+      return ComputeRange<double>(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<float>(), 
+        DataTypeImpl::GetTensorType<double>(),
+        DataTypeImpl::GetTensorType<int16_t>(), 
+        DataTypeImpl::GetTensorType<int32_t>(), 
+        DataTypeImpl::GetTensorType<int64_t>()}),
+    Range);
+
+
+}  // namespace contrib
+}  // namespace onnxruntime

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

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 <type_traits>
+
+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 <typename T>
+static T get_data(const TensorProto*) {
+    fail_shape_inference("Unsupported non-raw-data data type!");
+}
+
+template <>
+int32_t get_data<int32_t>(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<int64_t>(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<float>(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<double>(const TensorProto* shapeInitializer) {
+    if (shapeInitializer->double_data_size() > 0) return shapeInitializer->double_data(0);
+    fail_shape_inference("Can not get shape initializer data!");
+}
+
+template <typename T>
+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<const T*>(bytes.c_str());
+    } else {
+        return get_data<T>(shapeInitializer);
+    }
+}
+
+template <typename T>
+static int64_t CalcRangeDim(const TensorProto* startShapeInitializer,
+                            const TensorProto* limitShapeInitializer,
+                            const TensorProto* deltaShapeInitializer) {
+    T start = GetFirstElement<T>(startShapeInitializer);
+    T limit = GetFirstElement<T>(limitShapeInitializer);
+    T delta = GetFirstElement<T>(deltaShapeInitializer);
+    if (delta == T{0}) {
+        fail_shape_inference("delta in Range operator can not be zero!");
+    }
+    return static_cast<int64_t>(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<float>(startShapeInitializer, limitShapeInitializer, deltaShapeInitializer);
+    }
+    else if (dtype == TensorProto::INT32) {
+        dim = CalcRangeDim<int32_t>(startShapeInitializer, limitShapeInitializer, deltaShapeInitializer);
+    }
+    else if (dtype == TensorProto::INT64) {
+        dim = CalcRangeDim<int64_t>(startShapeInitializer, limitShapeInitializer, deltaShapeInitializer);
+    }
+    else if (dtype == TensorProto::INT16) {
+        dim = CalcRangeDim<int16_t>(startShapeInitializer, limitShapeInitializer, deltaShapeInitializer);
+    }
+    else if (dtype == TensorProto::DOUBLE) {
+        dim = CalcRangeDim<double>(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

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

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<int32_t> start = {0};
+  std::vector<int32_t> limit = {5};
+  std::vector<int32_t> expected_output = {0, 1, 2, 3, 4};
+
+  test.AddInput<int32_t>("start", {1}, start);
+  test.AddInput<int32_t>("limit", {1}, limit);
+  test.AddOutput<int32_t>("Y", {5LL}, expected_output);
+  test.Run();
+}
+
+TEST(RangeTest, PositiveInt32Delta_0) {
+  OpTester test("Range", 1, onnxruntime::kMSDomain);
+  std::vector<int32_t> start = {0};
+  std::vector<int32_t> limit = {10};
+  std::vector<int32_t> delta = {2};
+  std::vector<int32_t> expected_output = {0, 2, 4, 6, 8};
+
+  test.AddInput<int32_t>("start", {1}, start);
+  test.AddInput<int32_t>("limit", {1}, limit);
+  test.AddInput<int32_t>("delta", {1}, delta);
+  test.AddOutput<int32_t>("Y", {5LL}, expected_output);
+  test.Run();
+}
+
+TEST(RangeTest, PositiveInt32Delta_1) {
+  OpTester test("Range", 1, onnxruntime::kMSDomain);
+  std::vector<int32_t> start = {0};
+  std::vector<int32_t> limit = {9};
+  std::vector<int32_t> delta = {2};
+  std::vector<int32_t> expected_output = {0, 2, 4, 6, 8};
+
+  test.AddInput<int32_t>("start", {1}, start);
+  test.AddInput<int32_t>("limit", {1}, limit);
+  test.AddInput<int32_t>("delta", {1}, delta);
+  test.AddOutput<int32_t>("Y", {5LL}, expected_output);
+  test.Run();
+}
+
+TEST(RangeTest, Int32ScalarNegativeDelta_0) {
+  OpTester test("Range", 1, onnxruntime::kMSDomain);
+  std::vector<int32_t> start = {2};
+  std::vector<int32_t> limit = {-9};
+  std::vector<int32_t> delta = {-2};
+  std::vector<int32_t> expected_output = {2, 0, -2, -4, -6, -8};
+
+  test.AddInput<int32_t>("start", {}, start);
+  test.AddInput<int32_t>("limit", {}, limit);
+  test.AddInput<int32_t>("delta", {}, delta);
+  test.AddOutput<int32_t>("Y", {6LL}, expected_output);
+  test.Run();
+}
+
+
+TEST(RangeTest, Int32ScalarNegativeDelta_1) {
+  OpTester test("Range", 1, onnxruntime::kMSDomain);
+  std::vector<int32_t> start = {2};
+  std::vector<int32_t> limit = {-9};
+  std::vector<int32_t> delta = {-2};
+  std::vector<int32_t> expected_output = {2, 0, -2, -4, -6, -8};
+
+  test.AddInput<int32_t>("start", {}, start);
+  test.AddInput<int32_t>("limit", {}, limit);
+  test.AddInput<int32_t>("delta", {}, delta);
+  test.AddOutput<int32_t>("Y", {6LL}, expected_output);
+  test.Run();
+}
+
+TEST(RangeTest, ScalarFloatNegativeDelta) {
+  OpTester test("Range", 1, onnxruntime::kMSDomain);
+  std::vector<float> start = {2.0f};
+  std::vector<float> limit = {-8.1f};
+  std::vector<float> delta = {-2.0f};
+  std::vector<float> expected_output = {2.0f, 0.0f, -2.0f, -4.0f, -6.0f, -8.0f};
+  
+  test.AddInput<float>("start", {}, start);
+  test.AddInput<float>("limit", {}, limit);
+  test.AddInput<float>("delta", {}, delta);
+  test.AddOutput<float>("Y", {6LL}, expected_output);
+  test.Run();
+}
+
+}  // namespace test
+}  // namespace onnxruntime