From a71f7d33395afb0c50bae744e959e0d181ec06a9 Mon Sep 17 00:00:00 2001 From: Edward Chen <18449977+edgchen1@users.noreply.github.com> Date: Fri, 22 Jul 2022 11:21:45 -0700 Subject: [PATCH] [NNAPI EP] Add some support for MatMul with batch inputs (#12261) MatMul allows multiplying batches of matrices. This change enables limited support of batch inputs in the NNAPI EP. Some limitations: - Broadcasting is not supported. A and B must have the same leading dimensions. - Only float inputs are supported. QDQ MatMul or QLinearMatMul with batch inputs is not supported yet. Note that NNAPI's ANEURALNETWORKS_BATCH_MATMUL is pretty much what we need, but it is only available from NNAPI feature level 6. This change composes a bunch of NNAPI operations to achieve a similar result but this is not ideal. --- .../nnapi/nnapi_builtin/builders/helper.cc | 19 +- .../nnapi/nnapi_builtin/builders/helper.h | 6 + .../nnapi_builtin/builders/op_builder.cc | 57 +-- .../builders/op_builder_helpers.cc | 356 ++++++++++++++++++ .../builders/op_builder_helpers.h | 51 +++ .../builders/op_support_checker.cc | 13 +- .../nnapi/nnapi_builtin/builders/shaper.cc | 30 +- .../nnapi/nnapi_builtin/builders/shaper.h | 3 + .../test/providers/cpu/math/matmul_test.cc | 42 ++- 9 files changed, 516 insertions(+), 61 deletions(-) create mode 100644 onnxruntime/core/providers/nnapi/nnapi_builtin/builders/op_builder_helpers.cc create mode 100644 onnxruntime/core/providers/nnapi/nnapi_builtin/builders/op_builder_helpers.h diff --git a/onnxruntime/core/providers/nnapi/nnapi_builtin/builders/helper.cc b/onnxruntime/core/providers/nnapi/nnapi_builtin/builders/helper.cc index 0091a7cd7f..c5ddfce065 100644 --- a/onnxruntime/core/providers/nnapi/nnapi_builtin/builders/helper.cc +++ b/onnxruntime/core/providers/nnapi/nnapi_builtin/builders/helper.cc @@ -1,21 +1,23 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. +#include "core/providers/nnapi/nnapi_builtin/builders/helper.h" + +#include #include +#include #include #include -#include "helper.h" - -#include "core/common/safeint.h" #include "core/common/logging/logging.h" +#include "core/common/safeint.h" #include "core/framework/tensorprotoutils.h" -#include "core/graph/graph.h" #include "core/graph/graph_viewer.h" +#include "core/graph/graph.h" #include "core/providers/common.h" +#include "core/providers/nnapi/nnapi_builtin/builders/op_support_checker.h" #include "core/providers/shared/node_unit/node_unit.h" #include "core/providers/shared/utils/utils.h" -#include "op_support_checker.h" namespace onnxruntime { namespace nnapi { @@ -377,6 +379,13 @@ std::string Shape2String(const std::vector& shape) { return os.str(); } +uint32_t ShapeSize(const Shape& shape, size_t begin_idx, size_t end_idx) { + ORT_ENFORCE(begin_idx <= end_idx && begin_idx <= shape.size(), + "Invalid indices: begin [", begin_idx, "], end [", end_idx, "], shape size [", shape.size(), "]"); + return std::accumulate(shape.begin() + begin_idx, shape.begin() + end_idx, + SafeInt{1}, std::multiplies>{}); +} + bool CheckIsInitializer(const InitializedTensorSet& initializers, const NodeUnit& node_unit, const std::string& input_name, const char* input_description) { if (!Contains(initializers, input_name)) { diff --git a/onnxruntime/core/providers/nnapi/nnapi_builtin/builders/helper.h b/onnxruntime/core/providers/nnapi/nnapi_builtin/builders/helper.h index 252a875627..dce4c9088d 100644 --- a/onnxruntime/core/providers/nnapi/nnapi_builtin/builders/helper.h +++ b/onnxruntime/core/providers/nnapi/nnapi_builtin/builders/helper.h @@ -4,6 +4,7 @@ #pragma once #include +#include #include "core/graph/basic_types.h" #include "core/providers/nnapi/nnapi_builtin/nnapi_lib/NeuralNetworksTypes.h" @@ -158,6 +159,11 @@ bool IsValidSupportedNodeGroup(const std::vector& supported_node_gr // Get string representation of a Shape std::string Shape2String(const std::vector& shape); +uint32_t ShapeSize(const Shape& shape, size_t begin_idx, size_t end_idx); +inline uint32_t ShapeSize(const Shape& shape) { + return ShapeSize(shape, 0, shape.size()); +} + // Check the given input is an initializer tensor // input_name is the name of the initializer // input_description is the string describing the input in the output message (if any) diff --git a/onnxruntime/core/providers/nnapi/nnapi_builtin/builders/op_builder.cc b/onnxruntime/core/providers/nnapi/nnapi_builtin/builders/op_builder.cc index 7ef9607352..0b2336b13f 100644 --- a/onnxruntime/core/providers/nnapi/nnapi_builtin/builders/op_builder.cc +++ b/onnxruntime/core/providers/nnapi/nnapi_builtin/builders/op_builder.cc @@ -13,6 +13,7 @@ #include "core/providers/shared/utils/utils.h" #include "core/providers/shared/node_unit/node_unit.h" #include "core/providers/cpu/tensor/slice_helper.h" +#include "core/providers/nnapi/nnapi_builtin/builders/op_builder_helpers.h" #include "helper.h" #include "model_builder.h" #include "op_support_checker.h" @@ -29,37 +30,9 @@ struct OpBuilderRegistrations { std::unordered_map op_builder_map; }; -#define ADD_SCALAR_OPERAND(model_builder, input_indices, scalar_value) \ - do { \ - uint32_t _index = 0; \ - ORT_RETURN_IF_ERROR(model_builder.AddOperandFromScalar(scalar_value, _index)); \ - input_indices.push_back(_index); \ - } while (0) - -Status AddTransposeOperator(ModelBuilder& model_builder, - const std::string& input, - const std::string& perm_name, - std::vector perm, - const std::string& output) { - auto& shaper(model_builder.GetShaper()); - const auto& operand_indices(model_builder.GetOperandIndices()); - const auto& operand_types(model_builder.GetOperandTypes()); - - std::vector input_indices; - input_indices.push_back(operand_indices.at(input)); // input - - Shape perm_dimen = {SafeInt(perm.size())}; - OperandType perm_operand_type(Type::TENSOR_INT32, perm_dimen); - ORT_RETURN_IF_ERROR(model_builder.AddOperandFromPersistMemoryBuffer(perm_name, perm.data(), perm_operand_type)); - uint32_t perm_idx = operand_indices.at(perm_name); - - input_indices.push_back(perm_idx); // permutation - ORT_RETURN_IF_ERROR(shaper.Transpose(input, perm, output)); - OperandType output_operand_type = operand_types.at(input); - output_operand_type.SetDimensions(shaper[output]); - return model_builder.AddOperation(ANEURALNETWORKS_TRANSPOSE, input_indices, {output}, - {output_operand_type}); -} +#define ADD_SCALAR_OPERAND(model_builder, input_indices, scalar_value) \ + ORT_RETURN_IF_ERROR(onnxruntime::nnapi::op_builder_helpers::AddScalarOperand( \ + model_builder, input_indices, scalar_value)) static Status AddBinaryOperator(int32_t op_type, ModelBuilder& model_builder, @@ -776,7 +749,7 @@ Status TransposeOpBuilder::AddToModelBuilderImpl(ModelBuilder& model_builder, co std::string perm_name = model_builder.GetUniqueName(node_unit.Name() + input + "perm"); - ORT_RETURN_IF_ERROR(AddTransposeOperator(model_builder, input, perm_name, perm, output)); + ORT_RETURN_IF_ERROR(op_builder_helpers::AddNnapiTranspose(model_builder, input, perm_name, perm, output)); return Status::OK(); } @@ -905,16 +878,9 @@ bool ReshapeOpBuilder::IsQuantizedOp(const NodeUnit& node_unit) const { model_builder.RegisterOperand(output, operand_indices.at(input), output_operand_type); } else { // We still need to perform a reshape here - // Add input - std::vector input_indices; - input_indices.push_back(operand_indices.at(input)); - // Add new shape - Shape shape_dimen = {static_cast(shape.size())}; std::string shape_name = model_builder.GetUniqueName(node_unit.Name() + input + "newshape"); - OperandType shape_operand_type(Type::TENSOR_INT32, shape_dimen); - ORT_RETURN_IF_ERROR(model_builder.AddOperandFromPersistMemoryBuffer(shape_name, shape.data(), shape_operand_type)); - input_indices.push_back(operand_indices.at(shape_name)); - ORT_RETURN_IF_ERROR(model_builder.AddOperation(ANEURALNETWORKS_RESHAPE, input_indices, {output}, {output_operand_type})); + ORT_RETURN_IF_ERROR(op_builder_helpers::AddNnapiReshape(model_builder, input, shape_name, shape, output, + &shaper[output])); } return Status::OK(); @@ -1722,6 +1688,11 @@ bool GemmOpBuilder::IsQuantizedOp(const NodeUnit& node_unit) const { } void GemmOpBuilder::AddInitializersToSkip(ModelBuilder& model_builder, const NodeUnit& node_unit) const { + if (op_builder_helpers::IsSupportedBatchMatMul(node_unit, model_builder.GetNNAPIFeatureLevel())) { + // no initializers to skip for batch matmul + return; + } + const auto& inputs = node_unit.Inputs(); if (IsQuantizedOp(node_unit)) { if (node_unit.OpType() == "QLinearMatMul" || node_unit.OpType() == "MatMul") { // QLinearMatMul/QDQMatMul @@ -1758,6 +1729,10 @@ void GemmOpBuilder::AddInitializersToSkip(ModelBuilder& model_builder, const Nod } Status GemmOpBuilder::AddToModelBuilderImpl(ModelBuilder& model_builder, const NodeUnit& node_unit) const { + if (op_builder_helpers::IsSupportedBatchMatMul(node_unit, model_builder.GetNNAPIFeatureLevel())) { + return op_builder_helpers::BuildBatchMatMul(model_builder, node_unit); + } + auto& shaper(model_builder.GetShaper()); const auto& operand_indices(model_builder.GetOperandIndices()); const auto& operand_types(model_builder.GetOperandTypes()); diff --git a/onnxruntime/core/providers/nnapi/nnapi_builtin/builders/op_builder_helpers.cc b/onnxruntime/core/providers/nnapi/nnapi_builtin/builders/op_builder_helpers.cc new file mode 100644 index 0000000000..27e12ab649 --- /dev/null +++ b/onnxruntime/core/providers/nnapi/nnapi_builtin/builders/op_builder_helpers.cc @@ -0,0 +1,356 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "core/providers/nnapi/nnapi_builtin/builders/op_builder_helpers.h" + +#include + +#include "gsl/gsl" + +#include "core/common/safeint.h" +#include "core/graph/node_arg.h" +#include "core/providers/common.h" +#include "core/providers/nnapi/nnapi_builtin/builders/helper.h" +#include "core/providers/nnapi/nnapi_builtin/nnapi_lib/NeuralNetworksWrapper.h" + +namespace onnxruntime::nnapi::op_builder_helpers { + +using android::nn::wrapper::OperandType, android::nn::wrapper::Type; + +Status AddNnapiTranspose(ModelBuilder& model_builder, + const std::string& data_input, + const std::string& perm_input, const std::vector& perm, + const std::string& output) { + auto& shaper(model_builder.GetShaper()); + const auto& operand_indices(model_builder.GetOperandIndices()); + const auto& operand_types(model_builder.GetOperandTypes()); + + std::vector input_indices; + input_indices.push_back(operand_indices.at(data_input)); // input + + Shape perm_dimen = {SafeInt(perm.size())}; + OperandType perm_operand_type(Type::TENSOR_INT32, perm_dimen); + ORT_RETURN_IF_ERROR(model_builder.AddOperandFromPersistMemoryBuffer(perm_input, perm.data(), perm_operand_type)); + uint32_t perm_idx = operand_indices.at(perm_input); + + input_indices.push_back(perm_idx); // permutation + ORT_RETURN_IF_ERROR(shaper.Transpose(data_input, perm, output)); + OperandType output_operand_type = operand_types.at(data_input); + output_operand_type.SetDimensions(shaper[output]); + return model_builder.AddOperation(ANEURALNETWORKS_TRANSPOSE, input_indices, {output}, + {output_operand_type}); +} + +Status AddNnapiReshape(ModelBuilder& model_builder, + const std::string& data_input, + const std::string& shape_input, const std::vector& shape_value, + const std::string& output, const Shape* output_shape) { + if (output_shape == nullptr) { + auto& shaper = model_builder.GetShaper(); + ORT_RETURN_IF_ERROR(shaper.Reshape(data_input, shape_value, output)); + output_shape = &shaper[output]; + } + + const auto& operand_indices = model_builder.GetOperandIndices(); + const auto& operand_types = model_builder.GetOperandTypes(); + + // Add input + std::vector input_indices; + input_indices.push_back(operand_indices.at(data_input)); + + // Add new shape + const Shape shape_dimen{static_cast(shape_value.size())}; + const OperandType shape_operand_type{Type::TENSOR_INT32, shape_dimen}; + ORT_RETURN_IF_ERROR(model_builder.AddOperandFromPersistMemoryBuffer(shape_input, shape_value.data(), + shape_operand_type)); + input_indices.push_back(operand_indices.at(shape_input)); + + // For reshape, the output type should be the same as the input type except the shape is different + OperandType output_operand_type{operand_types.at(data_input)}; + output_operand_type.SetDimensions(*output_shape); + ORT_RETURN_IF_ERROR(model_builder.AddOperation(ANEURALNETWORKS_RESHAPE, + input_indices, {output}, {output_operand_type})); + + return Status::OK(); +} + +Status AddNnapiSplit(ModelBuilder& model_builder, + const std::string& input, + int32_t axis, + const std::vector& outputs) { + const auto& operand_indices = model_builder.GetOperandIndices(); + const auto& operand_types = model_builder.GetOperandTypes(); + auto& shaper = model_builder.GetShaper(); + + const auto input_rank = shaper[input].size(); + axis = static_cast(HandleNegativeAxis(axis, input_rank)); + + ORT_RETURN_IF_ERROR(shaper.Split(input, axis, outputs)); + + std::vector input_indices; + input_indices.push_back(operand_indices.at(input)); + ORT_RETURN_IF_ERROR(AddScalarOperand(model_builder, input_indices, axis)); + const auto count = gsl::narrow(outputs.size()); + ORT_RETURN_IF_ERROR(AddScalarOperand(model_builder, input_indices, count)); + + const OperandType& input_operand_type = operand_types.at(input); + std::vector output_operand_types; + output_operand_types.reserve(count); + std::transform(outputs.begin(), outputs.end(), std::back_inserter(output_operand_types), + [&](const std::string& output) { + OperandType output_operand_type = input_operand_type; + output_operand_type.SetDimensions(shaper[output]); + return output_operand_type; + }); + + ORT_RETURN_IF_ERROR(model_builder.AddOperation(ANEURALNETWORKS_SPLIT, + input_indices, outputs, output_operand_types)); + + return Status::OK(); +} + +bool IsSupportedBatchMatMul(const NodeUnit& node_unit, int32_t nnapi_feature_level) { + // Currently, batch MatMul is composed of various operations including ANEURALNETWORKS_SPLIT which requires + // ANEURALNETWORKS_FEATURE_LEVEL_3. + const auto min_nnapi_feature_level = ANEURALNETWORKS_FEATURE_LEVEL_3; + if (nnapi_feature_level < min_nnapi_feature_level) { + LOGS_DEFAULT(VERBOSE) << "Minimum NNAPI feature level required: " << min_nnapi_feature_level + << ", actual: " << nnapi_feature_level; + return false; + } + + // Only support non-QDQ MatMul for now. + // TODO could be expanded to support QLinearMatMul and QDQ MatMul + if (node_unit.UnitType() != NodeUnit::Type::SingleNode || + node_unit.OpType() != "MatMul") { + LOGS_DEFAULT(VERBOSE) << "Unsupported op type: " + << (node_unit.UnitType() == NodeUnit::Type::QDQGroup ? "QDQ " : "") << node_unit.OpType(); + return false; + } + + const auto& inputs = node_unit.Inputs(); + + // Verify shapes + // A and B should have at least three dimensions* and have the same leading dimensions except for the last two. + // [*] Having two dimensions is valid for a MatMul but for simplicity we don't support it in the current batch + // MatMul implementation. That case is handled by the regular Gemm/MatMul op building logic. + Shape a_shape; + if (!GetShape(inputs[0].node_arg, a_shape)) { + return false; + } + + Shape b_shape; + if (!GetShape(inputs[1].node_arg, b_shape)) { + return false; + } + + if (a_shape.size() < 3 || + a_shape.size() != b_shape.size() || + !std::equal(a_shape.begin(), a_shape.end() - 2, + b_shape.begin(), b_shape.end() - 2)) { + LOGS_DEFAULT(VERBOSE) + << "A and B must have at least three dimensions and have the same leading dimensions except for the last two. " + << "A shape: " << Shape2String(a_shape) << ", B shape: " << Shape2String(b_shape); + return false; + } + + // Verify type + int32_t a_type; + if (!GetType(inputs[0].node_arg, a_type)) { + return false; + } + + // Only support float for now. + // TODO could be expanded to support other types + if (a_type != ONNX_NAMESPACE::TensorProto_DataType_FLOAT) { + LOGS_DEFAULT(VERBOSE) << "Unsupported element data type: " << a_type; + return false; + } + + return true; +} + +Status BuildBatchMatMul(ModelBuilder& model_builder, const NodeUnit& node_unit) { + // we will implement batch MatMul by composing NNAPI operations + // this could be replaced with ANEURALNETWORKS_BATCH_MATMUL when that is more widely supported + + // assuming A and B have at least three dimensions and the same leading dimensions other than the last two + + const auto& inputs = node_unit.Inputs(); + + Shape a_shape; + ORT_RETURN_IF_NOT(GetShape(inputs[0].node_arg, a_shape), "Failed to get A's shape."); + + Shape b_shape; + ORT_RETURN_IF_NOT(GetShape(inputs[1].node_arg, b_shape), "Failed to get B's shape."); + + const std::string& a = inputs[0].node_arg.Name(); + const std::string& b = inputs[1].node_arg.Name(); + + const std::string& output = node_unit.Outputs()[0].node_arg.Name(); + + std::vector gemm_a_inputs{a}; + std::vector gemm_b_inputs{b}; + + const auto m = a_shape[a_shape.size() - 2], + k = a_shape[a_shape.size() - 1], + n = b_shape[b_shape.size() - 1]; + + const bool reshape_leading_dimensions = a_shape.size() > 3; + const auto batch_size = ShapeSize(a_shape, 0, a_shape.size() - 2); + + auto add_reshape = [&model_builder](const std::string& input, const Shape& new_shape, + const std::string& output) -> Status { + const std::string new_shape_name = model_builder.GetUniqueName(input + "/new_shape"); + std::vector new_shape_i32{}; + new_shape_i32.reserve(new_shape.size()); + std::transform(new_shape.begin(), new_shape.end(), std::back_inserter(new_shape_i32), + [](uint32_t d) { return gsl::narrow(d); }); + ORT_RETURN_IF_ERROR(AddNnapiReshape(model_builder, input, new_shape_name, new_shape_i32, output, nullptr)); + return Status::OK(); + }; + + auto add_reshape_generate_output = [&model_builder, &add_reshape](const std::string& input, const Shape& new_shape, + std::string& output) -> Status { + std::string reshaped = model_builder.GetUniqueName(input + "/reshaped"); + ORT_RETURN_IF_ERROR(add_reshape(input, new_shape, reshaped)); + output = std::move(reshaped); + return Status::OK(); + }; + + // collapse leading dimensions to a single one + if (reshape_leading_dimensions) { + const Shape a_new_shape_value = {batch_size, m, k}, + b_new_shape_value = {batch_size, k, n}; + std::string a_reshaped, b_reshaped; + + ORT_RETURN_IF_ERROR(add_reshape_generate_output(gemm_a_inputs.front(), a_new_shape_value, a_reshaped)); + gemm_a_inputs.front() = a_reshaped; + + ORT_RETURN_IF_ERROR(add_reshape_generate_output(gemm_b_inputs.front(), b_new_shape_value, b_reshaped)); + gemm_b_inputs.front() = b_reshaped; + } + + // transpose B + { + const std::string b_new_perm = model_builder.GetUniqueName(b + "/new_perm"), + b_transposed = model_builder.GetUniqueName(b + "/transposed"); + ORT_RETURN_IF_ERROR(AddNnapiTranspose(model_builder, gemm_b_inputs.front(), b_new_perm, {0, 2, 1}, b_transposed)); + gemm_b_inputs.front() = b_transposed; + } + + // split batch + { + auto add_split = [&model_builder, batch_size](const std::string& input, + std::vector& outputs_result) -> Status { + std::vector outputs; + outputs.reserve(batch_size); + for (size_t i = 0; i < batch_size; ++i) { + outputs.push_back(model_builder.GetUniqueName(MakeString(input, "/split_", i))); + } + ORT_RETURN_IF_ERROR(AddNnapiSplit(model_builder, input, 0, outputs)); + outputs_result = std::move(outputs); + return Status::OK(); + }; + + std::vector a_split_outputs; + ORT_RETURN_IF_ERROR(add_split(gemm_a_inputs.front(), a_split_outputs)); + gemm_a_inputs = std::move(a_split_outputs); + + std::vector b_split_outputs; + ORT_RETURN_IF_ERROR(add_split(gemm_b_inputs.front(), b_split_outputs)); + gemm_b_inputs = std::move(b_split_outputs); + } + + // GEMM per matrix pair + std::vector gemm_outputs; + gemm_outputs.reserve(batch_size); + { + const std::string bias = model_builder.GetUniqueName(node_unit.Name() + "/zero_bias"); + { + if (model_builder.GetOperandTypes().at(b).type != Type::TENSOR_FLOAT32) { + ORT_NOT_IMPLEMENTED("Only float input is supported now."); + } + const Shape bias_shape{n}; + const std::vector buffer(n, 0.0f); + const OperandType bias_operand_type(Type::TENSOR_FLOAT32, bias_shape); + ORT_RETURN_IF_ERROR(model_builder.AddOperandFromPersistMemoryBuffer(bias, buffer.data(), bias_operand_type)); + } + + auto add_fc = [&model_builder, &bias](const std::string& a, const std::string& b_transposed, + const std::string& output) -> Status { + const auto& operand_indices = model_builder.GetOperandIndices(); + const auto& operand_types = model_builder.GetOperandTypes(); + auto& shaper = model_builder.GetShaper(); + std::vector input_indices; + input_indices.push_back(operand_indices.at(a)); // A + input_indices.push_back(operand_indices.at(b_transposed)); // B' + input_indices.push_back(operand_indices.at(bias)); // C + int32_t fuse_code = ANEURALNETWORKS_FUSED_NONE; + ORT_RETURN_IF_ERROR(AddScalarOperand(model_builder, input_indices, fuse_code)); + + ORT_RETURN_IF_ERROR(shaper.FC(a, b_transposed, output)); + const OperandType output_operand_type(operand_types.at(a).type, shaper[output]); + ORT_RETURN_IF_ERROR(model_builder.AddOperation(ANEURALNETWORKS_FULLY_CONNECTED, input_indices, + {output}, {output_operand_type})); + return Status::OK(); + }; + + for (uint32_t i = 0; i < batch_size; ++i) { + const auto &gemm_a_input = gemm_a_inputs[i], + &gemm_b_input = gemm_b_inputs[i]; + + // make inputs 2D ([1, x, y] -> [x, y]) + std::string a_2d, b_transposed_2d; + ORT_RETURN_IF_ERROR(add_reshape_generate_output(gemm_a_input, Shape{m, k}, a_2d)); + ORT_RETURN_IF_ERROR(add_reshape_generate_output(gemm_b_input, Shape{n, k}, b_transposed_2d)); + + const std::string gemm_output = model_builder.GetUniqueName(MakeString(node_unit.Name(), "/gemm_", i)); + ORT_RETURN_IF_ERROR(add_fc(a_2d, b_transposed_2d, gemm_output)); + + // reshape output for concatenation ([x, y] -> [1, x, y]) + std::string gemm_output_3d; + ORT_RETURN_IF_ERROR(add_reshape_generate_output(gemm_output, Shape{1, m, n}, gemm_output_3d)); + + gemm_outputs.push_back(gemm_output_3d); + } + } + + // concat batch + const std::string joined_gemm_output = + reshape_leading_dimensions ? model_builder.GetUniqueName(node_unit.Name() + "/joined_gemm_output") : output; + { + auto add_concat = [&model_builder](const std::vector& inputs, + const std::string& output) -> Status { + const auto& operand_indices = model_builder.GetOperandIndices(); + const auto& operand_types = model_builder.GetOperandTypes(); + auto& shaper = model_builder.GetShaper(); + std::vector input_indices; + input_indices.reserve(inputs.size() + 1); + std::transform(inputs.begin(), inputs.end(), std::back_inserter(input_indices), + [&operand_indices](const std::string& input) { return operand_indices.at(input); }); + const int32_t axis = 0; + ORT_RETURN_IF_ERROR(AddScalarOperand(model_builder, input_indices, axis)); + ORT_RETURN_IF_ERROR(shaper.Concat(inputs, axis, output)); + OperandType output_operand_type = operand_types.at(inputs[0]); + output_operand_type.SetDimensions(shaper[output]); + ORT_RETURN_IF_ERROR(model_builder.AddOperation(ANEURALNETWORKS_CONCATENATION, + input_indices, {output}, {output_operand_type})); + return Status::OK(); + }; + + ORT_RETURN_IF_ERROR(add_concat(gemm_outputs, joined_gemm_output)); + } + + // reshape to original dimensions + if (reshape_leading_dimensions) { + Shape new_shape = a_shape; + new_shape[new_shape.size() - 2] = m; + new_shape[new_shape.size() - 1] = n; + ORT_RETURN_IF_ERROR(add_reshape(joined_gemm_output, new_shape, output)); + } + + return Status::OK(); +} + +} // namespace onnxruntime::nnapi::op_builder_helpers diff --git a/onnxruntime/core/providers/nnapi/nnapi_builtin/builders/op_builder_helpers.h b/onnxruntime/core/providers/nnapi/nnapi_builtin/builders/op_builder_helpers.h new file mode 100644 index 0000000000..20ffaf230d --- /dev/null +++ b/onnxruntime/core/providers/nnapi/nnapi_builtin/builders/op_builder_helpers.h @@ -0,0 +1,51 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once + +#include +#include + +#include "core/common/common.h" +#include "core/providers/nnapi/nnapi_builtin/builders/helper.h" +#include "core/providers/nnapi/nnapi_builtin/builders/model_builder.h" +#include "core/providers/shared/node_unit/node_unit.h" + +namespace onnxruntime::nnapi::op_builder_helpers { + +// adds a scalar operand to the NNAPI model and appends its index to `input_indices` +template +Status AddScalarOperand(ModelBuilder& model_builder, std::vector& input_indices, T scalar_value) { + uint32_t index = 0; + ORT_RETURN_IF_ERROR(model_builder.AddOperandFromScalar(std::move(scalar_value), index)); + input_indices.push_back(index); + return Status::OK(); +} + +// adds ANEURALNETWORKS_TRANSPOSE operation +Status AddNnapiTranspose(ModelBuilder& model_builder, + const std::string& data_input, + const std::string& perm_input, + const std::vector& perm, + const std::string& output); + +// adds ANEURALNETWORKS_RESHAPE operation +Status AddNnapiReshape(ModelBuilder& model_builder, + const std::string& data_input, + const std::string& shape_input, const std::vector& shape_value, + const std::string& output, const Shape* output_shape); + +// adds ANEURALNETWORKS_SPLIT operation +Status AddNnapiSplit(ModelBuilder& model_builder, + const std::string& input, + int32_t axis, + const std::vector& outputs); + +// checks whether batch MatMul in the given NodeUnit is supported by NNAPI EP +bool IsSupportedBatchMatMul(const NodeUnit& node_unit, int32_t nnapi_feature_level); + +// builds a batch MatMul in the NNAPI model from the given NodeUnit +// note: the pre-conditions of this function are checked in IsSupportedBatchMatMul() +Status BuildBatchMatMul(ModelBuilder& model_builder, const NodeUnit& node_unit); + +} // namespace onnxruntime::nnapi::op_builder_helpers diff --git a/onnxruntime/core/providers/nnapi/nnapi_builtin/builders/op_support_checker.cc b/onnxruntime/core/providers/nnapi/nnapi_builtin/builders/op_support_checker.cc index 092e5b451c..f5a033766a 100644 --- a/onnxruntime/core/providers/nnapi/nnapi_builtin/builders/op_support_checker.cc +++ b/onnxruntime/core/providers/nnapi/nnapi_builtin/builders/op_support_checker.cc @@ -8,9 +8,10 @@ #include "core/framework/tensorprotoutils.h" #include "core/graph/graph.h" #include "core/providers/common.h" +#include "core/providers/nnapi/nnapi_builtin/builders/helper.h" +#include "core/providers/nnapi/nnapi_builtin/builders/op_builder_helpers.h" #include "core/providers/shared/node_unit/node_unit.h" #include "core/providers/shared/utils/utils.h" -#include "helper.h" namespace onnxruntime { namespace nnapi { @@ -1431,6 +1432,16 @@ int GemmOpSupportChecker::GetMinSupportedOpSet(const NodeUnit& node_unit) const bool GemmOpSupportChecker::IsOpSupportedImpl(const InitializedTensorSet& initializers, const NodeUnit& node_unit, const OpSupportCheckParams& params) const { + // check batch matmul first, then fall back to checking single gemm/matmul + { + const bool is_supported_batch_matmul = + op_builder_helpers::IsSupportedBatchMatMul(node_unit, params.android_feature_level); + LOGS_DEFAULT(VERBOSE) << "Supported batch matmul: [" << is_supported_batch_matmul << "]"; + if (is_supported_batch_matmul) { + return true; + } + } + const auto& op_type = node_unit.OpType(); const auto& inputs = node_unit.Inputs(); const bool is_qlinear_matmul = op_type == "QLinearMatMul"; diff --git a/onnxruntime/core/providers/nnapi/nnapi_builtin/builders/shaper.cc b/onnxruntime/core/providers/nnapi/nnapi_builtin/builders/shaper.cc index 788cd769ac..7cf451123e 100644 --- a/onnxruntime/core/providers/nnapi/nnapi_builtin/builders/shaper.cc +++ b/onnxruntime/core/providers/nnapi/nnapi_builtin/builders/shaper.cc @@ -123,6 +123,11 @@ Status Shaper::Concat(const std::vector& input_names, SHAPER_FUNC(Concat, input_names, axis, output_name); } +Status Shaper::Split(const std::string& input_name, int32_t axis, + const std::vector& output_names) { + SHAPER_FUNC(Split, input_name, axis, output_names); +} + Status Shaper::Squeeze(const std::string& input_name, const std::vector& axes, const std::string& output_name) { @@ -362,13 +367,6 @@ Status Shaper::ConcatImpl(const std::vector& input_names, std::vector dimens; for (const auto& input_name : input_names) { const Shape& dimen = shape_map_.at(input_name); - if (!dimens.empty()) { - for (size_t i = 0; i < dimens[0].size(); i++) { - if ((int32_t)i == axis) - continue; - } - } - dimens.push_back(dimen); } @@ -388,6 +386,24 @@ Status Shaper::ConcatImpl(const std::vector& input_names, return Status::OK(); } +Status Shaper::SplitImpl(const std::string& input_name, int32_t axis, + const std::vector& output_names) { + const auto& input_shape = shape_map_.at(input_name); + const auto count = output_names.size(); + + ORT_RETURN_IF_NOT(input_shape[axis] % count == 0, + "count [", count, "] does not evenly divide dimension ", axis, " [", input_shape[axis], "]"); + + Shape output_shape = input_shape; + output_shape[axis] = input_shape[axis] / count; + + for (const auto& output_name : output_names) { + shape_map_[output_name] = output_shape; + } + + return Status::OK(); +} + Status Shaper::SqueezeImpl(const std::string& input_name, const std::vector& axes, const std::string& output_name) { diff --git a/onnxruntime/core/providers/nnapi/nnapi_builtin/builders/shaper.h b/onnxruntime/core/providers/nnapi/nnapi_builtin/builders/shaper.h index 58b4ef4398..31eb9aa9fa 100644 --- a/onnxruntime/core/providers/nnapi/nnapi_builtin/builders/shaper.h +++ b/onnxruntime/core/providers/nnapi/nnapi_builtin/builders/shaper.h @@ -57,6 +57,8 @@ class Shaper { common::Status Concat(const std::vector& input_names, const int32_t axis, const std::string& output_name); + common::Status Split(const std::string& input_name, int32_t axis, const std::vector& output_names); + common::Status Squeeze(const std::string& input_name, const std::vector& axes, const std::string& output_name); common::Status DepthToSpace(const std::string& input_name, @@ -122,6 +124,7 @@ class Shaper { common::Status IdentityImpl(const std::string& input_name, const std::string& output_name); common::Status FCImpl(const std::string& input1_name, const std::string& input2_name, const std::string& output_name); common::Status ConcatImpl(const std::vector& input_names, const int32_t axis, const std::string& output_name); + common::Status SplitImpl(const std::string& input_name, int32_t axis, const std::vector& output_names); common::Status SqueezeImpl(const std::string& input_names, const std::vector& axes, const std::string& output_name); common::Status DepthToSpaceImpl(const std::string& input_names, const int32_t blocksize, diff --git a/onnxruntime/test/providers/cpu/math/matmul_test.cc b/onnxruntime/test/providers/cpu/math/matmul_test.cc index 5633d12d75..31f32f932d 100644 --- a/onnxruntime/test/providers/cpu/math/matmul_test.cc +++ b/onnxruntime/test/providers/cpu/math/matmul_test.cc @@ -4,6 +4,7 @@ #include "gtest/gtest.h" #include "test/providers/provider_test_utils.h" #include "test/common/cuda_op_test_utils.h" +#include "test/common/tensor_op_test_utils.h" #include "default_providers.h" namespace onnxruntime { @@ -99,21 +100,49 @@ std::vector> GenerateTestCases() { {3, 0}, {}}); + test_cases.push_back( + {"test 3D batch", + {3, 1, 3}, + {3, 3, 2}, + {3, 1, 2}, + { + // clang-format off + 10, 13, + 100, 112, + 298, 319, + // clang-format on + }}); + + test_cases.push_back( + {"test 4D batch", + {2, 2, 1, 3}, + {2, 2, 3, 2}, + {2, 2, 1, 2}, + { + // clang-format off + 10, 13, + 100, 112, + 298, 319, + 604, 634, + // clang-format on + }}); + return test_cases; } template void RunMatMulTest(int32_t opset_version, bool is_a_constant, bool is_b_constant) { - std::vector common_input_vals{0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11}; for (auto t : GenerateTestCases()) { + SCOPED_TRACE("test case: " + t.name); + OpTester test("MatMul", opset_version); int64_t size0 = TensorShape::FromExistingBuffer(t.input0_dims).SizeHelper(0, t.input0_dims.size()); - std::vector input0_vals(common_input_vals.cbegin(), common_input_vals.cbegin() + size0); + std::vector input0_vals = ValueRange(size0); test.AddInput("A", t.input0_dims, input0_vals, is_a_constant); int64_t size1 = TensorShape::FromExistingBuffer(t.input1_dims).SizeHelper(0, t.input1_dims.size()); - std::vector input1_vals(common_input_vals.cbegin(), common_input_vals.cbegin() + size1); + std::vector input1_vals = ValueRange(size1); test.AddInput("B", t.input1_dims, input1_vals, is_b_constant); test.AddOutput("Y", t.expected_dims, t.expected_vals); @@ -121,7 +150,7 @@ void RunMatMulTest(int32_t opset_version, bool is_a_constant, bool is_b_constant // OpenVINO EP: Disabled temporarily matmul broadcasting not fully supported // Disable TensorRT because of unsupported data type std::unordered_set excluded_providers{kTensorrtExecutionProvider, kOpenVINOExecutionProvider}; - if (is_b_constant) { + if (t.name == "test 2D empty input") { // NNAPI: currently fails for the "test 2D empty input" case excluded_providers.insert(kNnapiExecutionProvider); } @@ -136,7 +165,6 @@ void RunMatMulTest(int32_t opset_version) { TEST(MathOpTest, MatMulFloatType) { RunMatMulTest(7, false, false); - RunMatMulTest(7, false, true); } TEST(MathOpTest, MatMulDoubleType) { @@ -190,7 +218,7 @@ TEST(MathOpTest, MatMul_Float16) { test.AddInput("A", {2, 4}, f_A); test.AddInput("B", {4, 3}, f_B); test.AddOutput("Y", {2, 3}, f_Y); - test.Run(OpTester::ExpectResult::kExpectSuccess, "", {kTensorrtExecutionProvider}); //TensorRT: fp16 is not supported + test.Run(OpTester::ExpectResult::kExpectSuccess, "", {kTensorrtExecutionProvider}); // TensorRT: fp16 is not supported } #endif @@ -213,7 +241,7 @@ TEST(MathOpTest, MatMul_BFloat16) { execution_providers.push_back(DefaultCudaExecutionProvider()); #elif USE_ROCM execution_providers.push_back(DefaultRocmExecutionProvider()); -#endif +#endif test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &execution_providers); } #endif