[CoreML EP] Add Identity, Ceil, Tile builders + drop trivial-only partitions

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

@@ -22,6 +22,12 @@
 
  public:
   bool SupportsMLProgram() const override { return true; }
+
+  // Cast is shape-only data movement from CoreML's perspective: per-element
+  // dtype conversion that the marshalling overhead dominates for small
+  // tensors. CoreML claims it but a partition consisting only of Casts
+  // doesn't earn its own marshalling cost.
+  bool IsTrivial(const Node& /*node*/) const override { return true; }
 };
 
 Status CastOpBuilder::AddToModelBuilderImpl([[maybe_unused]] ModelBuilder& model_builder,

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

@@ -17,6 +17,8 @@ class FlattenOpBuilder : public BaseOpBuilder {
 
   bool IsOpSupportedImpl(const Node& node, const OpBuilderInputParams& input_params,
                          const logging::Logger& logger) const override;
+
+  bool IsTrivial(const Node& /*node*/) const override { return true; }
 };
 
 Status FlattenOpBuilder::AddToModelBuilderImpl(ModelBuilder& model_builder,

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

@@ -0,0 +1,52 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+#include "core/providers/coreml/builders/helper.h"
+#include "core/providers/coreml/builders/impl/base_op_builder.h"
+#include "core/providers/coreml/builders/impl/builder_utils.h"
+#include "core/providers/coreml/builders/model_builder.h"
+#include "core/providers/coreml/builders/op_builder_factory.h"
+
+namespace onnxruntime {
+namespace coreml {
+
+class IdentityOpBuilder : public BaseOpBuilder {
+  Status AddToModelBuilderImpl(ModelBuilder& model_builder, const Node& node,
+                               const logging::Logger& logger) const override;
+
+  bool SupportsMLProgram() const override { return true; }
+
+  bool IsTrivial(const Node& /*node*/) const override { return true; }
+};
+
+Status IdentityOpBuilder::AddToModelBuilderImpl(ModelBuilder& model_builder, const Node& node,
+                                                const logging::Logger& /*logger*/) const {
+  const auto& input_defs = node.InputDefs();
+  const auto& output_def = *node.OutputDefs()[0];
+
+  if (model_builder.CreateMLProgram()) {
+    using namespace CoreML::Specification::MILSpec;
+    auto op = model_builder.CreateOperation(node, "identity");
+    AddOperationInput(*op, "x", input_defs[0]->Name());
+    AddOperationOutput(*op, output_def);
+    model_builder.AddOperation(std::move(op));
+  } else {
+    // NeuralNetwork: emulate via activation LINEAR(alpha=1, beta=0).
+    auto layer = model_builder.CreateNNLayer(node);
+    auto* linear = layer->mutable_activation()->mutable_linear();
+    linear->set_alpha(1.0f);
+    linear->set_beta(0.0f);
+    *layer->mutable_input()->Add() = input_defs[0]->Name();
+    *layer->mutable_output()->Add() = output_def.Name();
+    model_builder.AddLayer(std::move(layer));
+  }
+  return Status::OK();
+}
+
+void CreateIdentityOpBuilder(const std::string& op_type, OpBuilderRegistrations& op_registrations) {
+  op_registrations.builders.push_back(std::make_unique<IdentityOpBuilder>());
+  op_registrations.op_builder_map.emplace(op_type, op_registrations.builders.back().get());
+}
+
+}  // namespace coreml
+}  // namespace onnxruntime

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

@@ -27,6 +27,8 @@ class ReshapeOpBuilder : public BaseOpBuilder {
   int GetMinSupportedOpSet(const Node& /* node */) const override { return 5; }
 
   bool SupportsMLProgram() const override { return true; }
+
+  bool IsTrivial(const Node& /*node*/) const override { return true; }
 };
 
 void ReshapeOpBuilder::AddInitializersToSkip(ModelBuilder& model_builder, const Node& node) const {

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

@@ -29,6 +29,9 @@ class SqueezeOpBuilder : public BaseOpBuilder {
   bool IsOpSupportedImpl(const Node& node, const OpBuilderInputParams& input_params,
                          const logging::Logger& logger) const override;
   bool SupportsMLProgram() const override { return true; }
+
+  // SqueezeOpBuilder handles both Squeeze and Unsqueeze; both are shape-only.
+  bool IsTrivial(const Node& /*node*/) const override { return true; }
 };
 
 namespace {

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

@@ -0,0 +1,156 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+#include "core/optimizer/initializer.h"
+#include "core/providers/coreml/builders/helper.h"
+#include "core/providers/coreml/builders/impl/base_op_builder.h"
+#include "core/providers/coreml/builders/impl/builder_utils.h"
+#include "core/providers/coreml/builders/model_builder.h"
+#include "core/providers/coreml/builders/op_builder_factory.h"
+#include "core/providers/coreml/shape_utils.h"
+#include "core/providers/shared/utils/utils.h"
+
+namespace onnxruntime {
+namespace coreml {
+
+class TileOpBuilder : public BaseOpBuilder {
+  void AddInitializersToSkip(ModelBuilder& model_builder, const Node& node) const override;
+
+  Status AddToModelBuilderImpl(ModelBuilder& model_builder, const Node& node,
+                               const logging::Logger& logger) const override;
+
+  bool HasSupportedInputsImpl(const Node& node, const OpBuilderInputParams& input_params,
+                              const logging::Logger& logger) const override;
+
+  bool IsOpSupportedImpl(const Node& node, const OpBuilderInputParams& input_params,
+                         const logging::Logger& logger) const override;
+
+  bool SupportsMLProgram() const override { return true; }
+
+  bool IsTrivial(const Node& /*node*/) const override { return true; }
+};
+
+void TileOpBuilder::AddInitializersToSkip(ModelBuilder& model_builder, const Node& node) const {
+  // If 'repeats' is a constant initializer we bake it into the MIL constant
+  // and don't need the original to land in the model. If it's a runtime
+  // tensor the dynamic-shape MIL path consumes it directly.
+  if (model_builder.GetConstantInitializer(node.InputDefs()[1]->Name())) {
+    model_builder.AddInitializerToSkip(node.InputDefs()[1]->Name());
+  }
+}
+
+Status TileOpBuilder::AddToModelBuilderImpl(ModelBuilder& model_builder, const Node& node,
+                                            const logging::Logger& /*logger*/) const {
+  const auto& input_defs = node.InputDefs();
+  const auto& output_def = *node.OutputDefs()[0];
+  const auto* repeats_init = model_builder.GetConstantInitializer(input_defs[1]->Name());
+
+  if (model_builder.CreateMLProgram()) {
+    using namespace CoreML::Specification::MILSpec;
+    auto op = model_builder.CreateOperation(node, "tile");
+    AddOperationInput(*op, "x", input_defs[0]->Name());
+    if (repeats_init) {
+      Initializer unpacked(model_builder.GetGraphViewer().GetGraph(), *repeats_init);
+      auto repeats = unpacked.DataAsSpan<int64_t>();
+      AddOperationInput(*op, "reps", model_builder.AddConstant(op->type(), "reps", repeats));
+    } else {
+      // Runtime 'reps' (e.g. emitted by a Loop). Pass the tensor through.
+      AddOperationInput(*op, "reps", input_defs[1]->Name());
+    }
+    AddOperationOutput(*op, output_def);
+    model_builder.AddOperation(std::move(op));
+  } else {
+    if (!repeats_init) {
+      return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
+                             "TileOpBuilder NeuralNetwork path requires constant 'repeats'");
+    }
+    Initializer unpacked(model_builder.GetGraphViewer().GetGraph(), *repeats_init);
+    auto repeats = unpacked.DataAsSpan<int64_t>();
+    auto layer = model_builder.CreateNNLayer(node);
+    auto* tile_params = layer->mutable_tile();
+    for (int64_t r : repeats) {
+      tile_params->add_reps(r);
+    }
+    *layer->mutable_input()->Add() = input_defs[0]->Name();
+    *layer->mutable_output()->Add() = output_def.Name();
+    model_builder.AddLayer(std::move(layer));
+  }
+  return Status::OK();
+}
+
+bool TileOpBuilder::HasSupportedInputsImpl(const Node& node, const OpBuilderInputParams& input_params,
+                                           const logging::Logger& logger) const {
+  // Tile is shape-only data movement, so it can carry any element type CoreML
+  // can represent. ONNX Tile is commonly used in graph post-processing on
+  // INT32 grid-index tensors (e.g. YOLO anchor expansion), which the default
+  // base check (float-only) would reject.
+  int32_t input_type;
+  if (!GetType(*node.InputDefs()[0], input_type, logger)) {
+    return false;
+  }
+  switch (input_type) {
+    case ONNX_NAMESPACE::TensorProto_DataType_FLOAT:
+    case ONNX_NAMESPACE::TensorProto_DataType_INT32:
+    case ONNX_NAMESPACE::TensorProto_DataType_INT64:
+    case ONNX_NAMESPACE::TensorProto_DataType_BOOL:
+      return true;
+    case ONNX_NAMESPACE::TensorProto_DataType_FLOAT16:
+      if (input_params.create_mlprogram && input_params.coreml_version >= 6) {
+        return true;
+      }
+      [[fallthrough]];
+    default:
+      LOGS(logger, VERBOSE) << "[Tile] input type " << input_type << " is not supported";
+      return false;
+  }
+}
+
+bool TileOpBuilder::IsOpSupportedImpl(const Node& node, const OpBuilderInputParams& input_params,
+                                      const logging::Logger& logger) const {
+  const auto& input_defs = node.InputDefs();
+
+  // The NeuralNetwork emitter only supports constant 'repeats'; the MLProgram
+  // path also accepts a runtime 'reps' tensor.
+  const auto& repeats_name = input_defs[1]->Name();
+  const auto* repeats_tensor = input_params.graph_viewer.GetConstantInitializer(repeats_name);
+  if (!input_params.create_mlprogram && !repeats_tensor) {
+    LOGS(logger, VERBOSE) << "Tile NeuralNetwork path requires 'repeats' to be a constant initializer";
+    return false;
+  }
+
+  std::vector<int64_t> input_shape;
+  if (!GetShape(*input_defs[0], input_shape, logger)) {
+    return false;
+  }
+
+  if (input_shape.size() > 5) {
+    LOGS(logger, VERBOSE) << "Tile does not support input rank greater than 5. Input rank: " << input_shape.size();
+    return false;
+  }
+
+  if (repeats_tensor) {
+    Initializer unpacked(input_params.graph_viewer.GetGraph(), *repeats_tensor);
+    auto repeats = unpacked.DataAsSpan<int64_t>();
+    if (repeats.size() != input_shape.size()) {
+      LOGS(logger, VERBOSE) << "Tile 'repeats' length (" << repeats.size()
+                            << ") must match input rank (" << input_shape.size() << ")";
+      return false;
+    }
+    for (int64_t r : repeats) {
+      if (r < 1) {
+        LOGS(logger, VERBOSE) << "Tile 'repeats' values must be positive; got " << r;
+        return false;
+      }
+    }
+  }
+
+  return true;
+}
+
+void CreateTileOpBuilder(const std::string& op_type, OpBuilderRegistrations& op_registrations) {
+  op_registrations.builders.push_back(std::make_unique<TileOpBuilder>());
+  op_registrations.op_builder_map.emplace(op_type, op_registrations.builders.back().get());
+}
+
+}  // namespace coreml
+}  // namespace onnxruntime

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

@@ -17,6 +17,8 @@ class TransposeOpBuilder : public BaseOpBuilder {
                                const logging::Logger& logger) const override;
 
   bool SupportsMLProgram() const override { return true; }
+
+  bool IsTrivial(const Node& /*node*/) const override { return true; }
 };
 
 Status TransposeOpBuilder::AddToModelBuilderImpl(ModelBuilder& model_builder,

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

@@ -18,6 +18,11 @@
   bool SupportsMLProgram() const override { return true; }
   bool IsOpSupportedImpl(const Node& node, const OpBuilderInputParams& input_params,
                          const logging::Logger& logger) const override;
+
+  // Of the unary ops this builder handles, only Ceil is cheap enough to count
+  // as trivial. Erf/Round/Exp/Reciprocal/Sqrt are all transcendental or
+  // multi-cycle ops and earn their own marshalling cost.
+  bool IsTrivial(const Node& node) const override { return node.OpType() == "Ceil"; }
 };
 
 Status UnaryOpBuilder::AddToModelBuilderImpl(ModelBuilder& model_builder, const Node& node,
@@ -39,6 +44,8 @@
       coreml_op_type = "round";
     } else if (op_type == "Exp") {
       coreml_op_type = "exp";
+    } else if (op_type == "Ceil") {
+      coreml_op_type = "ceil";
     } else {
       return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
                              "UnaryOpBuilder::AddToModelBuilderImpl, unexpected op: ", op_type);
@@ -82,7 +89,8 @@
 bool UnaryOpBuilder::IsOpSupportedImpl(const Node& node, const OpBuilderInputParams& input_params,
                                        const logging::Logger& /*logger*/) const {
   if (!input_params.create_mlprogram) {
-    if (node.OpType() == "Erf" || node.OpType() == "Round" || node.OpType() == "Exp") {
+    if (node.OpType() == "Erf" || node.OpType() == "Round" || node.OpType() == "Exp" ||
+        node.OpType() == "Ceil") {
       return false;
     }
   }

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

@@ -44,6 +44,14 @@
 
   // Does the builder implementation support creating an ML Program?
   virtual bool SupportsMLProgram() const = 0;
+
+  // Is this op cheap enough that a CoreML partition consisting only of nodes
+  // like it isn't worth the marshalling cost? Used by the trivial-only
+  // partition heuristic in CoreMLExecutionProvider::GetCapability. Defaults
+  // to false; trivial-op builders override to true. Some builders dispatch
+  // multiple op types (e.g. UnaryOpBuilder), so the answer can depend on
+  // node.OpType().
+  virtual bool IsTrivial(const Node& /*node*/) const { return false; }
 };
 
 }  // namespace coreml

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

@@ -38,6 +38,7 @@ static OpBuilderRegistrations CreateOpBuilderRegistrations() {
   CreateUnaryOpBuilder("Round", op_registrations);
   CreateUnaryOpBuilder("Sqrt", op_registrations);
   CreateUnaryOpBuilder("Exp", op_registrations);
+  CreateUnaryOpBuilder("Ceil", op_registrations);
 
   // Binary elementwise ops
   CreateBinaryOpBuilder("Add", op_registrations);
@@ -77,6 +78,7 @@ static OpBuilderRegistrations CreateOpBuilderRegistrations() {
   CreateGatherOpBuilder("Gather", op_registrations);
   CreateGemmOpBuilder("Gemm", op_registrations);
   CreateGridSampleOpBuilder("GridSample", op_registrations);
+  CreateIdentityOpBuilder("Identity", op_registrations);
   CreateLRNOpBuilder("LRN", op_registrations);
   CreateGemmOpBuilder("MatMul", op_registrations);
   CreatePadOpBuilder("Pad", op_registrations);
@@ -87,6 +89,7 @@ static OpBuilderRegistrations CreateOpBuilderRegistrations() {
   CreateSplitOpBuilder("Split", op_registrations);
   CreateSoftmaxOpBuilder("Softmax", op_registrations);
   CreateSqueezeOpBuilder("Squeeze", op_registrations);
+  CreateTileOpBuilder("Tile", op_registrations);
   CreateTransposeOpBuilder("Transpose", op_registrations);
   CreateSqueezeOpBuilder("Unsqueeze", op_registrations);
 

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

@@ -31,6 +31,7 @@ void CreateFlattenOpBuilder(const std::string& op_type, OpBuilderRegistrations&
 void CreateGatherOpBuilder(const std::string& op_type, OpBuilderRegistrations& op_registrations);
 void CreateGemmOpBuilder(const std::string& op_type, OpBuilderRegistrations& op_registrations);
 void CreateGridSampleOpBuilder(const std::string& op_type, OpBuilderRegistrations& op_registrations);
+void CreateIdentityOpBuilder(const std::string& op_type, OpBuilderRegistrations& op_registrations);
 void CreateLRNOpBuilder(const std::string& op_type, OpBuilderRegistrations& op_registrations);
 void CreatePadOpBuilder(const std::string& op_type, OpBuilderRegistrations& op_registrations);
 void CreatePoolOpBuilder(const std::string& op_type, OpBuilderRegistrations& op_registrations);
@@ -42,6 +43,7 @@ void CreateSliceOpBuilder(const std::string& op_type, OpBuilderRegistrations& op
 void CreateSoftmaxOpBuilder(const std::string& op_type, OpBuilderRegistrations& op_registrations);
 void CreateSplitOpBuilder(const std::string& op_type, OpBuilderRegistrations& op_registrations);
 void CreateSqueezeOpBuilder(const std::string& op_type, OpBuilderRegistrations& op_registrations);
+void CreateTileOpBuilder(const std::string& op_type, OpBuilderRegistrations& op_registrations);
 void CreateTransposeOpBuilder(const std::string& op_type, OpBuilderRegistrations& op_registrations);
 void CreateUnaryOpBuilder(const std::string& op_type, OpBuilderRegistrations& op_registrations);
 void CreateQuickGeluOpBuilder(const std::string& op_type, OpBuilderRegistrations& op_registrations);

onnxruntime/core/providers/coreml/coreml_execution_provider.cc

@@ -11,6 +11,7 @@
 #include "core/framework/tensorprotoutils.h"
 #include "core/graph/graph_viewer.h"
 #include "core/providers/coreml/builders/helper.h"
+#include "core/providers/coreml/builders/op_builder_factory.h"
 #include "core/providers/partitioning_utils.h"
 #include "core/session/onnxruntime_cxx_api.h"
 
@@ -88,9 +89,32 @@
         return MakeString(user_provided_key, "_", COREML, "_", model_hash, "_", metadef_id);
       };
 
-  result = utils::CreateSupportedPartitions(graph_viewer, supported_nodes, {},
+  // Drop CoreML partitions that consist entirely of trivial shape / cheap-elementwise ops.
+  // These ops can each be claimed individually but the CPU↔CoreML round-trip cost
+  // (~50-100us marshalling) outweighs the saving when the partition has no compute-heavy
+  // op to amortise it over. Per-op CoreML dispatch cost is ~10-14us on M3 Max even for
+  // trivial ops (Identity/Ceil/Tile etc.), and CPU runs them in <1us each.
+  //
+  // The "trivial" marker lives on each op builder's IOpBuilder::IsTrivial(node)
+  // override rather than as a hardcoded set here, so adding a new trivial op
+  // builder doesn't risk drifting from a list maintained at the EP level.
+  const auto& op_builders = coreml::GetOpBuilders();
+  const auto is_node_trivial = [&](const Node* node) -> bool {
+    auto it = op_builders.find(node->OpType());
+    return it != op_builders.end() && it->second->IsTrivial(*node);
+  };
+  const auto is_node_supported = [&](const Node& node) -> bool {
+    return supported_nodes.find(&node) != supported_nodes.end();
+  };
+  const auto on_group_closed = [&](const std::vector<const Node*>& group) -> bool {
+    // Keep the partition only if at least one node is non-trivial.
+    return std::any_of(group.begin(), group.end(),
+                       [&](const Node* node) { return !is_node_trivial(node); });
+  };
+
+  result = utils::CreateSupportedPartitions(graph_viewer, is_node_supported, on_group_closed,
                                             gen_metadef_name, COREML, kCoreMLExecutionProvider,
-                                            nullptr,
+                                            /*node_unit_map*/ nullptr,
                                             /*drop_constant_initializers*/ true);
 
   const auto num_of_partitions = result.size();