[NNAPI EP] Update squeeze ops (#5946)

* [NNAPI EP] Update squeeze ops

Author: guoyu-wang

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 <iostream>
 #include <string>

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

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<int32_t> 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<int32_t> 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<int32_t>(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<uint32_t>(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<uint32_t> 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<int32_t> SqueezeOpBuilder::GetAxes(ModelBuilder& model_builder, const Node& node) {
   vector<int32_t> 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<int32_t> axes = GetAxes(model_builder, node);
-  const auto& input_shape(shaper[input]);
-  auto input_dims = input_shape.size();
-  for (auto& axis : axes) {
-    axis = static_cast<int32_t>(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<uint32_t>(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<uint32_t> 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

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<int32_t> 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<uint32_t> 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();
 }

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

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<float>("data", {1, 1, 1, 1}, std::vector<float>{1.0f});
+  test.AddOutput<float>("squeezed", {}, std::vector<float>{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<int64_t>{0});