Add ML Program support for transpose op (#21364)

### Description
Add support for transpose op



### Motivation and Context
Enable support for Autodesk model
This commit is contained in:
vraspar 2024-07-16 16:34:58 -07:00 committed by GitHub
parent 760a31c848
commit fa287042ca
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 23 additions and 6 deletions

View file

@ -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<COREML_SPEC::NeuralNetworkLayer> layer = model_builder.CreateNNLayer(node);
NodeAttrHelper helper(node);
std::vector<int64_t> perm = helper.Get("perm", std::vector<int64_t>());
std::vector<int64_t> 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;
*layer->mutable_input()->Add() = node.InputDefs()[0]->Name();
*layer->mutable_output()->Add() = node.OutputDefs()[0]->Name();
std::unique_ptr<Operation> 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));
model_builder.AddLayer(std::move(layer));
} else
#endif // defined(COREML_ENABLE_MLPROGRAM)
{
std::unique_ptr<COREML_SPEC::NeuralNetworkLayer> layer = model_builder.CreateNNLayer(node);
*layer->mutable_transpose()->mutable_axes() = {perm.cbegin(), perm.cend()};
*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();
}

View file

@ -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||