CoreML: Add ML Program Split Op

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

@@ -5,6 +5,7 @@
 #include "core/providers/common.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"
@@ -24,6 +25,8 @@
 
   // Split opset 13- uses "split" as attribute. Currently it's not supported.
   int GetMinSupportedOpSet(const Node& /* node */) const override { return 13; }
+
+  bool SupportsMLProgram() const override { return true; }
 };
 
 void SplitOpBuilder::AddInitializersToSkip(ModelBuilder& model_builder, const Node& node) const {
@@ -43,63 +46,105 @@
   ORT_RETURN_IF_NOT(GetShape(*node.InputDefs()[0], data_shape, logger), "Failed to get input shape.");
 
   NodeAttrHelper helper(node);
-  const auto axis = helper.Get("axis", 0);
+  int64_t axis = helper.Get("axis", 0);
 
-  // attribute introduced since opset 18
-  uint64_t num_outputs;
-
-  std::unique_ptr<COREML_SPEC::NeuralNetworkLayer> layer = model_builder.CreateNNLayer(node);
-  auto* coreml_splitnd = layer->mutable_splitnd();
-  coreml_splitnd->set_axis(axis);
-
-  if (input_defs.size() > 1) {
-    // if "split" is explicitly provided as an input
-    const auto& split_tensor = *model_builder.GetInitializerTensors().at(input_defs[1]->Name());
-    Initializer unpacked_tensor(split_tensor);
-    auto split_span = unpacked_tensor.DataAsSpan<uint64_t>();
-    auto split_sizes = split_span.size();
-    num_outputs = narrow<uint64_t>(split_sizes);
-    for (size_t i = 0; i < split_sizes; i++) {
-      coreml_splitnd->add_splitsizes(split_span[i]);
-    }
-  } else if (node.SinceVersion() < 18) {
-    num_outputs = narrow<uint64_t>(node.OutputDefs().size());
-    coreml_splitnd->set_numsplits(num_outputs);
-  } else {
-    // note: for opset 18+ 'num_outputs' is a required attribute
-    num_outputs = narrow<uint64_t>(helper.GetInt64("num_outputs").value());
+  auto calculate_remainder_and_chunk_size = [&](int32_t num_outputs) {
     // note: checked in IsOpSupportedImpl that ensures the dim value at splitting axis exists
     auto split_dim_size = data_shape[HandleNegativeAxis(axis, data_shape.size())];
-    uint64_t chunk_size = narrow<uint64_t>((split_dim_size + num_outputs - 1) / num_outputs);
+    uint64_t chunk_size = (split_dim_size + num_outputs - 1) / num_outputs;
     uint64_t remainder = split_dim_size % chunk_size;
-    if (remainder) {
-      // uneven
-      auto split_sizes = InlinedVector<uint64_t>(num_outputs, chunk_size);
-      split_sizes.back() = remainder;
-      for (size_t i = 0; i < split_sizes.size(); i++) {
-        coreml_splitnd->add_splitsizes(split_sizes[i]);
-      }
+    return std::make_tuple(remainder, chunk_size);
+  };
+
+#if defined(COREML_ENABLE_MLPROGRAM)
+  if (model_builder.CreateMLProgram()) {
+    using namespace CoreML::Specification::MILSpec;
+    std::unique_ptr<Operation> split_op = model_builder.CreateOperation(node, "split");
+    AddOperationInput(*split_op, "axis", model_builder.AddScalarConstant(split_op->type(), "axis", axis));
+
+    if (input_defs.size() > 1) {
+      // if "split" is explicitly provided as an input
+      Initializer unpacked_tensor(*model_builder.GetConstantInitializer(input_defs[1]->Name()));
+      auto split_span = unpacked_tensor.DataAsSpan<int64_t>();
+      AddOperationInput(*split_op, "split_sizes",
+                        model_builder.AddConstant(split_op->type(), "split_sizes", split_span));
+    } else if (node.SinceVersion() < 18) {
+      int64_t num_outputs = narrow<int64_t>(node.OutputDefs().size());
+      AddOperationInput(*split_op, "num_splits",
+                        model_builder.AddScalarConstant(split_op->type(), "num_splits", num_outputs));
     } else {
-      // even
+      // note: for opset 18+ 'num_outputs' is a required attribute
+      int64_t num_outputs = helper.GetInt64("num_outputs").value();
+      auto [remainder, chunk_size] = calculate_remainder_and_chunk_size(static_cast<int32_t>(num_outputs));
+      if (remainder) {
+        // uneven
+        std::vector<int64_t> split_sizes(num_outputs, chunk_size);
+        split_sizes.back() = remainder;
+        AddOperationInput(*split_op, "split_sizes",
+                          model_builder.AddConstant(split_op->type(), "split_sizes", split_sizes));
+      } else {
+        // even
+        AddOperationInput(*split_op, "num_splits",
+                          model_builder.AddScalarConstant(split_op->type(), "num_splits", num_outputs));
+      }
+    }
+
+    AddOperationInput(*split_op, "x", input_defs[0]->Name());
+    for (const auto& output_def : node.OutputDefs()) {
+      AddOperationOutput(*split_op, *output_def);
+    }
+    model_builder.AddOperation(std::move(split_op));
+
+  } else
+#endif
+  {
+    std::unique_ptr<COREML_SPEC::NeuralNetworkLayer> layer = model_builder.CreateNNLayer(node);
+    auto* coreml_splitnd = layer->mutable_splitnd();
+    coreml_splitnd->set_axis(axis);
+
+    if (input_defs.size() > 1) {
+      // if "split" is explicitly provided as an input
+      // const auto& split_tensor = *model_builder.GetInitializerTensors().at(input_defs[1]->Name());
+      Initializer unpacked_tensor(*model_builder.GetConstantInitializer(input_defs[1]->Name()));
+      auto split_span = unpacked_tensor.DataAsSpan<uint64_t>();
+      for (const auto& split_size : split_span) {
+        coreml_splitnd->add_splitsizes(split_size);
+      }
+    } else if (node.SinceVersion() < 18) {
+      uint64_t num_outputs = narrow<uint64_t>(node.OutputDefs().size());
       coreml_splitnd->set_numsplits(num_outputs);
+    } else {
+      // note: for opset 18+ 'num_outputs' is a required attribute
+      uint64_t num_outputs = narrow<uint64_t>(helper.GetInt64("num_outputs").value());
+      auto [remainder, chunk_size] = calculate_remainder_and_chunk_size(static_cast<int32_t>(num_outputs));
+      if (remainder) {
+        // uneven
+        auto split_sizes = InlinedVector<uint64_t>(num_outputs, chunk_size);
+        split_sizes.back() = remainder;
+        for (size_t i = 0; i < split_sizes.size(); i++) {
+          coreml_splitnd->add_splitsizes(split_sizes[i]);
+        }
+      } else {
+        // even
+        coreml_splitnd->set_numsplits(num_outputs);
+      }
     }
-  }
 
-  *layer->mutable_input()->Add() = node.InputDefs()[0]->Name();
-  // variadic number of outputs. Calculated based on the length of the given splitSizes if provided.
-  // Otherwise, uses attribute value 'num_outputs'.
-  for (uint64_t i = 0; i < num_outputs; i++) {
-    *layer->mutable_output()->Add() = node.OutputDefs()[i]->Name();
+    *layer->mutable_input()->Add() = node.InputDefs()[0]->Name();
+    // variadic number of outputs. Calculated based on the length of the given splitSizes if provided.
+    // Otherwise, uses attribute value 'num_outputs'.
+    for (const auto& output_def : node.OutputDefs()) {
+      *layer->mutable_output()->Add() = output_def->Name();
+    }
+    model_builder.AddLayer(std::move(layer));
   }
-  model_builder.AddLayer(std::move(layer));
 
   return Status::OK();
 }
 
 bool SplitOpBuilder::IsOpSupportedImpl(const Node& node, const OpBuilderInputParams& input_params,
                                        const logging::Logger& logger) const {
   const auto& input_defs = node.InputDefs();
-  const auto& initializers = input_params.graph_viewer.GetAllInitializedTensors();
 
   NodeAttrHelper helper(node);
   const auto axis = helper.Get("axis", 0);
@@ -110,16 +155,19 @@
 
   const auto split_dims_at_axis = input_shape[HandleNegativeAxis(axis, input_shape.size())];
   if (input_defs.size() > 1 && input_defs[1]->Exists()) {
-    if (!CheckIsConstantInitializer(*input_defs[1], input_params.graph_viewer, logger, "'split'")) {
+    const auto* splits_tensor = input_params.graph_viewer.GetConstantInitializer(input_defs[1]->Name());
+    if (!splits_tensor) {
+      LOGS(logger, VERBOSE) << "CoreML 'splits' input must be a constant initializer.";
       return false;
     }
+
     const auto split_shape = *input_defs[1]->Shape();
     if (split_shape.dim_size() < 2) {
-      LOGS(logger, VERBOSE) << "CoreML SplitND requires to produce at least 2 outputs.";
+      LOGS(logger, VERBOSE) << "CoreML Split must produce at least 2 outputs.";
       return false;
     }
-    const auto& splits_tensor = *initializers.at(input_defs[1]->Name());
-    Initializer unpacked_tensor(splits_tensor);
+
+    Initializer unpacked_tensor(*splits_tensor);
     auto splits_span = unpacked_tensor.DataAsSpan<int64_t>();
     int64_t sum_of_splits = std::accumulate(splits_span.begin(), splits_span.end(), int64_t{0});
     if (sum_of_splits != split_dims_at_axis) {

tools/ci_build/github/apple/coreml_supported_mlprogram_ops.md

@@ -24,6 +24,7 @@ Keep in sync with doco generated from /docs/execution-providers/CoreML-Execution
 |ai.onnx:Reshape||
 |ai.onnx:Resize|See [resize_op_builder.cc](https://github.com/microsoft/onnxruntime/blob/main/onnxruntime/core/providers/coreml/builders/impl/resize_op_builder.cc) implementation. There are too many permutations to describe the valid combinations.|
 |ai.onnx.Slice|starts/ends/axes/steps must be constant initializers.|
+|ai.onnx:Split||
 |ai.onnx:Sub||
 |ai.onnx:Sigmoid||
 |ai:onnx:Tanh||