From e18c9582a814b0d0d8875becd9198590be238286 Mon Sep 17 00:00:00 2001 From: KeDengMS Date: Wed, 30 Oct 2019 10:15:04 -0700 Subject: [PATCH] [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 --- onnxruntime/core/codegen/common/common.cc | 13 +++ onnxruntime/core/codegen/common/op_macro.h | 2 + onnxruntime/core/codegen/mti/tensor/expand.cc | 30 +++++ onnxruntime/core/codegen/mti/tensor/expand.h | 14 +++ onnxruntime/core/codegen/mti/tensor/gather.cc | 3 +- .../core/codegen/mti/tensor/shape_op.cc | 25 +++++ .../core/codegen/mti/tensor/shape_op.h | 14 +++ onnxruntime/core/codegen/mti/tensor/slice.cc | 79 +++++++++++--- onnxruntime/core/codegen/mti/tensor/slice.h | 6 +- .../passes/op_ir_creator/tensor/expand.cc | 26 +++++ .../passes/op_ir_creator/tensor/shape_op.cc | 26 +++++ .../passes/op_ir_creator/tensor/slice.cc | 55 ++-------- .../common/analysis/subgraph_codegen_stats.cc | 5 + .../core/providers/nuphar/common/utils.cc | 35 +++--- .../core/providers/nuphar/common/utils.h | 4 +- .../x86/op_ir_creator/tensor/slice.cc | 11 +- onnxruntime/core/providers/nuphar/kernel.h | 2 + .../nuphar/nuphar_execution_provider.cc | 37 ++++++- .../nuphar/partition/graph_partitioner.cc | 9 +- .../providers/nuphar/scripts/model_editor.py | 103 ++++++++++++++++++ .../providers/nuphar/scripts/node_factory.py | 4 +- .../python/onnxruntime_test_python_nuphar.py | 4 +- 22 files changed, 404 insertions(+), 103 deletions(-) create mode 100644 onnxruntime/core/codegen/mti/tensor/expand.cc create mode 100644 onnxruntime/core/codegen/mti/tensor/expand.h create mode 100644 onnxruntime/core/codegen/mti/tensor/shape_op.cc create mode 100644 onnxruntime/core/codegen/mti/tensor/shape_op.h create mode 100644 onnxruntime/core/codegen/passes/op_ir_creator/tensor/expand.cc create mode 100644 onnxruntime/core/codegen/passes/op_ir_creator/tensor/shape_op.cc diff --git a/onnxruntime/core/codegen/common/common.cc b/onnxruntime/core/codegen/common/common.cc index 05d0308635..a2934ac5fb 100644 --- a/onnxruntime/core/codegen/common/common.cc +++ b/onnxruntime/core/codegen/common/common.cc @@ -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"); } diff --git a/onnxruntime/core/codegen/common/op_macro.h b/onnxruntime/core/codegen/common/op_macro.h index c11285fea6..04305c4aa4 100644 --- a/onnxruntime/core/codegen/common/op_macro.h +++ b/onnxruntime/core/codegen/common/op_macro.h @@ -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) \ diff --git a/onnxruntime/core/codegen/mti/tensor/expand.cc b/onnxruntime/core/codegen/mti/tensor/expand.cc new file mode 100644 index 0000000000..cdac4f56e1 --- /dev/null +++ b/onnxruntime/core/codegen/mti/tensor/expand.cc @@ -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& new_shape, const std::string& name) { + MTI_ASSERT(new_shape.size() >= X->shape.size()); + return tvm::compute( + new_shape, + [&](const tvm::Array& out_indices) { + tvm::Array 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 diff --git a/onnxruntime/core/codegen/mti/tensor/expand.h b/onnxruntime/core/codegen/mti/tensor/expand.h new file mode 100644 index 0000000000..d66d41aeb0 --- /dev/null +++ b/onnxruntime/core/codegen/mti/tensor/expand.h @@ -0,0 +1,14 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once +#include +#include + +namespace onnxruntime { +namespace tvm_codegen { + +tvm::Tensor Expand(const tvm::Tensor& X, const tvm::Array& new_shape, const std::string& name = "expand"); + +} // namespace tvm_codegen +} // namespace onnxruntime diff --git a/onnxruntime/core/codegen/mti/tensor/gather.cc b/onnxruntime/core/codegen/mti/tensor/gather.cc index 18531d7f3d..6748b5913a 100644 --- a/onnxruntime/core/codegen/mti/tensor/gather.cc +++ b/onnxruntime/core/codegen/mti/tensor/gather.cc @@ -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); diff --git a/onnxruntime/core/codegen/mti/tensor/shape_op.cc b/onnxruntime/core/codegen/mti/tensor/shape_op.cc new file mode 100644 index 0000000000..b51bd67a8b --- /dev/null +++ b/onnxruntime/core/codegen/mti/tensor/shape_op.cc @@ -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(X->shape.size()); + tvm::Array out_shape{ndim}; + return tvm::compute( + out_shape, [&](const tvm::Array& 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 diff --git a/onnxruntime/core/codegen/mti/tensor/shape_op.h b/onnxruntime/core/codegen/mti/tensor/shape_op.h new file mode 100644 index 0000000000..67ee2de50e --- /dev/null +++ b/onnxruntime/core/codegen/mti/tensor/shape_op.h @@ -0,0 +1,14 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once +#include +#include + +namespace onnxruntime { +namespace tvm_codegen { + +tvm::Tensor Shape(const tvm::Tensor& X, const std::string& name = "shape"); + +} // namespace tvm_codegen +} // namespace onnxruntime diff --git a/onnxruntime/core/codegen/mti/tensor/slice.cc b/onnxruntime/core/codegen/mti/tensor/slice.cc index 4caf3946ce..6cbab43584 100644 --- a/onnxruntime/core/codegen/mti/tensor/slice.cc +++ b/onnxruntime/core/codegen/mti/tensor/slice.cc @@ -5,40 +5,83 @@ #include "core/codegen/mti/mti_tvm_utils.h" #include +#include #include #include 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& starts, - const tvm::Array& ends, + const std::vector& starts, + const std::vector& ends, + const std::vector& axes1, + const std::vector& steps, const std::string& name) { - tvm::Array 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 axes; + for (const auto& i : axes1) { + axes.push_back(HandleNegativeAxis(i, X->shape.size())); } + + tvm::Array output_shape; + bool empty = false; + for (int64_t i = 0; i < gsl::narrow(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& ovars) { tvm::Array 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); diff --git a/onnxruntime/core/codegen/mti/tensor/slice.h b/onnxruntime/core/codegen/mti/tensor/slice.h index 26f53650b1..ac5c943779 100644 --- a/onnxruntime/core/codegen/mti/tensor/slice.h +++ b/onnxruntime/core/codegen/mti/tensor/slice.h @@ -9,8 +9,10 @@ namespace onnxruntime { namespace tvm_codegen { tvm::Tensor Slice(const tvm::Tensor& X, - const tvm::Array& starts, - const tvm::Array& ends, + const std::vector& starts, + const std::vector& ends, + const std::vector& axes, + const std::vector& steps, const std::string& name = "slice"); } // namespace tvm_codegen diff --git a/onnxruntime/core/codegen/passes/op_ir_creator/tensor/expand.cc b/onnxruntime/core/codegen/passes/op_ir_creator/tensor/expand.cc new file mode 100644 index 0000000000..0f0e0cf098 --- /dev/null +++ b/onnxruntime/core/codegen/passes/op_ir_creator/tensor/expand.cc @@ -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& inputs, + const Node& node, + CodeGenContext& ctx_codegen, + tvm::Array& 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 diff --git a/onnxruntime/core/codegen/passes/op_ir_creator/tensor/shape_op.cc b/onnxruntime/core/codegen/passes/op_ir_creator/tensor/shape_op.cc new file mode 100644 index 0000000000..84761ecac1 --- /dev/null +++ b/onnxruntime/core/codegen/passes/op_ir_creator/tensor/shape_op.cc @@ -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& inputs, + const Node& node, + CodeGenContext& ctx_codegen, + tvm::Array& outputs) { + tvm::Tensor Y = Shape(inputs[0], node.Name() + "_Expand"); + outputs.push_back(Y); + return Status::OK(); +} + +} // namespace tvm_codegen +} // namespace onnxruntime diff --git a/onnxruntime/core/codegen/passes/op_ir_creator/tensor/slice.cc b/onnxruntime/core/codegen/passes/op_ir_creator/tensor/slice.cc index ecdd523616..2d9cac14fa 100644 --- a/onnxruntime/core/codegen/passes/op_ir_creator/tensor/slice.cc +++ b/onnxruntime/core/codegen/passes/op_ir_creator/tensor/slice.cc @@ -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& inputs, const Node& node, tvm::Array& outputs, const std::vector& starts, const std::vector& ends, - const std::vector& axes1) { + const std::vector& axes1, + const std::vector& steps1) { ORT_RETURN_IF_NOT(nullptr != node.InputDefs()[0]); - const ONNX_NAMESPACE::TensorShapeProto* shape_proto = node.InputDefs()[0]->Shape(); std::vector axes; if (axes1.size() == 0) { @@ -34,48 +31,14 @@ Status SliceCommon(const tvm::Array& inputs, axes = axes1; } - tvm::Array 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(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 shape; - for (size_t dim = 0; dim < gsl::narrow_cast(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 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 starts, ends; + std::vector starts, ends, steps; ORT_RETURN_IF_ERROR(info.GetAttrs("starts", starts)); ORT_RETURN_IF_ERROR(info.GetAttrs("ends", ends)); ORT_RETURN_IF_NOT(starts.size() == ends.size()); auto axes = info.GetAttrsOrDefault("axes"); - return SliceCommon(inputs, node, outputs, starts, ends, axes); + return SliceCommon(inputs, node, outputs, starts, ends, axes, steps); } } // namespace tvm_codegen diff --git a/onnxruntime/core/providers/nuphar/common/analysis/subgraph_codegen_stats.cc b/onnxruntime/core/providers/nuphar/common/analysis/subgraph_codegen_stats.cc index 836c3b9829..739b34086c 100644 --- a/onnxruntime/core/providers/nuphar/common/analysis/subgraph_codegen_stats.cc +++ b/onnxruntime/core/providers/nuphar/common/analysis/subgraph_codegen_stats.cc @@ -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(passes_[UseCountAnalysisOffset])->NodeUseCount(node) > CheapNodeTrueReuseCount; diff --git a/onnxruntime/core/providers/nuphar/common/utils.cc b/onnxruntime/core/providers/nuphar/common/utils.cc index d91cf1be97..a69a7547ee 100644 --- a/onnxruntime/core/providers/nuphar/common/utils.cc +++ b/onnxruntime/core/providers/nuphar/common/utils.cc @@ -35,38 +35,39 @@ bool HasUnknownShapeOnAxes(const NodeArg* def, std::vector& axes) { return false; } -Status GetSliceAxesFromTensorProto(std::vector& axes, - const ONNX_NAMESPACE::TensorProto& axes_tp) { +Status GetVectorInt64FromTensorProto(std::vector& 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 data(new char[tp_sz_in_bytes]); -#define UNPACK_TENSOR(T) \ - T* p = reinterpret_cast(data.get()); \ - ORT_RETURN_IF_ERROR(utils::UnpackTensor( \ - 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 tmp_axes(p, p + tp_sz_in_bytes / sizeof(T)); +#define UNPACK_TENSOR(T) \ + T* p = reinterpret_cast(data.get()); \ + ORT_RETURN_IF_ERROR(utils::UnpackTensor( \ + tp, \ + tp.raw_data().size() ? tp.raw_data().data() : nullptr, \ + tp.raw_data().size(), \ + p, \ + tp_sz_in_bytes / sizeof(T))); \ + std::vector 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(axis)); + for (auto axis : tmp_v) { + v.push_back(static_cast(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(); diff --git a/onnxruntime/core/providers/nuphar/common/utils.h b/onnxruntime/core/providers/nuphar/common/utils.h index 5077972fc3..1b7939b4b0 100644 --- a/onnxruntime/core/providers/nuphar/common/utils.h +++ b/onnxruntime/core/providers/nuphar/common/utils.h @@ -16,8 +16,8 @@ bool HasUnknownShapeOnAxis(const ConstPointerContainer>& d bool HasUnknownShapeOnAxes(const NodeArg* def, std::vector& axes); -Status GetSliceAxesFromTensorProto(std::vector& axes, - const ONNX_NAMESPACE::TensorProto& axes_tp); +Status GetVectorInt64FromTensorProto(std::vector& v, + const ONNX_NAMESPACE::TensorProto& tp); } // namespace nuphar } // namespace onnxruntime diff --git a/onnxruntime/core/providers/nuphar/compiler/x86/op_ir_creator/tensor/slice.cc b/onnxruntime/core/providers/nuphar/compiler/x86/op_ir_creator/tensor/slice.cc index d303f2f641..3eadeb0033 100644 --- a/onnxruntime/core/providers/nuphar/compiler/x86/op_ir_creator/tensor/slice.cc +++ b/onnxruntime/core/providers/nuphar/compiler/x86/op_ir_creator/tensor/slice.cc @@ -17,7 +17,8 @@ Status SliceCommon(const tvm::Array& inputs, tvm::Array& outputs, const std::vector& starts, const std::vector& ends, - const std::vector& axes); + const std::vector& axes, + const std::vector& steps); } // namespace tvm_codegen @@ -36,7 +37,7 @@ Status NUPHAR_TVM_X86_OP_IR_CREATOR_CLASS(Slice)::Evaluate( std::vector> slice_params; int version = ctx_codegen.GetCodeGenHandle()->domain_version_lookup_func(node.Domain()); if (version <= 9) { - std::vector starts, ends, axes; + std::vector starts, ends, axes, steps; ORT_RETURN_IF_ERROR(info.GetAttrs("starts", starts)); ORT_RETURN_IF_ERROR(info.GetAttrs("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()); } } - 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 diff --git a/onnxruntime/core/providers/nuphar/kernel.h b/onnxruntime/core/providers/nuphar/kernel.h index 1c28f8f3ae..d308e9dae7 100644 --- a/onnxruntime/core/providers/nuphar/kernel.h +++ b/onnxruntime/core/providers/nuphar/kernel.h @@ -96,6 +96,7 @@ class NupharKernelState { NUPHAR_OP(Equal, 11, DataTypeImpl::AllFixedSizeTensorTypes()) \ NUPHAR_OP(Erf, 9, DataTypeImpl::GetTensorType()) \ 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()) \ diff --git a/onnxruntime/core/providers/nuphar/nuphar_execution_provider.cc b/onnxruntime/core/providers/nuphar/nuphar_execution_provider.cc index 0679a1ca72..f0f803b30e 100644 --- a/onnxruntime/core/providers/nuphar/nuphar_execution_provider.cc +++ b/onnxruntime/core/providers/nuphar/nuphar_execution_provider.cc @@ -177,7 +177,8 @@ NupharExecutionProvider::GetCapability(const onnxruntime::GraphViewer& graph_vie std::vector> results; - auto is_supported_func = [&](const Node& node) { + typedef std::function 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 axes; + std::vector 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); diff --git a/onnxruntime/core/providers/nuphar/partition/graph_partitioner.cc b/onnxruntime/core/providers/nuphar/partition/graph_partitioner.cc index f7d32ff7fe..010b4a3bf1 100644 --- a/onnxruntime/core/providers/nuphar/partition/graph_partitioner.cc +++ b/onnxruntime/core/providers/nuphar/partition/graph_partitioner.cc @@ -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; } diff --git a/onnxruntime/core/providers/nuphar/scripts/model_editor.py b/onnxruntime/core/providers/nuphar/scripts/model_editor.py index 3beb47dd7f..9663aadba2 100644 --- a/onnxruntime/core/providers/nuphar/scripts/model_editor.py +++ b/onnxruntime/core/providers/nuphar/scripts/model_editor.py @@ -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) diff --git a/onnxruntime/core/providers/nuphar/scripts/node_factory.py b/onnxruntime/core/providers/nuphar/scripts/node_factory.py index 687c045f95..1c6aed54b6 100644 --- a/onnxruntime/core/providers/nuphar/scripts/node_factory.py +++ b/onnxruntime/core/providers/nuphar/scripts/node_factory.py @@ -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): diff --git a/onnxruntime/test/python/onnxruntime_test_python_nuphar.py b/onnxruntime/test/python/onnxruntime_test_python_nuphar.py index 81afbafa20..6bd7d33fa3 100644 --- a/onnxruntime/test/python/onnxruntime_test_python_nuphar.py +++ b/onnxruntime/test/python/onnxruntime_test_python_nuphar.py @@ -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