[CoreML EP] Add Squeeze Op support (#7730)

* add squeeze op builder initial

* fix mistakes

* modify

* enable UT passed and minor refine

* minor formatting

* address comment

Co-authored-by: rachguo <rachguo@rachguos-Mac-mini.local>

onnxruntime/core/providers/coreml/builders/impl/squeeze_op_builder.cc

@@ -0,0 +1,104 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+#include <core/common/safeint.h>
+
+#include "core/providers/common.h"
+#include "core/providers/shared/utils/utils.h"
+#include "core/providers/coreml/builders/model_builder.h"
+#include "core/providers/coreml/builders/op_builder_factory.h"
+
+#include "base_op_builder.h"
+
+namespace onnxruntime {
+namespace coreml {
+
+class SqueezeOpBuilder : public BaseOpBuilder {
+  // Add operator related
+ public:
+  void AddInitializersToSkip(ModelBuilder& model_builder, const Node& node) const override;
+
+ private:
+  Status AddToModelBuilderImpl(ModelBuilder& model_builder, const Node& node,
+                               const logging::Logger& logger) const override ORT_MUST_USE_RESULT;
+
+  // Operator support related
+ private:
+  bool IsOpSupportedImpl(const InitializedTensorSet& initializers, const Node& node,
+                         const logging::Logger& logger) const override;
+};
+
+// Add operator related
+
+void SqueezeOpBuilder::AddInitializersToSkip(ModelBuilder& model_builder, const Node& node) const {
+  if (node.SinceVersion() > 12 && node.InputDefs().size() > 1) {
+    model_builder.AddInitializerToSkip(node.InputDefs()[1]->Name());
+  }
+}
+
+/* static */ std::vector<int64_t> GetAxes(ModelBuilder& model_builder, const Node& node) {
+  std::vector<int64_t> axes;
+  // Squeeze opset 13 use input as axes
+  if (node.SinceVersion() > 12) {
+    // If axes is not provided, return an empty axes as default to squeeze all
+    if (node.InputDefs().size() > 1) {
+      const auto& initializers(model_builder.GetInitializerTensors());
+      const auto& axes_tensor = *initializers.at(node.InputDefs()[1]->Name());
+      const int64_t* raw_axes = GetTensorInt64Data(axes_tensor);
+      const auto size = SafeInt<size_t>(axes_tensor.dims()[0]);
+      axes.resize(size);
+      for (size_t i = 0; i < size; i++) {
+        axes[i] = raw_axes[i];
+      }
+    }
+  } else {
+    NodeAttrHelper helper(node);
+    axes = helper.Get("axes", std::vector<int64_t>());
+  }
+
+  return axes;
+}
+
+Status SqueezeOpBuilder::AddToModelBuilderImpl(ModelBuilder& model_builder,
+                                               const Node& node,
+                                               const logging::Logger& /* logger */) const {
+  std::unique_ptr<COREML_SPEC::NeuralNetworkLayer> layer = CreateNNLayer(node);
+
+  auto* coreml_squeeze = layer->mutable_squeeze();
+  std::vector<int64_t> axes = GetAxes(model_builder, node);
+  if (axes.empty()) {
+    coreml_squeeze->set_squeezeall(true);
+  } else {
+    *coreml_squeeze->mutable_axes() = {axes.cbegin(), axes.cend()};
+    coreml_squeeze->set_squeezeall(false);
+  }
+
+  *layer->mutable_input()->Add() = node.InputDefs()[0]->Name();
+  *layer->mutable_output()->Add() = node.OutputDefs()[0]->Name();
+
+  model_builder.AddLayer(std::move(layer));
+  return Status::OK();
+}
+
+// Operator support related
+
+bool SqueezeOpBuilder::IsOpSupportedImpl(const InitializedTensorSet& initializers, const Node& node,
+                                         const logging::Logger& /*logger*/) const {
+  // Squeeze opset 13 uses input 1 as axes, if we have input 1 then it needs to be an initializer
+  if (node.SinceVersion() > 12 && node.InputDefs().size() > 1) {
+    const auto& axes_name = node.InputDefs()[1]->Name();
+    if (!Contains(initializers, axes_name)) {
+      LOGS_DEFAULT(VERBOSE) << "Input axes of Squeeze must be known";
+      return false;
+    }
+  }
+
+  return true;
+}
+
+void CreateSqueezeOpBuilder(const std::string& op_type, OpBuilderRegistrations& op_registrations) {
+  op_registrations.builders.push_back(std::make_unique<SqueezeOpBuilder>());
+  op_registrations.op_builder_map.emplace(op_type, op_registrations.builders.back().get());
+}
+
+}  // namespace coreml
+}  // namespace onnxruntime

onnxruntime/core/providers/coreml/builders/op_builder_factory.cc

@@ -65,6 +65,10 @@ static OpBuilderRegistrations CreateOpBuilderRegistrations() {
     CreateClipOpBuilder("Clip", op_registrations);
   }
 
+  {  // Squeeze
+    CreateSqueezeOpBuilder("Squeeze", op_registrations);
+  }
+
   return op_registrations;
 }
 

onnxruntime/core/providers/coreml/builders/op_builder_factory.h

@@ -30,5 +30,6 @@ void CreateClipOpBuilder(const std::string& op_type, OpBuilderRegistrations& op_
 void CreateActivationOpBuilder(const std::string& op_type, OpBuilderRegistrations& op_registrations);
 void CreatePoolOpBuilder(const std::string& op_type, OpBuilderRegistrations& op_registrations);
 void CreateGemmOpBuilder(const std::string& op_type, OpBuilderRegistrations& op_registrations);
+void CreateSqueezeOpBuilder(const std::string& op_type, OpBuilderRegistrations& op_registrations);
 }  // namespace coreml
 }  // namespace onnxruntime

onnxruntime/test/providers/cpu/tensor/squeeze_op_test.cc

@@ -117,15 +117,20 @@ TEST(SqueezeOpTest, SqueezeNegAxis_2) {
 }
 
 TEST(SqueezeOpTest, Squeeze_2_axes_input) {
-  OpTester test("Squeeze", 13);
-  test.AddInput<float>("data", {1, 4, 1, 1, 2},
-                       std::vector<float>{1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f, 7.0f, 8.0f});
-  test.AddInput<int64_t>("axes", {3}, std::vector<int64_t>{0, 2, 3});
-  test.AddOutput<float>("squeezed", {4, 2},
-                        std::vector<float>{1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f, 7.0f, 8.0f});
-  // Incorrect precision for OpenVINO EP. Will be re-enabled after it's fixed
-  // TensorRT and OpenVINO dont support "axes" input in opset 13, re-enable after
-  test.Run(OpTester::ExpectResult::kExpectSuccess, "", {kOpenVINOExecutionProvider, kTensorrtExecutionProvider});
+  auto run_test = [](bool axes_is_initializer) {
+    OpTester test("Squeeze", 13);
+    test.AddInput<float>("data", {1, 4, 1, 1, 2},
+                         std::vector<float>{1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f, 7.0f, 8.0f});
+    test.AddInput<int64_t>("axes", {3}, std::vector<int64_t>{0, 2, 3}, axes_is_initializer);
+    test.AddOutput<float>("squeezed", {4, 2},
+                          std::vector<float>{1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f, 7.0f, 8.0f});
+    // Incorrect precision for OpenVINO EP. Will be re-enabled after it's fixed
+    // TensorRT and OpenVINO dont support "axes" input in opset 13, re-enable after
+    test.Run(OpTester::ExpectResult::kExpectSuccess, "", {kOpenVINOExecutionProvider, kTensorrtExecutionProvider});
+  };
+
+  run_test(false);
+  run_test(true);  // COREML EP will need axes as an initializer
 }
 
 TEST(SqueezeOpTest, Squeeze_Empty_Axes_opset13) {
@@ -137,17 +142,22 @@ TEST(SqueezeOpTest, Squeeze_Empty_Axes_opset13) {
 }
 
 TEST(SqueezeOpTest, SqueezeNegAxis_axes_input) {
-  OpTester test("Squeeze", 13);
-  test.AddInput<float>("data", {1, 4, 1, 1, 2},
-                       std::vector<float>{1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f, 7.0f, 8.0f});
+  auto run_test = [](bool axes_is_initializer) {
+    OpTester test("Squeeze", 13);
+    test.AddInput<float>("data", {1, 4, 1, 1, 2},
+                         std::vector<float>{1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f, 7.0f, 8.0f});
 
-  test.AddInput<int64_t>("axes", {3}, std::vector<int64_t>{0, -3, -2});
-  test.AddOutput<float>("squeezed", {4, 2},
-                        std::vector<float>{1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f, 7.0f, 8.0f});
+    test.AddInput<int64_t>("axes", {3}, std::vector<int64_t>{0, -3, -2}, axes_is_initializer);
+    test.AddOutput<float>("squeezed", {4, 2},
+                          std::vector<float>{1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f, 7.0f, 8.0f});
 
-  // OpenVINO EP Incorrect precision. Will be re-enabled after its fixed
-  // TensorRT and OpenVINO dont support "axes" input in opset 13, re-enable after
-  test.Run(OpTester::ExpectResult::kExpectSuccess, "", {kOpenVINOExecutionProvider, kTensorrtExecutionProvider});
+    // OpenVINO EP Incorrect precision. Will be re-enabled after its fixed
+    // TensorRT and OpenVINO dont support "axes" input in opset 13, re-enable after
+    test.Run(OpTester::ExpectResult::kExpectSuccess, "", {kOpenVINOExecutionProvider, kTensorrtExecutionProvider});
+  };
+
+  run_test(false);
+  run_test(true);  // COREML EP will need axes as an initializer
 }
 
 // Add 4d input shape test, since NNAPI supports up to 4d input shape