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