[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:
KeDengMS 2019-10-30 10:15:04 -07:00 committed by GitHub
parent 63e9961637
commit e18c9582a8
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
22 changed files with 404 additions and 103 deletions

View file

@ -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");
}

View file

@ -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) \

View 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

View 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

View file

@ -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);

View 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

View 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

View file

@ -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);

View file

@ -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

View file

@ -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

View file

@ -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

View file

@ -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

View file

@ -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;

View file

@ -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();

View file

@ -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

View file

@ -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

View file

@ -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()) \

View file

@ -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);

View file

@ -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;
}

View file

@ -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)

View file

@ -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):

View file

@ -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