diff --git a/onnxruntime/core/providers/nnapi/nnapi_builtin/builders/helper.cc b/onnxruntime/core/providers/nnapi/nnapi_builtin/builders/helper.cc index ef0fa3a98a..b4241b92d7 100644 --- a/onnxruntime/core/providers/nnapi/nnapi_builtin/builders/helper.cc +++ b/onnxruntime/core/providers/nnapi/nnapi_builtin/builders/helper.cc @@ -1,6 +1,5 @@ -// -// Created by daquexian on 8/3/18. -// +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. #include #include diff --git a/onnxruntime/core/providers/nnapi/nnapi_builtin/builders/helper.h b/onnxruntime/core/providers/nnapi/nnapi_builtin/builders/helper.h index 4bbe194741..44825599bc 100644 --- a/onnxruntime/core/providers/nnapi/nnapi_builtin/builders/helper.h +++ b/onnxruntime/core/providers/nnapi/nnapi_builtin/builders/helper.h @@ -1,6 +1,6 @@ -// -// Created by daquexian on 5/21/18. -// +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + #pragma once #include 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 d5e0118dba..e74add6a14 100644 --- a/onnxruntime/core/providers/nnapi/nnapi_builtin/builders/op_builder.cc +++ b/onnxruntime/core/providers/nnapi/nnapi_builtin/builders/op_builder.cc @@ -1767,6 +1767,10 @@ Status ConcatOpBuilder::AddToModelBuilderImpl(ModelBuilder& model_builder, const class SqueezeOpBuilder : public BaseOpBuilder { public: void AddInitializersToSkip(ModelBuilder& model_builder, const Node& node) const override; + static Status AddSqueezeOp(ModelBuilder& model_builder, + const std::string& node_name, + const std::string& input, const std::string& output, + vector axes) ORT_MUST_USE_RESULT; private: Status AddToModelBuilderImpl(ModelBuilder& model_builder, const Node& node) const override ORT_MUST_USE_RESULT; @@ -1779,6 +1783,49 @@ void SqueezeOpBuilder::AddInitializersToSkip(ModelBuilder& model_builder, const } } +/* static */ Status SqueezeOpBuilder::AddSqueezeOp(ModelBuilder& model_builder, + const std::string& node_name, + const std::string& input, const std::string& output, + vector axes) { + auto& shaper(model_builder.GetShaper()); + const auto& operand_indices(model_builder.GetOperandIndices()); + const auto& operand_types(model_builder.GetOperandTypes()); + + const auto& input_shape(shaper[input]); + auto input_dims = input_shape.size(); + for (auto& axis : axes) { + axis = static_cast(HandleNegativeAxis(axis, input_dims)); + } + + // Despite the spec of ANEURALNETWORKS_SQUEEZE at + // https://developer.android.com/ndk/reference/group/neural-networks + // states, that the axes (input 1 of ANEURALNETWORKS_SQUEEZE) is optional. + // + // The actual code of NNAPI requires the axes to be provided + // https://android.googlesource.com/platform/frameworks/ml/+/master/nn/common/operations/Squeeze.cpp#31 + if (axes.empty()) { // Squeeze all + for (size_t i = 0; i < input_dims; i++) { + if (input_shape[i] == 1) + axes.push_back(i); + } + } + + const auto axes_name = model_builder.GetUniqueName(node_name + input + "_axes"); + Shape axes_dimen = {static_cast(axes.size())}; + const OperandType axes_operand_type(Type::TENSOR_INT32, axes_dimen); + ORT_RETURN_IF_ERROR(model_builder.AddOperandFromPersistMemoryBuffer(axes_name, axes.data(), axes_operand_type)); + + std::vector input_indices; + input_indices.push_back(operand_indices.at(input)); // input + input_indices.push_back(operand_indices.at(axes_name)); // axes + + ORT_RETURN_IF_ERROR(shaper.Squeeze(input, axes, output)); + const OperandType output_operand_type(operand_types.at(input).type, shaper[output]); + ORT_RETURN_IF_ERROR(model_builder.AddOperation(ANEURALNETWORKS_SQUEEZE, input_indices, + {output}, {output_operand_type}, {false})); + return Status::OK(); +} + /* static */ vector SqueezeOpBuilder::GetAxes(ModelBuilder& model_builder, const Node& node) { vector axes; // Squeeze opset 13 use input as axes @@ -1804,47 +1851,13 @@ void SqueezeOpBuilder::AddInitializersToSkip(ModelBuilder& model_builder, const } Status SqueezeOpBuilder::AddToModelBuilderImpl(ModelBuilder& model_builder, const Node& node) const { - auto& shaper(model_builder.GetShaper()); - const auto& operand_indices(model_builder.GetOperandIndices()); - const auto& operand_types(model_builder.GetOperandTypes()); - auto input = node.InputDefs()[0]->Name(); if (model_builder.IsOperandNHWC(input)) { // We want to transpose nhwc operand back to nchw before squeeze ORT_RETURN_IF_ERROR(GetNCHWInput(model_builder, node, 0, input)); } - NodeAttrHelper helper(node); - vector axes = GetAxes(model_builder, node); - const auto& input_shape(shaper[input]); - auto input_dims = input_shape.size(); - for (auto& axis : axes) { - axis = static_cast(HandleNegativeAxis(axis, input_dims)); - } - - if (axes.empty()) { // Squeeze all - for (size_t i = 0; i < input_dims; i++) { - if (input_shape[i] == 1) - axes.push_back(i); - } - } - - const auto axes_name = model_builder.GetUniqueName(node.Name() + input + "_axes"); - Shape axes_dimen = {static_cast(axes.size())}; - shaper.AddShape(axes_name, axes_dimen); - const OperandType axes_operand_type(Type::TENSOR_INT32, axes_dimen); - ORT_RETURN_IF_ERROR(model_builder.AddOperandFromPersistMemoryBuffer(axes_name, axes.data(), axes_operand_type)); - - std::vector input_indices; - input_indices.push_back(operand_indices.at(input)); // input - input_indices.push_back(operand_indices.at(axes_name)); // axes - - const auto& output = node.OutputDefs()[0]->Name(); - ORT_RETURN_IF_ERROR(shaper.Squeeze(input, axes, output)); - const OperandType output_operand_type(operand_types.at(input).type, shaper[output]); - ORT_RETURN_IF_ERROR(model_builder.AddOperation(ANEURALNETWORKS_SQUEEZE, input_indices, - {output}, {output_operand_type}, {false})); - return Status::OK(); + return AddSqueezeOp(model_builder, node.Name(), input, node.OutputDefs()[0]->Name(), GetAxes(model_builder, node)); } #pragma endregion diff --git a/onnxruntime/core/providers/nnapi/nnapi_builtin/builders/shaper.cc b/onnxruntime/core/providers/nnapi/nnapi_builtin/builders/shaper.cc index a060ca421b..9b06fcfc8a 100644 --- a/onnxruntime/core/providers/nnapi/nnapi_builtin/builders/shaper.cc +++ b/onnxruntime/core/providers/nnapi/nnapi_builtin/builders/shaper.cc @@ -1,3 +1,6 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + #include "core/providers/common.h" #include "core/providers/nnapi/nnapi_builtin/nnapi_lib/NeuralNetworksWrapper.h" @@ -374,17 +377,12 @@ Status Shaper::SqueezeImpl(const std::string& input_name, const std::string& output_name) { const Shape& input_dimen = shape_map_.at(input_name); int32_t input_size = input_dimen.size(); - size_t axes_size = axes.size(); std::unordered_set axes_to_be_squeezed; - if (axes_size == 0) { - for (int32_t idx = 0; idx < input_size; ++idx) { - if (input_dimen[idx] == 1) - axes_to_be_squeezed.insert(idx); - } - } else { - for (const auto& axis : axes) - axes_to_be_squeezed.insert(axis); - } + + // If the Op is squeezing all by not specifying axes, the axes is pre-populate + // with axes of all single dimensions by the caller + for (const auto& axis : axes) + axes_to_be_squeezed.insert(axis); // Make output dimensions std::vector output_dimen; @@ -394,6 +392,11 @@ Status Shaper::SqueezeImpl(const std::string& input_name, output_dimen.push_back(input_dimen[i]); } + // In case of a tensor has all 1's in dimension such as {1,1,1,1} and gets squeezed all + // the output shape will be {1} + if (output_dimen.empty()) + output_dimen.push_back(1); + shape_map_[output_name] = output_dimen; return Status::OK(); } diff --git a/onnxruntime/core/providers/nnapi/nnapi_builtin/builders/shaper.h b/onnxruntime/core/providers/nnapi/nnapi_builtin/builders/shaper.h index ee3ab17ea9..b9299454dc 100644 --- a/onnxruntime/core/providers/nnapi/nnapi_builtin/builders/shaper.h +++ b/onnxruntime/core/providers/nnapi/nnapi_builtin/builders/shaper.h @@ -1,3 +1,6 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + #pragma once #include diff --git a/onnxruntime/test/providers/cpu/tensor/squeeze_op_test.cc b/onnxruntime/test/providers/cpu/tensor/squeeze_op_test.cc index e07d9845a9..39561a6ba5 100644 --- a/onnxruntime/test/providers/cpu/tensor/squeeze_op_test.cc +++ b/onnxruntime/test/providers/cpu/tensor/squeeze_op_test.cc @@ -35,6 +35,15 @@ TEST(SqueezeOpTest, Squeeze_Empty_Axes_2) { test.Run(OpTester::ExpectResult::kExpectSuccess, "", {kTensorrtExecutionProvider}); } +TEST(SqueezeOpTest, Squeeze_Empty_Axes_3) { + OpTester test("Squeeze"); + // Squeeze all for all 1's shape will end up as a scalar + test.AddInput("data", {1, 1, 1, 1}, std::vector{1.0f}); + test.AddOutput("squeezed", {}, std::vector{1.0f}); + // TensorRT doesn't seem to support missing 'axes' + test.Run(OpTester::ExpectResult::kExpectSuccess, "", {kTensorrtExecutionProvider}); +} + TEST(SqueezeOpTest, Squeeze_1_int32) { OpTester test("Squeeze"); test.AddAttribute("axes", std::vector{0});