CoreML: Add ML Program Split Op (#21456)
### Description Add support for Split Op ### Motivation and Context Address operator gaps in high priority model. --------- Co-authored-by: Scott McKay <skottmckay@gmail.com> Co-authored-by: Edward Chen <18449977+edgchen1@users.noreply.github.com>
This commit is contained in:
vraspar
GitHub
Родитель
5d78b9a17b
Коммит
07d3be5b0e
|
|
@ -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 @@ class SplitOpBuilder : public BaseOpBuilder {
|
|||
|
||||
// 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,55 +46,98 @@ Status SplitOpBuilder::AddToModelBuilderImpl(ModelBuilder& model_builder,
|
|||
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]);
|
||||
}
|
||||
} else {
|
||||
// even
|
||||
coreml_splitnd->set_numsplits(num_outputs);
|
||||
}
|
||||
}
|
||||
return std::make_tuple(remainder, chunk_size);
|
||||
};
|
||||
|
||||
*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();
|
||||
#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 {
|
||||
// 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 (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();
|
||||
}
|
||||
|
|
@ -99,7 +145,6 @@ Status SplitOpBuilder::AddToModelBuilderImpl(ModelBuilder& model_builder,
|
|||
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 @@ bool SplitOpBuilder::IsOpSupportedImpl(const Node& node, const OpBuilderInputPar
|
|||
|
||||
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) {
|
||||
|
|
|
|||
|
|
@ -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||
|
||||
|
|
|
|||
Загрузка…
Ссылка в новой задаче