From fa287042ca05debb35b07f65147a9c3a39330231 Mon Sep 17 00:00:00 2001 From: vraspar Date: Tue, 16 Jul 2024 16:34:58 -0700 Subject: [PATCH] Add ML Program support for transpose op (#21364) ### Description Add support for transpose op ### Motivation and Context Enable support for Autodesk model --- .../builders/impl/transpose_op_builder.cc | 28 +++++++++++++++---- .../apple/coreml_supported_mlprogram_ops.md | 1 + 2 files changed, 23 insertions(+), 6 deletions(-) diff --git a/onnxruntime/core/providers/coreml/builders/impl/transpose_op_builder.cc b/onnxruntime/core/providers/coreml/builders/impl/transpose_op_builder.cc index f6a61d55a3d63..831c4cf4d08ba 100644 --- a/onnxruntime/core/providers/coreml/builders/impl/transpose_op_builder.cc +++ b/onnxruntime/core/providers/coreml/builders/impl/transpose_op_builder.cc @@ -3,6 +3,7 @@ #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" @@ -14,13 +15,13 @@ namespace coreml { class TransposeOpBuilder : public BaseOpBuilder { Status AddToModelBuilderImpl(ModelBuilder& model_builder, const Node& node, const logging::Logger& logger) const override; + + bool SupportsMLProgram() const override { return true; } }; Status TransposeOpBuilder::AddToModelBuilderImpl(ModelBuilder& model_builder, const Node& node, const logging::Logger& logger) const { - std::unique_ptr layer = model_builder.CreateNNLayer(node); - NodeAttrHelper helper(node); std::vector perm = helper.Get("perm", std::vector()); std::vector input_shape; @@ -33,12 +34,27 @@ Status TransposeOpBuilder::AddToModelBuilderImpl(ModelBuilder& model_builder, ORT_RETURN_IF_NOT(perm.size() == input_dims, "Perm and input should have same dimension"); } - *layer->mutable_transpose()->mutable_axes() = {perm.cbegin(), perm.cend()}; +#if defined(COREML_ENABLE_MLPROGRAM) + if (model_builder.CreateMLProgram()) { + using namespace CoreML::Specification::MILSpec; + + std::unique_ptr op = model_builder.CreateOperation(node, "transpose"); + AddOperationInput(*op, "x", node.InputDefs()[0]->Name()); + AddOperationInput(*op, "perm", model_builder.AddConstant(op->type(), "perm", perm)); + AddOperationOutput(*op, *node.OutputDefs()[0]); + model_builder.AddOperation(std::move(op)); - *layer->mutable_input()->Add() = node.InputDefs()[0]->Name(); - *layer->mutable_output()->Add() = node.OutputDefs()[0]->Name(); + } else +#endif // defined(COREML_ENABLE_MLPROGRAM) + { + std::unique_ptr layer = model_builder.CreateNNLayer(node); + *layer->mutable_transpose()->mutable_axes() = {perm.cbegin(), perm.cend()}; - model_builder.AddLayer(std::move(layer)); + *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(); } diff --git a/tools/ci_build/github/apple/coreml_supported_mlprogram_ops.md b/tools/ci_build/github/apple/coreml_supported_mlprogram_ops.md index 9961d2668e6f5..1bbb933f66ba4 100644 --- a/tools/ci_build/github/apple/coreml_supported_mlprogram_ops.md +++ b/tools/ci_build/github/apple/coreml_supported_mlprogram_ops.md @@ -20,3 +20,4 @@ Keep in sync with doco generated from /docs/execution-providers/CoreML-Execution |ai.onnx:Sub|| |ai.onnx:Sigmoid|| |ai:onnx:Tanh|| +|ai:onnx:Transpose||