mirror of
https://github.com/saymrwulf/onnxruntime.git
synced 2026-07-29 20:14:01 +00:00
[NupharEP] performance improvements (#2283)
* [Nuphar EP] performance improvements 1. Add new ops: Shape, Expand 2. Add support for steps in Slice 3. Simplify Gather 4. Always inline alias nodes 5. Transpose nodes with inner loop being symbolic falls back to CPU provider when vectorization is not possible 6. Add opt_inproj option to model_editor to extract MatMuls inside Scan for input projection to outside
This commit is contained in:
parent
63e9961637
commit
e18c9582a8
22 changed files with 404 additions and 103 deletions
|
|
@ -36,6 +36,19 @@ bool IsRecurrentNode(const onnxruntime::Node& node) {
|
|||
|
||||
bool IsAliasNode(const onnxruntime::Node& node) {
|
||||
auto op_type = node.OpType();
|
||||
if (op_type == "Transpose") {
|
||||
// Treat Transpose (1,N) -> (N,1) as Alias
|
||||
const auto shape = node.OutputDefs()[0]->Shape();
|
||||
if (shape != nullptr && shape->dim_size() == 2) {
|
||||
for (int i = 0; i < 2; ++i) {
|
||||
if (shape->dim(i).has_dim_value() && shape->dim(i).dim_value() == 1) {
|
||||
return true;
|
||||
}
|
||||
}
|
||||
}
|
||||
return false;
|
||||
}
|
||||
|
||||
return (op_type == "Flatten" || op_type == "Identity" || op_type == "Reshape" ||
|
||||
op_type == "Squeeze" || op_type == "Unsqueeze");
|
||||
}
|
||||
|
|
|
|||
|
|
@ -77,6 +77,7 @@ namespace onnxruntime {
|
|||
ADD_OP_ITEM(Conv) \
|
||||
ADD_OP_ITEM(Crop) \
|
||||
ADD_OP_ITEM(Dropout) \
|
||||
ADD_OP_ITEM(Expand) \
|
||||
ADD_OP_ITEM(Flatten) \
|
||||
ADD_OP_ITEM(Gather) \
|
||||
ADD_OP_ITEM(GatherElements) \
|
||||
|
|
@ -88,6 +89,7 @@ namespace onnxruntime {
|
|||
ADD_OP_ITEM(MatMulInteger) \
|
||||
ADD_OP_ITEM(Pad) \
|
||||
ADD_OP_ITEM(Reshape) \
|
||||
ADD_OP_ITEM(Shape) \
|
||||
ADD_OP_ITEM(Slice) \
|
||||
ADD_OP_ITEM(Softmax) \
|
||||
ADD_OP_ITEM(Split) \
|
||||
|
|
|
|||
30
onnxruntime/core/codegen/mti/tensor/expand.cc
Normal file
30
onnxruntime/core/codegen/mti/tensor/expand.cc
Normal file
|
|
@ -0,0 +1,30 @@
|
|||
// Copyright (c) Microsoft Corporation. All rights reserved.
|
||||
// Licensed under the MIT License.
|
||||
|
||||
#include "core/codegen/mti/tensor/expand.h"
|
||||
#include "core/codegen/mti/common.h"
|
||||
|
||||
namespace onnxruntime {
|
||||
namespace tvm_codegen {
|
||||
|
||||
tvm::Tensor Expand(const tvm::Tensor& X, const tvm::Array<tvm::Expr>& new_shape, const std::string& name) {
|
||||
MTI_ASSERT(new_shape.size() >= X->shape.size());
|
||||
return tvm::compute(
|
||||
new_shape,
|
||||
[&](const tvm::Array<tvm::Var>& out_indices) {
|
||||
tvm::Array<tvm::Expr> indices;
|
||||
size_t broadcasted_rank = new_shape.size() - X->shape.size();
|
||||
for (size_t d = broadcasted_rank; d < new_shape.size(); ++d) {
|
||||
if (tvm::is_const_int(X->shape[d - broadcasted_rank], 1)) {
|
||||
indices.push_back(tvm::make_zero(HalideIR::Int(32)));
|
||||
} else {
|
||||
indices.push_back(out_indices[d]);
|
||||
}
|
||||
}
|
||||
return X(indices);
|
||||
},
|
||||
name);
|
||||
}
|
||||
|
||||
} // namespace tvm_codegen
|
||||
} // namespace onnxruntime
|
||||
14
onnxruntime/core/codegen/mti/tensor/expand.h
Normal file
14
onnxruntime/core/codegen/mti/tensor/expand.h
Normal file
|
|
@ -0,0 +1,14 @@
|
|||
// 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 Expand(const tvm::Tensor& X, const tvm::Array<tvm::Expr>& new_shape, const std::string& name = "expand");
|
||||
|
||||
} // namespace tvm_codegen
|
||||
} // namespace onnxruntime
|
||||
|
|
@ -45,8 +45,7 @@ tvm::Tensor Gather(const tvm::Tensor& t,
|
|||
ivars.push_back(ovars[i - 1 + indices->shape.size()]);
|
||||
}
|
||||
}
|
||||
return tvm::ir::Select::make((ivars[axis_t] >= 0) && (ivars[axis_t] < t->shape[axis_t]),
|
||||
t(ivars), tvm::make_zero(t->dtype));
|
||||
return t(ivars);
|
||||
};
|
||||
|
||||
return tvm::compute(output_shape, l, name);
|
||||
|
|
|
|||
25
onnxruntime/core/codegen/mti/tensor/shape_op.cc
Normal file
25
onnxruntime/core/codegen/mti/tensor/shape_op.cc
Normal file
|
|
@ -0,0 +1,25 @@
|
|||
// Copyright (c) Microsoft Corporation. All rights reserved.
|
||||
// Licensed under the MIT License.
|
||||
|
||||
#include "core/codegen/mti/tensor/shape_op.h"
|
||||
|
||||
namespace onnxruntime {
|
||||
namespace tvm_codegen {
|
||||
|
||||
tvm::Tensor Shape(const tvm::Tensor& X, const std::string& name) {
|
||||
int ndim = static_cast<int>(X->shape.size());
|
||||
tvm::Array<tvm::Expr> out_shape{ndim};
|
||||
return tvm::compute(
|
||||
out_shape, [&](const tvm::Array<tvm::Var>& indices) {
|
||||
auto idx = indices[0];
|
||||
tvm::Expr ret = 0;
|
||||
for (int i = 0; i < ndim; ++i) {
|
||||
ret = tvm::ir::Select::make(idx == i, X->shape[i], ret);
|
||||
}
|
||||
return tvm::cast(HalideIR::Int(64), ret);
|
||||
},
|
||||
name);
|
||||
}
|
||||
|
||||
} // namespace tvm_codegen
|
||||
} // namespace onnxruntime
|
||||
14
onnxruntime/core/codegen/mti/tensor/shape_op.h
Normal file
14
onnxruntime/core/codegen/mti/tensor/shape_op.h
Normal file
|
|
@ -0,0 +1,14 @@
|
|||
// 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 Shape(const tvm::Tensor& X, const std::string& name = "shape");
|
||||
|
||||
} // namespace tvm_codegen
|
||||
} // namespace onnxruntime
|
||||
|
|
@ -5,40 +5,83 @@
|
|||
|
||||
#include "core/codegen/mti/mti_tvm_utils.h"
|
||||
#include <climits>
|
||||
#include <gsl/gsl>
|
||||
#include <topi/transform.h>
|
||||
#include <tvm/ir_pass.h>
|
||||
|
||||
namespace onnxruntime {
|
||||
namespace tvm_codegen {
|
||||
|
||||
static const int64_t max_range = INT_MAX;
|
||||
// local constexpr for INT_MAX
|
||||
constexpr int64_t max_range = INT_MAX;
|
||||
|
||||
tvm::Expr position(const tvm::Expr& dim, const tvm::Integer& offset) {
|
||||
if (offset->value >= max_range)
|
||||
return dim;
|
||||
else if (offset->value < 0)
|
||||
return dim + offset;
|
||||
else
|
||||
return offset;
|
||||
tvm::Expr position(const tvm::Expr& dim, const tvm::Integer& offset, bool allow_out_of_bound = false) {
|
||||
if (offset->value >= max_range) {
|
||||
return allow_out_of_bound ? dim : dim - 1;
|
||||
} else if (offset->value <= -max_range) {
|
||||
return tvm::make_const(HalideIR::Int(32), allow_out_of_bound ? -1 : 0);
|
||||
} else {
|
||||
if (offset->value >= 0) {
|
||||
return tvm::ir::Simplify(tvm::ir::Min::make(offset, dim + (allow_out_of_bound ? 0 : -1)));
|
||||
} else {
|
||||
return tvm::ir::Simplify(dim + tvm::ir::Max::make(offset, -dim + (allow_out_of_bound ? -1 : 0)));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
tvm::Tensor Slice(const tvm::Tensor& X,
|
||||
const tvm::Array<tvm::Integer>& starts,
|
||||
const tvm::Array<tvm::Integer>& ends,
|
||||
const std::vector<int64_t>& starts,
|
||||
const std::vector<int64_t>& ends,
|
||||
const std::vector<int64_t>& axes1,
|
||||
const std::vector<int64_t>& steps,
|
||||
const std::string& name) {
|
||||
tvm::Array<tvm::Expr> output_shape;
|
||||
for (size_t i = 0; i < X->shape.size(); ++i) {
|
||||
tvm::Expr start = position(X->shape[i], starts[i]);
|
||||
tvm::Expr end = position(X->shape[i], ends[i]);
|
||||
output_shape.push_back(tvm::ir::Simplify(end - start));
|
||||
MTI_ASSERT(starts.size() == ends.size());
|
||||
MTI_ASSERT(starts.size() == axes1.size());
|
||||
MTI_ASSERT(starts.size() == steps.size());
|
||||
|
||||
std::vector<int64_t> axes;
|
||||
for (const auto& i : axes1) {
|
||||
axes.push_back(HandleNegativeAxis(i, X->shape.size()));
|
||||
}
|
||||
|
||||
tvm::Array<tvm::Expr> output_shape;
|
||||
bool empty = false;
|
||||
for (int64_t i = 0; i < gsl::narrow<int64_t>(X->shape.size()); ++i) {
|
||||
auto axes_iter = std::find(axes.begin(), axes.end(), i);
|
||||
if (axes_iter != axes.end()) {
|
||||
auto axis = axes_iter - axes.begin();
|
||||
tvm::Expr start = position(X->shape[i], starts[axis]);
|
||||
tvm::Expr end = position(X->shape[i], ends[axis], /*allow_out_of_bound*/ true);
|
||||
auto dim = tvm::ir::Simplify((end - start + tvm::Integer(steps[axis] + (steps[axis] < 0 ? 1 : -1))) / tvm::Integer(steps[axis]));
|
||||
auto int_dim = tvm::as_const_int(dim);
|
||||
if (int_dim && *int_dim <= 0) {
|
||||
output_shape.push_back(0);
|
||||
empty = true;
|
||||
} else {
|
||||
output_shape.push_back(dim);
|
||||
}
|
||||
} else {
|
||||
output_shape.push_back(X->shape[i]);
|
||||
}
|
||||
}
|
||||
|
||||
if (empty) {
|
||||
return MakeZeroTensor(output_shape, X->dtype, name);
|
||||
}
|
||||
|
||||
return tvm::compute(
|
||||
output_shape,
|
||||
[&](const tvm::Array<tvm::Var>& ovars) {
|
||||
tvm::Array<tvm::Expr> ivars;
|
||||
for (size_t i = 0; i < X->shape.size(); ++i)
|
||||
ivars.push_back(ovars[i] + tvm::ir::Simplify(position(X->shape[i], starts[i])));
|
||||
|
||||
for (size_t i = 0; i < X->shape.size(); ++i) {
|
||||
auto axes_iter = std::find(axes.begin(), axes.end(), i);
|
||||
if (axes_iter != axes.end()) {
|
||||
auto axis = axes_iter - axes.begin();
|
||||
ivars.push_back(tvm::ir::Simplify(ovars[i] * tvm::Integer(steps[axis]) + position(X->shape[i], starts[axis])));
|
||||
} else {
|
||||
ivars.push_back(ovars[i]);
|
||||
}
|
||||
}
|
||||
return X(ivars);
|
||||
},
|
||||
name);
|
||||
|
|
|
|||
|
|
@ -9,8 +9,10 @@ namespace onnxruntime {
|
|||
namespace tvm_codegen {
|
||||
|
||||
tvm::Tensor Slice(const tvm::Tensor& X,
|
||||
const tvm::Array<tvm::Integer>& starts,
|
||||
const tvm::Array<tvm::Integer>& ends,
|
||||
const std::vector<int64_t>& starts,
|
||||
const std::vector<int64_t>& ends,
|
||||
const std::vector<int64_t>& axes,
|
||||
const std::vector<int64_t>& steps,
|
||||
const std::string& name = "slice");
|
||||
|
||||
} // namespace tvm_codegen
|
||||
|
|
|
|||
|
|
@ -0,0 +1,26 @@
|
|||
// 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/mti_tvm_utils.h"
|
||||
#include "core/codegen/mti/tensor/expand.h"
|
||||
#include "core/codegen/passes/utils/ort_tvm_utils.h"
|
||||
#include "core/framework/op_kernel_info.h"
|
||||
|
||||
namespace onnxruntime {
|
||||
namespace tvm_codegen {
|
||||
|
||||
// Evaluate of Expand OpIRCreator
|
||||
Status GENERIC_OP_IR_CREATOR_CLASS(Expand)::Evaluate(
|
||||
const tvm::Array<tvm::Tensor>& inputs,
|
||||
const Node& node,
|
||||
CodeGenContext& ctx_codegen,
|
||||
tvm::Array<tvm::Tensor>& outputs) {
|
||||
tvm::Tensor Y = Expand(inputs[0], ShapeToTvmArray(node.OutputDefs()[0], ctx_codegen), node.Name() + "_Expand");
|
||||
outputs.push_back(Y);
|
||||
return Status::OK();
|
||||
}
|
||||
|
||||
} // namespace tvm_codegen
|
||||
} // namespace onnxruntime
|
||||
|
|
@ -0,0 +1,26 @@
|
|||
// 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/mti_tvm_utils.h"
|
||||
#include "core/codegen/mti/tensor/shape_op.h"
|
||||
#include "core/codegen/passes/utils/ort_tvm_utils.h"
|
||||
#include "core/framework/op_kernel_info.h"
|
||||
|
||||
namespace onnxruntime {
|
||||
namespace tvm_codegen {
|
||||
|
||||
// Evaluate of Expand OpIRCreator
|
||||
Status GENERIC_OP_IR_CREATOR_CLASS(Shape)::Evaluate(
|
||||
const tvm::Array<tvm::Tensor>& inputs,
|
||||
const Node& node,
|
||||
CodeGenContext& ctx_codegen,
|
||||
tvm::Array<tvm::Tensor>& outputs) {
|
||||
tvm::Tensor Y = Shape(inputs[0], node.Name() + "_Expand");
|
||||
outputs.push_back(Y);
|
||||
return Status::OK();
|
||||
}
|
||||
|
||||
} // namespace tvm_codegen
|
||||
} // namespace onnxruntime
|
||||
|
|
@ -13,17 +13,14 @@
|
|||
namespace onnxruntime {
|
||||
namespace tvm_codegen {
|
||||
|
||||
// local constexpr for INT_MAX
|
||||
constexpr int64_t max_range = INT_MAX;
|
||||
|
||||
Status SliceCommon(const tvm::Array<tvm::Tensor>& inputs,
|
||||
const Node& node,
|
||||
tvm::Array<tvm::Tensor>& outputs,
|
||||
const std::vector<int64_t>& starts,
|
||||
const std::vector<int64_t>& ends,
|
||||
const std::vector<int64_t>& axes1) {
|
||||
const std::vector<int64_t>& axes1,
|
||||
const std::vector<int64_t>& steps1) {
|
||||
ORT_RETURN_IF_NOT(nullptr != node.InputDefs()[0]);
|
||||
const ONNX_NAMESPACE::TensorShapeProto* shape_proto = node.InputDefs()[0]->Shape();
|
||||
|
||||
std::vector<int64_t> axes;
|
||||
if (axes1.size() == 0) {
|
||||
|
|
@ -34,48 +31,14 @@ Status SliceCommon(const tvm::Array<tvm::Tensor>& inputs,
|
|||
axes = axes1;
|
||||
}
|
||||
|
||||
tvm::Array<tvm::Integer> tvm_starts, tvm_ends;
|
||||
bool empty = false;
|
||||
|
||||
for (int dim = 0; dim < shape_proto->dim_size(); ++dim) {
|
||||
auto axes_iter = std::find(axes.begin(), axes.end(), dim);
|
||||
const ONNX_NAMESPACE::TensorShapeProto_Dimension& proto_dim = shape_proto->dim(dim);
|
||||
bool found_in_axes = (axes_iter != axes.end());
|
||||
if (!found_in_axes) {
|
||||
tvm_starts.push_back(0);
|
||||
if (utils::HasDimValue(proto_dim)) {
|
||||
tvm_ends.push_back(proto_dim.dim_value());
|
||||
} else {
|
||||
tvm_ends.push_back(max_range);
|
||||
}
|
||||
} else {
|
||||
auto axes_index = axes_iter - axes.begin();
|
||||
int64_t start = starts[axes_index];
|
||||
int64_t end = ends[axes_index];
|
||||
if (utils::HasDimValue(proto_dim)) {
|
||||
int64_t dim_max = proto_dim.dim_value();
|
||||
if (start < 0) start += dim_max;
|
||||
if (end < 0) end += dim_max;
|
||||
start = std::min(dim_max, std::max(static_cast<int64_t>(0), start));
|
||||
end = std::min(dim_max, std::max(start, end));
|
||||
}
|
||||
tvm_starts.push_back(start);
|
||||
tvm_ends.push_back(end);
|
||||
empty = empty || (start == end);
|
||||
}
|
||||
}
|
||||
|
||||
tvm::Tensor Y;
|
||||
if (empty) {
|
||||
tvm::Array<tvm::Expr> shape;
|
||||
for (size_t dim = 0; dim < gsl::narrow_cast<size_t>(shape_proto->dim_size()); ++dim) {
|
||||
shape.push_back(tvm::ir::Simplify(tvm_ends[dim] - tvm_starts[dim]));
|
||||
}
|
||||
Y = MakeZeroTensor(shape, inputs[0]->dtype, node.Name() + "_zeros");
|
||||
std::vector<int64_t> steps;
|
||||
if (steps1.size() == 0) {
|
||||
steps.resize(starts.size(), 1);
|
||||
} else {
|
||||
Y = Slice(inputs[0], tvm_starts, tvm_ends, node.Name() + "_Slice");
|
||||
steps = steps1;
|
||||
}
|
||||
|
||||
tvm::Tensor Y = Slice(inputs[0], starts, ends, axes, steps, node.Name() + "_Slice");
|
||||
outputs.push_back(Y);
|
||||
return Status::OK();
|
||||
}
|
||||
|
|
@ -94,14 +57,14 @@ Status GENERIC_OP_IR_CREATOR_CLASS(Slice)::Evaluate(
|
|||
int version = ctx_codegen.GetCodeGenHandle()->domain_version_lookup_func(node.Domain());
|
||||
ORT_RETURN_IF_NOT(version <= 9, "Dynamic Slice is not supported yet");
|
||||
|
||||
std::vector<int64_t> starts, ends;
|
||||
std::vector<int64_t> starts, ends, steps;
|
||||
ORT_RETURN_IF_ERROR(info.GetAttrs<int64_t>("starts", starts));
|
||||
ORT_RETURN_IF_ERROR(info.GetAttrs<int64_t>("ends", ends));
|
||||
ORT_RETURN_IF_NOT(starts.size() == ends.size());
|
||||
|
||||
auto axes = info.GetAttrsOrDefault<int64_t>("axes");
|
||||
|
||||
return SliceCommon(inputs, node, outputs, starts, ends, axes);
|
||||
return SliceCommon(inputs, node, outputs, starts, ends, axes, steps);
|
||||
}
|
||||
|
||||
} // namespace tvm_codegen
|
||||
|
|
|
|||
|
|
@ -36,6 +36,11 @@ int CodeGenUnitStats::NodeUseCount(const onnxruntime::Node* node) const {
|
|||
|
||||
bool CodeGenUnitStats::IsCheapNodeReuse(const onnxruntime::Node* node) const {
|
||||
ORT_ENFORCE(passes_.size() > UseCountAnalysisOffset);
|
||||
|
||||
// always inline(not reuse) alias node
|
||||
if (IsAliasNode(*node))
|
||||
return false;
|
||||
|
||||
// Define cheap nodes include Add / Sub / Mul
|
||||
if (node->OpType() == "Add" || node->OpType() == "Sub" || node->OpType() == "Mul")
|
||||
return Promote<NupharUseCountAnalysis>(passes_[UseCountAnalysisOffset])->NodeUseCount(node) > CheapNodeTrueReuseCount;
|
||||
|
|
|
|||
|
|
@ -35,38 +35,39 @@ bool HasUnknownShapeOnAxes(const NodeArg* def, std::vector<int64_t>& axes) {
|
|||
return false;
|
||||
}
|
||||
|
||||
Status GetSliceAxesFromTensorProto(std::vector<int64_t>& axes,
|
||||
const ONNX_NAMESPACE::TensorProto& axes_tp) {
|
||||
Status GetVectorInt64FromTensorProto(std::vector<int64_t>& v,
|
||||
const ONNX_NAMESPACE::TensorProto& tp) {
|
||||
size_t tp_sz_in_bytes;
|
||||
ORT_RETURN_IF_ERROR(utils::GetSizeInBytesFromTensorProto<0>(axes_tp, &tp_sz_in_bytes));
|
||||
ORT_RETURN_IF_ERROR(utils::GetSizeInBytesFromTensorProto<0>(tp, &tp_sz_in_bytes));
|
||||
OrtValue ort_value;
|
||||
std::unique_ptr<char[]> data(new char[tp_sz_in_bytes]);
|
||||
|
||||
#define UNPACK_TENSOR(T) \
|
||||
T* p = reinterpret_cast<T*>(data.get()); \
|
||||
ORT_RETURN_IF_ERROR(utils::UnpackTensor<T>( \
|
||||
axes_tp, \
|
||||
axes_tp.raw_data().size() ? axes_tp.raw_data().data() : nullptr, \
|
||||
axes_tp.raw_data().size(), \
|
||||
p, \
|
||||
tp_sz_in_bytes / sizeof(T))); \
|
||||
std::vector<T> tmp_axes(p, p + tp_sz_in_bytes / sizeof(T));
|
||||
#define UNPACK_TENSOR(T) \
|
||||
T* p = reinterpret_cast<T*>(data.get()); \
|
||||
ORT_RETURN_IF_ERROR(utils::UnpackTensor<T>( \
|
||||
tp, \
|
||||
tp.raw_data().size() ? tp.raw_data().data() : nullptr, \
|
||||
tp.raw_data().size(), \
|
||||
p, \
|
||||
tp_sz_in_bytes / sizeof(T))); \
|
||||
std::vector<T> tmp_v(p, p + tp_sz_in_bytes / sizeof(T));
|
||||
v.clear();
|
||||
|
||||
switch (axes_tp.data_type()) {
|
||||
switch (tp.data_type()) {
|
||||
case ONNX_NAMESPACE::TensorProto_DataType_INT32: {
|
||||
UNPACK_TENSOR(int32_t);
|
||||
for (auto axis : tmp_axes) {
|
||||
axes.push_back(static_cast<int64_t>(axis));
|
||||
for (auto axis : tmp_v) {
|
||||
v.push_back(static_cast<int64_t>(axis));
|
||||
}
|
||||
break;
|
||||
}
|
||||
case ONNX_NAMESPACE::TensorProto_DataType_INT64: {
|
||||
UNPACK_TENSOR(int64_t);
|
||||
axes.insert(axes.end(), tmp_axes.begin(), tmp_axes.end());
|
||||
v.insert(v.end(), tmp_v.begin(), tmp_v.end());
|
||||
break;
|
||||
}
|
||||
default:
|
||||
ORT_NOT_IMPLEMENTED("Unimplemented type: ", axes_tp.data_type());
|
||||
ORT_NOT_IMPLEMENTED("Unimplemented type: ", tp.data_type());
|
||||
}
|
||||
|
||||
return Status::OK();
|
||||
|
|
|
|||
|
|
@ -16,8 +16,8 @@ bool HasUnknownShapeOnAxis(const ConstPointerContainer<std::vector<NodeArg*>>& d
|
|||
|
||||
bool HasUnknownShapeOnAxes(const NodeArg* def, std::vector<int64_t>& axes);
|
||||
|
||||
Status GetSliceAxesFromTensorProto(std::vector<int64_t>& axes,
|
||||
const ONNX_NAMESPACE::TensorProto& axes_tp);
|
||||
Status GetVectorInt64FromTensorProto(std::vector<int64_t>& v,
|
||||
const ONNX_NAMESPACE::TensorProto& tp);
|
||||
|
||||
} // namespace nuphar
|
||||
} // namespace onnxruntime
|
||||
|
|
|
|||
|
|
@ -17,7 +17,8 @@ Status SliceCommon(const tvm::Array<tvm::Tensor>& inputs,
|
|||
tvm::Array<tvm::Tensor>& outputs,
|
||||
const std::vector<int64_t>& starts,
|
||||
const std::vector<int64_t>& ends,
|
||||
const std::vector<int64_t>& axes);
|
||||
const std::vector<int64_t>& axes,
|
||||
const std::vector<int64_t>& steps);
|
||||
|
||||
} // namespace tvm_codegen
|
||||
|
||||
|
|
@ -36,7 +37,7 @@ Status NUPHAR_TVM_X86_OP_IR_CREATOR_CLASS(Slice)::Evaluate(
|
|||
std::vector<std::vector<int64_t>> slice_params;
|
||||
int version = ctx_codegen.GetCodeGenHandle()->domain_version_lookup_func(node.Domain());
|
||||
if (version <= 9) {
|
||||
std::vector<int64_t> starts, ends, axes;
|
||||
std::vector<int64_t> starts, ends, axes, steps;
|
||||
ORT_RETURN_IF_ERROR(info.GetAttrs<int64_t>("starts", starts));
|
||||
ORT_RETURN_IF_ERROR(info.GetAttrs<int64_t>("ends", ends));
|
||||
ORT_RETURN_IF_NOT(starts.size() == ends.size());
|
||||
|
|
@ -44,11 +45,11 @@ Status NUPHAR_TVM_X86_OP_IR_CREATOR_CLASS(Slice)::Evaluate(
|
|||
slice_params.push_back(starts);
|
||||
slice_params.push_back(ends);
|
||||
slice_params.push_back(axes);
|
||||
slice_params.push_back(steps);
|
||||
} else {
|
||||
// for opset 10 Slice, input 1/2/3/4 are starts/ends/axes/steps
|
||||
// while axes and steps are optional
|
||||
ORT_ENFORCE(node.InputDefs().size() < 5, "Slice opset 10: steps is not supported yet");
|
||||
for (size_t i = 1; i < 4; ++i) {
|
||||
for (size_t i = 1; i < 5; ++i) {
|
||||
if (i < node.InputDefs().size()) {
|
||||
const auto* tensor = ctx_nuphar->GetOrtInitializerTensor(node.InputDefs()[i]->Name());
|
||||
if (tensor) {
|
||||
|
|
@ -65,7 +66,7 @@ Status NUPHAR_TVM_X86_OP_IR_CREATOR_CLASS(Slice)::Evaluate(
|
|||
slice_params.push_back(std::vector<int64_t>());
|
||||
}
|
||||
}
|
||||
return tvm_codegen::SliceCommon(inputs, node, outputs, slice_params[0], slice_params[1], slice_params[2]);
|
||||
return tvm_codegen::SliceCommon(inputs, node, outputs, slice_params[0], slice_params[1], slice_params[2], slice_params[3]);
|
||||
}
|
||||
|
||||
} // namespace nuphar
|
||||
|
|
|
|||
|
|
@ -96,6 +96,7 @@ class NupharKernelState {
|
|||
NUPHAR_OP(Equal, 11, DataTypeImpl::AllFixedSizeTensorTypes()) \
|
||||
NUPHAR_OP(Erf, 9, DataTypeImpl::GetTensorType<float>()) \
|
||||
NUPHAR_OP(Exp, 6, DataTypeImpl::AllIEEEFloatTensorTypes()) \
|
||||
NUPHAR_OP(Expand, 8, DataTypeImpl::AllFixedSizeTensorTypes()) \
|
||||
NUPHAR_VERSIONED_OP(Flatten, 1, 8, DataTypeImpl::AllIEEEFloatTensorTypes()) \
|
||||
NUPHAR_VERSIONED_OP(Flatten, 9, 10, DataTypeImpl::AllIEEEFloatTensorTypes()) \
|
||||
NUPHAR_OP(Flatten, 11, DataTypeImpl::AllIEEEFloatTensorTypes()) \
|
||||
|
|
@ -151,6 +152,7 @@ class NupharKernelState {
|
|||
NUPHAR_OP(Reshape, 5, DataTypeImpl::AllFixedSizeTensorTypes()) \
|
||||
NUPHAR_OP(ScaledTanh, 1, DataTypeImpl::AllIEEEFloatTensorTypes()) \
|
||||
NUPHAR_OP(Selu, 6, DataTypeImpl::AllIEEEFloatTensorTypes()) \
|
||||
NUPHAR_OP(Shape, 1, DataTypeImpl::AllFixedSizeTensorTypes()) \
|
||||
NUPHAR_OP(Sigmoid, 6, DataTypeImpl::AllIEEEFloatTensorTypes()) \
|
||||
NUPHAR_VERSIONED_OP(Slice, 1, 9, DataTypeImpl::AllFixedSizeTensorTypes()) \
|
||||
NUPHAR_OP(Slice, 10, DataTypeImpl::AllFixedSizeTensorTypes()) \
|
||||
|
|
|
|||
|
|
@ -177,7 +177,8 @@ NupharExecutionProvider::GetCapability(const onnxruntime::GraphViewer& graph_vie
|
|||
|
||||
std::vector<std::unique_ptr<ComputeCapability>> results;
|
||||
|
||||
auto is_supported_func = [&](const Node& node) {
|
||||
typedef std::function<bool(const Node&)> IsSupportedFunc;
|
||||
IsSupportedFunc is_supported_func = [&](const Node& node) {
|
||||
bool all_shape_defined = true;
|
||||
node.ForEachDef([&all_shape_defined](const NodeArg& def, bool /*is_input*/) {
|
||||
if (def.Shape() == nullptr) {
|
||||
|
|
@ -218,21 +219,26 @@ NupharExecutionProvider::GetCapability(const onnxruntime::GraphViewer& graph_vie
|
|||
auto num_inputs = inputs.size();
|
||||
ORT_ENFORCE(num_inputs > 0);
|
||||
std::vector<int64_t> axes;
|
||||
std::vector<int64_t> steps;
|
||||
if (num_inputs > 1) {
|
||||
// Slice-10
|
||||
bool is_starts_dynamic = !graph_viewer.IsConstantInitializer(inputs[1]->Name(), true);
|
||||
bool is_ends_dynamic = !graph_viewer.IsConstantInitializer(inputs[2]->Name(), true);
|
||||
|
||||
bool is_axes_dynamic = inputs.size() > 3 && !graph_viewer.IsConstantInitializer(inputs[3]->Name(), true);
|
||||
|
||||
bool has_steps = inputs.size() > 4;
|
||||
if (is_starts_dynamic || is_ends_dynamic || is_axes_dynamic || has_steps)
|
||||
bool is_steps_dynamic = inputs.size() > 4 && !graph_viewer.IsConstantInitializer(inputs[4]->Name(), true);
|
||||
if (is_starts_dynamic || is_ends_dynamic || is_axes_dynamic || is_steps_dynamic)
|
||||
return false;
|
||||
|
||||
const ONNX_NAMESPACE::TensorProto* steps_tp = nullptr;
|
||||
bool found_steps = inputs.size() > 4 && graph_viewer.GetInitializedTensor(inputs[4]->Name(), steps_tp);
|
||||
if (found_steps) {
|
||||
GetVectorInt64FromTensorProto(steps, *steps_tp);
|
||||
}
|
||||
|
||||
const ONNX_NAMESPACE::TensorProto* axes_tp = nullptr;
|
||||
bool found_axes = inputs.size() > 3 && graph_viewer.GetInitializedTensor(inputs[3]->Name(), axes_tp);
|
||||
if (found_axes) {
|
||||
GetSliceAxesFromTensorProto(axes, *axes_tp);
|
||||
GetVectorInt64FromTensorProto(axes, *axes_tp);
|
||||
}
|
||||
} else {
|
||||
const onnxruntime::NodeAttributes& attrs = node.GetAttributes();
|
||||
|
|
@ -247,6 +253,25 @@ NupharExecutionProvider::GetCapability(const onnxruntime::GraphViewer& graph_vie
|
|||
if (HasUnknownShapeOnAxes(inputs[0], axes))
|
||||
return false;
|
||||
}
|
||||
|
||||
if (IsAliasNode(node)) {
|
||||
// for AliasNode as final output, skip them to avoid potential copy
|
||||
for (auto iter = node.OutputEdgesBegin(); iter != node.OutputEdgesEnd(); ++iter) {
|
||||
if (!is_supported_func(iter->GetNode())) {
|
||||
return false;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if (node.OpType() == "Transpose") {
|
||||
// When there's symbolic dim in last dim of Transpose output
|
||||
// reject it since it was not able to vectorize inlined ops
|
||||
const auto output_shape = node.OutputDefs()[0]->Shape();
|
||||
const auto output_rank = output_shape->dim_size();
|
||||
if (output_rank > 0 && output_shape->dim(output_rank - 1).has_dim_param()) {
|
||||
return false;
|
||||
}
|
||||
}
|
||||
return true;
|
||||
};
|
||||
GraphPartitioner graph_partitioner(is_supported_func);
|
||||
|
|
|
|||
|
|
@ -51,11 +51,16 @@ bool GraphPartitioner::IsNodeSupported(const Node& node) const {
|
|||
});
|
||||
return std::move(symbolic_dimensions);
|
||||
};
|
||||
// if there are any output symbols not in input symbols, fallback to CPU
|
||||
auto input_sym = get_symbolic_dimensions(node, true);
|
||||
auto output_sym = get_symbolic_dimensions(node, false);
|
||||
if (input_sym != output_sym && output_sym.size() > 0)
|
||||
if (std::count_if(output_sym.begin(),
|
||||
output_sym.end(),
|
||||
[&input_sym](const std::string& name) {
|
||||
return input_sym.count(name) == 0;
|
||||
})) {
|
||||
return false;
|
||||
|
||||
}
|
||||
return true;
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -609,10 +609,110 @@ def remove_initializers_from_inputs(input_model, output_model, remain_inputs=[])
|
|||
mp.graph.input.extend(new_inputs)
|
||||
onnx.save(mp, output_model)
|
||||
|
||||
def optimize_input_projection(input_model, output_model):
|
||||
in_mp = onnx.load(input_model)
|
||||
out_mp = onnx.ModelProto()
|
||||
out_mp.CopyFrom(in_mp)
|
||||
out_mp.ir_version = 5 # update ir version to avoid requirement of initializer in graph input
|
||||
out_mp.graph.ClearField('node')
|
||||
nf = NodeFactory(out_mp.graph, prefix='opt_inproj_')
|
||||
initializers = dict([(i.name, i) for i in in_mp.graph.initializer])
|
||||
# first find possible fused SVD and do constant folding on MatMul of initializers
|
||||
const_matmuls = [n for n in in_mp.graph.node if n.op_type == 'MatMul' and all([i in initializers for i in n.input])]
|
||||
for mm in const_matmuls:
|
||||
lhs = numpy_helper.to_array(initializers[mm.input[0]])
|
||||
rhs = numpy_helper.to_array(initializers[mm.input[1]])
|
||||
val = np.matmul(lhs, rhs)
|
||||
new_initializer = out_mp.graph.initializer.add()
|
||||
new_initializer.CopyFrom(numpy_helper.from_array(val, mm.output[0]))
|
||||
if not [n for n in in_mp.graph.node if n != mm and mm.input[0] in n.input]:
|
||||
nf.remove_initializer(mm.input[0])
|
||||
if not [n for n in in_mp.graph.node if n != mm and mm.input[1] in n.input]:
|
||||
nf.remove_initializer(mm.input[1])
|
||||
|
||||
initializers = dict([(i.name,i) for i in out_mp.graph.initializer])
|
||||
|
||||
# remove const_matmul output from graph outputs
|
||||
new_outputs = [i for i in out_mp.graph.output if not [m for m in const_matmuls if m.output[0] == i.name]]
|
||||
out_mp.graph.ClearField('output')
|
||||
out_mp.graph.output.extend(new_outputs)
|
||||
|
||||
for in_n in in_mp.graph.node:
|
||||
if in_n in const_matmuls:
|
||||
continue
|
||||
|
||||
optimize_scan = False
|
||||
if in_n.op_type == 'Scan':
|
||||
in_sg = NodeFactory.get_attribute(in_n, 'body')
|
||||
num_scan_inputs = NodeFactory.get_attribute(in_n, 'num_scan_inputs')
|
||||
# only support 1 scan input
|
||||
if num_scan_inputs == 1:
|
||||
optimize_scan = True
|
||||
|
||||
# copy the node if it's not the scan node that is supported at the moment
|
||||
if not optimize_scan:
|
||||
out_n = out_mp.graph.node.add()
|
||||
out_n.CopyFrom(in_n)
|
||||
continue
|
||||
|
||||
scan_input_directions = NodeFactory.get_attribute(in_n, 'scan_input_directions')
|
||||
scan_output_directions = NodeFactory.get_attribute(in_n, 'scan_output_directions')
|
||||
out_sg = onnx.GraphProto()
|
||||
out_sg.CopyFrom(in_sg)
|
||||
out_sg.ClearField('node')
|
||||
nf_subgraph = NodeFactory(out_mp.graph, out_sg, prefix='opt_inproj_sg_' + in_n.name + '_')
|
||||
new_inputs = list(in_n.input)
|
||||
in_sg_inputs = [i.name for i in in_sg.input]
|
||||
replaced_matmul = None
|
||||
for in_sn in in_sg.node:
|
||||
if in_sn.op_type == 'Concat' and len(in_sn.input) == 2 and all([i in in_sg_inputs for i in in_sn.input]):
|
||||
# make sure the concat's inputs are scan input and scan state
|
||||
if NodeFactory.get_attribute(in_sn, 'axis') != len(in_sg.input[-1].type.tensor_type.shape.dim) - 1:
|
||||
continue # must concat last dim
|
||||
matmul_node = [nn for nn in in_sg.node if nn.op_type == 'MatMul' and in_sn.output[0] in nn.input]
|
||||
if not matmul_node:
|
||||
continue
|
||||
replaced_matmul = matmul_node[0]
|
||||
assert replaced_matmul.input[1] in initializers
|
||||
aa = nf.get_initializer(replaced_matmul.input[1])
|
||||
input_size = in_sg.input[-1].type.tensor_type.shape.dim[-1].dim_value
|
||||
if in_sg_inputs[-1] == in_sn.input[0]:
|
||||
hidden_idx = 1
|
||||
input_proj_weights, hidden_proj_weights = np.vsplit(aa, [input_size])
|
||||
else:
|
||||
hidden_idx = 0
|
||||
hidden_proj_weights, input_proj_weights = np.vsplit(aa, [aa.shape[-1] - input_size])
|
||||
# add matmul for input_proj outside of Scan
|
||||
input_proj = nf.make_node('MatMul', [new_inputs[-1], input_proj_weights])
|
||||
input_proj.doc_string = replaced_matmul.doc_string
|
||||
new_inputs[-1] = input_proj.name
|
||||
out_sg.input[-1].type.tensor_type.shape.dim[-1].dim_value = input_proj_weights.shape[-1]
|
||||
# add matmul for hidden_proj inside Scan
|
||||
hidden_proj = nf_subgraph.make_node('MatMul', [in_sn.input[hidden_idx], hidden_proj_weights])
|
||||
hidden_proj.doc_string = replaced_matmul.doc_string
|
||||
nf_subgraph.make_node('Add', [out_sg.input[-1].name, hidden_proj], output_names=replaced_matmul.output[0])
|
||||
# remove initializer of concat matmul
|
||||
if not [n for n in in_mp.graph.node if n != in_n and replaced_matmul.input[1] in n.input]:
|
||||
nf.remove_initializer(replaced_matmul.input[1])
|
||||
elif in_sn != replaced_matmul:
|
||||
out_sg.node.add().CopyFrom(in_sn)
|
||||
|
||||
scan = nf.make_node('Scan', new_inputs,
|
||||
{'body':out_sg,
|
||||
'scan_input_directions':scan_input_directions,
|
||||
'scan_output_directions':scan_output_directions,
|
||||
'num_scan_inputs':num_scan_inputs},
|
||||
output_names=list(in_n.output))
|
||||
scan.name = in_n.name
|
||||
scan.doc_string = in_n.doc_string
|
||||
|
||||
onnx.save(out_mp, output_model)
|
||||
|
||||
def parse_arguments():
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument('--mode', help='The modification mode',
|
||||
choices=['to_scan',
|
||||
'opt_inproj',
|
||||
'remove_initializers_from_inputs'])
|
||||
parser.add_argument('--input', help='The input model file', default=None)
|
||||
parser.add_argument('--output', help='The output model file', default=None)
|
||||
|
|
@ -625,6 +725,9 @@ if __name__ == '__main__':
|
|||
if args.mode == 'to_scan':
|
||||
print('Convert LSTM/GRU/RNN to Scan...')
|
||||
convert_to_scan_model(args.input, args.output)
|
||||
elif args.mode == 'opt_inproj':
|
||||
print('Optimize input projection in Scan...')
|
||||
optimize_input_projection(args.input, args.output)
|
||||
elif args.mode == 'remove_initializers_from_inputs':
|
||||
print('Remove all initializers from input for model with IR version >= 4...')
|
||||
remove_initializers_from_inputs(args.input, args.output)
|
||||
|
|
|
|||
|
|
@ -13,10 +13,10 @@ class NodeFactory:
|
|||
node_count_ = 0
|
||||
const_count_ = 0
|
||||
|
||||
def __init__(self, main_graph, sub_graph=None):
|
||||
def __init__(self, main_graph, sub_graph=None, prefix=''):
|
||||
self.graph_ = sub_graph if sub_graph else main_graph
|
||||
self.main_graph_ = main_graph
|
||||
self.name_prefix_ = ''
|
||||
self.name_prefix_ = prefix
|
||||
|
||||
class ScopedPrefix:
|
||||
def __init__(self, nf, name):
|
||||
|
|
|
|||
|
|
@ -32,9 +32,11 @@ class TestNuphar(unittest.TestCase):
|
|||
bidaf_dir = os.path.join(cwd, 'bidaf')
|
||||
bidaf_model = os.path.join(bidaf_dir, 'bidaf.onnx')
|
||||
bidaf_scan_model = os.path.join(bidaf_dir, 'bidaf_scan.onnx')
|
||||
bidaf_opt_scan_model = os.path.join(bidaf_dir, 'bidaf_opt_scan.onnx')
|
||||
bidaf_int8_scan_only_model = os.path.join(bidaf_dir, 'bidaf_int8_scan_only.onnx')
|
||||
subprocess.run([sys.executable, '-m', 'onnxruntime.nuphar.model_editor', '--input', bidaf_model, '--output', bidaf_scan_model, '--mode', 'to_scan'], check=True, cwd=cwd)
|
||||
subprocess.run([sys.executable, '-m', 'onnxruntime.nuphar.model_quantizer', '--input', bidaf_scan_model, '--output', bidaf_int8_scan_only_model, '--only_for_scan'], check=True, cwd=cwd)
|
||||
subprocess.run([sys.executable, '-m', 'onnxruntime.nuphar.model_editor', '--input', bidaf_scan_model, '--output', bidaf_opt_scan_model, '--mode', 'opt_inproj'], check=True, cwd=cwd)
|
||||
subprocess.run([sys.executable, '-m', 'onnxruntime.nuphar.model_quantizer', '--input', bidaf_opt_scan_model, '--output', bidaf_int8_scan_only_model, '--only_for_scan'], check=True, cwd=cwd)
|
||||
|
||||
# run onnx_test_runner to verify results
|
||||
# use -M to disable memory pattern
|
||||
|
|
|
|||
Loading…
Reference in a new issue