saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-30 20:18:08 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Added GatherElements to Nuphar (#2016)

Browse source 
* Added GatherElements to Nuphar

This change added GatherElements (op_ver 11) to the Nuphar provider.

* address CR feedback

* create a utilify function for accessing index safely

* address more CR

* SafeIndex -> ClampIndex
This commit is contained in:
Yang Chen 2019-10-04 23:53:02 -07:00 committed by GitHub
parent 1ba76c5f74
commit e8285a7996
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
9 changed files with 127 additions and 2 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

1
onnxruntime/core/codegen/common/op_macro.h
View file

@ -79,6 +79,7 @@ namespace onnxruntime {
ADD_OP_ITEM(Dropout) \
ADD_OP_ITEM(Flatten) \
ADD_OP_ITEM(Gather) \
ADD_OP_ITEM(GatherElements) \
ADD_OP_ITEM(Gemm) \
ADD_OP_ITEM(Identity) \
ADD_OP_ITEM(LogSoftmax) \

8
onnxruntime/core/codegen/mti/mti_tvm_utils.cc
View file

@ -192,5 +192,13 @@ tvm::Array<tvm::Tensor> MakeInputsForExtern(const tvm::Array<tvm::Tensor>& input
return fixed_inputs;
}
// Make sure idx is clamped in the range of [-bound, bound - 1]
tvm::Expr ClampIndex(const tvm::Expr& idx, const tvm::Expr& bound) {
// when idx >= 0, we take tvm::max(..., 0), because (idx < 0) is 0
// when idx < 0, we take bound + tvm::max(...), because tvm::max(idx, 0) is 0
return tvm::max(tvm::min(idx, bound - 1), 0) +
(idx < 0) * (bound + tvm::max(idx, -bound));
}
} // namespace tvm_codegen
} // namespace onnxruntime

3
onnxruntime/core/codegen/mti/mti_tvm_utils.h
View file

@ -60,6 +60,9 @@ inline int64_t HandleNegativeAxis(int64_t axis, int64_t rank) {
return axis = axis < 0 ? (axis + rank) : axis;
}
// Make sure idx is clamped in the range of [-bound, bound - 1]
tvm::Expr ClampIndex(const tvm::Expr& idx, const tvm::Expr& bound);
// Helper function to workaround tvm ExternOp issue when input has symbolic dimensions
tvm::Array<tvm::Tensor> MakeInputsForExtern(const tvm::Array<tvm::Tensor>& inputs, const std::string& name = "make_inputs_for_extern");

45
onnxruntime/core/codegen/mti/tensor/gather_elements.cc Normal file
View file

@ -0,0 +1,45 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
#include "core/codegen/mti/tensor/gather_elements.h"
#include "core/codegen/mti/mti_tvm_utils.h"
#include <topi/transform.h>
namespace onnxruntime {
namespace tvm_codegen {
tvm::Tensor GatherElements(const tvm::Tensor& t,
int64_t axis,
const tvm::Tensor& indices,
const std::string& name) {
tvm::Array<tvm::Expr> output_shape;
int64_t indices_rank = static_cast<int64_t>(indices->shape.size());
// output shape is the same as indices
for (int64_t i = 0; i < indices_rank; ++i)
output_shape.push_back(indices->shape[i]);
tvm::Expr idx_upper_bound = t->shape[axis];
auto l = [&](const tvm::Array<tvm::Var>& ovars) {
tvm::Array<tvm::Expr> ivars;
for (int64_t i = 0; i < indices_rank; i++) {
if (i == axis) {
tvm::Array<tvm::Expr> idx_vars;
for (int64_t j = 0; j < indices_rank; j++)
idx_vars.push_back(ovars[j]);
// make sure idx is clamped in the range of [-idx_upper_bound, idx_upper_bound - 1]
tvm::Expr real_idx = tvm_codegen::ClampIndex(indices(idx_vars), idx_upper_bound);
// tvm idx must be of Int(32)
ivars.push_back(tvm::cast(tvm::Int(32), real_idx));
} else {
ivars.push_back(ovars[i]);
}
}
return t(ivars);
};
return tvm::compute(output_shape, l, name);
}
} // namespace tvm_codegen
} // namespace onnxruntime

17
onnxruntime/core/codegen/mti/tensor/gather_elements.h Normal file
View file

@ -0,0 +1,17 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
#pragma once
#include <string>
#include <tvm/tvm.h>
namespace onnxruntime {
namespace tvm_codegen {
tvm::Tensor GatherElements(const tvm::Tensor& t,
int64_t axis,
const tvm::Tensor& indices,
const std::string& name = "gather_elements");
} // namespace tvm_codegen
} // namespace onnxruntime

32
onnxruntime/core/codegen/passes/op_ir_creator/tensor/gather_elements.cc Normal file
View file

@ -0,0 +1,32 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
#include "core/codegen/passes/op_ir_creator/all_ops.h"
#include "core/codegen/mti/tensor/gather_elements.h"
#include "core/framework/op_kernel_info.h"
#include "core/providers/common.h"
namespace onnxruntime {
namespace tvm_codegen {
// Evaluate of GatherElements OpIRCreator
Status GENERIC_OP_IR_CREATOR_CLASS(GatherElements)::Evaluate(
const tvm::Array<tvm::Tensor>& inputs,
const Node& node,
CodeGenContext&,
tvm::Array<tvm::Tensor>& outputs) {
ProtoHelperNodeContext ctx(node);
OpNodeProtoHelper<ProtoHelperNodeContext> attrs(&ctx);
int64_t axis;
ORT_ENFORCE(attrs.GetAttr<int64_t>("axis", &axis).IsOK());
axis = HandleNegativeAxis(axis, gsl::narrow_cast<int64_t>(inputs[0]->shape.size()));
tvm::Tensor Y = GatherElements(inputs[0], axis, inputs[1], node.Name() + "_GatherElements");
outputs.push_back(Y);
return Status::OK();
}
} // namespace tvm_codegen
} // namespace onnxruntime

11
onnxruntime/core/providers/nuphar/kernel.cc
View file

@ -193,6 +193,17 @@ ONNX_OPERATOR_KERNEL_EX(
DataTypeImpl::GetTensorType<int64_t>()}),
nuphar::NupharKernel);
ONNX_OPERATOR_KERNEL_EX(
GatherElements,
kOnnxDomain,
11,
kNupharExecutionProvider,
KernelDefBuilder()
.TypeConstraint("T", DataTypeImpl::AllFixedSizeTensorTypes())
.TypeConstraint("Tind", std::vector<MLDataType>{DataTypeImpl::GetTensorType<int32_t>(),
DataTypeImpl::GetTensorType<int64_t>()}),
nuphar::NupharKernel);
ONNX_OPERATOR_KERNEL_EX(
MatMulInteger,
kOnnxDomain,

2
onnxruntime/core/providers/nuphar/nuphar_execution_provider.cc
View file

@ -391,6 +391,7 @@ LIST_NUPHAR_OPS()
class ONNX_OPERATOR_VERSIONED_KERNEL_CLASS_NAME(kNupharExecutionProvider, kOnnxDomain, 6, 8, Cast);
class ONNX_OPERATOR_KERNEL_CLASS_NAME(kNupharExecutionProvider, kOnnxDomain, 9, Cast);
class ONNX_OPERATOR_KERNEL_CLASS_NAME(kNupharExecutionProvider, kOnnxDomain, 1, Gather);
class ONNX_OPERATOR_KERNEL_CLASS_NAME(kNupharExecutionProvider, kOnnxDomain, 11, GatherElements);
class ONNX_OPERATOR_KERNEL_CLASS_NAME(kNupharExecutionProvider, kOnnxDomain, 10, MatMulInteger);
class ONNX_OPERATOR_KERNEL_CLASS_NAME(kNupharExecutionProvider, kMSDomain, 1, MatMulInteger16);
class ONNX_OPERATOR_KERNEL_CLASS_NAME(kNupharExecutionProvider, kOnnxDomain, 9, Scan);
@ -413,6 +414,7 @@ static void RegisterStandaloneNupharKernels(KernelRegistry& kernel_registry) {
kernel_registry.Register(BuildKernelCreateInfo<ONNX_OPERATOR_VERSIONED_KERNEL_CLASS_NAME(kNupharExecutionProvider, kOnnxDomain, 6, 8, Cast)>());
kernel_registry.Register(BuildKernelCreateInfo<ONNX_OPERATOR_KERNEL_CLASS_NAME(kNupharExecutionProvider, kOnnxDomain, 9, Cast)>());
kernel_registry.Register(BuildKernelCreateInfo<ONNX_OPERATOR_KERNEL_CLASS_NAME(kNupharExecutionProvider, kOnnxDomain, 1, Gather)>());
kernel_registry.Register(BuildKernelCreateInfo<ONNX_OPERATOR_KERNEL_CLASS_NAME(kNupharExecutionProvider, kOnnxDomain, 11, GatherElements)>());
kernel_registry.Register(BuildKernelCreateInfo<ONNX_OPERATOR_KERNEL_CLASS_NAME(kNupharExecutionProvider, kOnnxDomain, 10, MatMulInteger)>());
kernel_registry.Register(BuildKernelCreateInfo<ONNX_OPERATOR_KERNEL_CLASS_NAME(kNupharExecutionProvider, kMSDomain, 1, MatMulInteger16)>());
kernel_registry.Register(BuildKernelCreateInfo<ONNX_OPERATOR_KERNEL_CLASS_NAME(kNupharExecutionProvider, kOnnxDomain, 9, Scan)>());

10
onnxruntime/test/providers/cpu/tensor/gather_elements_op_test.cc
View file

@ -76,7 +76,10 @@ void RunTypedTest()
test5.AddOutput<T>("output", {2, 2},
{1, 1,
4, 4});
test5.Run(OpTester::ExpectResult::kExpectFailure, "GatherElements op: Value in indices must be within bounds [-2 , 1]. Actual value is 2");
// skip nuphar, which will not throw error message but will ensure no out-of-bound access
test5.Run(OpTester::ExpectResult::kExpectFailure,
"GatherElements op: Value in indices must be within bounds [-2 , 1]. Actual value is 2",
{kNupharExecutionProvider});
// 3D input - axis 1
OpTester test6("GatherElements", 11);
@ -158,7 +161,10 @@ void RunTypedTest<std::string>() {
test4.AddOutput<std::string>("output", {2, 2},
{"a", "a",
"d", "d"});
test4.Run(OpTester::ExpectResult::kExpectFailure, "GatherElements op: Value in indices must be within bounds [-2 , 1]. Actual value is -3");
// skip nuphar, which will not throw error message but will ensure no out-of-bound access
test4.Run(OpTester::ExpectResult::kExpectFailure,
"GatherElements op: Value in indices must be within bounds [-2 , 1]. Actual value is -3",
{kNupharExecutionProvider});
// 3D input - axis 1
OpTester test5("GatherElements", 11);