address comments

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

@@ -106,7 +106,7 @@ bool BaseOpBuilder::IsInputDtypeSupport(const Node& node, size_t idx,
     return true;
   }
 
-  // only support MLProgram for FP16
+  // only MLProgram support FP16
   if (!input_params.create_mlprogram && input_type == ONNX_NAMESPACE::TensorProto_DataType_FLOAT16) {
     return false;
   }

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

@@ -56,6 +56,44 @@ bool CheckIfBothInputShapesMatch(const Node& node, const logging::Logger& logger
 }
 }  // namespace
 
+// Add variadic inputs to the model builder
+// in onnx spec, some node allows variadic inputs, such as max(x, y, z, ...)
+// while in coreml, maximum op only allows two inputs maximum(x, y)
+// the conversion is doing the following:
+// max(x, y, z, ...) -> max(max(x, y), z, ...)
+#if defined(COREML_ENABLE_MLPROGRAM)
+static void AddVariadicInputsToMB(std::unique_ptr<CoreML::Specification::MILSpec::Operation>* op, ModelBuilder& model_builder, const Node& node,
+                                  const logging::Logger& logger) {
+  using namespace CoreML::Specification::MILSpec;
+  const auto& input_defs(node.InputDefs());
+  std::string_view layer_input_name_x = model_builder.GetUniqueName(node, "variadic");
+  auto input_dtype = input_defs[0]->TypeAsProto()->tensor_type().elem_type();
+  const int32_t elem_type = static_cast<int32_t>(input_dtype);
+  std::vector<int64_t> x0_shape, x1_shape;
+  GetShape(*input_defs[0], x0_shape, logger);
+  GetShape(*input_defs[1], x1_shape, logger);
+  if (x0_shape.size() < x1_shape.size()) {
+    std::swap(x0_shape, x1_shape);
+  }
+  std::fill(x0_shape.begin(), x0_shape.end(), -1);
+  std::unique_ptr<Operation> op_prev = std::move(*op);
+  for (size_t i = 2; i < input_defs.size(); i++) {
+    AddIntermediateOperationOutput(*op_prev, layer_input_name_x, elem_type, x0_shape);
+    std::unique_ptr<Operation> op_cur = model_builder.CreateOperation(node, op_prev->type());
+    AddOperationInput(*op_cur, "x", layer_input_name_x);
+    AddOperationInput(*op_cur, "y", input_defs[i]->Name());
+    model_builder.AddOperation(std::move(op_prev));
+    op_prev = std::move(op_cur);
+    layer_input_name_x = model_builder.GetUniqueName(node, "variadic");
+    GetShape(*input_defs[i], x1_shape, logger);
+    if (x0_shape.size() < x1_shape.size()) {
+      x0_shape.resize(x1_shape.size(), -1);
+    }
+  }
+  *op = std::move(op_prev);
+}
+#endif
+
 Status BinaryOpBuilder::AddToModelBuilderImpl(ModelBuilder& model_builder, const Node& node,
                                               const logging::Logger& logger) const {
   const auto& op_type(node.OpType());
@@ -90,30 +128,8 @@ Status BinaryOpBuilder::AddToModelBuilderImpl(ModelBuilder& model_builder, const
     AddOperationInput(*op, "x", input_defs[0]->Name());
     AddOperationInput(*op, "y", input_defs[1]->Name());
     if (input_defs.size() > 2) {
-      std::string_view layer_input_name_x = model_builder.GetUniqueName(node, "variadic");
-      auto input_dtype = input_defs[0]->TypeAsProto()->tensor_type().elem_type();
-      const int32_t elem_type = static_cast<int32_t>(input_dtype);
-      std::vector<int64_t> x0_shape, x1_shape;
-      GetShape(*input_defs[0], x0_shape, logger);
-      GetShape(*input_defs[1], x1_shape, logger);
-      for (size_t i = 0; i < x0_shape.size(); i++) {
-        x0_shape[i] = std::max(x0_shape[i], x1_shape[i]);
-      }
-      std::unique_ptr<Operation> op_prev = std::move(op);
-      for (size_t i = 2; i < input_defs.size(); i++) {
-        AddIntermediateOperationOutput(*op_prev, layer_input_name_x, elem_type, x0_shape);
-        std::unique_ptr<Operation> op_cur = model_builder.CreateOperation(node, coreml_op_type);
-        AddOperationInput(*op_cur, "x", layer_input_name_x);
-        AddOperationInput(*op_cur, "y", input_defs[i]->Name());
-        model_builder.AddOperation(std::move(op_prev));
-        op_prev = std::move(op_cur);
-        layer_input_name_x = model_builder.GetUniqueName(node, "variadic");
-        GetShape(*input_defs[i], x1_shape, logger);
-        for (size_t i = 0; i < x0_shape.size(); i++) {
-          x0_shape[i] = std::max(x0_shape[i], x1_shape[i]);
-        }
-      }
-      op = std::move(op_prev);
+      // "max" node may have variadic inputs
+      AddVariadicInputsToMB(&op, model_builder, node, logger);
     }
     AddOperationOutput(*op, *node.OutputDefs()[0]);
     model_builder.AddOperation(std::move(op));

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

@@ -80,6 +80,10 @@ Status ReductionOpBuilder::AddToModelBuilderImpl(ModelBuilder& model_builder, co
       coreml_op_type = "reduce_mean";
     } else if (op_type == "ReduceMax") {
       coreml_op_type = "reduce_max";
+    } else if (op_type == "ReduceMin") {
+      coreml_op_type = "reduce_min";
+    } else if (op_type == "ReduceProd") {
+      coreml_op_type = "reduce_prod";
     } else {
       return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
                              "ReductionOpBuilder::AddToModelBuilderImpl, unexpected op: ", op_type);
@@ -120,7 +124,10 @@ Status ReductionOpBuilder::AddToModelBuilderImpl(ModelBuilder& model_builder, co
 bool ReductionOpBuilder::IsOpSupportedImpl(const Node& node, const OpBuilderInputParams& input_params,
                                            const logging::Logger& logger) const {
   const auto& input_defs = node.InputDefs();
-
+  if (!input_params.create_mlprogram &&
+      (node.OpType() == "ReduceMax" || node.OpType() == "ReduceMin" || node.OpType() == "ReduceProd")) {
+    return false;
+  }
   NodeAttrHelper helper(node);
 
   // noop_with_empty_axes defaults to false and is only available in newer opsets where 'axes' is an optional input
@@ -140,12 +147,10 @@ bool ReductionOpBuilder::IsOpSupportedImpl(const Node& node, const OpBuilderInpu
     empty_axes = axes->int64_data_size() == 0;
   }
   if (empty_axes && noop_with_empty_axes && !input_params.create_mlprogram) {
-    LOGS(logger, VERBOSE) << "CoreML doesn't support noop on empty axes for reduction layers";
-    return false;
-  }
-  if (!input_params.create_mlprogram && node.OpType() == "ReduceMax") {
+    LOGS(logger, VERBOSE) << "NeuralNetwork doesn't support noop on empty axes for reduction layers";
     return false;
   }
+
   return true;
 }
 

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

@@ -74,9 +74,6 @@ Status ShapeOpBuilder::AddToModelBuilderImpl(ModelBuilder& model_builder, const
 bool ShapeOpBuilder::IsOpSupportedImpl(const Node& node, const OpBuilderInputParams& input_params,
                                        const logging::Logger& logger) const {
   const auto* tensor_shape = node.InputDefs()[0]->Shape();
-  if (tensor_shape == nullptr) {
-    return false;
-  }
 
   NodeAttrHelper node_attr_helper{node};
   if (!input_params.create_mlprogram) {
@@ -96,6 +93,10 @@ bool ShapeOpBuilder::IsOpSupportedImpl(const Node& node, const OpBuilderInputPar
     int64_t start = HandleNegativeAxis(node_attr_helper.Get("start", 0), tensor_shape->dim_size());
     size = size - start;
     if (size == 0) {
+      LOGS(logger, VERBOSE) << "Shape does not support slicing when size is 0";
+      return false;
+    } else if (size != tensor_shape->dim_size() && tensor_shape == nullptr) {
+      LOGS(logger, VERBOSE) << "Shape does not support slicing when tensor_shape is not available";
       return false;
     }
   }

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

@@ -38,46 +38,40 @@ Status SoftmaxOpBuilder::AddToModelBuilderImpl(ModelBuilder& model_builder,
   auto axis_nonnegative = HandleNegativeAxis(axis, data_shape.size());
 
 #if defined(COREML_ENABLE_MLPROGRAM)
+  // CoreML's softmax match onnx's softmax behavior since opset 13.
+  // For opset < 13, we need to reshape to 2D and set axis to -1 to simulate onnx softmax behavior.
+  // [B,D,...](onnx softmax opset 12, axis=1)->[B,D*...](CoreML softmax, axis=-1)->[B,D,...](reshape back)
   if (model_builder.CreateMLProgram()) {
     using namespace CoreML::Specification::MILSpec;
     auto input_dtype = node.InputDefs()[0]->TypeAsProto()->tensor_type().elem_type();
     const int32_t elem_type = static_cast<int32_t>(input_dtype);
 
     std::string_view layer_input_name_x = node.InputDefs()[0]->Name();
-    std::vector<int64_t> shape1(2, 1);
-    if (node.SinceVersion() < 13 && axis_nonnegative != static_cast<int64_t>(data_shape.size()) - 1) {
+    const bool need_reshape = node.SinceVersion() < 13 && axis_nonnegative != static_cast<int64_t>(data_shape.size()) - 1;
+    std::vector<int64_t> target_shape;
+    if (need_reshape) {
       // reshape to 2D to simulate onnx softmax behavior
       auto reshape1 = model_builder.CreateOperation(node, "reshape", "pre");
-      for (size_t i = 0, shape1_i = 0; i < data_shape.size(); i++) {
-        if (static_cast<int64_t>(i) == axis_nonnegative) {
-          shape1_i++;
-        }
-        shape1[shape1_i] *= data_shape[i];
-      }
-      if (axis_nonnegative == 0) {
-        shape1[0] = shape1[1];
-        shape1.pop_back();
-      }
-      axis_nonnegative = -1;
+      TensorShape input_shape(data_shape);
+      target_shape.push_back(input_shape.SizeToDimension(axis_nonnegative));
+      target_shape.push_back(input_shape.SizeFromDimension(axis_nonnegative));
+      axis_nonnegative = 1;
       AddOperationInput(*reshape1, "x", layer_input_name_x);
-      AddOperationInput(*reshape1, "shape", model_builder.AddConstant(reshape1->type(), "shape1", shape1));
+      AddOperationInput(*reshape1, "shape", model_builder.AddConstant(reshape1->type(), "shape1", target_shape));
       layer_input_name_x = model_builder.GetUniqueName(node, "ln_reshape1_");
-      AddIntermediateOperationOutput(*reshape1, layer_input_name_x, elem_type, shape1);
+      AddIntermediateOperationOutput(*reshape1, layer_input_name_x, elem_type, target_shape);
       model_builder.AddOperation(std::move(reshape1));
     }
     std::unique_ptr<Operation> op = model_builder.CreateOperation(node, "softmax");
     AddOperationInput(*op, "x", layer_input_name_x);
     AddOperationInput(*op, "axis", model_builder.AddScalarConstant(op->type(), "axis", axis_nonnegative));
-    std::string_view ln_output_name;
-    if (node.SinceVersion() < 13 && axis_nonnegative != static_cast<int64_t>(data_shape.size()) - 1) {
-      ln_output_name = model_builder.GetUniqueName(node, "ln_reshape1_");
-      AddIntermediateOperationOutput(*op, ln_output_name, elem_type, shape1);
-    } else {
+    if (!need_reshape) {
       AddOperationOutput(*op, *node.OutputDefs()[0]);
-    }
-    model_builder.AddOperation(std::move(op));
-
-    if (node.SinceVersion() < 13 && axis_nonnegative != static_cast<int64_t>(data_shape.size()) - 1) {
+      model_builder.AddOperation(std::move(op));
+    } else {
+      std::string_view ln_output_name = model_builder.GetUniqueName(node, "ln_reshape1_");
+      AddIntermediateOperationOutput(*op, ln_output_name, elem_type, target_shape);
+      model_builder.AddOperation(std::move(op));
       auto reshape2 = model_builder.CreateOperation(node, "reshape", "post");
       AddOperationInput(*reshape2, "x", ln_output_name);
       AddOperationInput(*reshape2, "shape", model_builder.AddConstant(reshape2->type(), "shape2", data_shape));

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

@@ -11,6 +11,7 @@
 #include "core/providers/coreml/shape_utils.h"
 #include "core/providers/shared/utils/utils.h"
 #include "core/optimizer/initializer.h"
+#include "core/providers/cpu/tensor/unsqueeze.h"
 
 namespace onnxruntime {
 namespace coreml {
@@ -27,7 +28,7 @@ class SqueezeOpBuilder : public BaseOpBuilder {
 };
 
 namespace {
-void GetAxes(ModelBuilder& model_builder, const Node& node, std::vector<int64_t>& axes) {
+void GetAxes(ModelBuilder& model_builder, const Node& node, TensorShapeVector& 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
@@ -41,7 +42,8 @@ void GetAxes(ModelBuilder& model_builder, const Node& node, std::vector<int64_t>
     }
   } else {
     NodeAttrHelper helper(node);
-    axes = helper.Get("axes", std::vector<int64_t>());
+    auto axes_attr = helper.Get("axes", std::vector<int64_t>());
+    axes.assign(axes_attr.begin(), axes_attr.end());
   }
 }
 }  // namespace
@@ -58,7 +60,7 @@ Status SqueezeOpBuilder::AddToModelBuilderImpl(ModelBuilder& model_builder,
   std::unique_ptr<COREML_SPEC::NeuralNetworkLayer> layer = model_builder.CreateNNLayer(node);
   const auto& input_defs(node.InputDefs());
   auto* coreml_squeeze = layer->mutable_squeeze();
-  std::vector<int64_t> axes;
+  TensorShapeVector axes;
   GetAxes(model_builder, node, axes);
   std::vector<int64_t> input_shape;
   GetShape(*input_defs[0], input_shape, logger);
@@ -72,24 +74,12 @@ Status SqueezeOpBuilder::AddToModelBuilderImpl(ModelBuilder& model_builder,
 
     if (coreml_op_type == "squeeze") {
       if (!axes.empty()) {
-        AddOperationInput(*op, "axes", model_builder.AddConstant(op->type(), "axes", axes));
+        // coreml squeeze op does support negative axes
+        AddOperationInput(*op, "axes", model_builder.AddConstant(op->type(), "axes", AsSpan(axes)));
       }
     } else {
-      for (auto& axis : axes) {
-        axis = HandleNegativeAxis(axis, input_shape.size() + axes.size());
-      }
-      std::vector<int64_t> new_shape(axes.size() + input_shape.size(), 1);
-      std::sort(axes.begin(), axes.end());
-      // For example: Given an input tensor (data) of shape [3, 4, 5],
-      // then Unsqueeze(data, axes=[0, 4]) outputs a tensor (expanded) containing same data as data but with shape [1, 3, 4, 5, 1].
-      for (size_t i = 0, ori_i = 0, axes_i = 0; i < new_shape.size(); i++) {
-        if ((axes_i >= axes.size() || static_cast<int64_t>(i) != axes[axes_i]) && input_shape.size() >= ori_i) {
-          new_shape[i] = input_shape[ori_i++];
-        } else {
-          axes_i++;
-        }
-      }
-      AddOperationInput(*op, "shape", model_builder.AddConstant(op->type(), "shape", new_shape));
+      TensorShapeVector output_shape = UnsqueezeBase::ComputeOutputShape(TensorShape(input_shape), axes);
+      AddOperationInput(*op, "shape", model_builder.AddConstant(op->type(), "shape", AsSpan(output_shape)));
     }
     AddOperationOutput(*op, *node.OutputDefs()[0]);
     model_builder.AddOperation(std::move(op));
@@ -118,7 +108,7 @@ bool SqueezeOpBuilder::IsOpSupportedImpl(const Node& node, const OpBuilderInputP
   if (node.SinceVersion() > 12 && input_defs.size() > 1) {
     const auto& axes_name = input_defs[1]->Name();
     if (!input_params.graph_viewer.GetConstantInitializer(axes_name)) {
-      LOGS(logger, VERBOSE) << "Input axes of Squeeze must be known";
+      LOGS(logger, VERBOSE) << "Input axes must be known";
       return false;
     }
   }
@@ -127,19 +117,19 @@ bool SqueezeOpBuilder::IsOpSupportedImpl(const Node& node, const OpBuilderInputP
     if (!input_params.create_mlprogram) {
       return false;
     }
-    int64_t rank = -1;
+
+    int64_t num_of_new_dims = 0;
     if (node.SinceVersion() > 12) {
-      const auto& axes_tensor = *input_params.graph_viewer.GetConstantInitializer(input_defs[1]->Name());
-      Initializer unpacked_tensor(axes_tensor);
-      rank = unpacked_tensor.size();
+      num_of_new_dims = node.InputDefs()[1]->Shape()->dim(0).dim_value();
     } else {
       NodeAttrHelper helper(node);
       auto axes = helper.Get("axes", std::vector<int64_t>());
-      rank = static_cast<int64_t>(axes.size());
+      num_of_new_dims = static_cast<int64_t>(axes.size());
     }
+
     std::vector<int64_t> input_shape;
-    if (!GetShape(*input_defs[0], input_shape, logger) || input_shape.size() + rank > 5) {
-      LOGS(logger, VERBOSE) << "Unsqueeze with rank > 5 is not supported";
+    if (!GetShape(*input_defs[0], input_shape, logger) || input_shape.size() + num_of_new_dims > 5) {
+      LOGS(logger, VERBOSE) << "Unsqueeze with num_of_new_dims > 5 is not supported";
       return false;
     }
   }

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

@@ -45,7 +45,9 @@ static OpBuilderRegistrations CreateOpBuilderRegistrations() {
 
   // Reduction ops
   CreateReductionOpBuilder("ReduceMean", op_registrations);
+  CreateReductionOpBuilder("ReduceMin", op_registrations);
   CreateReductionOpBuilder("ReduceMax", op_registrations);
+  CreateReductionOpBuilder("ReduceProd", op_registrations);
   CreateReductionOpBuilder("ReduceSum", op_registrations);
 
   // Normalization ops

onnxruntime/core/providers/cpu/tensor/unsqueeze.h

@@ -20,15 +20,6 @@ class UnsqueezeBase {
   };
 
   Status PrepareCompute(OpKernelContext* context, Prepare& p) const;
-
- protected:
-  UnsqueezeBase(const OpKernelInfo& info) {
-    size_t num_inputs = info.GetInputCount();
-    if (num_inputs == 1) {  // axes must be a valid attribute
-      ORT_ENFORCE(info.GetAttrs("axes", axes_).IsOK(), "Missing/Invalid 'axes' attribute value");
-    }
-  }
-
   static TensorShapeVector ComputeOutputShape(
       const TensorShape& input_shape,
       const TensorShapeVector& axes) {
@@ -59,6 +50,14 @@ class UnsqueezeBase {
     return output_shape;
   }
 
+ protected:
+  UnsqueezeBase(const OpKernelInfo& info) {
+    size_t num_inputs = info.GetInputCount();
+    if (num_inputs == 1) {  // axes must be a valid attribute
+      ORT_ENFORCE(info.GetAttrs("axes", axes_).IsOK(), "Missing/Invalid 'axes' attribute value");
+    }
+  }
+
   TensorShapeVector axes_;
 };